From f09b9b0ab706cfe25994a7ed35066e61e5af5d8f Mon Sep 17 00:00:00 2001 From: nikochiko Date: Mon, 18 Nov 2019 23:36:54 +0530 Subject: [PATCH 0001/1390] Revert "Update save.py" This reverts commit fdadd0e5e524df6488cd763c4ab7595d469ed1ef. --- tensorflow/python/keras/saving/save.py | 44 +++++--------------------- 1 file changed, 8 insertions(+), 36 deletions(-) diff --git a/tensorflow/python/keras/saving/save.py b/tensorflow/python/keras/saving/save.py index 9f7f5778afe..4be3aa0bbda 100644 --- a/tensorflow/python/keras/saving/save.py +++ b/tensorflow/python/keras/saving/save.py @@ -23,7 +23,6 @@ import os import six from tensorflow.python import tf2 -from tensorflow.python.keras.engine.network import _is_hdf5_filepath from tensorflow.python.keras.saving import hdf5_format from tensorflow.python.keras.saving.saved_model import load as saved_model_load from tensorflow.python.keras.saving.saved_model import save as saved_model_save @@ -37,6 +36,9 @@ except ImportError: h5py = None # pylint: enable=g-import-not-at-top +_HDF5_EXTENSIONS = ['.h5', '.hdf5', '.keras'] + + # TODO(kathywu): Remove this when Keras SavedModel is not experimental. _KERAS_SAVED_MODEL_STILL_EXPERIMENTAL = True @@ -90,42 +92,12 @@ def save_model(model, """ from tensorflow.python.keras.engine import sequential # pylint: disable=g-import-not-at-top - if type(filepath) != str and not isinstance(filepath, h5py.File): - raise ValueError( - 'Expected `filepath` to be a String or `h5py.File` object. Got' - 'unsupported value %s of type %s' - % (filepath, type(filepath))) + default_format = 'tf' if tf2.enabled() else 'h5' + save_format = save_format or default_format - filepath_is_h5py_file = h5py is not None and isinstance(filepath, h5py.File) - filepath_is_h5 = type(filepath) == str and _is_hdf5_filepath(filepath) - if save_format is None: - if (filepath_is_h5 or - (filepath_is_h5py_file)): - save_format = 'h5' - else: - save_format = 'tf' if tf2.enabled() else 'h5' - else: - user_format = save_format.lower().strip() - if user_format in ('tensorflow', 'tf'): - save_format = 'tf' - elif user_format in ('hdf5', 'h5', 'keras'): - save_format = 'h5' - else: - raise ValueError( - 'Unknown format "%s". Was expecting one of {"tf", "h5"}.' % ( - save_format,)) - if save_format == 'tf' and filepath_is_h5: - raise ValueError( - ('`save` got save_format="tf"/"tensorflow", but the ' - 'filepath ("%s") looks like an HDF5 file. Omit the ".h5"/".keras" ' - 'when saving in TensorFlow format.') - % filepath) - if save_format == 'tf' and filepath_is_h5py_file: - raise ValueError( - '`save` got save_format="tf"/"tensorflow", but the given `filepath`' - 'is an `h5py.File` object.') - - if save_format == 'h5': + if (save_format == 'h5' or + (h5py is not None and isinstance(filepath, h5py.File)) or + os.path.splitext(filepath)[1] in _HDF5_EXTENSIONS): # TODO(b/130258301): add utility method for detecting model type. if (not model._is_graph_network and # pylint:disable=protected-access not isinstance(model, sequential.Sequential)): From ae86fd4b96ee934a8af239b5d01381bba8a470f1 Mon Sep 17 00:00:00 2001 From: nikochiko Date: Wed, 20 Nov 2019 15:48:54 +0530 Subject: [PATCH 0002/1390] Update docstring for tf.ensure_shape Update docstring Enhance example. --- tensorflow/python/ops/check_ops.py | 42 ++++++++++++++++++------------ 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index 34106f61fd8..68546d792a8 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -144,7 +144,7 @@ def _unary_assert_doc(sym, sym_name): Raises: InvalidArgumentError: if the check can be performed immediately and - `x {sym}` is False. The check can be performed immediately during + `x {sym}` is False. The check can be performed immediately during eager execution or if `x` is statically known. """.format( sym=sym, sym_name=cap_sym_name, opname=opname) @@ -207,7 +207,7 @@ def _binary_assert_doc(sym): Raises: InvalidArgumentError: if the check can be performed immediately and - `x {sym} y` is False. The check can be performed immediately during + `x {sym} y` is False. The check can be performed immediately during eager execution or if `x` and `y` are statically known. """.format( sym=sym, opname=opname) @@ -2107,24 +2107,30 @@ def ensure_shape(x, shape, name=None): """Updates the shape of a tensor and checks at runtime that the shape holds. For example: - ```python - x = tf.compat.v1.placeholder(tf.int32) - print(x.shape) - ==> TensorShape(None) - y = x * 2 - print(y.shape) - ==> TensorShape(None) + + >>> # tf.placeholder() is not compatible with eager execution + ... + >>> tf.compat.v1.disable_eager_execution() + >>> x = tf.compat.v1.placeholder(tf.int32) + >>> print(x.shape) + TensorShape(None) + >>> y = x * 2 + >>> print(y.shape) + TensorShape(None) y = tf.ensure_shape(y, (None, 3, 3)) print(y.shape) - ==> TensorShape([Dimension(None), Dimension(3), Dimension(3)]) + TensorShape([None, 3, 3]) - with tf.compat.v1.Session() as sess: - # Raises tf.errors.InvalidArgumentError, because the shape (3,) is not - # compatible with the shape (None, 3, 3) - sess.run(y, feed_dict={x: [1, 2, 3]}) + >>> with tf.compat.v1.Session() as sess: + >>> sess.run(y, feed_dict={x: [1, 2, 3]}) + Traceback (most recent call last): + ... + InvalidArgumentError: Shape of tensor mul [3] is not compatible with + expected shape [?,3,3]. - ``` + The above example raises `tf.errors.InvalidArgumentError`, + because the shape (3,) is not compatible with the shape (None, 3, 3) NOTE: This differs from `Tensor.set_shape` in that it sets the static shape of the resulting tensor and enforces it at runtime, raising an error if the @@ -2140,8 +2146,10 @@ def ensure_shape(x, shape, name=None): name: A name for this operation (optional). Defaults to "EnsureShape". Returns: - A `Tensor`. Has the same type and contents as `x`. At runtime, raises a - `tf.errors.InvalidArgumentError` if `shape` is incompatible with the shape + A `Tensor`. Has the same type and contents as `x`. + + Raises: + tf.errors.InvalidArgumentError: If `shape` is incompatible with the shape of `x`. """ if not isinstance(shape, tensor_shape.TensorShape): From c6281f660b3f593054737de0bc9388113181aa4d Mon Sep 17 00:00:00 2001 From: nikochiko Date: Thu, 21 Nov 2019 00:06:38 +0530 Subject: [PATCH 0003/1390] Make requested changes --- tensorflow/python/ops/array_ops.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 6a18d08f22f..b31c28deb28 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -444,7 +444,7 @@ def broadcast_dynamic_shape(shape_x, shape_y): >>> shape_x = [1, 2, 3] >>> shape_y = [5, 1, 3] - >>> broadcast_dynamic_shape(shape_x, shape_y) + >>> tf.broadcast_dynamic_shape(shape_x, shape_y) Args: @@ -479,7 +479,7 @@ def broadcast_static_shape(shape_x, shape_y): >>> shape_x = tf.TensorShape([1, 2, 3]) >>> shape_y = tf.TensorShape([5, 1 ,3]) - >>> broadcast_static_shape(shape_x, shape_y) + >>> tf.broadcast_static_shape(shape_x, shape_y) TensorShape([Dimension(5), Dimension(2), Dimension(3)]) Args: @@ -1556,12 +1556,12 @@ def boolean_mask(tensor, mask, name="boolean_mask", axis=None): # 1-D example tensor = [0, 1, 2, 3] mask = np.array([True, False, True, False]) - boolean_mask(tensor, mask) # [0, 2] + tf.boolean_mask(tensor, mask) # [0, 2] # 2-D example tensor = [[1, 2], [3, 4], [5, 6]] mask = np.array([True, False, True]) - boolean_mask(tensor, mask) # [[1, 2], [5, 6]] + tf.boolean_mask(tensor, mask) # [[1, 2], [5, 6]] ``` Args: @@ -1640,12 +1640,12 @@ def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"): >>> tensor = [0, 1, 2, 3] # 1-D example >>> mask = np.array([True, False, True, False]) - >>> boolean_mask(tensor, mask) + >>> tf.boolean_mask(tensor, mask) >>> tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example >>> mask = np.array([True, False, True]) - >>> boolean_mask(tensor, mask) + >>> tf.boolean_mask(tensor, mask) Args: @@ -3175,7 +3175,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): ... [1, 1, 0]], ... ["a", "b", "c", "a"], ... (2, 2, 2)) - >>> edit_distance(hypothesis, truth, normalize=True) + >>> tf.edit_distance(hypothesis, truth, normalize=True) From d6cff22e271692d770ebc94a8aec6b83e2d533b9 Mon Sep 17 00:00:00 2001 From: nikochiko Date: Thu, 21 Nov 2019 14:47:02 +0530 Subject: [PATCH 0004/1390] Make requested changes - Shapes from lists to tuples - TensorShapes to `TensorShape([...])` --- tensorflow/python/ops/array_ops.py | 21 +++++++++++---------- tensorflow/python/ops/check_ops.py | 10 ++++------ 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index b31c28deb28..a345d290f69 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -442,8 +442,8 @@ def broadcast_dynamic_shape(shape_x, shape_y): Example: - >>> shape_x = [1, 2, 3] - >>> shape_y = [5, 1, 3] + >>> shape_x = (1, 2, 3) + >>> shape_y = (5, 1, 3) >>> tf.broadcast_dynamic_shape(shape_x, shape_y) @@ -469,8 +469,9 @@ def broadcast_static_shape(shape_x, shape_y): `TensorShape` which is the shape of the result of a broadcasting op applied in tensors of shapes `shape_x` and `shape_y`. - For example, if shape_x is [1, 2, 3] and shape_y is [5, 1, 3], the result is a - TensorShape whose value is [5, 2, 3]. + For example, if shape_x is `TensorShape([1, 2, 3])` and shape_y is + `TensorShape([5, 1, 3])`, the result is a TensorShape whose value is + `TensorShape([5, 2, 3])`. This is useful when validating the result of a broadcasting operation when the tensors have statically known shapes. @@ -3160,8 +3161,8 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): For example: Given the following input, - * `hypothesis` is a `tf.SparseTensor` of shape `[2, 1, 1]` - * `truth` is a `tf.SparseTensor` of shape `[2, 2, 2]` + * `hypothesis` is a `tf.SparseTensor` of shape `(2, 1, 1)` + * `truth` is a `tf.SparseTensor` of shape `(2, 2, 2)` >>> hypothesis = tf.SparseTensor( ... [[0, 0, 0], @@ -3180,7 +3181,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): array([[inf, 1. ], [0.5, 1. ]], dtype=float32)> - The operaton returns a dense Tensor of shape `[2, 2]` with + The operaton returns a dense Tensor of shape `(2, 2)` with edit distances normalized by `truth` lengths. **Note**: It is possible to calculate edit distance between two @@ -3190,14 +3191,14 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): For the following inputs, ```python - # 'hypothesis' is a tensor of shape `[2, 1]` with variable-length values: + # 'hypothesis' is a tensor of shape `(2, 1)` with variable-length values: hypothesis = tf.SparseTensor( [[0, 0], [1,0]], ["a", "b"], (2, 1)) - # 'truth' is a tensor of shape `[2, 2]` with variable-length values: + # 'truth' is a tensor of shape `(2, 2)` with variable-length values: truth = tf.SparseTensor( [[0, 1], [1, 0], @@ -3207,7 +3208,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): normalize = True - # The output would be a dense Tensor of shape `[2,]`, with edit distances + # The output would be a dense Tensor of shape `(2,)`, with edit distances noramlized by 'truth' lengths. # output => array([0., 0.5], dtype=float32) ``` diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index 68546d792a8..c76f3939f1e 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -2117,13 +2117,11 @@ def ensure_shape(x, shape, name=None): >>> y = x * 2 >>> print(y.shape) TensorShape(None) - - y = tf.ensure_shape(y, (None, 3, 3)) - print(y.shape) - TensorShape([None, 3, 3]) - + >>> y = tf.ensure_shape(y, (None, 3, 3)) + >>> print(y.shape) + TensorShape([Dimension(None), Dimension(3), Dimension(3)]) >>> with tf.compat.v1.Session() as sess: - >>> sess.run(y, feed_dict={x: [1, 2, 3]}) + >>> sess.run(y, feed_dict={x: [1, 2, 3]}) Traceback (most recent call last): ... InvalidArgumentError: Shape of tensor mul [3] is not compatible with From 9ccf3973b0af7232e83d18a8ea6c6fc27632f0fc Mon Sep 17 00:00:00 2001 From: nikochiko Date: Tue, 26 Nov 2019 21:47:35 +0530 Subject: [PATCH 0005/1390] Minor changes Made few more minor changes to documentation of `tf.extract_volume_patches` and `tf.fill` --- .../api_def_ExtractVolumePatches.pbtxt | 8 +++---- tensorflow/python/ops/array_ops.py | 21 ++++++++++--------- 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/tensorflow/core/api_def/base_api/api_def_ExtractVolumePatches.pbtxt b/tensorflow/core/api_def/base_api/api_def_ExtractVolumePatches.pbtxt index 9c4015eaa4c..32a3c33dc35 100644 --- a/tensorflow/core/api_def/base_api/api_def_ExtractVolumePatches.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_ExtractVolumePatches.pbtxt @@ -34,13 +34,13 @@ END description: < Date: Sat, 14 Dec 2019 22:23:21 +0530 Subject: [PATCH 0006/1390] Update example to use tf.function --- tensorflow/python/ops/check_ops.py | 27 +++++++++------------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index e53a00bc232..ed28b1b093a 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -2121,25 +2121,16 @@ def ensure_shape(x, shape, name=None): For example: - >>> # tf.placeholder() is not compatible with eager execution - ... - >>> tf.compat.v1.disable_eager_execution() - >>> x = tf.compat.v1.placeholder(tf.int32) - >>> print(x.shape) - TensorShape(None) - >>> y = x * 2 - >>> print(y.shape) - TensorShape(None) - >>> y = tf.ensure_shape(y, (None, 3, 3)) - >>> print(y.shape) - TensorShape([Dimension(None), Dimension(3), Dimension(3)]) - >>> with tf.compat.v1.Session() as sess: - >>> sess.run(y, feed_dict={x: [1, 2, 3]}) + >>> @tf.function(input_signature=[tf.TensorSpec(dtype=tf.float32, shape=None)]) + >>> def f(tensor): + >>> return tf.ensure_shape(x, [3, 3]) + >>> + >>> f(tf.zeros([3, 3])) # Passes + >>> f([1, 2, 3]) # fails Traceback (most recent call last): - ... - InvalidArgumentError: Shape of tensor mul [3] is not compatible with - expected shape [?,3,3]. - + ... + InvalidArgumentError: Shape of tensor x [3] is not compatible with expected shape [3,3]. + The above example raises `tf.errors.InvalidArgumentError`, because the shape (3,) is not compatible with the shape (None, 3, 3) From e96d18ac0fa564b666db80bba1e044fb88d7251a Mon Sep 17 00:00:00 2001 From: Stephan Uphoff Date: Sun, 3 Nov 2019 17:11:03 -0700 Subject: [PATCH 0007/1390] Set g_pdm_dma_error_reporter pointer to the error_reporter passed to InitAudioRecording() --- .../micro/examples/micro_speech/apollo3evb/audio_provider.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/micro/examples/micro_speech/apollo3evb/audio_provider.cc b/tensorflow/lite/micro/examples/micro_speech/apollo3evb/audio_provider.cc index 0f9a91a9dba..cf10785071b 100644 --- a/tensorflow/lite/micro/examples/micro_speech/apollo3evb/audio_provider.cc +++ b/tensorflow/lite/micro/examples/micro_speech/apollo3evb/audio_provider.cc @@ -247,7 +247,6 @@ void pdm_start_dma(tflite::ErrorReporter* error_reporter) { // Reset the PDM DMA flags. g_pdm_dma_error = false; - g_pdm_dma_error_reporter = error_reporter; } #if USE_MAYA @@ -460,11 +459,12 @@ TfLiteStatus InitAudioRecording(tflite::ErrorReporter* error_reporter) { #endif // USE_TIME_STAMP // Configure, turn on PDM + g_pdm_dma_error_reporter = error_reporter; pdm_init(); am_hal_interrupt_master_enable(); am_hal_pdm_fifo_flush(g_pdm_handle); // Trigger the PDM DMA for the first time manually. - pdm_start_dma(g_pdm_dma_error_reporter); + pdm_start_dma(error_reporter); error_reporter->Report("\nPDM DMA Threshold = %d", PDMn(0)->FIFOTHR); From 6d6411b339ef2b14f38c57e38c05b2f24122364a Mon Sep 17 00:00:00 2001 From: nikochiko Date: Tue, 17 Dec 2019 16:50:45 +0530 Subject: [PATCH 0008/1390] Fix sanity build errors --- tensorflow/python/ops/array_ops.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 52c265608fe..65a55b5b2bd 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -490,10 +490,10 @@ setdiff1d.__doc__ = gen_array_ops.list_diff.__doc__ def broadcast_dynamic_shape(shape_x, shape_y): """Computes the shape of a broadcast given symbolic shapes. - When `shape_x` and `shape_y` are Tensors representing shapes (i.e. the result of - calling tf.shape on another Tensor) this computes a Tensor which is the shape - of the result of a broadcasting op applied in tensors of shapes `shape_x` and - `shape_y`. + When `shape_x` and `shape_y` are Tensors representing shapes (i.e. the result + of calling tf.shape on another Tensor) this computes a Tensor which is the + shape of the result of a broadcasting op applied in tensors of shapes + `shape_x` and `shape_y`. This is useful when validating the result of a broadcasting operation when the tensors do not have statically known shapes. @@ -503,7 +503,8 @@ def broadcast_dynamic_shape(shape_x, shape_y): >>> shape_x = (1, 2, 3) >>> shape_y = (5, 1, 3) >>> tf.broadcast_dynamic_shape(shape_x, shape_y) - + Args: shape_x: A rank 1 integer `Tensor`, representing the shape of x. @@ -1732,7 +1733,8 @@ def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"): >>> tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example >>> mask = np.array([True, False, True]) >>> tf.boolean_mask(tensor, mask) - + Args: tensor: N-D Tensor. From 2f3f3ee15bbc837a22b33c6177e064b5b92d5fec Mon Sep 17 00:00:00 2001 From: nikochiko Date: Tue, 17 Dec 2019 16:58:07 +0530 Subject: [PATCH 0009/1390] Fix pylint, wrap at 80 chars --- tensorflow/python/ops/array_ops.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 65a55b5b2bd..16dda049420 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -503,8 +503,8 @@ def broadcast_dynamic_shape(shape_x, shape_y): >>> shape_x = (1, 2, 3) >>> shape_y = (5, 1, 3) >>> tf.broadcast_dynamic_shape(shape_x, shape_y) - + Args: shape_x: A rank 1 integer `Tensor`, representing the shape of x. @@ -1733,8 +1733,8 @@ def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"): >>> tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example >>> mask = np.array([True, False, True]) >>> tf.boolean_mask(tensor, mask) - + Args: tensor: N-D Tensor. From d25cff86db264b870439dbf2b6f37ae30096af0b Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Sat, 21 Dec 2019 20:25:12 +0530 Subject: [PATCH 0010/1390] Update array_ops.py --- tensorflow/python/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 16dda049420..d980babf99d 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -503,7 +503,7 @@ def broadcast_dynamic_shape(shape_x, shape_y): >>> shape_x = (1, 2, 3) >>> shape_y = (5, 1, 3) >>> tf.broadcast_dynamic_shape(shape_x, shape_y) - Args: From 6746a6b8317313d952c8c8083be2476c75d7cc23 Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Sat, 21 Dec 2019 20:26:41 +0530 Subject: [PATCH 0011/1390] Update array_ops.py --- tensorflow/python/ops/array_ops.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index d980babf99d..57d953481a9 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -1728,12 +1728,12 @@ def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"): >>> tensor = [0, 1, 2, 3] # 1-D example >>> mask = np.array([True, False, True, False]) >>> tf.boolean_mask(tensor, mask) - + >>> tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example >>> mask = np.array([True, False, True]) >>> tf.boolean_mask(tensor, mask) - Args: @@ -3459,7 +3459,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): ... ["a", "b", "c", "a"], ... (2, 2, 2)) >>> tf.edit_distance(hypothesis, truth, normalize=True) - From 6a89b01c85893ecba281447e4d54c77889cd4d83 Mon Sep 17 00:00:00 2001 From: nikochiko Date: Sat, 4 Jan 2020 13:01:11 +0530 Subject: [PATCH 0012/1390] Fix docstrings --- tensorflow/python/ops/array_ops.py | 7 ++++--- tensorflow/python/ops/check_ops.py | 14 +++++++------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 57d953481a9..a208272a7e2 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -540,7 +540,7 @@ def broadcast_static_shape(shape_x, shape_y): >>> shape_x = tf.TensorShape([1, 2, 3]) >>> shape_y = tf.TensorShape([5, 1 ,3]) >>> tf.broadcast_static_shape(shape_x, shape_y) - TensorShape([Dimension(5), Dimension(2), Dimension(3)]) + TensorShape([5, 2, 3]) Args: shape_x: A `TensorShape` @@ -1733,8 +1733,9 @@ def boolean_mask_v2(tensor, mask, axis=None, name="boolean_mask"): >>> tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example >>> mask = np.array([True, False, True]) >>> tf.boolean_mask(tensor, mask) - + Args: tensor: N-D Tensor. diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index ed28b1b093a..2088dea7c12 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -1595,7 +1595,7 @@ def assert_shapes_v2(shapes, data=None, summarize=None, message=None, >>> n = 10 >>> q = 3 >>> d = 7 - >>> x = tf.zeros([n,q]) + >>> x = tf.zeros([n,q]) >>> y = tf.ones([n,d]) >>> param = tf.Variable([1.0, 2.0, 3.0]) >>> scalar = 1.0 @@ -1605,9 +1605,9 @@ def assert_shapes_v2(shapes, data=None, summarize=None, message=None, ... (param, ('Q',)), ... (scalar, ()), ... ]) - + >>> tf.debugging.assert_shapes([ - ... (x, ('N', 'D')), + ... (x, ('N', 'D')), ... (y, ('N', 'D')) ... ]) Traceback (most recent call last): @@ -2121,16 +2121,16 @@ def ensure_shape(x, shape, name=None): For example: - >>> @tf.function(input_signature=[tf.TensorSpec(dtype=tf.float32, shape=None)]) - >>> def f(tensor): - >>> return tf.ensure_shape(x, [3, 3]) + >>> @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)]) + ... def f(tensor): + ... return tf.ensure_shape(tensor, [3, 3]) >>> >>> f(tf.zeros([3, 3])) # Passes >>> f([1, 2, 3]) # fails Traceback (most recent call last): ... InvalidArgumentError: Shape of tensor x [3] is not compatible with expected shape [3,3]. - + The above example raises `tf.errors.InvalidArgumentError`, because the shape (3,) is not compatible with the shape (None, 3, 3) From aee9e604872ecb6bdd15cd6a1864b2d1108b671b Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Sun, 12 Jan 2020 10:41:37 +0530 Subject: [PATCH 0013/1390] Fix doctest for ensure_shape --- tensorflow/python/ops/check_ops.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py index 2088dea7c12..d38b3958205 100644 --- a/tensorflow/python/ops/check_ops.py +++ b/tensorflow/python/ops/check_ops.py @@ -2126,6 +2126,10 @@ def ensure_shape(x, shape, name=None): ... return tf.ensure_shape(tensor, [3, 3]) >>> >>> f(tf.zeros([3, 3])) # Passes + >>> f([1, 2, 3]) # fails Traceback (most recent call last): ... From a7899d7544230fce8dae4895733d82623af2b934 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 21 Jan 2020 13:18:55 +0000 Subject: [PATCH 0014/1390] Added an option TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 to enable sym quantization with activations in 16-bit and weigths in 8-bit. --- tensorflow/lite/python/convert.py | 6 + tensorflow/lite/python/lite.py | 13 +- tensorflow/lite/python/lite_constants.py | 3 + tensorflow/lite/python/lite_test.py | 14 +- .../python/optimize/calibration_wrapper.cc | 8 +- .../python/optimize/calibration_wrapper.h | 3 +- tensorflow/lite/python/optimize/calibrator.py | 6 +- .../lite/python/optimize/calibrator_test.py | 39 ++- .../lite/tools/optimize/operator_property.cc | 17 +- .../lite/tools/optimize/operator_property.h | 10 +- .../lite/tools/optimize/quantization_utils.cc | 102 +++++-- .../lite/tools/optimize/quantization_utils.h | 10 +- .../tools/optimize/quantization_utils_test.cc | 4 +- .../tools/optimize/quantization_wrapper.cc | 4 +- .../lite/tools/optimize/quantize_model.cc | 175 +++++++----- .../lite/tools/optimize/quantize_model.h | 7 +- .../tools/optimize/quantize_model_test.cc | 258 ++++++++++++------ 17 files changed, 477 insertions(+), 202 deletions(-) diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 2fe4d172487..494f32a515c 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -93,6 +93,12 @@ class OpsSet(enum.Enum): # quantized implementations. TFLITE_BUILTINS_INT8 = "TFLITE_BUILTINS_INT8" + # Convert model using only TensorFlow Lite operations with quantized int8 weights + # and int16 activations. + # Specifying this will throw an error for operations that do not yet have + # quantized implementations. + TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = "TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" + def __str__(self): return self.value diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 657cfea1bb8..fc9c064faf0 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -224,6 +224,10 @@ class TFLiteConverterBase(object): self.target_spec.supported_ops) or self._smallest_supported_type() == constants.INT8) + def _is_int16x8_target_required(self): + return (set([OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]) == + set(self.target_spec.supported_ops)) + def _smallest_supported_type(self): if self.target_spec.supported_types: return min(self.target_spec.supported_types, key=lambda x: x.size) @@ -238,7 +242,9 @@ class TFLiteConverterBase(object): ])) def _is_post_training_optimize(self): - return self._is_int8_target_required() or self._any_optimization_enabled() + return self._is_int8_target_required() or \ + self._is_int16x8_target_required() or \ + self._any_optimization_enabled() def _is_int8_weight_only_quantize(self): return (self._is_post_training_optimize() and @@ -255,11 +261,12 @@ class TFLiteConverterBase(object): def _calibrate_quantize_model(self, result, inference_input_type, inference_output_type, enable_mlir_quantizer): - allow_float = not self._is_int8_target_required() + allow_float = not self._is_int8_target_required() and not self._is_int16x8_target_required() calibrate_quantize = _calibrator.Calibrator(result) + activations_type = constants.INT16 if self._is_int16x8_target_required() else constants.INT8 return calibrate_quantize.calibrate_and_quantize( self.representative_dataset.input_gen, inference_input_type, - inference_output_type, allow_float, enable_mlir_quantizer) + inference_output_type, allow_float, activations_type, enable_mlir_quantizer) def _get_base_converter_args(self): """Returns the base converter args. diff --git a/tensorflow/lite/python/lite_constants.py b/tensorflow/lite/python/lite_constants.py index d43452c775b..4902f23795e 100644 --- a/tensorflow/lite/python/lite_constants.py +++ b/tensorflow/lite/python/lite_constants.py @@ -30,6 +30,7 @@ INT64 = dtypes.int64 STRING = dtypes.string QUANTIZED_UINT8 = dtypes.uint8 INT8 = dtypes.int8 +INT16 = dtypes.int16 COMPLEX64 = dtypes.complex64 TENSORFLOW_GRAPHDEF = _toco_flags_pb2.TENSORFLOW_GRAPHDEF TFLITE = _toco_flags_pb2.TFLITE @@ -43,6 +44,7 @@ _tf_export(v1=["lite.constants.STRING"]).export_constant(__name__, "STRING") _tf_export(v1=["lite.constants.QUANTIZED_UINT8"]).export_constant( __name__, "QUANTIZED_UINT8") _tf_export(v1=["lite.constants.INT8"]).export_constant(__name__, "INT8") +_tf_export(v1=["lite.constants.INT16"]).export_constant(__name__, "INT16") _tf_export(v1=["lite.constants.TFLITE"]).export_constant(__name__, "TFLITE") _tf_export(v1=["lite.constants.GRAPHVIZ_DOT"]).export_constant( __name__, "GRAPHVIZ_DOT") @@ -62,6 +64,7 @@ _allowed_symbols = [ "STRING", "QUANTIZED_UINT8", "INT8", + "INT16", "COMPLEX64", "TENSORFLOW_GRAPHDEF", "TFLITE", diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index 16959c84146..ef5e5d1cdf4 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -769,9 +769,13 @@ class FromSessionTest(TestModels, parameterized.TestCase): self.assertLess(len(quantized_tflite), len(float_tflite)) @parameterized.named_parameters( - ('EnableMlirConverter', True), # enable mlir - ('DisableMlirConverter', False)) # disable mlir - def testCalibrateAndQuantizeBuiltinInt8(self, enable_mlir): + # Quantize model to Int8: with enable mlir + ('UseTfliteBuiltinsIntEnableMLIR', [lite.OpsSet.TFLITE_BUILTINS_INT8], True), + # Quantize model to Int8: with disable mlir + ('UseTfliteBuiltinsIntDisableMLIR', [lite.OpsSet.TFLITE_BUILTINS_INT8], False), + # Quantize model to Int16: with disable mlir + ('UseTfliteBuiltinsInt16DisableMLIR', [lite.OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], False)) + def testCalibrateAndQuantizeBuiltinInt(self, supported_ops, enable_mlir): with ops.Graph().as_default(): inp, output, calibration_gen = self._getCalibrationQuantizeModel() sess = session.Session() @@ -787,9 +791,7 @@ class FromSessionTest(TestModels, parameterized.TestCase): quantized_converter = lite.TFLiteConverter.from_session( sess, [inp], [output]) quantized_converter.experimental_new_converter = enable_mlir - quantized_converter.target_spec.supported_ops = [ - lite.OpsSet.TFLITE_BUILTINS_INT8 - ] + quantized_converter.target_spec.supported_ops = supported_ops quantized_converter.representative_dataset = calibration_gen quantized_tflite = quantized_converter.convert() self.assertTrue(quantized_tflite) diff --git a/tensorflow/lite/python/optimize/calibration_wrapper.cc b/tensorflow/lite/python/optimize/calibration_wrapper.cc index 89ffb3430ea..88995136726 100644 --- a/tensorflow/lite/python/optimize/calibration_wrapper.cc +++ b/tensorflow/lite/python/optimize/calibration_wrapper.cc @@ -204,6 +204,7 @@ PyObject* CalibrationWrapper::SetTensor(int index, PyObject* value) { PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, int output_py_type, bool allow_float, + int activations_py_type, bool enable_mlir_quantizer) { if (NoOpModel(*model_)) { return python_utils::ConvertToPyString(model_str_->data(), @@ -212,6 +213,9 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, TfLiteType input_type = python_utils::TfLiteTypeFromPyType(input_py_type); TfLiteType output_type = python_utils::TfLiteTypeFromPyType(output_py_type); + TfLiteType activations_type = + python_utils::TfLiteTypeFromPyType(activations_py_type); + if (input_type == kTfLiteNoType || output_type == kTfLiteNoType) { PyErr_SetString(PyExc_ValueError, "Input/output type cannot be kTfLiteNoType"); @@ -230,7 +234,7 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, status = tflite::optimize::QuantizeModel( &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), TfLiteTypeToSchemaType(output_type), allow_float, - error_reporter_.get()); + TfLiteTypeToSchemaType(activations_type), error_reporter_.get()); } if (status != kTfLiteOk) { @@ -262,7 +266,7 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, auto status = tflite::optimize::QuantizeModel( &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), TfLiteTypeToSchemaType(output_type), allow_float, {op_name}, - error_reporter_.get()); + TensorType_INT8, error_reporter_.get()); if (status != kTfLiteOk) { error_reporter_->exception(); return nullptr; diff --git a/tensorflow/lite/python/optimize/calibration_wrapper.h b/tensorflow/lite/python/optimize/calibration_wrapper.h index 0fefc29dd81..e72fe15e958 100644 --- a/tensorflow/lite/python/optimize/calibration_wrapper.h +++ b/tensorflow/lite/python/optimize/calibration_wrapper.h @@ -60,7 +60,8 @@ class CalibrationWrapper { PyObject* FeedTensor(PyObject* input_value); PyObject* QuantizeModel(int input_py_type, int output_py_type, - bool allow_float, bool enable_mlir_quantizer = false); + bool allow_float, int activations_py_type, + bool enable_mlir_quantizer = false); // Allows quantizing only the operator that produces the tensor with name // operator_output_name. (This can be used to help debug.). diff --git a/tensorflow/lite/python/optimize/calibrator.py b/tensorflow/lite/python/optimize/calibrator.py index 6d9a29236f0..1f962917551 100644 --- a/tensorflow/lite/python/optimize/calibrator.py +++ b/tensorflow/lite/python/optimize/calibrator.py @@ -19,6 +19,7 @@ from __future__ import print_function import numpy as np from tensorflow.python.util.lazy_loader import LazyLoader +from tensorflow.lite.python import lite_constants # Lazy load since some of the performance benchmark skylark rules # break dependencies. Must use double quotes to match code internal rewrite @@ -55,7 +56,8 @@ class Calibrator(object): raise ValueError("Failed to parse the model.") def calibrate_and_quantize(self, dataset_gen, input_type, output_type, - allow_float, enable_mlir_quantizer=False): + allow_float, activations_type = lite_constants.INT8, + enable_mlir_quantizer=False): """Calibrates the model with specified generator and then quantizes it. Returns: @@ -69,6 +71,7 @@ class Calibrator(object): computation, useful when targeting an integer-only backend. If False, an error will be thrown if an operation cannot be quantized, otherwise the model will fallback to float ops. + activations_type: A tf.dtype representing the desired type for activations enable_mlir_quantizer: A boolean. True if wants to use mlir quantizer to quantize the calibrated model. """ @@ -78,6 +81,7 @@ class Calibrator(object): return self._calibrator.QuantizeModel( np.dtype(input_type.as_numpy_dtype()).num, np.dtype(output_type.as_numpy_dtype()).num, allow_float, + np.dtype(activations_type.as_numpy_dtype()).num, enable_mlir_quantizer) def calibrate_and_quantize_single(self, dataset_gen, input_type, output_type, diff --git a/tensorflow/lite/python/optimize/calibrator_test.py b/tensorflow/lite/python/optimize/calibrator_test.py index 28e8723f23d..7ec5f8f526c 100644 --- a/tensorflow/lite/python/optimize/calibrator_test.py +++ b/tensorflow/lite/python/optimize/calibrator_test.py @@ -33,9 +33,13 @@ from tensorflow.python.platform import test class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): @parameterized.named_parameters( - ('EnableMlirQuantizer', True), # enable mlir quantizer - ('DisableMlirQuantizer', False)) # disable mlir quantizer - def test_calibration_with_quantization(self, enable_mlir): + # Activation type Int8 - enable mlir quantizer + ('UseActivationTypeInt8EnabledMlir', constants.INT8, True), + # Activation type Int8 - disable mlir quantizer + ('UseActivationTypeInt8DisabledMlir', constants.INT8, False), + # Activation type Int16 + ('UseActivationTypeInt16', constants.INT16, False)) + def test_calibration_with_quantization(self, activations_type, enable_mlir): model_path = resource_loader.get_path_to_datafile( 'test_data/mobilenet_like_model.bin') float_model = open(model_path, 'rb').read() @@ -49,13 +53,18 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): quantized_model = quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, constants.FLOAT, False, + activations_type, enable_mlir) self.assertIsNotNone(quantized_model) @parameterized.named_parameters( - ('EnableMlirQuantizer', True), # enable mlir quantizer - ('DisableMlirQuantizer', False)) # disable mlir quantizer - def test_calibration_with_quantization_allow_float(self, enable_mlir): + # Activation type Int8 - enable mlir quantizer + ('UseActivationTypeInt8EnabledMlir', constants.INT8, True), + # Activation type Int8 - disable mlir quantizer + ('UseActivationTypeInt8DisableMlir', constants.INT8, False), + # Activation type Int16 - disable mlir quantizer + ('UseActivationTypeInt16', constants.INT16, False)) + def test_calibration_with_quantization_allow_float(self, activations_type, enable_mlir): model_path = resource_loader.get_path_to_datafile( 'test_data/mobilenet_like_model.bin') float_model = open(model_path, 'rb').read() @@ -69,6 +78,7 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): quantized_model = quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, constants.FLOAT, True, + activations_type, enable_mlir) self.assertIsNotNone(quantized_model) @@ -88,9 +98,13 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): self.assertIsNotNone(quantized_model) @parameterized.named_parameters( - ('EnableMlirQuantizer', True), # enable mlir quantizer - ('DisableMlirQuantizer', False)) # disable mlir quantizer - def test_calibration_with_quantization_multiple_inputs(self, enable_mlir): + # Activation type Int8 - enable mlir quantizer + ('UseActivationTypeInt8 - EnableMlirQuantizer', constants.INT8, True), + # Activation type Int8 - disable mlir quantizer + ('UseActivationTypeInt8 - DisableMlirQuantizer', constants.INT8, False), + # Activation type Int16 - disable mlir quantizer + ('UseActivationTypeInt16 - DisableEnableMlirQuantizer', constants.INT16, False)) + def test_calibration_with_quantization_multiple_inputs(self, activations_type, enable_mlir): # Load multi add model from test data. # This model has 4 inputs of size (1, 8, 8, 3). model_path = resource_loader.get_path_to_datafile( @@ -106,6 +120,7 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): quantized_model = quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, constants.FLOAT, False, + activations_type, enable_mlir) self.assertIsNotNone(quantized_model) @@ -148,7 +163,8 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): with self.assertRaisesRegex(ValueError, 'Size mismatch'): quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, - constants.FLOAT, False, enable_mlir) + constants.FLOAT, False, + enable_mlir) @parameterized.named_parameters( ('EnableMlirQuantizer', True), # enable mlir quantizer @@ -166,7 +182,8 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): with self.assertRaises(ValueError): quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, - constants.FLOAT, False, enable_mlir) + constants.FLOAT, False, + constants.INT8, enable_mlir) if __name__ == '__main__': diff --git a/tensorflow/lite/tools/optimize/operator_property.cc b/tensorflow/lite/tools/optimize/operator_property.cc index 13f63092761..1f2d8bb4a4d 100644 --- a/tensorflow/lite/tools/optimize/operator_property.cc +++ b/tensorflow/lite/tools/optimize/operator_property.cc @@ -64,6 +64,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}, {1, {}}}; property.outputs = {{0, {}}}; property.version = 2; + property.quantize_input_as_activations = true; break; case BuiltinOperator_ARG_MAX: property.inputs = {{0, {}}}; @@ -176,7 +177,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, // LogSoftmax requires output with 16/256 as scale and 127 as zero point. TensorProperty tensor_property; tensor_property.restriction = true; - tensor_property.restricted_value = {16.0 / 256.0, 127}; + tensor_property.restricted_value_int8 = {16.0 / 256.0, 127}; property.outputs = {{0, tensor_property}}; property.version = 2; break; @@ -186,7 +187,8 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, // Logistic requires output with 1/256 as scale and -128 as zero point. TensorProperty tensor_property; tensor_property.restriction = true; - tensor_property.restricted_value = {1 / 256.0, -128}; + tensor_property.restricted_value_int8 = {1 / 256.0, -128}; + tensor_property.restricted_value_int16 = {1 / 32768.0, 0}; property.outputs = {{0, tensor_property}}; property.version = 2; break; @@ -741,7 +743,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, // L2 Norm requires output with 1/128 as scale and 0 as zero point. TensorProperty tensor_property; tensor_property.restriction = true; - tensor_property.restricted_value = {1 / 128.0, 0}; + tensor_property.restricted_value_int8 = {1 / 128.0, 0}; property.outputs = {{0, tensor_property}}; property.version = 2; break; @@ -756,6 +758,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.arbitrary_inputs = true; property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; + property.quantize_input_as_activations = true; property.version = 2; break; case BuiltinOperator_MEAN: @@ -767,6 +770,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.arbitrary_inputs = true; property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; + property.quantize_input_as_activations = true; property.version = 2; break; case BuiltinOperator_MUL: @@ -778,6 +782,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.arbitrary_inputs = true; property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; + property.restrict_same_input_output_scale = true; property.version = 2; break; case BuiltinOperator_PAD: @@ -840,7 +845,8 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, // Softmax requires output with 1/256 as scale and -128 as zero point. TensorProperty tensor_property; tensor_property.restriction = true; - tensor_property.restricted_value = {1 / 256.0, -128}; + tensor_property.restricted_value_int8 = {1 / 256.0, -128}; + tensor_property.restricted_value_int16 = {1 / 32768.0, 0}; property.outputs = {{0, tensor_property}}; property.version = 2; break; @@ -866,7 +872,8 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, // Tanh requires output with 1/128 as scale and 0 as zero point. TensorProperty tensor_property; tensor_property.restriction = true; - tensor_property.restricted_value = {1 / 128.0, 0}; + tensor_property.restricted_value_int8 = {1 / 128.0, 0}; + tensor_property.restricted_value_int16 = {1 / 32768.0, 0}; property.outputs = {{0, tensor_property}}; property.version = 2; break; diff --git a/tensorflow/lite/tools/optimize/operator_property.h b/tensorflow/lite/tools/optimize/operator_property.h index 5d37aa304e5..23052308568 100644 --- a/tensorflow/lite/tools/optimize/operator_property.h +++ b/tensorflow/lite/tools/optimize/operator_property.h @@ -43,7 +43,8 @@ struct TensorProperty { // Constraints. bool restriction = false; // scale/zero_point hardcoded. - std::pair restricted_value = {0.0, 0}; + std::pair restricted_value_int8 = {0.0, 0}; + std::pair restricted_value_int16 = {0.0, 0}; // Use derived scale. bool use_derived_scale = false; @@ -93,6 +94,13 @@ struct OperatorProperty { // Op version. int version = 1; + + // When we quantize activations into 16 bit and weights into 8 bit, + // we want to quantize all inputs, including constant tensors, + // for the operators like Add, Mul into 16-bit as well. The constant + // inputs are quantized as weights and this variable indicates + // that we want to do quantizations of these tensors as activations. + bool quantize_input_as_activations = false; }; OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, diff --git a/tensorflow/lite/tools/optimize/quantization_utils.cc b/tensorflow/lite/tools/optimize/quantization_utils.cc index 10680758d72..4bc9686ec2c 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils.cc +++ b/tensorflow/lite/tools/optimize/quantization_utils.cc @@ -20,7 +20,6 @@ limitations under the License. #include #include "absl/memory/memory.h" -#include "third_party/eigen3/Eigen/Core" #include "tensorflow/lite/c/common.h" #include "tensorflow/lite/core/api/error_reporter.h" #include "tensorflow/lite/kernels/internal/quantization_util.h" @@ -30,6 +29,7 @@ limitations under the License. #include "tensorflow/lite/minimal_logging.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/tools/optimize/model_utils.h" +#include "third_party/eigen3/Eigen/Core" namespace tflite { namespace optimize { @@ -85,6 +85,46 @@ void GetAsymmetricQuantizationParams( quantization_params->zero_point = std::vector(1, zero_point); } +void GetSymmetricQuantizationParams( + float min, float max, const int half_quant_range, + QuantizationParametersT* quantization_params) { + // Adjust the boundaries to guarantee 0 is included. + min = std::min(min, 0.0f); + max = std::max(max, 0.0f); + const float scale = std::max(std::abs(max), std::abs(min)) / half_quant_range; + int64_t zero_point = 0; + quantization_params->min = std::vector(1, min); + quantization_params->max = std::vector(1, max); + quantization_params->scale = std::vector(1, scale); + quantization_params->zero_point = std::vector(1, 0); +} + +TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type, + QuantizationParametersT* quantization_params, + ErrorReporter* error_reporter) { + if (activations_type == TensorType_INT8) { + GetAsymmetricQuantizationParams( + tensor->quantization->min[0], tensor->quantization->max[0], + std::numeric_limits::min(), std::numeric_limits::max(), + quantization_params); + } else if (activations_type == TensorType_INT16) { + float range = std::max(std::abs(tensor->quantization->min[0]), + std::abs(tensor->quantization->max[0])); + const float quantized_range = 32767.0; + const float scale = range / quantized_range; + quantization_params->min = std::vector(1, -range); + quantization_params->max = std::vector(1, range); + quantization_params->scale = std::vector(1, scale); + quantization_params->zero_point = std::vector(1, 0); + } else { + error_reporter->Report( + "Unsupported activation type for quantize-activation: %s", + activations_type); + return kTfLiteError; + } + return kTfLiteOk; +} + // Set the max and min quantization parameter for a single tensor given its // values. void FillSingleMinMax(const float* const input, const uint64_t input_size, @@ -536,6 +576,7 @@ TfLiteStatus SymmetricQuantizeTensorPerChannel(ModelT* model, TensorT* tensor, model, tensor, error_reporter); } +template TfLiteStatus SymmetricPerLayerBiasQuantize(ModelT* model, TensorT* tensor, float scaling_factor, ErrorReporter* error_reporter) { @@ -548,25 +589,38 @@ TfLiteStatus SymmetricPerLayerBiasQuantize(ModelT* model, TensorT* tensor, uint64_t num_elements; TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements)); - std::vector final_buffer(num_elements); - const int32_t kScale = std::numeric_limits::max(); + std::vector final_buffer(num_elements); + const BiasType kScale = std::numeric_limits::max(); for (size_t i = 0; i < num_elements; i++) { - const int32_t quantized_value = tflite::SafeCast( + const BiasType quantized_value = tflite::SafeCast( TfLiteRound(float_data[i] * scaling_factor_inv)); final_buffer[i] = std::min(kScale, std::max(-kScale, quantized_value)); } // Set the buffers and output type. uint8_t* uint8_buffer = reinterpret_cast(final_buffer.data()); - size_t buffer_size = num_elements * sizeof(int32_t); + size_t buffer_size = num_elements * sizeof(BiasType); std::vector scales(1, scaling_factor); std::vector zero_points(1, 0); + + auto output_type = std::is_same::value + ? TensorType_INT32 + : TensorType_INT64; return AddQuantizationParams(scales, zero_points, 0, uint8_buffer, - buffer_size, TensorType_INT32, model, tensor, + buffer_size, output_type, model, tensor, error_reporter); } +template TfLiteStatus SymmetricPerLayerBiasQuantize( + ModelT* model, TensorT* tensor, float scaling_factor, + ErrorReporter* error_reporter); + +template TfLiteStatus SymmetricPerLayerBiasQuantize( + ModelT* model, TensorT* tensor, float scaling_factor, + ErrorReporter* error_reporter); + +template TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor, float input_scale, const float* weight_scales, @@ -583,14 +637,14 @@ TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor, uint64_t num_elements; TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements)); - std::vector final_buffer(num_elements); - const int32_t kScale = std::numeric_limits::max(); + std::vector final_buffer(num_elements); + const BiasType kScale = std::numeric_limits::max(); for (int32_t channel_idx = 0; channel_idx < number_of_dimension; channel_idx++) { float scaling_factor = scales[channel_idx]; float scaling_factor_inv = (scaling_factor == 0) ? 0 : 1.0 / scaling_factor; - const int32_t quantized_value = tflite::SafeCast( + const BiasType quantized_value = tflite::SafeCast( TfLiteRound(float_data[channel_idx] * scaling_factor_inv)); final_buffer[channel_idx] = std::min(kScale, std::max(-kScale, quantized_value)); @@ -598,12 +652,26 @@ TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor, // Set the buffers and output type. uint8_t* uint8_buffer = reinterpret_cast(final_buffer.data()); - size_t buffer_size = num_elements * sizeof(int32_t); + size_t buffer_size = num_elements * sizeof(BiasType); std::vector zero_point(scales.size(), 0); + + auto output_type = std::is_same::value + ? TensorType_INT32 + : TensorType_INT64; return AddQuantizationParams(scales, zero_point, 0, uint8_buffer, buffer_size, - TensorType_INT32, model, tensor, error_reporter); + output_type, model, tensor, error_reporter); } +template TfLiteStatus SymmetricPerChannelBiasQuantize( + ModelT* model, TensorT* tensor, float input_scale, + const float* weight_scales, int number_of_dimension, + ErrorReporter* error_reporter); + +template TfLiteStatus SymmetricPerChannelBiasQuantize( + ModelT* model, TensorT* tensor, float input_scale, + const float* weight_scales, int number_of_dimension, + ErrorReporter* error_reporter); + TfLiteStatus QuantizeWeight(ModelT* model, TensorT* tensor, bool per_channel, int per_axis_index, ErrorReporter* error_reporter) { // TODO(suharshs): Currently we conflate quantizing weights and constants. Its @@ -645,12 +713,12 @@ float GetEffectiveScale(ModelT* model, SubGraphT* subgraph, int op_idx, return scale; } -void QuantizeActivation(TensorT* tensor) { - GetAsymmetricQuantizationParams( - tensor->quantization->min[0], tensor->quantization->max[0], - std::numeric_limits::min(), std::numeric_limits::max(), - tensor->quantization.get()); - tensor->type = TensorType_INT8; +TfLiteStatus QuantizeActivation(TensorT* tensor, TensorType activations_type, + ErrorReporter* error_reporter) { + TF_LITE_ENSURE_STATUS(GetQuantizationParams( + tensor, activations_type, tensor->quantization.get(), error_reporter)); + tensor->type = activations_type; + return kTfLiteOk; } TfLiteStatus QuantizeActivationToInt16(TensorT* tensor, float scale) { diff --git a/tensorflow/lite/tools/optimize/quantization_utils.h b/tensorflow/lite/tools/optimize/quantization_utils.h index 18ed707e175..752b4253250 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils.h +++ b/tensorflow/lite/tools/optimize/quantization_utils.h @@ -113,12 +113,14 @@ TfLiteStatus SymmetricQuantizeFloatsToInt16(ModelT* model, TensorT* tensor, ErrorReporter* error_reporter); // Symmetrically quantized the bias for per-layer ops (i.e. FullyConnected). +template TfLiteStatus SymmetricPerLayerBiasQuantize(ModelT* model, TensorT* tensor, float scaling_factor, ErrorReporter* error_reporter); // Symmetrically quantizes the bias for ops like Conv and DepthwiseConv. // The scale of bias if weight_per_channel_scale[channel] * input_scale. +template TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor, float input_scale, const float* weight_scales, @@ -135,8 +137,14 @@ float GetEffectiveScale(ModelT* model, SubGraphT* subgraph, int op_idx, std::vector intermediate_index, std::vector factors); +// Return quantization parameters depending on activations type. +TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type, + QuantizationParametersT* quantization_params, + ErrorReporter* error_reporter); + // Quantize activation. -void QuantizeActivation(TensorT* tensor); +TfLiteStatus QuantizeActivation(TensorT* tensor, TensorType activations_type, + ErrorReporter* error_reporter); // Quantize activation to 16bit. TfLiteStatus QuantizeActivationToInt16(TensorT* tensor, float scale); diff --git a/tensorflow/lite/tools/optimize/quantization_utils_test.cc b/tensorflow/lite/tools/optimize/quantization_utils_test.cc index ece0123d166..49009e49600 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils_test.cc +++ b/tensorflow/lite/tools/optimize/quantization_utils_test.cc @@ -701,7 +701,7 @@ TEST_F(QuantizationUtilsTest, SymmetricPerLayerBiasQuantize) { model->buffers.push_back(std::move(buffer)); // Call and verify. - EXPECT_EQ(SymmetricPerLayerBiasQuantize( + EXPECT_EQ(SymmetricPerLayerBiasQuantize( model.get(), model->subgraphs[0]->tensors[0].get(), input_scale * weight_scale, &error_reporter_), kTfLiteOk); @@ -759,7 +759,7 @@ TEST_F(QuantizationUtilsTest, SymmetricPerChannelBiasQuantize) { model->buffers.push_back(std::move(buffer)); // Call and verify. - EXPECT_EQ(SymmetricPerChannelBiasQuantize( + EXPECT_EQ(SymmetricPerChannelBiasQuantize( model.get(), model->subgraphs[0]->tensors[0].get(), input_scale, weight_scales.data(), 2, &error_reporter_), kTfLiteOk); diff --git a/tensorflow/lite/tools/optimize/quantization_wrapper.cc b/tensorflow/lite/tools/optimize/quantization_wrapper.cc index bd3331da6bf..5002c382bc7 100644 --- a/tensorflow/lite/tools/optimize/quantization_wrapper.cc +++ b/tensorflow/lite/tools/optimize/quantization_wrapper.cc @@ -42,7 +42,9 @@ bool CreateQuantizedModel(const std::string& path) { tflite::StderrReporter error_reporter; if (tflite::optimize::QuantizeModel( &builder, &model, tflite::TensorType_FLOAT32, - tflite::TensorType_FLOAT32, &error_reporter) != kTfLiteOk) { + tflite::TensorType_FLOAT32, + // TODO: Pass required activation type if needed + tflite::TensorType_INT8, &error_reporter) != kTfLiteOk) { return false; } return WriteFile(path, builder.GetBufferPointer(), builder.GetSize()); diff --git a/tensorflow/lite/tools/optimize/quantize_model.cc b/tensorflow/lite/tools/optimize/quantize_model.cc index 6fc19ff2a56..ee562fe9c4c 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.cc +++ b/tensorflow/lite/tools/optimize/quantize_model.cc @@ -64,6 +64,7 @@ operator_property::OperatorProperty GetOperatorProperty( TfLiteStatus QuantizeBias(ModelT* model, const TensorT* input_tensor, const TensorT* weight_tensor, TensorT* bias_tensor, bool is_per_channel, int channel_dim_index, + const TensorType& activations_type, ErrorReporter* error_reporter) { if (bias_tensor->shape.size() != 1) { error_reporter->Report("Expected bias tensor shape to be 1."); @@ -92,9 +93,15 @@ TfLiteStatus QuantizeBias(ModelT* model, const TensorT* input_tensor, weight_scales.size()); return kTfLiteError; } - return utils::SymmetricPerChannelBiasQuantize( - model, bias_tensor, input_tensor->quantization->scale[0], - weight_scales.data(), channel_dim_size, error_reporter); + if (activations_type == tflite::TensorType_INT16) { + return utils::SymmetricPerChannelBiasQuantize( + model, bias_tensor, input_tensor->quantization->scale[0], + weight_scales.data(), channel_dim_size, error_reporter); + } else { + return utils::SymmetricPerChannelBiasQuantize( + model, bias_tensor, input_tensor->quantization->scale[0], + weight_scales.data(), channel_dim_size, error_reporter); + } } else { if (weight_scales.size() != 1) { error_reporter->Report( @@ -102,40 +109,54 @@ TfLiteStatus QuantizeBias(ModelT* model, const TensorT* input_tensor, weight_scales.size()); return kTfLiteError; } - return utils::SymmetricPerLayerBiasQuantize( - model, bias_tensor, - input_tensor->quantization->scale[0] * weight_scales[0], - error_reporter); + if (activations_type == tflite::TensorType_INT16) { + return utils::SymmetricPerLayerBiasQuantize( + model, bias_tensor, + input_tensor->quantization->scale[0] * weight_scales[0], + error_reporter); + } else { + return utils::SymmetricPerLayerBiasQuantize( + model, bias_tensor, + input_tensor->quantization->scale[0] * weight_scales[0], + error_reporter); + } } return kTfLiteError; } // True if the tensor type has to be modified. bool TensorTypeChangeRequired(const TensorT* tensor, const TensorType& type) { - // The quantized model is type INT8, so if the user provided type is INT8, we - // do not have to do any custom logic. Additionally, if the current tensor - // isn't INT8 quantized, the custom type doesn't apply. - return (type != TensorType_INT8 && tensor->type == TensorType_INT8 && - !tensor->quantization->scale.empty()); + // The quantized model is type INT8/INT16, so if the user provided type is + // INT8/INT16, we do not have to do any custom logic. Additionally, if the + // current tensor isn't INT8/INT16 quantized, the custom type doesn't apply. + bool int8check = type != TensorType_INT8 && tensor->type == TensorType_INT8 && + !tensor->quantization->scale.empty(); + bool int16check = type != TensorType_INT16 && + tensor->type == TensorType_INT16 && + !tensor->quantization->scale.empty(); + return (int8check || int16check); } // Sets the input type, adding a Leading Op node at the start of the model if // necessary. // Returns the new input tensor index. int32_t SetInputType(ModelT* model, SubGraphT* subgraph, - const int32_t tensor_idx, const TensorType& input_type) { + const int32_t tensor_idx, const TensorType& input_type, + const TensorType& activations_type) { TensorT* tensor = subgraph->tensors[tensor_idx].get(); if (!TensorTypeChangeRequired(tensor, input_type)) { return -1; } if (input_type == TensorType_FLOAT32 || input_type == TensorType_UINT8) { + std::string type_string = + activations_type == TensorType_INT16 ? "int16" : "int8"; // Create a new tensor to be the input of the leading Op. std::unique_ptr leading_op_input; if (input_type == TensorType_FLOAT32) { // Add tensor for quantize operator. Scales and zero points are not // needed. const string leading_op_name = tensor->name; - const string new_name_original_input = tensor->name + "_int8"; + const string new_name_original_input = tensor->name + "_" + type_string; tensor->name = new_name_original_input; utils::MakeTensor(leading_op_name, tensor->shape, input_type, &leading_op_input); @@ -150,7 +171,7 @@ int32_t SetInputType(ModelT* model, SubGraphT* subgraph, TFLITE_DCHECK_GE(zero_point, -128); TFLITE_DCHECK_LE(zero_point, 127); const string leading_op_name = tensor->name; - const string new_name_original_input = tensor->name + "_int8"; + const string new_name_original_input = tensor->name + "_" + type_string; tensor->name = new_name_original_input; utils::MakeTensorWithQuantParam(leading_op_name, tensor->shape, input_type, scale, zero_point + 128, @@ -177,17 +198,20 @@ int32_t SetInputType(ModelT* model, SubGraphT* subgraph, // necessary. // Returns the new output tensor index. int32_t SetOutputType(ModelT* model, SubGraphT* subgraph, - const int32_t tensor_idx, const TensorType& output_type) { + const int32_t tensor_idx, const TensorType& output_type, + const TensorType& activations_type) { TensorT* tensor = subgraph->tensors[tensor_idx].get(); if (!TensorTypeChangeRequired(tensor, output_type)) { return -1; } if (output_type == TensorType_FLOAT32 || output_type == TensorType_UINT8) { + std::string type_string = + activations_type == TensorType_INT16 ? "int16" : "int8"; // Create a new tensor to be the output of the tailing op. std::unique_ptr tailing_op_output; if (output_type == TensorType_FLOAT32) { const string tailing_op_name = tensor->name; - const string new_name_original_output = tensor->name + "_int8"; + const string new_name_original_output = tensor->name + "_" + type_string; tensor->name = new_name_original_output; utils::MakeTensor(tailing_op_name, tensor->shape, output_type, &tailing_op_output); @@ -202,7 +226,7 @@ int32_t SetOutputType(ModelT* model, SubGraphT* subgraph, TFLITE_DCHECK_GE(zero_point, -128); TFLITE_DCHECK_LE(zero_point, 127); const string tailing_op_name = tensor->name; - const string new_name_original_output = tensor->name + "_int8"; + const string new_name_original_output = tensor->name + "_" + type_string; tensor->name = new_name_original_output; utils::MakeTensorWithQuantParam(tailing_op_name, tensor->shape, output_type, scale, zero_point + 128, @@ -238,6 +262,7 @@ int32_t SetOutputType(ModelT* model, SubGraphT* subgraph, // uint8, can be thought as "requant"). TfLiteStatus SetInputAndOutputTypes(ModelT* model, const TensorType& input_type, const TensorType& output_type, + const TensorType& activations_type, ErrorReporter* error_reporter) { for (int subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { @@ -253,8 +278,8 @@ TfLiteStatus SetInputAndOutputTypes(ModelT* model, const TensorType& input_type, EnumNameTensorType(tensor->type)); return kTfLiteError; } - const int32_t input_idx = - SetInputType(model, subgraph, subgraph->inputs[i], input_type); + const int32_t input_idx = SetInputType( + model, subgraph, subgraph->inputs[i], input_type, activations_type); if (input_idx < 0) { continue; } @@ -270,8 +295,8 @@ TfLiteStatus SetInputAndOutputTypes(ModelT* model, const TensorType& input_type, EnumNameTensorType(tensor->type)); return kTfLiteError; } - const int32_t output_idx = - SetOutputType(model, subgraph, subgraph->outputs[i], output_type); + const int32_t output_idx = SetOutputType( + model, subgraph, subgraph->outputs[i], output_type, activations_type); if (output_idx < 0) { continue; } @@ -287,6 +312,7 @@ TfLiteStatus SetInputAndOutputTypes(ModelT* model, const TensorType& input_type, // The other ones with constraints are handled in QuantizeWeightsAndInput. TfLiteStatus ApplyConstraints(ModelT* model, const std::unordered_set& operator_names, + TensorType activations_type, ErrorReporter* error_reporter) { for (int subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { @@ -332,7 +358,7 @@ TfLiteStatus ApplyConstraints(ModelT* model, std::unique_ptr additional_tensor; const string requant_tensor_name = input_tensor->name + "_requantized"; utils::MakeTensorWithQuantParam( - requant_tensor_name, input_tensor->shape, TensorType_INT8, + requant_tensor_name, input_tensor->shape, activations_type, output_scale, output_zp, &additional_tensor); const int32_t additional_tensor_idx = subgraph->tensors.size(); subgraph->tensors.push_back(std::move(additional_tensor)); @@ -382,7 +408,8 @@ std::vector> GetOutputs( bool ShouldRestrictSameInputOutputScale( operator_property::OperatorProperty property) { - // Ops with multiple inputs (i.e. concat) gets restricted in ApplyConstraints. + // Ops with multiple inputs (i.e. concat, max and min) gets restricted in + // ApplyConstraints. return (!property.arbitrary_inputs && property.restrict_same_input_output_scale); } @@ -401,7 +428,7 @@ TfLiteStatus QuantizeOpInput( ModelT* model, int32_t subgraph_idx, size_t* op_idx, operator_property::OperatorProperty property, const std::pair& input, - ErrorReporter* error_reporter) { + const TensorType& activations_type, ErrorReporter* error_reporter) { int32_t input_idx = input.first; operator_property::TensorProperty tensor_property = input.second; SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get(); @@ -429,7 +456,9 @@ TfLiteStatus QuantizeOpInput( if (utils::HasBuffer(model, subgraph, tensor_idx)) { // TODO(suharshs): Look at consumers, throw error if one consumer is // per-channel and one per-layer. - if (tensor_property.number_of_bits == 8) { + bool quantize_const_input = property.quantize_input_as_activations && + activations_type == TensorType_INT16; + if (tensor_property.number_of_bits == 8 && !quantize_const_input) { if (tensor_property.use_derived_scale) { // Currently 8bit tensors in input do not accept derived scale. return kTfLiteError; @@ -444,7 +473,7 @@ TfLiteStatus QuantizeOpInput( *op_idx); return kTfLiteError; } - } else if (tensor_property.number_of_bits == 16) { + } else if (tensor_property.number_of_bits == 16 || quantize_const_input) { if (tensor_property.use_derived_scale) { // Currently 16bit tensors in input do not accept derived scale. return kTfLiteError; @@ -476,8 +505,8 @@ TfLiteStatus QuantizeOpInput( tensor_property.derived_scale.input_tensors, tensor_property.derived_scale.intermediate_tensors, tensor_property.derived_scale.factors); - return utils::SymmetricPerLayerBiasQuantize(model, tensor, scale, - error_reporter); + return utils::SymmetricPerLayerBiasQuantize( + model, tensor, scale, error_reporter); } else if (tensor_property.number_of_bits == 10) { // When the number of bits is 10 (instead of 16), quantize the tensor to @@ -514,7 +543,8 @@ TfLiteStatus QuantizeOpInput( // Currently 8bit tensors in input do not accept derived scale. return kTfLiteError; } - utils::QuantizeActivation(tensor); + TF_LITE_ENSURE_STATUS(utils::QuantizeActivation( + tensor, activations_type, error_reporter)); } else if (tensor_property.number_of_bits == 16) { TensorT* tensor = subgraph->tensors[tensor_idx].get(); float quantized_range = 32767.0; @@ -532,13 +562,16 @@ TfLiteStatus QuantizeOpInput( } else { // If the tensor is not a model input, we need to add a Quantize // operation since the preceding op may require a float output. + std::string type_string = + activations_type == TensorType_INT16 ? "int16" : "int8"; std::unique_ptr op_output; - utils::MakeTensor(tensor->name + "_int8", tensor->shape, - TensorType_INT8, &op_output); + utils::MakeTensor(tensor->name + "_" + type_string, tensor->shape, + activations_type, &op_output); op_output->quantization = absl::make_unique(); op_output->quantization->min.push_back(tensor->quantization->min[0]); op_output->quantization->max.push_back(tensor->quantization->max[0]); - utils::QuantizeActivation(op_output.get()); + TF_LITE_ENSURE_STATUS(utils::QuantizeActivation( + op_output.get(), activations_type, error_reporter)); const int32_t quant_op_output_idx = subgraph->tensors.size(); subgraph->tensors.push_back(std::move(op_output)); std::unique_ptr quant_op; @@ -580,7 +613,7 @@ TfLiteStatus QuantizeOpOutput( ModelT* model, int32_t subgraph_idx, int32_t op_idx, operator_property::OperatorProperty property, const std::pair& output, - ErrorReporter* error_reporter) { + TensorType activations_type, ErrorReporter* error_reporter) { int32_t output_idx = output.first; operator_property::TensorProperty tensor_property = output.second; // If the operator is not quantizable, we don't need to do anything for the @@ -644,18 +677,22 @@ TfLiteStatus QuantizeOpOutput( const float max = input_tensor->quantization->max[0]; output_tensor->quantization->max = {max}; } - output_tensor->type = TensorType_INT8; + output_tensor->type = activations_type; } else if (tensor_property.restriction) { - const auto scale_and_zp = tensor_property.restricted_value; + const auto scale_and_zp = activations_type == TensorType_INT16 + ? tensor_property.restricted_value_int16 + : tensor_property.restricted_value_int8; + // Apply to output. output_tensor->quantization = absl::make_unique(); output_tensor->quantization->scale.push_back(scale_and_zp.first); output_tensor->quantization->zero_point.push_back(scale_and_zp.second); - output_tensor->type = TensorType_INT8; + output_tensor->type = activations_type; } else { // Process regular output that doesn't have any restrictions. if (utils::HasMinMax(output_tensor)) { - utils::QuantizeActivation(output_tensor); + utils::QuantizeActivation(output_tensor, activations_type, + error_reporter); } else { error_reporter->Report( "Unable to find min/max value for output %d in %s in " @@ -668,6 +705,7 @@ TfLiteStatus QuantizeOpOutput( } TfLiteStatus QuantizeIntemediateTensors(ModelT* model, + TensorType activations_type, ErrorReporter* error_reporter) { for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { @@ -691,7 +729,8 @@ TfLiteStatus QuantizeIntemediateTensors(ModelT* model, input.second.symmetric == false) { TensorT* tensor = subgraph->tensors[index_global].get(); if (utils::HasMinMax(tensor)) { - utils::QuantizeActivation(tensor); + utils::QuantizeActivation(tensor, activations_type, + error_reporter); } else { error_reporter->Report( "Unable to find min/max value for output %d in %s in " @@ -793,7 +832,7 @@ TfLiteStatus QuantizeSharedRange(ModelT* model, ErrorReporter* error_reporter) { TfLiteStatus QuantizeWeightsInputOutput( ModelT* model, bool allow_float, const std::unordered_set& operator_names, - ErrorReporter* error_reporter) { + const TensorType& activations_type, ErrorReporter* error_reporter) { for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get(); @@ -815,14 +854,16 @@ TfLiteStatus QuantizeWeightsInputOutput( for (const std::pair& input : GetInputs(op, property)) { TF_LITE_ENSURE_STATUS(QuantizeOpInput(model, subgraph_idx, &op_idx, - property, input, error_reporter)); + property, input, activations_type, + error_reporter)); } // Quantize operator outputs. for (const std::pair& output : GetOutputs(op, property)) { - TF_LITE_ENSURE_STATUS(QuantizeOpOutput( - model, subgraph_idx, op_idx, property, output, error_reporter)); + TF_LITE_ENSURE_STATUS( + QuantizeOpOutput(model, subgraph_idx, op_idx, property, output, + activations_type, error_reporter)); } } } @@ -832,6 +873,7 @@ TfLiteStatus QuantizeWeightsInputOutput( // Quantize bias. TfLiteStatus QuantizeBiases(ModelT* model, const std::unordered_set& operator_names, + const TensorType& activations_type, ErrorReporter* error_reporter) { for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { @@ -877,10 +919,10 @@ TfLiteStatus QuantizeBiases(ModelT* model, subgraph->tensors[op->inputs[property.inputs[1].first]].get(); operator_property::TensorProperty weight_property = property.inputs[1].second; - TF_LITE_ENSURE_STATUS( - QuantizeBias(model, input_tensor, weight_tensor, bias_tensor, - weight_property.per_axis, - weight_property.per_axis_index, error_reporter)); + TF_LITE_ENSURE_STATUS(QuantizeBias( + model, input_tensor, weight_tensor, bias_tensor, + weight_property.per_axis, weight_property.per_axis_index, + activations_type, error_reporter)); } } } @@ -1000,7 +1042,7 @@ TfLiteStatus FillQuantizationParams( // Check compatibility of activation, weight and bias scales. Adjust if needed. TfLiteStatus EnsureBiasScaleCompatibility( ModelT* model, const std::unordered_set& operator_names, - ErrorReporter* error_reporter) { + TensorType activations_type, ErrorReporter* error_reporter) { for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get(); @@ -1049,11 +1091,9 @@ TfLiteStatus EnsureBiasScaleCompatibility( // Get input scale for assymmetric quantization. QuantizationParametersT temp_quant_params = QuantizationParametersT(); - utils::GetAsymmetricQuantizationParams( - input_tensor->quantization->min[0], - input_tensor->quantization->max[0], - std::numeric_limits::min(), - std::numeric_limits::max(), &temp_quant_params); + TF_LITE_ENSURE_STATUS( + utils::GetQuantizationParams(input_tensor, activations_type, + &temp_quant_params, error_reporter)); if (temp_quant_params.scale.size() != 1) { error_reporter->Report("Unexpected input quantization scale size."); return kTfLiteError; @@ -1132,21 +1172,24 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* model, const TensorType& input_type, const TensorType& output_type, bool allow_float, const std::unordered_set& operator_names, + const TensorType& activations_type, ErrorReporter* error_reporter) { TF_LITE_ENSURE_STATUS( FillQuantizationParams(model, operator_names, error_reporter)); + TF_LITE_ENSURE_STATUS(EnsureBiasScaleCompatibility( + model, operator_names, activations_type, error_reporter)); TF_LITE_ENSURE_STATUS( - EnsureBiasScaleCompatibility(model, operator_names, error_reporter)); - TF_LITE_ENSURE_STATUS(QuantizeIntemediateTensors(model, error_reporter)); + QuantizeIntemediateTensors(model, activations_type, error_reporter)); TF_LITE_ENSURE_STATUS(QuantizeSharedRange(model, error_reporter)); TF_LITE_ENSURE_STATUS(QuantizeWeightsInputOutput( - model, allow_float, operator_names, error_reporter)); + model, allow_float, operator_names, activations_type, error_reporter)); + TF_LITE_ENSURE_STATUS(ApplyConstraints(model, operator_names, + activations_type, error_reporter)); TF_LITE_ENSURE_STATUS( - ApplyConstraints(model, operator_names, error_reporter)); - TF_LITE_ENSURE_STATUS(QuantizeBiases(model, operator_names, error_reporter)); + QuantizeBiases(model, operator_names, activations_type, error_reporter)); utils::SetOperatorCodeVersion(model); - TF_LITE_ENSURE_STATUS( - SetInputAndOutputTypes(model, input_type, output_type, error_reporter)); + TF_LITE_ENSURE_STATUS(SetInputAndOutputTypes( + model, input_type, output_type, activations_type, error_reporter)); flatbuffers::Offset output_model_location = Model::Pack(*builder, model); @@ -1158,23 +1201,27 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* model, const TensorType& input_type, const TensorType& output_type, bool allow_float, + const TensorType& activations_type, ErrorReporter* error_reporter) { return QuantizeModel(builder, model, input_type, output_type, allow_float, - GetAllOperatorOutputs(model), error_reporter); + GetAllOperatorOutputs(model), activations_type, + error_reporter); } TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* model, const TensorType& input_type, const TensorType& output_type, + const TensorType& activations_type, ErrorReporter* error_reporter) { return QuantizeModel(builder, model, input_type, output_type, - /*allow_float=*/false, error_reporter); + /*allow_float=*/false, activations_type, error_reporter); } TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* model, ErrorReporter* error_reporter) { + ModelT* model, const TensorType& activations_type, + ErrorReporter* error_reporter) { return QuantizeModel(builder, model, TensorType_FLOAT32, TensorType_FLOAT32, - /*allow_float=*/false, error_reporter); + /*allow_float=*/false, activations_type, error_reporter); } } // namespace optimize diff --git a/tensorflow/lite/tools/optimize/quantize_model.h b/tensorflow/lite/tools/optimize/quantize_model.h index 9b0353f6b6b..cc801ec9870 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.h +++ b/tensorflow/lite/tools/optimize/quantize_model.h @@ -35,7 +35,9 @@ namespace optimize { // // Note: This is a private API, subject to change. TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* input_model, ErrorReporter* error_reporter); + ModelT* input_model, + const TensorType& activations_type, + ErrorReporter* error_reporter); // Same as above, but the types of quantized inputs and outputs are // configurable. @@ -44,6 +46,7 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, + const TensorType& activations_type, ErrorReporter* error_reporter); // Same as above, but can enable allowing float intermediate operations for ops @@ -53,6 +56,7 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, bool allow_float, + const TensorType& activations_type, ErrorReporter* error_reporter); // Same as above, but enables only quantizing a whitelist of operations, @@ -63,6 +67,7 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, bool allow_float, const std::unordered_set& operator_names, + const TensorType& activations_type, ErrorReporter* error_reporter); } // namespace optimize diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc index da1b293c84b..166d60ecc66 100644 --- a/tensorflow/lite/tools/optimize/quantize_model_test.cc +++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc @@ -80,28 +80,35 @@ class QuantizeModelTest : public testing::Test { internal::FailOnErrorReporter error_reporter_; }; -class QuantizeConvModelTest : public QuantizeModelTest { +class QuantizeConvModelTest : public QuantizeModelTest, + public testing::WithParamInterface { protected: QuantizeConvModelTest() { + tensor_type_ = GetParam(); input_model_ = ReadModel(internal::kConvModelWith0Plus10Weights); readonly_model_ = input_model_->GetModel(); readonly_model_->UnPackTo(&model_); } + TensorType tensor_type_; }; -TEST_F(QuantizeConvModelTest, QuantizationSucceeds) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); +INSTANTIATE_TEST_SUITE_P(QuantizeConvModelTestInst, QuantizeConvModelTest, + testing::ValuesIn({TensorType_INT8, + TensorType_INT16})); + +TEST_P(QuantizeConvModelTest, QuantizationSucceeds) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); const uint8_t* buffer = builder_.GetBufferPointer(); const Model* output_model = GetModel(buffer); ASSERT_TRUE(output_model); } -TEST_F(QuantizeConvModelTest, SkipUnspecifiedLayer) { +TEST_P(QuantizeConvModelTest, SkipUnspecifiedLayer) { auto status = QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, - /*allow_float=*/true, {}, &error_reporter_); + /*allow_float=*/true, {}, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); ASSERT_EQ(model_.subgraphs.size(), readonly_model_->subgraphs()->size()); // The resulting model should be the same. @@ -123,9 +130,9 @@ TEST_F(QuantizeConvModelTest, SkipUnspecifiedLayer) { } } -TEST_F(QuantizeConvModelTest, TensorShapesAndStructureIsUnchanged) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); +TEST_P(QuantizeConvModelTest, TensorShapesAndStructureIsUnchanged) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); ASSERT_EQ(model_.subgraphs.size(), readonly_model_->subgraphs()->size()); for (size_t subgraph_idx = 0; subgraph_idx < model_.subgraphs.size(); @@ -148,9 +155,9 @@ TEST_F(QuantizeConvModelTest, TensorShapesAndStructureIsUnchanged) { EXPECT_EQ(model_.operator_codes[0]->version, 3); } -TEST_F(QuantizeConvModelTest, OperatorsAreUnchanged) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); +TEST_P(QuantizeConvModelTest, OperatorsAreUnchanged) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); ASSERT_EQ(model_.operator_codes.size(), readonly_model_->operator_codes()->size()); @@ -182,20 +189,28 @@ TEST_F(QuantizeConvModelTest, OperatorsAreUnchanged) { } } -TEST_F(QuantizeConvModelTest, GraphIsFullyQuantized) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); +TEST_P(QuantizeConvModelTest, GraphIsFullyQuantized) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (const auto& subgraph : model_.subgraphs) { for (const auto& tensor : subgraph->tensors) { - EXPECT_TRUE(tensor->type == TensorType_INT32 || - tensor->type == TensorType_INT8); + if (tensor_type_ == TensorType_INT8) { + EXPECT_TRUE(tensor->type == TensorType_INT32 || + tensor->type == TensorType_INT8); + } else if (tensor_type_ == TensorType_INT16) { + EXPECT_TRUE(tensor->type == TensorType_INT64 || // bias + tensor->type == TensorType_INT8 || // weights + tensor->type == TensorType_INT16); // activations + } } } } -TEST_F(QuantizeConvModelTest, FloatInputAndOutput) { - auto status = QuantizeModel(&builder_, &model_, &error_reporter_); +TEST_P(QuantizeConvModelTest, FloatInputAndOutput) { + auto status = + QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (int32_t subgraph_idx = 0; subgraph_idx < model_.subgraphs.size(); @@ -234,22 +249,33 @@ TEST_F(QuantizeConvModelTest, FloatInputAndOutput) { EXPECT_EQ(subgraph->tensors[output_idx]->type, TensorType_FLOAT32); EXPECT_EQ(subgraph->tensors[output_idx]->name, "output"); // The original input and output has been renamed. - EXPECT_EQ(subgraph->tensors[quant_op->outputs[0]]->name, "input_int8"); - EXPECT_EQ(subgraph->tensors[dequant_op->inputs[0]]->name, "output_int8"); + std::string control_suffix = + (tensor_type_ == TensorType_INT16) ? "int16" : "int8"; + EXPECT_EQ(subgraph->tensors[quant_op->outputs[0]]->name, + "input_" + control_suffix); + EXPECT_EQ(subgraph->tensors[dequant_op->inputs[0]]->name, + "output_" + control_suffix); for (int tensor_idx = 0; tensor_idx < subgraph->tensors.size(); ++tensor_idx) { const auto& tensor = subgraph->tensors[tensor_idx]; if (input_idx != tensor_idx && output_idx != tensor_idx) { - EXPECT_TRUE(tensor->type == TensorType_INT32 || - tensor->type == TensorType_INT8); + if (tensor_type_ == TensorType_INT8) { + EXPECT_TRUE(tensor->type == TensorType_INT32 || + tensor->type == TensorType_INT8); + } else if (tensor_type_ == TensorType_INT16) { + EXPECT_TRUE(tensor->type == TensorType_INT64 || // bias + tensor->type == TensorType_INT8 || // weights + tensor->type == TensorType_INT16); // activations + } } } } } -TEST_F(QuantizeConvModelTest, Uint8InputAndOutput) { - auto status = QuantizeModel(&builder_, &model_, TensorType_UINT8, - TensorType_UINT8, &error_reporter_); +TEST_P(QuantizeConvModelTest, Uint8InputAndOutput) { + auto status = + QuantizeModel(&builder_, &model_, TensorType_UINT8, TensorType_UINT8, + TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (int32_t subgraph_idx = 0; subgraph_idx < model_.subgraphs.size(); @@ -326,21 +352,25 @@ class QuantizeConvNoBiasModelTest : public QuantizeModelTest { }; TEST_F(QuantizeConvNoBiasModelTest, QuantizationSucceeds) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); const uint8_t* buffer = builder_.GetBufferPointer(); const Model* output_model = GetModel(buffer); ASSERT_TRUE(output_model); } -class QuantizeConcatModelTest : public QuantizeModelTest { +class QuantizeConcatModelTest : public QuantizeModelTest, + public testing::WithParamInterface { protected: QuantizeConcatModelTest() { input_model_ = ReadModel(internal::kFloatConcatMax5Max10Max10); readonly_model_ = input_model_->GetModel(); readonly_model_->UnPackTo(&model_); } + + TensorType tensor_type_; }; // There are two inputs for concat, "input0" and "input1". "input0" has [0, 5] @@ -352,9 +382,9 @@ class QuantizeConcatModelTest : public QuantizeModelTest { // input0 -> requant -> input0_requant \ // concat - output // input1 / -TEST_F(QuantizeConcatModelTest, AddRequantBeforeConcat) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); +TEST_P(QuantizeConcatModelTest, AddRequantBeforeConcat) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); // There is only one subgraph. @@ -373,32 +403,51 @@ TEST_F(QuantizeConcatModelTest, AddRequantBeforeConcat) { EXPECT_EQ(model_.operator_codes[concat->opcode_index]->builtin_code, BuiltinOperator_CONCATENATION); + auto zero_point_control = tensor_type_ == TensorType_INT8 ? -128 : 0; + /* + input0_scale_control + INT8: (5-0) / (2^8 - 1) + INT16: (5-0) / (2^16 / 2 - 1) + input1_scale + INT8: (10-0) / (2^8 - 1) + INT16: (10-0) / (2^16 / 2 - 1) + */ + auto input0_scale_control = + tensor_type_ == TensorType_INT8 ? 0.019607844 : 0.00015259254; + auto input1_scale = + tensor_type_ == TensorType_INT8 ? 0.039215688 : 0.00030518509; + // There should be 4 tensors: input0, input1, input0_requantized, output. EXPECT_EQ(subgraph->tensors.size(), 4); - EXPECT_EQ(subgraph->tensors[0]->type, TensorType_INT8); + EXPECT_EQ(subgraph->tensors[0]->type, tensor_type_); EXPECT_EQ(subgraph->tensors[0]->name, "input0"); EXPECT_EQ(subgraph->tensors[0]->quantization->scale.size(), 1); EXPECT_EQ(subgraph->tensors[0]->quantization->zero_point.size(), 1); - EXPECT_FLOAT_EQ(subgraph->tensors[0]->quantization->scale[0], 0.019607844); - EXPECT_FLOAT_EQ(subgraph->tensors[0]->quantization->zero_point[0], -128); - EXPECT_EQ(subgraph->tensors[1]->type, TensorType_INT8); + EXPECT_FLOAT_EQ(subgraph->tensors[0]->quantization->scale[0], + input0_scale_control); + EXPECT_FLOAT_EQ(subgraph->tensors[0]->quantization->zero_point[0], + zero_point_control); + EXPECT_EQ(subgraph->tensors[1]->type, tensor_type_); EXPECT_EQ(subgraph->tensors[1]->name, "input1"); EXPECT_EQ(subgraph->tensors[1]->quantization->scale.size(), 1); EXPECT_EQ(subgraph->tensors[1]->quantization->zero_point.size(), 1); - EXPECT_FLOAT_EQ(subgraph->tensors[1]->quantization->scale[0], 0.039215688); - EXPECT_FLOAT_EQ(subgraph->tensors[1]->quantization->zero_point[0], -128); - EXPECT_EQ(subgraph->tensors[2]->type, TensorType_INT8); + EXPECT_FLOAT_EQ(subgraph->tensors[1]->quantization->scale[0], input1_scale); + EXPECT_FLOAT_EQ(subgraph->tensors[1]->quantization->zero_point[0], + zero_point_control); + EXPECT_EQ(subgraph->tensors[2]->type, tensor_type_); EXPECT_EQ(subgraph->tensors[2]->name, "output"); EXPECT_EQ(subgraph->tensors[2]->quantization->scale.size(), 1); EXPECT_EQ(subgraph->tensors[2]->quantization->zero_point.size(), 1); - EXPECT_FLOAT_EQ(subgraph->tensors[2]->quantization->scale[0], 0.039215688); - EXPECT_FLOAT_EQ(subgraph->tensors[2]->quantization->zero_point[0], -128); - EXPECT_EQ(subgraph->tensors[3]->type, TensorType_INT8); + EXPECT_FLOAT_EQ(subgraph->tensors[2]->quantization->scale[0], input1_scale); + EXPECT_FLOAT_EQ(subgraph->tensors[2]->quantization->zero_point[0], + zero_point_control); + EXPECT_EQ(subgraph->tensors[3]->type, tensor_type_); EXPECT_EQ(subgraph->tensors[3]->name, "input0_requantized"); EXPECT_EQ(subgraph->tensors[3]->quantization->scale.size(), 1); EXPECT_EQ(subgraph->tensors[3]->quantization->zero_point.size(), 1); - EXPECT_FLOAT_EQ(subgraph->tensors[3]->quantization->scale[0], 0.039215688); - EXPECT_FLOAT_EQ(subgraph->tensors[3]->quantization->zero_point[0], -128); + EXPECT_FLOAT_EQ(subgraph->tensors[3]->quantization->scale[0], input1_scale); + EXPECT_FLOAT_EQ(subgraph->tensors[3]->quantization->zero_point[0], + zero_point_control); // The connection should be what is described in the comment. EXPECT_EQ(requant->inputs.size(), 1); @@ -419,7 +468,9 @@ TEST_F(QuantizeConcatModelTest, AddRequantBeforeConcat) { EXPECT_EQ(model_.operator_codes[1]->builtin_code, BuiltinOperator_QUANTIZE); EXPECT_EQ(model_.operator_codes[1]->version, 2); } - +INSTANTIATE_TEST_SUITE_P(QuantizeConcatModelInst, QuantizeConcatModelTest, + testing::ValuesIn({TensorType_INT8, + TensorType_INT16})); class QuantizeSplitModelTest : public QuantizeModelTest { protected: QuantizeSplitModelTest() { @@ -432,8 +483,9 @@ class QuantizeSplitModelTest : public QuantizeModelTest { // There are two outputs for split with different scales, the resulting model // should have the scales be hardcodes to the input scale value. TEST_F(QuantizeSplitModelTest, QuantizeSplit) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); // There is only one subgraph. @@ -496,8 +548,9 @@ class QuantizeConvModel1Test : public QuantizeModelTest { }; TEST_F(QuantizeConvModel1Test, VerifyConvQuantizationWithUnitScale) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); const auto& subgraph = model_.subgraphs[0]; @@ -587,18 +640,25 @@ TEST_F(QuantizeConvModel1Test, VerifyConvQuantizationWithUnitScale) { EXPECT_EQ(model_.operator_codes[0]->version, 3); } -class QuantizeConvModel2Test : public QuantizeModelTest { +class QuantizeConvModel2Test : public QuantizeModelTest, + public testing::WithParamInterface { protected: QuantizeConvModel2Test() { + tensor_type_ = GetParam(); input_model_ = ReadModel(internal::kConvModelWith0Plus10Weights); readonly_model_ = input_model_->GetModel(); readonly_model_->UnPackTo(&model_); } -}; -TEST_F(QuantizeConvModel2Test, VerifyConvQuantization) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + TensorType tensor_type_; +}; +INSTANTIATE_TEST_SUITE_P(QuantizeConvModel2TestInst, QuantizeConvModel2Test, + testing::ValuesIn({TensorType_INT8, + TensorType_INT16})); + +TEST_P(QuantizeConvModel2Test, VerifyConvQuantization) { + auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + tensor_type_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto conv_op = subgraph->operators[0].get(); @@ -615,8 +675,10 @@ TEST_F(QuantizeConvModel2Test, VerifyConvQuantization) { const auto output_tensor = subgraph->tensors[conv_op->outputs[output_tensor_idx]].get(); - EXPECT_EQ(bias_tensor->type, TensorType_INT32); - EXPECT_EQ(input_tensor->type, TensorType_INT8); + EXPECT_EQ(bias_tensor->type, tensor_type_ == TensorType_INT8 + ? TensorType_INT32 + : TensorType_INT64); + EXPECT_EQ(input_tensor->type, tensor_type_); EXPECT_EQ(weights_tensor->type, TensorType_INT8); ASSERT_TRUE(weights_tensor->quantization); @@ -644,17 +706,28 @@ TEST_F(QuantizeConvModel2Test, VerifyConvQuantization) { } const auto bias_buffer = model_.buffers[bias_tensor->buffer].get(); - ASSERT_EQ(bias_buffer->data.size(), sizeof(int32_t) * bias_tensor->shape[0]); - const int32_t* bias_values = - reinterpret_cast(bias_buffer->data.data()); + auto control_size = tensor_type_ == TensorType_INT8 + ? sizeof(int32_t) * bias_tensor->shape[0] + : sizeof(int64_t) * bias_tensor->shape[0]; + + ASSERT_EQ(bias_buffer->data.size(), control_size); const auto original_bias_buffer = readonly_model_->buffers()->Get(bias_tensor->buffer); const float* bias_float_buffer = reinterpret_cast(original_bias_buffer->data()->data()); - for (size_t i = 0; i < out_channel_size; i++) { - auto dequantized_value = bias_values[i] * bias_scales[i]; - EXPECT_NEAR(dequantized_value, bias_float_buffer[i], bias_scales[i] / 2); + if (tensor_type_ == TensorType_INT8) { + int32_t* bias_values = reinterpret_cast(bias_buffer->data.data()); + for (size_t i = 0; i < out_channel_size; i++) { + auto dequantized_value = bias_values[i] * bias_scales[i]; + EXPECT_NEAR(dequantized_value, bias_float_buffer[i], bias_scales[i] / 2); + } + } else if (tensor_type_ == TensorType_INT16) { + int64_t* bias_values = reinterpret_cast(bias_buffer->data.data()); + for (size_t i = 0; i < out_channel_size; i++) { + auto dequantized_value = bias_values[i] * bias_scales[i]; + EXPECT_NEAR(dequantized_value, bias_float_buffer[i], bias_scales[i] / 2); + } } const auto weights_buffer = model_.buffers[weights_tensor->buffer].get(); @@ -695,8 +768,9 @@ class QuantizeSoftmaxTest : public QuantizeModelTest { }; TEST_F(QuantizeSoftmaxTest, VerifySoftmaxQuantization) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -755,8 +829,9 @@ class QuantizeAvgPoolTest : public QuantizeModelTest { }; TEST_F(QuantizeAvgPoolTest, VerifyAvgPoolQuantization) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -816,8 +891,9 @@ class QuantizeMultiInputAddWithReshapeTest : public QuantizeModelTest { }; TEST_F(QuantizeMultiInputAddWithReshapeTest, VerifyReshapeQuantization) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Verify Reshape is quantized. @@ -863,8 +939,9 @@ TEST_F(QuantizeMultiInputAddWithReshapeTest, VerifyReshapeQuantization) { } TEST_F(QuantizeMultiInputAddWithReshapeTest, VerifyAddQuantization) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Verify ADD is quantized. @@ -923,8 +1000,9 @@ class QuantizeConstInputTest : public QuantizeModelTest { }; TEST_F(QuantizeConstInputTest, VerifyConstOpInput) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Verify ConstOp is quantized. @@ -965,8 +1043,9 @@ class QuantizeArgMaxTest : public QuantizeModelTest { }; TEST_F(QuantizeArgMaxTest, VerifyArgMax) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -1008,8 +1087,9 @@ class QuantizeLSTMTest : public QuantizeModelTest { TEST_F(QuantizeLSTMTest, VerifyLSTM) { // Quantize model. - auto status = QuantizeModel(&builder_, &model_, TensorType_FLOAT32, - TensorType_FLOAT32, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Read expected model. @@ -1067,8 +1147,9 @@ class QuantizeLSTM2Test : public QuantizeModelTest { TEST_F(QuantizeLSTM2Test, VerifyLSTM) { // Quantize model. - auto status = QuantizeModel(&builder_, &model_, TensorType_FLOAT32, - TensorType_FLOAT32, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Read expected model. @@ -1126,8 +1207,9 @@ class QuantizeSVDFTest : public QuantizeModelTest { TEST_F(QuantizeSVDFTest, VerifySVDF) { // Quantize model. - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); // Read expected model. @@ -1184,8 +1266,9 @@ class QuantizeFCTest : public QuantizeModelTest { }; TEST_F(QuantizeFCTest, VerifyFC) { - auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, - TensorType_INT8, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -1236,7 +1319,7 @@ class QuantizeCustomOpTest : public QuantizeModelTest { TEST_F(QuantizeCustomOpTest, VerifyMixedQuantization) { auto status = QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, - /*allow_float=*/true, &error_reporter_); + /*allow_float=*/true, TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto float_graph = readonly_model_->subgraphs()->Get(0); @@ -1270,7 +1353,8 @@ class QuantizePackTest : public QuantizeModelTest { }; TEST_F(QuantizePackTest, VerifyPack) { - auto status = QuantizeModel(&builder_, &model_, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); @@ -1334,7 +1418,8 @@ class QuantizeMinimumMaximumTest }; TEST_P(QuantizeMinimumMaximumTest, VerifyMinimumMaximum) { - auto status = QuantizeModel(&builder_, &model_, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -1415,7 +1500,8 @@ class QuantizeUnpackTest : public QuantizeModelTest { } }; TEST_F(QuantizeUnpackTest, VerifyUnpack) { - auto status = QuantizeModel(&builder_, &model_, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); From 3130f64c4783d987ef32ae6853f61d1c60983d34 Mon Sep 17 00:00:00 2001 From: jmsmdy Date: Fri, 7 Feb 2020 21:11:21 -0500 Subject: [PATCH 0015/1390] Corrected docstring for tf.signal.frame --- tensorflow/python/ops/signal/shape_ops.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/ops/signal/shape_ops.py b/tensorflow/python/ops/signal/shape_ops.py index b0e8537c2be..bb4ac7ec33c 100644 --- a/tensorflow/python/ops/signal/shape_ops.py +++ b/tensorflow/python/ops/signal/shape_ops.py @@ -73,15 +73,25 @@ def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, axis=-1, audio = tf.random.normal([3, 9152]) # Compute overlapping frames of length 512 with a step of 180 (frames overlap - # by 332 samples). By default, only 50 frames are generated since the last - # 152 samples do not form a full frame. + # by 332 samples). By default, only 49 frames are generated since a frame + # with start position j*180 for j > 48 would overhang the end. frames = tf.signal.frame(audio, 512, 180) - frames.shape.assert_is_compatible_with([3, 50, 512]) + frames.shape.assert_is_compatible_with([3, 49, 512]) - # When pad_end is enabled, the final frame is kept (padded with zeros). + # When pad_end is enabled, the final two frames are kept (padded with zeros). frames = tf.signal.frame(audio, 512, 180, pad_end=True) frames.shape.assert_is_compatible_with([3, 51, 512]) ``` + + If the dimension along `axis` is N, and `pad_end=False`, the number of frames + can be computed by: + ```python + frames = 1 + (N - frame_size) // frame_step + ``` + If `pad_end=True`, the number of frames can be computed by: + ```python + frames = -(-N // frame_step) # ceiling division + ``` Args: signal: A `[..., samples, ...]` `Tensor`. The rank and dimensions From de6afc5d6b509c8f3d709bf1e275373864ec0936 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 12 Feb 2020 13:51:24 +0000 Subject: [PATCH 0016/1390] Changed per reviewer comments. --- .../python/optimize/calibration_wrapper.cc | 3 +- .../lite/tools/optimize/quantization_utils.cc | 10 ++----- .../lite/tools/optimize/quantize_model.cc | 19 +++++++++---- .../lite/tools/optimize/quantize_model.h | 25 +++++++++++++---- .../tools/optimize/quantize_model_test.cc | 28 +++++++++---------- 5 files changed, 51 insertions(+), 34 deletions(-) diff --git a/tensorflow/lite/python/optimize/calibration_wrapper.cc b/tensorflow/lite/python/optimize/calibration_wrapper.cc index 88995136726..cdc8adaaf2b 100644 --- a/tensorflow/lite/python/optimize/calibration_wrapper.cc +++ b/tensorflow/lite/python/optimize/calibration_wrapper.cc @@ -266,7 +266,8 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, auto status = tflite::optimize::QuantizeModel( &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), TfLiteTypeToSchemaType(output_type), allow_float, {op_name}, - TensorType_INT8, error_reporter_.get()); + TensorType_INT8, + error_reporter_.get()); if (status != kTfLiteOk) { error_reporter_->exception(); return nullptr; diff --git a/tensorflow/lite/tools/optimize/quantization_utils.cc b/tensorflow/lite/tools/optimize/quantization_utils.cc index 4bc9686ec2c..ba43416cf04 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils.cc +++ b/tensorflow/lite/tools/optimize/quantization_utils.cc @@ -108,14 +108,10 @@ TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type, std::numeric_limits::min(), std::numeric_limits::max(), quantization_params); } else if (activations_type == TensorType_INT16) { - float range = std::max(std::abs(tensor->quantization->min[0]), - std::abs(tensor->quantization->max[0])); const float quantized_range = 32767.0; - const float scale = range / quantized_range; - quantization_params->min = std::vector(1, -range); - quantization_params->max = std::vector(1, range); - quantization_params->scale = std::vector(1, scale); - quantization_params->zero_point = std::vector(1, 0); + GetSymmetricQuantizationParams(tensor->quantization->min[0], + tensor->quantization->max[0], + quantized_range, quantization_params); } else { error_reporter->Report( "Unsupported activation type for quantize-activation: %s", diff --git a/tensorflow/lite/tools/optimize/quantize_model.cc b/tensorflow/lite/tools/optimize/quantize_model.cc index ee562fe9c4c..bbb40080fbc 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.cc +++ b/tensorflow/lite/tools/optimize/quantize_model.cc @@ -1210,18 +1210,25 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* model, const TensorType& input_type, - const TensorType& output_type, - const TensorType& activations_type, + const TensorType& output_type, bool allow_float, ErrorReporter* error_reporter) { - return QuantizeModel(builder, model, input_type, output_type, - /*allow_float=*/false, activations_type, error_reporter); + return QuantizeModel(builder, model, input_type, output_type, allow_float, + GetAllOperatorOutputs(model), TensorType_INT8, + error_reporter); } TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* model, const TensorType& activations_type, + ModelT* model, const TensorType& input_type, + const TensorType& output_type, ErrorReporter* error_reporter) { + return QuantizeModel(builder, model, input_type, output_type, + /*allow_float=*/false, error_reporter); +} + +TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, + ModelT* model, ErrorReporter* error_reporter) { return QuantizeModel(builder, model, TensorType_FLOAT32, TensorType_FLOAT32, - /*allow_float=*/false, activations_type, error_reporter); + /*allow_float=*/false, error_reporter); } } // namespace optimize diff --git a/tensorflow/lite/tools/optimize/quantize_model.h b/tensorflow/lite/tools/optimize/quantize_model.h index cc801ec9870..06c30b88fd0 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.h +++ b/tensorflow/lite/tools/optimize/quantize_model.h @@ -35,9 +35,7 @@ namespace optimize { // // Note: This is a private API, subject to change. TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* input_model, - const TensorType& activations_type, - ErrorReporter* error_reporter); + ModelT* input_model, ErrorReporter* error_reporter); // Same as above, but the types of quantized inputs and outputs are // configurable. @@ -46,7 +44,6 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, - const TensorType& activations_type, ErrorReporter* error_reporter); // Same as above, but can enable allowing float intermediate operations for ops @@ -56,7 +53,6 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, bool allow_float, - const TensorType& activations_type, ErrorReporter* error_reporter); // Same as above, but enables only quantizing a whitelist of operations, @@ -67,6 +63,25 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, ModelT* input_model, const TensorType& input_type, const TensorType& output_type, bool allow_float, const std::unordered_set& operator_names, + ErrorReporter* error_reporter); + +// Same as above, but enables to provide activation type, which +// could be TensorType_INT16 or TensorType_INT8. +// +// Note: This is a private API, subject to change. +TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, + ModelT* model, const TensorType& input_type, + const TensorType& output_type, bool allow_float, + const TensorType& activations_type, + ErrorReporter* error_reporter); + +// Quantizes input_model and populates the provided builder with the new model +// with all possible input parameters. +// All functions above call this function underneath. +TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, + ModelT* model, const TensorType& input_type, + const TensorType& output_type, bool allow_float, + const std::unordered_set& operator_names, const TensorType& activations_type, ErrorReporter* error_reporter); diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc index 166d60ecc66..ef46b3fbd5d 100644 --- a/tensorflow/lite/tools/optimize/quantize_model_test.cc +++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc @@ -106,9 +106,9 @@ TEST_P(QuantizeConvModelTest, QuantizationSucceeds) { } TEST_P(QuantizeConvModelTest, SkipUnspecifiedLayer) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, - /*allow_float=*/true, {}, tensor_type_, &error_reporter_); + auto status = QuantizeModel( + &builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, + /*allow_float=*/true, {}, TensorType_FLOAT32, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); ASSERT_EQ(model_.subgraphs.size(), readonly_model_->subgraphs()->size()); // The resulting model should be the same. @@ -190,8 +190,9 @@ TEST_P(QuantizeConvModelTest, OperatorsAreUnchanged) { } TEST_P(QuantizeConvModelTest, GraphIsFullyQuantized) { - auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - tensor_type_, &error_reporter_); + auto status = + QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, + /*allow_float*/ false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (const auto& subgraph : model_.subgraphs) { for (const auto& tensor : subgraph->tensors) { @@ -210,7 +211,7 @@ TEST_P(QuantizeConvModelTest, GraphIsFullyQuantized) { TEST_P(QuantizeConvModelTest, FloatInputAndOutput) { auto status = QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, - tensor_type_, &error_reporter_); + /*allow_float*/ false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (int32_t subgraph_idx = 0; subgraph_idx < model_.subgraphs.size(); @@ -384,7 +385,7 @@ class QuantizeConcatModelTest : public QuantizeModelTest, // input1 / TEST_P(QuantizeConcatModelTest, AddRequantBeforeConcat) { auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - tensor_type_, &error_reporter_); + false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); // There is only one subgraph. @@ -549,7 +550,7 @@ class QuantizeConvModel1Test : public QuantizeModelTest { TEST_F(QuantizeConvModel1Test, VerifyConvQuantizationWithUnitScale) { auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, + QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, false, TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); const auto& subgraph = model_.subgraphs[0]; @@ -658,7 +659,7 @@ INSTANTIATE_TEST_SUITE_P(QuantizeConvModel2TestInst, QuantizeConvModel2Test, TEST_P(QuantizeConvModel2Test, VerifyConvQuantization) { auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - tensor_type_, &error_reporter_); + false, tensor_type_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto conv_op = subgraph->operators[0].get(); @@ -1353,8 +1354,7 @@ class QuantizePackTest : public QuantizeModelTest { }; TEST_F(QuantizePackTest, VerifyPack) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); + auto status = QuantizeModel(&builder_, &model_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); @@ -1418,8 +1418,7 @@ class QuantizeMinimumMaximumTest }; TEST_P(QuantizeMinimumMaximumTest, VerifyMinimumMaximum) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); + auto status = QuantizeModel(&builder_, &model_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; @@ -1500,8 +1499,7 @@ class QuantizeUnpackTest : public QuantizeModelTest { } }; TEST_F(QuantizeUnpackTest, VerifyUnpack) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, &error_reporter_); + auto status = QuantizeModel(&builder_, &model_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); From 792f553fd078a425d66c81567ca8f3588d44fcdc Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 5 Feb 2020 11:55:27 +0000 Subject: [PATCH 0017/1390] Added non-strict mode for 16x8 quantization --- tensorflow/lite/python/lite.py | 16 ++++-- tensorflow/lite/tools/optimize/BUILD | 1 + .../lite/tools/optimize/operator_property.cc | 16 ++++++ .../lite/tools/optimize/operator_property.h | 3 +- .../lite/tools/optimize/quantize_model.cc | 50 +++++++++++------- .../tools/optimize/quantize_model_test.cc | 50 +++++++++++++++--- tensorflow/lite/tools/optimize/test_util.cc | 1 + tensorflow/lite/tools/optimize/test_util.h | 5 ++ .../tools/optimize/testdata/mixed16x8.bin | Bin 0 -> 1184 bytes 9 files changed, 111 insertions(+), 31 deletions(-) create mode 100644 tensorflow/lite/tools/optimize/testdata/mixed16x8.bin diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index fc9c064faf0..1e0c89d3aa5 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -220,13 +220,16 @@ class TFLiteConverterBase(object): "type to be INT8.") def _is_int8_target_required(self): - return (set([OpsSet.TFLITE_BUILTINS_INT8]) == set( + return ((set([OpsSet.TFLITE_BUILTINS_INT8]) == set( self.target_spec.supported_ops) or - self._smallest_supported_type() == constants.INT8) + self._smallest_supported_type() == constants.INT8) and + not self._is_int16x8_target_required()) def _is_int16x8_target_required(self): - return (set([OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]) == - set(self.target_spec.supported_ops)) + return bool( + set(self.target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 + ])) def _smallest_supported_type(self): if self.target_spec.supported_types: @@ -262,6 +265,11 @@ class TFLiteConverterBase(object): def _calibrate_quantize_model(self, result, inference_input_type, inference_output_type, enable_mlir_quantizer): allow_float = not self._is_int8_target_required() and not self._is_int16x8_target_required() + if (self._is_int16x8_target_required()): + allow_float = bool( + set(self.target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS + ])) calibrate_quantize = _calibrator.Calibrator(result) activations_type = constants.INT16 if self._is_int16x8_target_required() else constants.INT8 return calibrate_quantize.calibrate_and_quantize( diff --git a/tensorflow/lite/tools/optimize/BUILD b/tensorflow/lite/tools/optimize/BUILD index 27be0f829ba..ee5e845b96b 100644 --- a/tensorflow/lite/tools/optimize/BUILD +++ b/tensorflow/lite/tools/optimize/BUILD @@ -245,6 +245,7 @@ tf_cc_test( "//tensorflow/lite/tools/optimize:testdata/maximum.bin", "//tensorflow/lite/tools/optimize:testdata/minimum.bin", "//tensorflow/lite/tools/optimize:testdata/mixed.bin", + "//tensorflow/lite/tools/optimize:testdata/mixed16x8.bin", "//tensorflow/lite/tools/optimize:testdata/multi_input_add_reshape.bin", "//tensorflow/lite/tools/optimize:testdata/pack.bin", "//tensorflow/lite/tools/optimize:testdata/single_avg_pool_min_minus_5_max_plus_5.bin", diff --git a/tensorflow/lite/tools/optimize/operator_property.cc b/tensorflow/lite/tools/optimize/operator_property.cc index 1f2d8bb4a4d..c31ad9dbb1e 100644 --- a/tensorflow/lite/tools/optimize/operator_property.cc +++ b/tensorflow/lite/tools/optimize/operator_property.cc @@ -70,6 +70,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}}; // ArgMax has no quantizable output. property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_AVERAGE_POOL_2D: property.inputs = {{0, {}}}; @@ -85,6 +86,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_SPLIT: // We skip input 0 since it is the split dim which is not real valued. @@ -143,6 +145,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}, {1, {}}}; // Comparisons have no quantizable outputs. property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_EXPAND_DIMS: // We skip input 1 as it is not real valued (it's the index of axis) and @@ -165,11 +168,13 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_HARD_SWISH: { property.inputs = {{0, {}}}; property.outputs = {{0, {}}}; property.version = 1; + property.quantizable_int16 = false; break; } case BuiltinOperator_LOG_SOFTMAX: { @@ -180,6 +185,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, tensor_property.restricted_value_int8 = {16.0 / 256.0, 127}; property.outputs = {{0, tensor_property}}; property.version = 2; + property.quantizable_int16 = false; break; } case BuiltinOperator_LOGISTIC: { @@ -736,6 +742,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.restrict_scale = {{18, 0}}; property.version = 2; } + property.quantizable_int16 = false; break; } case BuiltinOperator_L2_NORMALIZATION: { @@ -746,6 +753,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, tensor_property.restricted_value_int8 = {1 / 128.0, 0}; property.outputs = {{0, tensor_property}}; property.version = 2; + property.quantizable_int16 = false; break; } case BuiltinOperator_MAX_POOL_2D: @@ -765,6 +773,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}}; property.outputs = {{0, {}}}; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_MINIMUM: property.arbitrary_inputs = true; @@ -791,6 +800,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_QUANTIZE: property.inputs = {{0, {}}}; @@ -802,11 +812,13 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}}; property.outputs = {{0, {}}}; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_RELU_N1_TO_1: property.inputs = {{0, {}}}; property.outputs = {{0, {}}}; property.version = 1; + property.quantizable_int16 = false; break; case BuiltinOperator_RESHAPE: property.inputs = {{0, {}}}; @@ -820,6 +832,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.outputs = {{0, {}}}; property.restrict_same_input_output_scale = true; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_SHAPE: property.inputs = {{0, {}}}; @@ -866,6 +879,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, property.inputs = {{0, {}}}; property.outputs = {{0, {}}}; property.version = 2; + property.quantizable_int16 = false; break; case BuiltinOperator_TANH: { property.inputs = {{0, {}}}; @@ -899,6 +913,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, {3, tensor_property_bias}}; property.outputs = {{0, {}}}; property.version = 3; + property.quantizable_int16 = false; break; } case BuiltinOperator_TRANSPOSE: @@ -916,6 +931,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, default: // No quantized implementation exists for this operation. property.quantizable = false; + property.quantizable_int16 = false; } return property; } diff --git a/tensorflow/lite/tools/optimize/operator_property.h b/tensorflow/lite/tools/optimize/operator_property.h index 23052308568..151e314f335 100644 --- a/tensorflow/lite/tools/optimize/operator_property.h +++ b/tensorflow/lite/tools/optimize/operator_property.h @@ -65,7 +65,8 @@ struct TensorProperty { struct OperatorProperty { // Is a quantized operations currently supported. bool quantizable = true; - + // Is a quantized operations currently supported for 16x8 + bool quantizable_int16 = true; // Op has arbitrary number of inputs, such as concat. bool arbitrary_inputs = false; // Op has arbitrary number of outputs, such as slice. diff --git a/tensorflow/lite/tools/optimize/quantize_model.cc b/tensorflow/lite/tools/optimize/quantize_model.cc index bbb40080fbc..ceae3c29d9e 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.cc +++ b/tensorflow/lite/tools/optimize/quantize_model.cc @@ -43,13 +43,17 @@ namespace { // operator_names. operator_property::OperatorProperty GetOperatorProperty( const std::unordered_set& operator_names, const ModelT* model, - int subgraph_index, int op_idx, const string& operator_name) { + int subgraph_index, int op_idx, const string& operator_name, + const TensorType& activations_type) { operator_property::OperatorProperty property = operator_property::GetOperatorProperty(model, subgraph_index, op_idx); const OperatorT* op = model->subgraphs[subgraph_index]->operators[op_idx].get(); const BuiltinOperator op_code = model->operator_codes[op->opcode_index]->builtin_code; + if (activations_type == TensorType_INT16 && !property.quantizable_int16) { + property.quantizable = false; + } // The algorithm adds Dequantize and Quantize, so we don't require them to be // in the operator_names. if (op_code != BuiltinOperator_DEQUANTIZE && @@ -320,9 +324,9 @@ TfLiteStatus ApplyConstraints(ModelT* model, // Iterate backward to avoid messing with index. for (int op_idx = subgraph->operators.size() - 1; op_idx >= 0; op_idx--) { OperatorT* op = subgraph->operators[op_idx].get(); - operator_property::OperatorProperty property = - GetOperatorProperty(operator_names, model, subgraph_idx, op_idx, - subgraph->tensors[op->outputs[0]]->name); + operator_property::OperatorProperty property = GetOperatorProperty( + operator_names, model, subgraph_idx, op_idx, + subgraph->tensors[op->outputs[0]]->name, activations_type); if (!property.quantizable) { continue; } @@ -840,11 +844,17 @@ TfLiteStatus QuantizeWeightsInputOutput( OperatorT* op = subgraph->operators[op_idx].get(); const BuiltinOperator op_code = model->operator_codes[op->opcode_index]->builtin_code; - operator_property::OperatorProperty property = - GetOperatorProperty(operator_names, model, subgraph_idx, op_idx, - subgraph->tensors[op->outputs[0]]->name); + operator_property::OperatorProperty property = GetOperatorProperty( + operator_names, model, subgraph_idx, op_idx, + subgraph->tensors[op->outputs[0]]->name, activations_type); - if (!property.quantizable && !allow_float) { + if (activations_type == TensorType_INT16 && !property.quantizable && + !allow_float) { + error_reporter->Report( + "Quantization to 16x8-bit not yet supported for op: %s", + EnumNameBuiltinOperator(op_code)); + return kTfLiteError; + } else if (!property.quantizable && !allow_float) { error_reporter->Report("Quantization not yet supported for op: %s", EnumNameBuiltinOperator(op_code)); return kTfLiteError; @@ -882,9 +892,9 @@ TfLiteStatus QuantizeBiases(ModelT* model, OperatorT* op = subgraph->operators[op_idx].get(); const BuiltinOperator op_code = model->operator_codes[op->opcode_index]->builtin_code; - operator_property::OperatorProperty property = - GetOperatorProperty(operator_names, model, subgraph_idx, op_idx, - subgraph->tensors[op->outputs[0]]->name); + operator_property::OperatorProperty property = GetOperatorProperty( + operator_names, model, subgraph_idx, op_idx, + subgraph->tensors[op->outputs[0]]->name, activations_type); if (!property.quantizable) { continue; } @@ -951,15 +961,15 @@ std::unordered_set GetAllOperatorOutputs(ModelT* model) { // will not be filled by this function. TfLiteStatus FillQuantizationParams( ModelT* model, const std::unordered_set& operator_names, - ErrorReporter* error_reporter) { + const TensorType& activations_type, ErrorReporter* error_reporter) { for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size(); subgraph_idx++) { SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get(); for (size_t op_idx = 0; op_idx < subgraph->operators.size(); op_idx++) { OperatorT* op = subgraph->operators[op_idx].get(); - operator_property::OperatorProperty property = - GetOperatorProperty(operator_names, model, subgraph_idx, op_idx, - subgraph->tensors[op->outputs[0]]->name); + operator_property::OperatorProperty property = GetOperatorProperty( + operator_names, model, subgraph_idx, op_idx, + subgraph->tensors[op->outputs[0]]->name, activations_type); // Populate max, min for each input tensor. for (const std::pair& input : @@ -1048,9 +1058,9 @@ TfLiteStatus EnsureBiasScaleCompatibility( SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get(); for (size_t op_idx = 0; op_idx < subgraph->operators.size(); op_idx++) { OperatorT* op = subgraph->operators[op_idx].get(); - operator_property::OperatorProperty property = - GetOperatorProperty(operator_names, model, subgraph_idx, op_idx, - subgraph->tensors[op->outputs[0]]->name); + operator_property::OperatorProperty property = GetOperatorProperty( + operator_names, model, subgraph_idx, op_idx, + subgraph->tensors[op->outputs[0]]->name, activations_type); // Loop over all bias tensors. for (const int bias_idx : property.biases) { @@ -1174,8 +1184,8 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, const std::unordered_set& operator_names, const TensorType& activations_type, ErrorReporter* error_reporter) { - TF_LITE_ENSURE_STATUS( - FillQuantizationParams(model, operator_names, error_reporter)); + TF_LITE_ENSURE_STATUS(FillQuantizationParams( + model, operator_names, activations_type, error_reporter)); TF_LITE_ENSURE_STATUS(EnsureBiasScaleCompatibility( model, operator_names, activations_type, error_reporter)); TF_LITE_ENSURE_STATUS( diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc index ef46b3fbd5d..b73cb9a79ca 100644 --- a/tensorflow/lite/tools/optimize/quantize_model_test.cc +++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc @@ -1308,7 +1308,8 @@ TEST_F(QuantizeFCTest, VerifyFC) { EXPECT_EQ(model_.operator_codes[1]->version, 1); } -class QuantizeCustomOpTest : public QuantizeModelTest { +class QuantizeCustomOpTest : public QuantizeModelTest, + public ::testing::WithParamInterface { protected: QuantizeCustomOpTest() { input_model_ = ReadModel(internal::kModelMixed); @@ -1317,10 +1318,10 @@ class QuantizeCustomOpTest : public QuantizeModelTest { } }; -TEST_F(QuantizeCustomOpTest, VerifyMixedQuantization) { +TEST_P(QuantizeCustomOpTest, VerifyMixedQuantization) { auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, - /*allow_float=*/true, TensorType_INT8, &error_reporter_); + QuantizeModel(&builder_, &model_, GetParam(), GetParam(), + /*allow_float=*/true, GetParam(), &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto float_graph = readonly_model_->subgraphs()->Get(0); @@ -1334,8 +1335,45 @@ TEST_F(QuantizeCustomOpTest, VerifyMixedQuantization) { BuiltinOperator_CUSTOM, BuiltinOperator_CUSTOM, BuiltinOperator_QUANTIZE, BuiltinOperator_SQUEEZE}; const std::vector op_input_types = { - TensorType_INT8, TensorType_INT8, TensorType_FLOAT32, - TensorType_FLOAT32, TensorType_FLOAT32, TensorType_INT8}; + GetParam(), GetParam(), TensorType_FLOAT32, + TensorType_FLOAT32, TensorType_FLOAT32, GetParam()}; + for (int i = 0; i < subgraph->operators.size(); ++i) { + OperatorT* op = subgraph->operators[i].get(); + ASSERT_EQ(model_.operator_codes[op->opcode_index]->builtin_code, + op_codes[i]); + ASSERT_EQ(subgraph->tensors[op->inputs[0]]->type, op_input_types[i]); + } +} + +INSTANTIATE_TEST_SUITE_P(QuantizeCustomOpTest, QuantizeCustomOpTest, + ::testing::Values(TensorType_INT8, TensorType_INT16)); + +class QuantizeOp16x8Test : public QuantizeModelTest { + protected: + QuantizeOp16x8Test() { + input_model_ = ReadModel(internal::kModelMixed16x8); + readonly_model_ = input_model_->GetModel(); + readonly_model_->UnPackTo(&model_); + } +}; + +TEST_F(QuantizeOp16x8Test, VerifyMixedQuantization16x8) { + auto status = + QuantizeModel(&builder_, &model_, TensorType_INT16, TensorType_FLOAT32, + /*allow_float=*/true, TensorType_INT16, &error_reporter_); + ASSERT_EQ(kTfLiteOk, status); + const auto& subgraph = model_.subgraphs[0]; + auto float_graph = readonly_model_->subgraphs()->Get(0); + // The original model conv_2d->log_softmax + ASSERT_EQ(float_graph->operators()->size(), 2); + // The resulting model should be: + // conv_2d->dequantize->log_softmax + ASSERT_EQ(subgraph->operators.size(), 3); + const std::vector op_codes = { + BuiltinOperator_CONV_2D, BuiltinOperator_DEQUANTIZE, + BuiltinOperator_LOG_SOFTMAX}; + const std::vector op_input_types = { + TensorType_INT16, TensorType_INT16, TensorType_FLOAT32}; for (int i = 0; i < subgraph->operators.size(); ++i) { OperatorT* op = subgraph->operators[i].get(); ASSERT_EQ(model_.operator_codes[op->opcode_index]->builtin_code, diff --git a/tensorflow/lite/tools/optimize/test_util.cc b/tensorflow/lite/tools/optimize/test_util.cc index 7d5e9d65f06..379be64059f 100644 --- a/tensorflow/lite/tools/optimize/test_util.cc +++ b/tensorflow/lite/tools/optimize/test_util.cc @@ -48,6 +48,7 @@ const char* kModelWithArgMaxOp = "argmax.bin"; const char* kModelWithFCOp = "fc.bin"; const char* kModelMixed = "mixed.bin"; +const char* kModelMixed16x8 = "mixed16x8.bin"; const char* kModelSplit = "split.bin"; diff --git a/tensorflow/lite/tools/optimize/test_util.h b/tensorflow/lite/tools/optimize/test_util.h index abcdbc21d36..a49f3500288 100644 --- a/tensorflow/lite/tools/optimize/test_util.h +++ b/tensorflow/lite/tools/optimize/test_util.h @@ -76,6 +76,11 @@ extern const char* kModelWithFCOp; // reshape->custom->custom->squeeze. extern const char* kModelMixed; +// Test model with mixed quantizable and +// and un-quantizable ops for +// activations in 16-bit. +extern const char* kModelMixed16x8; + // Test model with split op. extern const char* kModelSplit; diff --git a/tensorflow/lite/tools/optimize/testdata/mixed16x8.bin b/tensorflow/lite/tools/optimize/testdata/mixed16x8.bin new file mode 100644 index 0000000000000000000000000000000000000000..c1f615e966eb164341d8ba4d56e1e359ef516388 GIT binary patch literal 1184 zcmb1PU|hJk@Whk=1Xg@J(qWH-pHr2qf_gJ_T$hI)o3dIk&(3IG59Px=4cJI#2x6#}#XYV(I z+0J^akG-Ncll|PMa(3Mu8TRL0SnTgR%h)@mnb}(lm)YNIw6YgxVzg)bKG|+RlZpKt zJw5xz<5l)73~XS3f_%oo0OCRX4RR*~*zYW0d45p%_&fV6IOpe;r52T>rs#pgMu35V zfrmi^9-=%93=C%&7#I#PFfeQZ>t|qKVqjp10EaXKg9uzL$aD~f!~!UsCBSAu><7t! z#9%atEy4g&0g5?x29UiF(?BX17%iaT4N?JfI|Bn3g8&0NIIg+C$`}|y<{)EGyo1z( z+z4WW zD#KjsMVOIGd{Sm&G1PpJ-Aw=g{|BWtklir* zzTmOXnne5bgHjVy!V-%z6O(dM{R`lJ2H6M0ZU6uOPhenR0EGd}4Hf@!hao6Gk>Y1i zJi_7vW+%w+3=DkWRLB5M83GJ^3=9k|43KmRvICTcL8%N>N`Ue_t~dqBd8VZ1m1LGw i;<5{ru0Z)8WF{z&u`#fM%Mb=5HU Date: Tue, 25 Feb 2020 16:20:43 -0500 Subject: [PATCH 0018/1390] added doctest to tf.signal.frame --- tensorflow/python/ops/signal/shape_ops.py | 32 +++++++++++------------ 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/tensorflow/python/ops/signal/shape_ops.py b/tensorflow/python/ops/signal/shape_ops.py index bb4ac7ec33c..89b35014005 100644 --- a/tensorflow/python/ops/signal/shape_ops.py +++ b/tensorflow/python/ops/signal/shape_ops.py @@ -68,29 +68,27 @@ def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, axis=-1, For example: - ```python - # A batch size 3 tensor of 9152 audio samples. - audio = tf.random.normal([3, 9152]) - - # Compute overlapping frames of length 512 with a step of 180 (frames overlap - # by 332 samples). By default, only 49 frames are generated since a frame - # with start position j*180 for j > 48 would overhang the end. - frames = tf.signal.frame(audio, 512, 180) - frames.shape.assert_is_compatible_with([3, 49, 512]) - - # When pad_end is enabled, the final two frames are kept (padded with zeros). - frames = tf.signal.frame(audio, 512, 180, pad_end=True) - frames.shape.assert_is_compatible_with([3, 51, 512]) - ``` + >>> # A batch size 3 tensor of 9152 audio samples. + >>> audio = tf.random.normal([3, 9152]) + >>> + >>> # Compute overlapping frames of length 512 with a step of 180 (frames overlap + >>> # by 332 samples). By default, only 49 frames are generated since a frame + >>> # with start position j*180 for j > 48 would overhang the end. + >>> frames = tf.signal.frame(audio, 512, 180) + >>> frames.shape.assert_is_compatible_with([3, 49, 512]) + >>> + >>> # When pad_end is enabled, the final two frames are kept (padded with zeros). + >>> frames = tf.signal.frame(audio, 512, 180, pad_end=True) + >>> frames.shape.assert_is_compatible_with([3, 51, 512]) If the dimension along `axis` is N, and `pad_end=False`, the number of frames can be computed by: ```python - frames = 1 + (N - frame_size) // frame_step + num_frames = 1 + (N - frame_size) // frame_step ``` If `pad_end=True`, the number of frames can be computed by: ```python - frames = -(-N // frame_step) # ceiling division + num_frames = -(-N // frame_step) # ceiling division ``` Args: @@ -106,7 +104,7 @@ def frame(signal, frame_length, frame_step, pad_end=False, pad_value=0, axis=-1, name: An optional name for the operation. Returns: - A `Tensor` of frames with shape `[..., frames, frame_length, ...]`. + A `Tensor` of frames with shape `[..., num_frames, frame_length, ...]`. Raises: ValueError: If `frame_length`, `frame_step`, `pad_value`, or `axis` are not From db0d4681210c831f505a455a0951fa37482184d6 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 17 Mar 2020 16:30:32 +0000 Subject: [PATCH 0019/1390] Corrected after merge with master. Tested: strict mode and non-strict mode. Change-Id: I7e03d08133f39cc65a18875e65ce5cdddaf2d6a4 --- tensorflow/lite/python/lite.py | 62 ++++++++++++++++--- .../optimize/calibration_wrapper_pybind11.cc | 8 +-- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 7e5f8ce704f..900398d7a6f 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -180,11 +180,13 @@ class QuantizationMode(object): def post_training_int8_no_float(self): """Post training int8 quantize, disallow float fallback.""" return (self._is_int8_target_required() and + not self._is_int16x8_target_required() and self._representative_dataset is not None) def post_training_int8_allow_float(self): """Post training int8 quantize, allow float fallback.""" return (self._any_optimization_enabled() and + not self._is_int16x8_target_required() and self._representative_dataset is not None and self._smallest_supported_type() == constants.INT8) @@ -193,6 +195,18 @@ class QuantizationMode(object): return (self._any_optimization_enabled() and self._contains_training_quant_op()) + def post_training_int16x8_no_float(self): + """Post training int16x8 quantize, disallow float fallback.""" + return (not self._is_int8_target_required() and + self._is_int16x8_target_required() and + not self._is_allow_float() and + self._representative_dataset is not None) + + def post_training_int16x8_allow_float(self): + """Post training int16x8 quantize, allow float fallback.""" + return (self._is_int16x8_target_required() and + self._is_allow_float()) + def post_training_dynamic_range_int8(self): """Post training int8 const, on-the-fly int8 quantize of dynamic tensors.""" # Post-training dynamic range quantization is only enabled if post-training @@ -212,9 +226,14 @@ class QuantizationMode(object): return not (self.post_training_int8_no_float() or self.post_training_int8_allow_float() or self.training_time_int8_allow_float() or + self.post_training_int16x8_no_float() or + self.post_training_int16x8_allow_float() or self.post_training_dynamic_range_int8() or self.post_training_fp16()) + def activations_type(self): + return constants.INT16 if self._is_int16x8_target_required() else constants.INT8 + # Below are helpers for the above functions. def _validate_int8_required(self): @@ -244,6 +263,18 @@ class QuantizationMode(object): self._target_spec.supported_ops) or set(self._target_spec.supported_types) == set([constants.INT8])) + def _is_int16x8_target_required(self): + return bool( + set(self._target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 + ])) + + def _is_allow_float(self): + return bool( + set(self._target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS + ])) + def _any_optimization_enabled(self): return bool( set(self._optimizations).intersection([ @@ -309,13 +340,13 @@ class TFLiteConverterBase(object): return _get_grappler_config(optimizers) def _calibrate_quantize_model(self, result, inference_input_type, - inference_output_type, allow_float): + inference_output_type, activations_type, allow_float): if not isinstance(self.representative_dataset, RepresentativeDataset): self.representative_dataset = RepresentativeDataset( self.representative_dataset) calibrate_quantize = _calibrator.Calibrator(result) - activations_type = constants.INT16 if self._is_int16x8_target_required() else constants.INT8 + if (self.experimental_calibrate_only: return calibrate_quantize.calibrate(self.representative_dataset.input_gen) else: @@ -608,12 +639,20 @@ class TFLiteConverterV2(TFLiteConverterBase): output_tensors=output_tensors, **converter_kwargs) + activations_type = quant_mode.activations_type() + if quant_mode.post_training_int8_no_float(): result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, False) + constants.FLOAT, activations_type, False) elif quant_mode.post_training_int8_allow_float(): result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, True) + constants.FLOAT, activations_type, True) + elif quant_mode.post_training_int16x8_no_float(): + result = self._calibrate_quantize_model(result, constants.FLOAT, + constants.FLOAT, activations_type, False) + elif quant_mode.post_training_int16x8_allow_float(): + result = self._calibrate_quantize_model(result, constants.FLOAT, + constants.FLOAT, activations_type, True) return result @@ -1114,6 +1153,8 @@ class TFLiteConverter(TFLiteConverterBase): quant_mode.post_training_int8_no_float() or quant_mode.post_training_int8_allow_float() or quant_mode.post_training_dynamic_range_int8() or + quant_mode.post_training_int16x8_no_float() or + quant_mode.post_training_int16x8_allow_float() or quant_mode.post_training_fp16()) if post_training_optimize: # Post training optimizations require that TOCO outputs a float model. @@ -1223,12 +1264,20 @@ class TFLiteConverter(TFLiteConverterBase): output_arrays=self._output_arrays, **converter_kwargs) + activations_type = quant_mode.activations_type() + if quant_mode.post_training_int8_no_float(): result = self._calibrate_quantize_model(result, inference_input_type, - inference_output_type, False) + inference_output_type, activations_type, False) elif quant_mode.post_training_int8_allow_float(): result = self._calibrate_quantize_model(result, inference_input_type, - inference_output_type, True) + inference_output_type, activations_type, True) + elif quant_mode.post_training_int16x8_no_float(): + result = self._calibrate_quantize_model(result, inference_input_type, + inference_output_type, activations_type, False) + elif quant_mode.post_training_int16x8_allow_float(): + result = self._calibrate_quantize_model(result, inference_input_type, + inference_output_type, activations_type, True) return result @@ -1334,7 +1383,6 @@ class TocoConverter(object): @classmethod @_deprecation.deprecated( - None, "Use `lite.TFLiteConverter.from_keras_model_file` instead.") def from_keras_model_file(cls, model_file, input_arrays=None, diff --git a/tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc b/tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc index f56b23090b9..9a8fea5d1f6 100644 --- a/tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc +++ b/tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc @@ -40,17 +40,17 @@ PYBIND11_MODULE(_pywrap_tensorflow_lite_calibration_wrapper, m) { }) .def("QuantizeModel", [](CalibrationWrapper& self, int input_py_type, int output_py_type, - bool allow_float, bool enable_mlir_quantizer) { + bool allow_float, int activations_py_type, bool enable_mlir_quantizer) { return tensorflow::pyo_or_throw( self.QuantizeModel(input_py_type, output_py_type, allow_float, - enable_mlir_quantizer)); + activations_py_type, enable_mlir_quantizer)); }) .def("QuantizeModel", [](CalibrationWrapper& self, int input_py_type, int output_py_type, - bool allow_float) { + bool allow_float, int activations_py_type) { return tensorflow::pyo_or_throw( self.QuantizeModel(input_py_type, output_py_type, allow_float, - /*enable_mlir_quantizer=*/false)); + activations_py_type, /*enable_mlir_quantizer=*/false)); }) .def("QuantizeModel", [](CalibrationWrapper& self, int input_py_type, int output_py_type, bool allow_float, From fea3433cfc42d1a7ed050779092f3c92b077cd48 Mon Sep 17 00:00:00 2001 From: Tomohiro Ubukata Date: Sun, 15 Mar 2020 05:22:41 +0000 Subject: [PATCH 0020/1390] Add error checks --- .../core/platform/cloud/curl_http_request.cc | 8 +++-- .../platform/default/posix_file_system.cc | 14 ++++++-- .../core/platform/default/subprocess.cc | 36 ++++++++++++++----- tensorflow/core/platform/path.cc | 4 ++- tensorflow/core/platform/platform_strings.cc | 4 ++- .../android_armv7a_cpu_utils_helper.cc | 8 +++-- 6 files changed, 56 insertions(+), 18 deletions(-) diff --git a/tensorflow/core/platform/cloud/curl_http_request.cc b/tensorflow/core/platform/cloud/curl_http_request.cc index a227edb1fb0..39a1f7a35b2 100644 --- a/tensorflow/core/platform/cloud/curl_http_request.cc +++ b/tensorflow/core/platform/cloud/curl_http_request.cc @@ -166,7 +166,9 @@ CurlHttpRequest::~CurlHttpRequest() { libcurl_->curl_slist_free_all(resolve_list_); } if (put_body_) { - fclose(put_body_); + if (fclose(put_body_) != 0) { + LOG(ERROR) << "fclose() failed: " << strerror(errno); + } } if (curl_) { libcurl_->curl_easy_cleanup(curl_); @@ -237,7 +239,9 @@ Status CurlHttpRequest::SetPutFromFile(const string& body_filepath, is_method_set_ = true; method_ = RequestMethod::kPut; if (put_body_) { - fclose(put_body_); + if (fclose(put_body_) != 0) { + LOG(ERROR) << "fclose() failed: " << strerror(errno); + } } put_body_ = fopen(body_filepath.c_str(), "r"); if (!put_body_) { diff --git a/tensorflow/core/platform/default/posix_file_system.cc b/tensorflow/core/platform/default/posix_file_system.cc index 106a0412fb7..05c2b2762d4 100644 --- a/tensorflow/core/platform/default/posix_file_system.cc +++ b/tensorflow/core/platform/default/posix_file_system.cc @@ -51,7 +51,11 @@ class PosixRandomAccessFile : public RandomAccessFile { public: PosixRandomAccessFile(const string& fname, int fd) : filename_(fname), fd_(fd) {} - ~PosixRandomAccessFile() override { close(fd_); } + ~PosixRandomAccessFile() override { + if (close(fd_) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } + } Status Name(StringPiece* result) const override { *result = filename_; @@ -229,7 +233,9 @@ Status PosixFileSystem::NewReadOnlyMemoryRegionFromFile( } else { result->reset(new PosixReadOnlyMemoryRegion(address, st.st_size)); } - close(fd); + if (close(fd) < 0) { + s = IOError(fname, errno); + } } return s; } @@ -256,7 +262,9 @@ Status PosixFileSystem::GetChildren(const string& dir, result->push_back(entry->d_name); } } - closedir(d); + if (closedir(d) < 0) { + return IOError(dir, errno); + } return Status::OK(); } diff --git a/tensorflow/core/platform/default/subprocess.cc b/tensorflow/core/platform/default/subprocess.cc index 562f4cd2d0c..acf7073b9a4 100644 --- a/tensorflow/core/platform/default/subprocess.cc +++ b/tensorflow/core/platform/default/subprocess.cc @@ -102,11 +102,15 @@ void SubProcess::FreeArgs() { void SubProcess::ClosePipes() { for (int i = 0; i < kNFds; i++) { if (parent_pipe_[i] >= 0) { - close(parent_pipe_[i]); + if (close(parent_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } parent_pipe_[i] = -1; } if (child_pipe_[i] >= 0) { - close(child_pipe_[i]); + if (close(child_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } child_pipe_[i] = -1; } } @@ -215,7 +219,9 @@ bool SubProcess::Start() { running_ = true; for (int i = 0; i < kNFds; i++) { if (child_pipe_[i] >= 0) { - close(child_pipe_[i]); + if (close(child_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } child_pipe_[i] = -1; } } @@ -227,7 +233,9 @@ bool SubProcess::Start() { int devnull_fd = -1; for (int i = 0; i < kNFds; i++) { if (parent_pipe_[i] >= 0) { - close(parent_pipe_[i]); + if (close(parent_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } parent_pipe_[i] = -1; } @@ -242,7 +250,9 @@ bool SubProcess::Start() { _exit(1); } } - close(child_pipe_[i]); + if (close(child_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } child_pipe_[i] = -1; break; @@ -264,14 +274,18 @@ bool SubProcess::Start() { } } } else { - close(i); + if (close(i) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } } break; } } if (devnull_fd >= 0) { - close(devnull_fd); + if (close(devnull_fd) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } } // Execute the child program. @@ -379,7 +393,9 @@ int SubProcess::Communicate(const string* stdin_input, string* stdout_output, // Special case: if no data is given to send to the child process, // close the pipe to unblock the child, and skip the file descriptor. if (stdin_input == nullptr) { - close(parent_pipe_[i]); + if (close(parent_pipe_[i]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } parent_pipe_[i] = -1; continue; } @@ -441,7 +457,9 @@ int SubProcess::Communicate(const string* stdin_input, string* stdout_output, fds[i].fd = -1; fd_remain--; // Close the child's stdin pipe to unblock the process. - close(parent_pipe_[CHAN_STDIN]); + if (close(parent_pipe_[CHAN_STDIN]) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } parent_pipe_[CHAN_STDIN] = -1; } } else if (!retry(errno)) { diff --git a/tensorflow/core/platform/path.cc b/tensorflow/core/platform/path.cc index 1e88328aace..a041ac67d72 100644 --- a/tensorflow/core/platform/path.cc +++ b/tensorflow/core/platform/path.cc @@ -320,7 +320,9 @@ string GetTempFilename(const string& extension) { if (fd < 0) { LOG(FATAL) << "Failed to create temp file."; } else { - close(fd); + if (close(fd) < 0) { + LOG(ERROR) << "close() failed: " << strerror(errno); + } return tmp_filepath; } } diff --git a/tensorflow/core/platform/platform_strings.cc b/tensorflow/core/platform/platform_strings.cc index 489a211ccf7..af8787f4fbc 100644 --- a/tensorflow/core/platform/platform_strings.cc +++ b/tensorflow/core/platform/platform_strings.cc @@ -52,7 +52,9 @@ int GetPlatformStrings(const std::string& path, } result = (ferror(ifp) == 0) ? 0 : errno; - fclose(ifp); + if (fclose(ifp) != 0) { + result = errno; + } } else { result = errno; } diff --git a/tensorflow/core/platform/profile_utils/android_armv7a_cpu_utils_helper.cc b/tensorflow/core/platform/profile_utils/android_armv7a_cpu_utils_helper.cc index 0534443d17c..f75d3533d17 100644 --- a/tensorflow/core/platform/profile_utils/android_armv7a_cpu_utils_helper.cc +++ b/tensorflow/core/platform/profile_utils/android_armv7a_cpu_utils_helper.cc @@ -118,10 +118,14 @@ int64 AndroidArmV7ACpuUtilsHelper::ReadCpuFrequencyFile( const int retval = fscanf(fp, "%lld", &freq_in_khz); if (retval < 0) { LOG(WARNING) << "Failed to \"" << file_path << "\""; - fclose(fp); + if (fclose(fp) != 0) { + LOG(WARNING) << "fclose() failed: " << strerror(errno); + } return INVALID_CPU_FREQUENCY; } - fclose(fp); + if (fclose(fp) != 0) { + LOG(WARNING) << "fclose() failed: " << strerror(errno); + } return freq_in_khz * 1000; // The file contains cpu frequency in khz } From e5e81a3b52cabd767ce504a209f3b1f1bf9fa58b Mon Sep 17 00:00:00 2001 From: khaled besrour Date: Sat, 21 Mar 2020 18:13:50 +0100 Subject: [PATCH 0021/1390] Bug Fix : resolve cast error by adding cleaner type for MKL --- tensorflow/core/kernels/mkl_conv_ops.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/mkl_conv_ops.h b/tensorflow/core/kernels/mkl_conv_ops.h index 7d1e19566ee..2ee2a621067 100644 --- a/tensorflow/core/kernels/mkl_conv_ops.h +++ b/tensorflow/core/kernels/mkl_conv_ops.h @@ -51,7 +51,7 @@ using mkldnn::stream; namespace tensorflow { #ifdef ENABLE_MKLDNN_V1 -#define MKLDNN_SIZE_DTYPE long int +#define MKLDNN_SIZE_DTYPE memory::dim #else #define MKLDNN_SIZE_DTYPE int #endif // ENABLE_MKLDNN_V1 From ce81212fdfe7c5e84e354b1a52b24299a20899e9 Mon Sep 17 00:00:00 2001 From: khaled besrour Date: Sat, 21 Mar 2020 18:19:39 +0100 Subject: [PATCH 0022/1390] Bug Fix : correct compilation erreur by using dynamic allocation --- tensorflow/core/util/mkl_util.h | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h index e0a399f2d6c..a1b6cc758f8 100644 --- a/tensorflow/core/util/mkl_util.h +++ b/tensorflow/core/util/mkl_util.h @@ -1260,8 +1260,8 @@ inline Status CreateBlockedMemDescHelper(const memory::dims& dim, DCHECK_EQ(dim.size(), strides.size()); #ifdef ENABLE_MKLDNN_V1 const int kNumDims = dim.size(); - mkldnn_dim_t input_dims[kNumDims]; - mkldnn_dim_t input_strides[kNumDims]; + mkldnn_dim_t* input_dims = new mkldnn_dim_t[kNumDims]; + mkldnn_dim_t* input_strides = new mkldnn_dim_t[kNumDims]; for (int i = 0; i < kNumDims; ++i) { input_dims[i] = dim[i]; input_strides[i] = strides[i]; @@ -1270,7 +1270,11 @@ inline Status CreateBlockedMemDescHelper(const memory::dims& dim, mkldnn_memory_desc_init_by_strides(blocked_md, kNumDims, input_dims, memory::convert_to_c(dtype), input_strides); + delete[] input_dims; + delete[] input_strides; } catch (mkldnn::error& e) { + delete[] input_dims; + delete[] input_strides; return Status(error::Code::INTERNAL, tensorflow::strings::StrCat( "Failed to create blocked memory descriptor.", From 6da9fa6111ab1be251b1138a0a591bc022b58cdf Mon Sep 17 00:00:00 2001 From: Peng Meng Date: Thu, 26 Mar 2020 00:05:02 +0800 Subject: [PATCH 0023/1390] fixOpFuse --- .../core/grappler/optimizers/remapper.cc | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index 5b41ad38089..b47c71fd5c8 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -1476,6 +1476,13 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return true; }; + const auto is_conv2d_candidate = [&]() -> bool { + if (!IsConv2D(*node_def)) return false; + if (GetDataTypeFromAttr(*node_def, "T") != DT_FLOAT) return false; + + return true; + }; + // Candidate for a FusedBatchNorm fusion. const auto is_batch_norm_fusion_candidate = [&]() -> bool { if (!IsRelu(*node_def)) return false; @@ -1506,7 +1513,7 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return false; }; - return is_batch_norm_candidate() || is_batch_norm_fusion_candidate(); + return is_conv2d_candidate() || is_batch_norm_candidate() || is_batch_norm_fusion_candidate(); } } // namespace @@ -1564,6 +1571,17 @@ Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item, } #endif //! INTEL_MKL + // Infer properties lazily in case they are not needed. + if (!ctx.inferred_graph_properties && RequiresInferredShapes(ctx, i)) { + const bool assume_valid_feeds = opt_level_ == RewriterConfig::AGGRESSIVE; + TF_RETURN_IF_ERROR(ctx.graph_properties.InferStatically( + assume_valid_feeds, + /*aggressive_shape_inference=*/false, + /*include_input_tensor_values=*/true, + /*include_output_tensor_values=*/false)); + ctx.inferred_graph_properties = true; + } + // Remap {Conv2D,MatMul}+BiasAdd into the _Fused{Conv2D,MatMul} ContractionWithBiasAdd contract_with_bias; if (allow_non_differentiable_rewrites && @@ -1592,6 +1610,7 @@ Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item, // Remove this once TF-MKL supports _FusedConv2D with these operations. #ifndef INTEL_MKL // Remap Conv2D+Squeeze+BiasAdd into the _FusedConv2D+Squeeze. + ContractionWithSqueezeAndBiasAdd contract_with_squeeze_and_bias; if (allow_non_differentiable_rewrites && FindConv2DWithSqueezeAndBias(ctx, i, &contract_with_squeeze_and_bias)) { @@ -1624,16 +1643,6 @@ Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item, } #endif // !INTEL_MKL - // Infer properties lazily in case they are not needed. - if (!ctx.inferred_graph_properties && RequiresInferredShapes(ctx, i)) { - const bool assume_valid_feeds = opt_level_ == RewriterConfig::AGGRESSIVE; - TF_RETURN_IF_ERROR(ctx.graph_properties.InferStatically( - assume_valid_feeds, - /*aggressive_shape_inference=*/false, - /*include_input_tensor_values=*/true, - /*include_output_tensor_values=*/false)); - ctx.inferred_graph_properties = true; - } // Remap FusedBatchNorm++ into the _FusedBatchNormEx. FusedBatchNormEx fused_batch_norm_ex; From 5673a09ff777ee08660ce9f71d90958e18bd995c Mon Sep 17 00:00:00 2001 From: Peng Meng Date: Thu, 26 Mar 2020 00:30:54 +0800 Subject: [PATCH 0024/1390] fix code format --- tensorflow/core/grappler/optimizers/remapper.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index b47c71fd5c8..fdfb8e379b2 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -1513,7 +1513,8 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return false; }; - return is_conv2d_candidate() || is_batch_norm_candidate() || is_batch_norm_fusion_candidate(); + return is_conv2d_candidate() || is_batch_norm_candidate() || + is_batch_norm_fusion_candidate(); } } // namespace From 69ee4de053a14bdf883a0e6726bb2b374b71c973 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Fri, 27 Mar 2020 15:59:37 +0000 Subject: [PATCH 0025/1390] Fix for the broken 16-bit interface after latest changes to master. --- .../python/optimize/calibration_wrapper.cc | 14 ++--- .../lite/tools/optimize/quantize_model.cc | 12 ++-- .../lite/tools/optimize/quantize_model.h | 12 ++-- .../tools/optimize/quantize_model_test.cc | 55 ++++++++++--------- 4 files changed, 49 insertions(+), 44 deletions(-) diff --git a/tensorflow/lite/python/optimize/calibration_wrapper.cc b/tensorflow/lite/python/optimize/calibration_wrapper.cc index 660ea8d2d1b..ad82581bcba 100644 --- a/tensorflow/lite/python/optimize/calibration_wrapper.cc +++ b/tensorflow/lite/python/optimize/calibration_wrapper.cc @@ -233,12 +233,11 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false); flatbuffers::FlatBufferBuilder builder; auto status = kTfLiteOk; - - status = tflite::optimize::QuantizeModel( - &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), - TfLiteTypeToSchemaType(output_type), allow_float, - TfLiteTypeToSchemaType(activations_type), error_reporter_.get()); - } + + status = tflite::optimize::QuantizeModelAllOperators( + &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), + TfLiteTypeToSchemaType(output_type), allow_float, + TfLiteTypeToSchemaType(activations_type), error_reporter_.get()); if (status != kTfLiteOk) { error_reporter_->exception(); @@ -269,8 +268,7 @@ PyObject* CalibrationWrapper::QuantizeModel(int input_py_type, auto status = tflite::optimize::QuantizeModel( &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type), TfLiteTypeToSchemaType(output_type), allow_float, {op_name}, - TensorType_INT8, - error_reporter_.get()); + TensorType_INT8, error_reporter_.get()); if (status != kTfLiteOk) { error_reporter_->exception(); return nullptr; diff --git a/tensorflow/lite/tools/optimize/quantize_model.cc b/tensorflow/lite/tools/optimize/quantize_model.cc index 32cd2b8c25a..0892e7ae52a 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.cc +++ b/tensorflow/lite/tools/optimize/quantize_model.cc @@ -1240,11 +1240,13 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, return kTfLiteOk; } -TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* model, const TensorType& input_type, - const TensorType& output_type, bool allow_float, - const TensorType& activations_type, - ErrorReporter* error_reporter) { +TfLiteStatus QuantizeModelAllOperators(flatbuffers::FlatBufferBuilder* builder, + ModelT* model, + const TensorType& input_type, + const TensorType& output_type, + bool allow_float, + const TensorType& activations_type, + ErrorReporter* error_reporter) { return QuantizeModel(builder, model, input_type, output_type, allow_float, GetAllOperatorOutputs(model), activations_type, error_reporter); diff --git a/tensorflow/lite/tools/optimize/quantize_model.h b/tensorflow/lite/tools/optimize/quantize_model.h index 06c30b88fd0..29f581d2b35 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.h +++ b/tensorflow/lite/tools/optimize/quantize_model.h @@ -69,11 +69,13 @@ TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, // could be TensorType_INT16 or TensorType_INT8. // // Note: This is a private API, subject to change. -TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder, - ModelT* model, const TensorType& input_type, - const TensorType& output_type, bool allow_float, - const TensorType& activations_type, - ErrorReporter* error_reporter); +TfLiteStatus QuantizeModelAllOperators(flatbuffers::FlatBufferBuilder* builder, + ModelT* model, + const TensorType& input_type, + const TensorType& output_type, + bool allow_float, + const TensorType& activations_type, + ErrorReporter* error_reporter); // Quantizes input_model and populates the provided builder with the new model // with all possible input parameters. diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc index 0f780e4d3da..885fa98992c 100644 --- a/tensorflow/lite/tools/optimize/quantize_model_test.cc +++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc @@ -190,9 +190,9 @@ TEST_P(QuantizeConvModelTest, OperatorsAreUnchanged) { } TEST_P(QuantizeConvModelTest, GraphIsFullyQuantized) { - auto status = - QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - /*allow_float*/ false, tensor_type_, &error_reporter_); + auto status = QuantizeModelAllOperators( + &builder_, &model_, tensor_type_, tensor_type_, + /*allow_float*/ false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (const auto& subgraph : model_.subgraphs) { for (const auto& tensor : subgraph->tensors) { @@ -209,9 +209,9 @@ TEST_P(QuantizeConvModelTest, GraphIsFullyQuantized) { } TEST_P(QuantizeConvModelTest, FloatInputAndOutput) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, - /*allow_float*/ false, tensor_type_, &error_reporter_); + auto status = QuantizeModelAllOperators( + &builder_, &model_, TensorType_FLOAT32, TensorType_FLOAT32, + /*allow_float*/ false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); for (int32_t subgraph_idx = 0; subgraph_idx < model_.subgraphs.size(); @@ -384,8 +384,9 @@ class QuantizeConcatModelTest : public QuantizeModelTest, // concat - output // input1 / TEST_P(QuantizeConcatModelTest, AddRequantBeforeConcat) { - auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - false, tensor_type_, &error_reporter_); + auto status = + QuantizeModelAllOperators(&builder_, &model_, tensor_type_, tensor_type_, + false, tensor_type_, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); // There is only one subgraph. @@ -549,9 +550,9 @@ class QuantizeConvModel1Test : public QuantizeModelTest { }; TEST_F(QuantizeConvModel1Test, VerifyConvQuantizationWithUnitScale) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_INT8, TensorType_INT8, false, - TensorType_INT8, &error_reporter_); + auto status = QuantizeModelAllOperators(&builder_, &model_, TensorType_INT8, + TensorType_INT8, false, + TensorType_INT8, &error_reporter_); EXPECT_EQ(status, kTfLiteOk); const auto& subgraph = model_.subgraphs[0]; @@ -658,8 +659,9 @@ INSTANTIATE_TEST_SUITE_P(QuantizeConvModel2TestInst, QuantizeConvModel2Test, TensorType_INT16})); TEST_P(QuantizeConvModel2Test, VerifyConvQuantization) { - auto status = QuantizeModel(&builder_, &model_, tensor_type_, tensor_type_, - false, tensor_type_, &error_reporter_); + auto status = + QuantizeModelAllOperators(&builder_, &model_, tensor_type_, tensor_type_, + false, tensor_type_, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto conv_op = subgraph->operators[0].get(); @@ -1308,8 +1310,9 @@ TEST_F(QuantizeFCTest, VerifyFC) { EXPECT_EQ(model_.operator_codes[1]->version, 1); } -class QuantizeCustomOpTest : public QuantizeModelTest, - public ::testing::WithParamInterface { +class QuantizeCustomOpTest + : public QuantizeModelTest, + public ::testing::WithParamInterface { protected: QuantizeCustomOpTest() { input_model_ = ReadModel(internal::kModelMixed); @@ -1319,9 +1322,9 @@ class QuantizeCustomOpTest : public QuantizeModelTest, }; TEST_P(QuantizeCustomOpTest, VerifyMixedQuantization) { - auto status = - QuantizeModel(&builder_, &model_, GetParam(), GetParam(), - /*allow_float=*/true, GetParam(), &error_reporter_); + auto status = QuantizeModelAllOperators( + &builder_, &model_, GetParam(), GetParam(), + /*allow_float=*/true, GetParam(), &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto float_graph = readonly_model_->subgraphs()->Get(0); @@ -1335,7 +1338,7 @@ TEST_P(QuantizeCustomOpTest, VerifyMixedQuantization) { BuiltinOperator_CUSTOM, BuiltinOperator_CUSTOM, BuiltinOperator_QUANTIZE, BuiltinOperator_SQUEEZE}; const std::vector op_input_types = { - GetParam(), GetParam(), TensorType_FLOAT32, + GetParam(), GetParam(), TensorType_FLOAT32, TensorType_FLOAT32, TensorType_FLOAT32, GetParam()}; for (int i = 0; i < subgraph->operators.size(); ++i) { OperatorT* op = subgraph->operators[i].get(); @@ -1358,9 +1361,9 @@ class QuantizeOp16x8Test : public QuantizeModelTest { }; TEST_F(QuantizeOp16x8Test, VerifyMixedQuantization16x8) { - auto status = - QuantizeModel(&builder_, &model_, TensorType_INT16, TensorType_FLOAT32, - /*allow_float=*/true, TensorType_INT16, &error_reporter_); + auto status = QuantizeModelAllOperators( + &builder_, &model_, TensorType_INT16, TensorType_FLOAT32, + /*allow_float=*/true, TensorType_INT16, &error_reporter_); ASSERT_EQ(kTfLiteOk, status); const auto& subgraph = model_.subgraphs[0]; auto float_graph = readonly_model_->subgraphs()->Get(0); @@ -1369,11 +1372,11 @@ TEST_F(QuantizeOp16x8Test, VerifyMixedQuantization16x8) { // The resulting model should be: // conv_2d->dequantize->log_softmax ASSERT_EQ(subgraph->operators.size(), 3); - const std::vector op_codes = { - BuiltinOperator_CONV_2D, BuiltinOperator_DEQUANTIZE, - BuiltinOperator_LOG_SOFTMAX}; + const std::vector op_codes = {BuiltinOperator_CONV_2D, + BuiltinOperator_DEQUANTIZE, + BuiltinOperator_LOG_SOFTMAX}; const std::vector op_input_types = { - TensorType_INT16, TensorType_INT16, TensorType_FLOAT32}; + TensorType_INT16, TensorType_INT16, TensorType_FLOAT32}; for (int i = 0; i < subgraph->operators.size(); ++i) { OperatorT* op = subgraph->operators[i].get(); ASSERT_EQ(model_.operator_codes[op->opcode_index]->builtin_code, From e6fd34c57c4ad7402c2c127a3380cfba65ecdb55 Mon Sep 17 00:00:00 2001 From: Fredrik Knutsson Date: Wed, 11 Sep 2019 09:43:52 +0200 Subject: [PATCH 0026/1390] Add support for offline planned tensor allocations By adding metadata to the model, it is possible to set arena offset for each tensor. Change-Id: Idd646c00a6e34e0c2603896d748cd5680a57f015 --- .../memory_planner/greedy_memory_planner.cc | 181 +++++++++++------- .../memory_planner/greedy_memory_planner.h | 19 +- tensorflow/lite/micro/micro_allocator.cc | 82 +++++++- 3 files changed, 210 insertions(+), 72 deletions(-) diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc index faea73e9169..7763afa2075 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc @@ -42,8 +42,8 @@ GreedyMemoryPlanner::GreedyMemoryPlanner(unsigned char* scratch_buffer, int scratch_buffer_size) : buffer_count_(0), need_to_calculate_offsets_(true) { const int per_buffer_size = sizeof(BufferRequirements) + // requirements_ - sizeof(int) + // buffer_sizes_sorted_by_size_ - sizeof(int) + // buffer_ids_sorted_by_size_ + sizeof(int) + // buffer_sizes_sorted_ + sizeof(int) + // buffer_ids_sorted_ sizeof(ListEntry) + // buffers_sorted_by_offset_ sizeof(int); // buffer_offsets_; // Allocate the arrays we need within the scratch buffer arena. @@ -53,10 +53,10 @@ GreedyMemoryPlanner::GreedyMemoryPlanner(unsigned char* scratch_buffer, requirements_ = reinterpret_cast(next_free); next_free += sizeof(BufferRequirements) * max_buffer_count_; - buffer_sizes_sorted_by_size_ = reinterpret_cast(next_free); + buffer_sizes_sorted_ = reinterpret_cast(next_free); next_free += sizeof(int) * max_buffer_count_; - buffer_ids_sorted_by_size_ = reinterpret_cast(next_free); + buffer_ids_sorted_ = reinterpret_cast(next_free); next_free += sizeof(int) * max_buffer_count_; buffers_sorted_by_offset_ = reinterpret_cast(next_free); @@ -81,11 +81,24 @@ TfLiteStatus GreedyMemoryPlanner::AddBuffer( current->size = size; current->first_time_used = first_time_used; current->last_time_used = last_time_used; + current->offline_offset = kOnlinePlannedBuffer; ++buffer_count_; need_to_calculate_offsets_ = true; return kTfLiteOk; } +TfLiteStatus GreedyMemoryPlanner::AddBuffer( + tflite::ErrorReporter* error_reporter, int size, int first_time_used, + int last_time_used, int offline_offset) { + BufferRequirements* current = &requirements_[buffer_count_]; + if (AddBuffer(error_reporter, size, first_time_used, last_time_used) != + kTfLiteOk) { + return kTfLiteError; + } + current->offline_offset = offline_offset; + return kTfLiteOk; +} + bool GreedyMemoryPlanner::DoesEntryOverlapInTime( const GreedyMemoryPlanner::ListEntry* entry, const int first_time_used, const int last_time_used) const { @@ -107,7 +120,7 @@ GreedyMemoryPlanner::NextSimultaneouslyActiveBuffer( ListEntry* result = nullptr; ListEntry* candidate_next_entry; if (start == nullptr) { - candidate_next_entry = &buffers_sorted_by_offset_[0]; + candidate_next_entry = &buffers_sorted_by_offset_[first_entry_index_]; } else { if (start->next_entry_index == -1) { return nullptr; @@ -139,29 +152,51 @@ void GreedyMemoryPlanner::CalculateOffsetsIfNeeded() { // This helps find a more compact layout. Intuitively, you can think // about putting the large buffers in place first, and then the // smaller buffers can fit in the gaps, rather than fragmenting the - // gaps with small buffers at the beginning. + // gaps with small buffers at the beginning. Add offline planned offsets + // first in the list, since they have a predetermined offset. + int idx_from_tail = buffer_count_; + int idx_from_head = 0; for (int i = 0; i < buffer_count_; ++i) { - buffer_sizes_sorted_by_size_[i] = requirements_[i].size; - buffer_ids_sorted_by_size_[i] = i; - buffer_offsets_[i] = -1; + if (requirements_[i].offline_offset == kOnlinePlannedBuffer) { + idx_from_tail--; + buffer_sizes_sorted_[idx_from_tail] = requirements_[i].size; + buffer_ids_sorted_[idx_from_tail] = i; + buffer_offsets_[i] = -1; + } else { + buffer_sizes_sorted_[idx_from_head] = requirements_[i].size; + buffer_ids_sorted_[idx_from_head] = i; + buffer_offsets_[i] = requirements_[i].offline_offset; + idx_from_head++; + } } - // This sorting algorithm is naive, and may end up taking a very long time - // with hundreds of buffers. - ReverseSortInPlace(buffer_sizes_sorted_by_size_, buffer_ids_sorted_by_size_, - buffer_count_); - // Put the largest buffer at offset zero to start the process. - ListEntry* first_entry = &buffers_sorted_by_offset_[0]; - first_entry->offset = 0; - first_entry->requirements_index = buffer_ids_sorted_by_size_[0]; - first_entry->next_entry_index = -1; + // This sorting algorithm is naive, and may end up taking a very long time + // with hundreds of buffers. Do not sort the offline planned offsets. + ReverseSortInPlace(&buffer_sizes_sorted_[idx_from_head], + &buffer_ids_sorted_[idx_from_head], + buffer_count_ - idx_from_head); + + // Initialize the first entry to the first buffer in + // buffer_ids_sorted_. + // - If there are no offline planned offsets, the largest buffer will be + // first, and the buffers will be handled in size order. + // - If offline offsets are present, these will be handled first in order + // for the greedy algorithm to utilized gaps in the offline plan. + first_entry_index_ = 0; next_free_entry_ = 1; - buffer_offsets_[buffer_ids_sorted_by_size_[0]] = 0; + ListEntry* first_entry = &buffers_sorted_by_offset_[first_entry_index_]; + first_entry->next_entry_index = -1; // to mark the entry as end of list + int buffer_id = buffer_ids_sorted_[0]; + first_entry->requirements_index = buffer_id; + if (requirements_[buffer_id].offline_offset == kOnlinePlannedBuffer) { + buffer_offsets_[buffer_id] = 0; + } + first_entry->offset = buffer_offsets_[buffer_id]; // Work through the rest of the buffers to find a good gap to place each one. for (int i = 1; i < buffer_count_; ++i) { // The id is the order the buffer was originally added by the client. - const int buffer_id = buffer_ids_sorted_by_size_[i]; + const int buffer_id = buffer_ids_sorted_[i]; // Look at what size and time range the buffer needs to be active. BufferRequirements* wanted_requirements = &requirements_[buffer_id]; const int wanted_size = wanted_requirements->size; @@ -173,37 +208,43 @@ void GreedyMemoryPlanner::CalculateOffsetsIfNeeded() { // so that it's easy to find the next buffer in memory, and so the gap. // The candidate_entry variable holds the buffer that we're considering // placing the current buffer after. - ListEntry* prior_entry = nullptr; + int candidate_offset = 0; // Loop through the offset-ordered list of buffers, looking for gaps. - while (true) { - // Find out what the next active buffer is. - ListEntry* next_entry = NextSimultaneouslyActiveBuffer( - prior_entry, wanted_first_time_used, wanted_last_time_used); + if (wanted_requirements->offline_offset == kOnlinePlannedBuffer) { + ListEntry* prior_entry = nullptr; + while (true) { + // Find out what the next active buffer is. + ListEntry* next_entry = NextSimultaneouslyActiveBuffer( + prior_entry, wanted_first_time_used, wanted_last_time_used); - if (prior_entry) { - BufferRequirements* candidate_requirements = - &requirements_[prior_entry->requirements_index]; - const int prior_entry_offset = - prior_entry->offset + candidate_requirements->size; - if (prior_entry_offset > candidate_offset) { - candidate_offset = prior_entry_offset; + if (prior_entry) { + BufferRequirements* candidate_requirements = + &requirements_[prior_entry->requirements_index]; + const int prior_entry_offset = + prior_entry->offset + candidate_requirements->size; + if (prior_entry_offset > candidate_offset) { + candidate_offset = prior_entry_offset; + } } + if (next_entry == nullptr) { + // We're at the end of the list, so we can always append the buffer + // here. + break; + } + // Find out how much space there is between us and the next buffer. + const int gap = next_entry->offset - candidate_offset; + if (gap >= wanted_size) { + // This entry has a big enough gap between it and the next, so + // use it! + break; + } + // The gap wasn't big enough, so move on to another candidate. + prior_entry = next_entry; } - if (next_entry == nullptr) { - // We're at the end of the list, so we can always append the buffer - // here. - break; - } - // Find out how much space there is between us and the next buffer. - const int gap = next_entry->offset - candidate_offset; - if (gap >= wanted_size) { - // This entry has a big enough gap between it and the next, so - // use it! - break; - } - // The gap wasn't big enough, so move on to another candidate. - prior_entry = next_entry; + } else { + // Offline planned offset are to be considered constant + candidate_offset = wanted_requirements->offline_offset; } // At this point, we've either found a gap (possibly at the end of the // list) and want to place the buffer there, or there are no other active @@ -217,26 +258,36 @@ void GreedyMemoryPlanner::CalculateOffsetsIfNeeded() { new_entry->requirements_index = buffer_id; const int new_entry_index = next_free_entry_; ++next_free_entry_; - ListEntry* current_entry = first_entry; - // Make sure that we insert the buffer at the correct place in the ordered - // list. - while (true) { - const int next_entry_index = current_entry->next_entry_index; - if (next_entry_index == -1) { - // We're at the end of the list, so just add the new entry here. - current_entry->next_entry_index = new_entry_index; - new_entry->next_entry_index = -1; - break; + + if (first_entry->offset > candidate_offset) { + // The new entry offset is smaller than the first entry offset => + // replace the first entry + first_entry = new_entry; + first_entry->next_entry_index = first_entry_index_; + first_entry_index_ = new_entry_index; + } else { + ListEntry* current_entry = first_entry; + // Make sure that we insert the buffer at the correct place in the + // buffer-offset-ordered list + while (true) { + const int next_entry_index = current_entry->next_entry_index; + if (next_entry_index == -1) { + // We're at the end of the list, so just add the new entry here. + current_entry->next_entry_index = new_entry_index; + new_entry->next_entry_index = -1; + break; + } + // not at the end of the list -> take a look at next entry + ListEntry* next_entry = &buffers_sorted_by_offset_[next_entry_index]; + if (next_entry->offset > candidate_offset) { + // We're at the right spot to do an insertion and retain the sorting + // order, so place the new entry here. + new_entry->next_entry_index = current_entry->next_entry_index; + current_entry->next_entry_index = new_entry_index; + break; + } + current_entry = next_entry; } - ListEntry* next_entry = &buffers_sorted_by_offset_[next_entry_index]; - if (next_entry->offset > candidate_offset) { - // We're at the right spot to do an insertion and retain the sorting - // order, so place the new entry here. - new_entry->next_entry_index = current_entry->next_entry_index; - current_entry->next_entry_index = new_entry_index; - break; - } - current_entry = next_entry; } } } diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h index f2c77ed94f3..d874b70e732 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h @@ -21,6 +21,8 @@ limitations under the License. namespace tflite { +constexpr int kOnlinePlannedBuffer = -1; + // A memory planner that uses a greedy algorithm to arrange buffers in memory // to minimize the overall arena size needed. // @@ -59,6 +61,12 @@ class GreedyMemoryPlanner : public MemoryPlanner { TfLiteStatus AddBuffer(ErrorReporter* error_reporter, int size, int first_time_used, int last_time_used) override; + // Record details of an offline planned buffer offset we want to place. + // offline_offset is the buffer offset from the start of the arena. + TfLiteStatus AddBuffer(ErrorReporter* error_reporter, int size, + int first_time_used, int last_time_used, + int offline_offset); + // Returns the high-water mark of used memory. This is the minimum size of a // memory arena you'd need to allocate to hold these buffers. size_t GetMaximumMemorySize() override; @@ -110,16 +118,23 @@ class GreedyMemoryPlanner : public MemoryPlanner { // Records the client-provided information about each buffer. struct BufferRequirements { int size; + int offline_offset; int first_time_used; int last_time_used; }; // Working arrays used during the layout algorithm. BufferRequirements* requirements_; - int* buffer_sizes_sorted_by_size_; - int* buffer_ids_sorted_by_size_; + // buffer_sizes_sorted_ and buffer_ids_sorted_ are sorted according to: + // { + // offline planned buffers, + // online planned buffers sorted by size + // } + int* buffer_sizes_sorted_; + int* buffer_ids_sorted_; ListEntry* buffers_sorted_by_offset_; int next_free_entry_; + int first_entry_index_; // Stores the outcome of the plan, the location of each buffer in the arena. int* buffer_offsets_; diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index c3044a0351f..46edec2bf43 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -37,6 +37,7 @@ struct AllocationInfo { int last_used; bool needs_allocating; void** output_ptr; + int offline_offset; }; // We align tensor buffers to 16-byte boundaries, since this is a common @@ -112,9 +113,17 @@ class AllocationInfoBuilder { return Allocate(); } + // Check if model contains offline planned buffer offsets. + // - If there's no metadata available, offline_planner_offsets is not set + // - If there's metadata available, offline_planner_offsets will point to the + // first offset in the metadata buffer list. + TfLiteStatus GetOfflinePlannedOffsets(const Model* model, + int** offline_planner_offsets); + // Add allocaiton information for the tensors. - TfLiteStatus AddTensors(const SubGraph* subgraph, + TfLiteStatus AddTensors(const SubGraph* subgraph, int* offline_offsets, TfLiteTensor* runtime_tensors); + // Add allocation information for the scratch buffers. TfLiteStatus AddScratchBuffers(internal::ScratchBufferHandle* buffer_handles); @@ -148,6 +157,7 @@ TfLiteStatus AllocationInfoBuilder::Allocate() { } TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, + int* offline_offsets, TfLiteTensor* runtime_tensors) { // Set up allocation info for all tensors. for (size_t i = 0; i < tensor_count_; ++i) { @@ -159,6 +169,11 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, current->last_used = -1; current->needs_allocating = (runtime_tensors[i].data.raw == nullptr) && (!subgraph->tensors()->Get(i)->is_variable()); + if (offline_offsets) { + current->offline_offset = offline_offsets[i]; + } else { + current->offline_offset = kOnlinePlannedBuffer; + } } for (size_t i = 0; i < subgraph->inputs()->size(); ++i) { @@ -216,6 +231,51 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, return kTfLiteOk; } +// The tensor offsets will be encoded in the metadata:[Metadata] field of the +// Model. The following encoding applies: +// +// | Metadata component | Value | +// | name:string | “OfflineMemoryAllocation” | +// | buffer:unit | Index of buffer containing memory allocation data | +// +// The buffer contents for the memory allocation is a list of 32-bit integers of +// the following format: +// +// | Offset | Value | +// | 0 | Offline allocation format version – set to 0 | +// | 1 | Subgraph index to which this allocation applies | +// | 2 | Number offsets following: n | +// | 3 | Arena byte offset of tensor #0 or -1 to allocate at runtime | +// | 4 | Arena byte offset of tensor #1 or -1 to allocate at runtime | +// | 3+(n-1) | Arena byte offset of tensor #(n-1) or -1 to allocate at runtime | +TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets( + const Model* model, int** offline_planner_offsets) { + if (model->metadata()) { + for (int i = 0; i < model->metadata()->size(); ++i) { + auto metadata = model->metadata()->Get(i); + if (strncmp(metadata->name()->c_str(), "OfflineMemoryAllocation", + strlen("OfflineMemoryAllocation")) == 0) { + const flatbuffers::Vector>* buffers = + model->buffers(); + auto* buffer = (*buffers)[metadata->buffer()]; + auto* array = buffer->data(); + const uint32_t* metadata_buffer = (uint32_t*)array->data(); + const int32_t nbr_tensors = metadata_buffer[2]; + *offline_planner_offsets = (int32_t*)&metadata_buffer[3]; + + if (tensor_count_ != nbr_tensors) { + TF_LITE_REPORT_ERROR(reporter_, + "Nbr of offline buffer offsets (%d) in metadata " + "not equal nbr tensors (%d)\n", + nbr_tensors, tensor_count_); + return kTfLiteError; + } + } + } + } + return kTfLiteOk; +} + TfLiteStatus AllocationInfoBuilder::AddScratchBuffers( internal::ScratchBufferHandle* buffer_handles) { // Set up allocation info for buffers. @@ -241,9 +301,17 @@ TfLiteStatus CreatePlan(ErrorReporter* error_reporter, MemoryPlanner* planner, if (current->needs_allocating) { size_t aligned_bytes_required = AlignSizeUp(current->bytes, kBufferAlignment); - TF_LITE_ENSURE_STATUS( - planner->AddBuffer(error_reporter, aligned_bytes_required, - current->first_created, current->last_used)); + if (current->offline_offset == kOnlinePlannedBuffer) { + TF_LITE_ENSURE_STATUS( + planner->AddBuffer(error_reporter, aligned_bytes_required, + current->first_created, current->last_used)); + } else { + TF_LITE_ENSURE_STATUS( + (static_cast(planner)) + ->AddBuffer(error_reporter, aligned_bytes_required, + current->first_created, current->last_used, + current->offline_offset)); + } } } return kTfLiteOk; @@ -546,7 +614,11 @@ TfLiteStatus MicroAllocator::FinishTensorAllocation() { AllocationInfoBuilder builder(error_reporter_, &tmp_allocator); TF_LITE_ENSURE_STATUS( builder.Init(tensors_->size(), scratch_buffer_count_)); - TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph_, context_->tensors)); + int* offline_planner_offsets = nullptr; + TF_LITE_ENSURE_STATUS( + builder.GetOfflinePlannedOffsets(model_, &offline_planner_offsets)); + TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph_, offline_planner_offsets, + context_->tensors)); TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_handles_)); const AllocationInfo* allocation_info = builder.Finish(); From 2ac92c48cd31721ec882b110e2af95edc505dc66 Mon Sep 17 00:00:00 2001 From: Fredrik Knutsson Date: Wed, 25 Mar 2020 12:30:25 +0100 Subject: [PATCH 0027/1390] Add helper functions in micro allocator One function to check metadata correctness and another to print model data. Change-Id: I2500fbbac25b376d068e3d9a1d190249da461eef --- tensorflow/lite/micro/micro_allocator.cc | 88 ++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 46edec2bf43..a9324eb35cd 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -96,6 +96,94 @@ TfLiteStatus AllocateVariables( return kTfLiteOk; } +// Helper function to print model flatbuffer data. This function is not called +// by default. Hence it's not linked in to the final binary code. +void PrintModelData(const Model* model, ErrorReporter* error_reporter) { + auto* subgraphs = model->subgraphs(); + const SubGraph* subgraph = (*subgraphs)[0]; + const flatbuffers::Vector>* tensors = + subgraph->tensors(); + const flatbuffers::Vector>* buffers = + model->buffers(); + TF_LITE_REPORT_ERROR(error_reporter, "==== Model info: ====="); + for (int i = 0; i < tensors->size(); ++i) { + const tflite::Tensor& flatbuffer_tensor = *tensors->Get(i); + auto* quantization = flatbuffer_tensor.quantization(); + size_t type_size, tensor_size; + auto* buffer = (*buffers)[flatbuffer_tensor.buffer()]; + auto* array = buffer->data(); + int array_size = 0; + if (array) { + array_size = array->size(); + } + BytesRequiredForTensor(flatbuffer_tensor, &tensor_size, &type_size, + error_reporter); + TF_LITE_REPORT_ERROR( + error_reporter, "Tensor index: %d arena tensor %d size %d", i, + !array_size && !flatbuffer_tensor.is_variable(), tensor_size); + } +} + +// Helper function to check flatbuffer metadata correctness. This function is +// not called by default. Hence it's not linked in to the final binary code. +TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, + ErrorReporter* error_reporter) { + if (model->metadata()) { + for (int i = 0; i < model->metadata()->size(); ++i) { + auto metadata = model->metadata()->Get(i); + if (strncmp(metadata->name()->c_str(), "OfflineMemoryAllocation", + strlen("OfflineMemoryAllocation")) == 0) { + auto* subgraphs = model->subgraphs(); + const SubGraph* subgraph = (*subgraphs)[0]; + const flatbuffers::Vector>* tensors = + subgraph->tensors(); + const flatbuffers::Vector>* buffers = + model->buffers(); + int nbr_tflite_tensors = tensors->size(); + auto* buffer = (*buffers)[metadata->buffer()]; + auto* array = buffer->data(); + const uint32_t* metadata_buffer = (uint32_t*)array->data(); + int version = metadata_buffer[0]; + int subgraph_idx = metadata_buffer[1]; + const int nbr_offline_offsets = metadata_buffer[2]; + int* offline_planner_offsets = (int*)&metadata_buffer[3]; + + TF_LITE_REPORT_ERROR(error_reporter, "==== Model metadata info: ====="); + TF_LITE_REPORT_ERROR(error_reporter, + "Offline planner metadata found, version %d, " + "subgraph %d, nbr offline offsets %d", + version, subgraph_idx, nbr_offline_offsets); + for (int i = 0; i < nbr_offline_offsets; ++i) { + TF_LITE_REPORT_ERROR( + error_reporter, + "Offline planner tensor index %d, offline offset: %d", i, + offline_planner_offsets[i]); + } + + if (version != 1) { + TF_LITE_REPORT_ERROR(error_reporter, "Version not supported! (%d)\n", + version); + return kTfLiteError; + } + if (subgraph_idx != 0) { + TF_LITE_REPORT_ERROR(error_reporter, + "Only 1 subgraph supported! Subgraph idx (%d)\n", + subgraph_idx); + return kTfLiteError; + } + if (nbr_tflite_tensors != nbr_offline_offsets) { + TF_LITE_REPORT_ERROR(error_reporter, + "Nbr of offline buffer offsets (%d) in metadata " + "not equal nbr tensors (%d)\n", + nbr_offline_offsets, nbr_tflite_tensors); + return kTfLiteError; + } + } + } + } + return kTfLiteOk; +} + // A helper class to construct AllocationInfo array. This array contains the // lifetime of tensors / scratch_buffer and will be used to calculate the memory // plan. Methods need to be called in order from `Init`, `Add*`, to `Finish`. From dd673e306a0fec5b87f9c979d4c29524c575b766 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Fri, 6 Mar 2020 14:04:46 +0100 Subject: [PATCH 0028/1390] OfflinePlanner: Swap offsets in TestAllocationForModelsWithBranches The offline planner sorts the tensors in reverse order, so the testcase have to be updated accordingly. Change-Id: Ic3a1193489d6ad5f592db1c9a289b01083ad9c62 --- tensorflow/lite/micro/micro_allocator_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 47eefff90b5..765a447c044 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -185,10 +185,10 @@ TF_LITE_MICRO_TEST(TestAllocationForModelsWithBranches) { // bytes = 2 * 2 * 3 * sizeof(float32) = 48, same for other tensors. TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[0].bytes); // t1 can't reuse any memory, as n0 requires both t0 and t1. - TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(96, context.tensors[1].data.uint8 - start); // t2 can't reuse any memory, as n1 requires both t0 and t2. Also n2 requires // both t1 and t2. - TF_LITE_MICRO_EXPECT_EQ(96, context.tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[2].data.uint8 - start); // t3 reuses the same memory from t0 as t0 is not an input to any node. TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[3].data.uint8 - start); } From 7737c6222fafdb6425d0d2395a1be5c6f4205780 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Wed, 11 Mar 2020 12:49:23 +0100 Subject: [PATCH 0029/1390] Testcases for offline planned tensors Added testcases. Expanded ModelBuilder for have metadata support. Change-Id: I7e8bf3b3537d126086aef52bb3d6fe572aa8e7a0 --- tensorflow/lite/micro/micro_allocator_test.cc | 183 ++++++++++++++++++ tensorflow/lite/micro/test_helpers.cc | 106 +++++++++- tensorflow/lite/micro/test_helpers.h | 13 ++ 3 files changed, 292 insertions(+), 10 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 765a447c044..69fb82910b0 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -227,4 +227,187 @@ TF_LITE_MICRO_TEST(TestFinishComplexTensorAllocation) { tflite::testing::EnsureUniqueVariableTensorBuffer(&context, 7); } +TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { + int version = 1; + int subgraph = 0; + int nbr_tensors = 4; + const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + + nbr_tensors] = {version, subgraph, + nbr_tensors, // header + // memory offsets: + -1, -1, -1, -1}; + + // The structure is identical to the one in + // TestAllocationForModelsWithBranches + std::vector node_list = { + { + {0}, // input + {1} // output + }, + { + {0}, // input + {2} // output + }, + { + {1, 2}, // input1, input2 + {3} // output + }}; + + const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( + nbr_tensors, metadata_buffer, node_list); + + TfLiteContext context; + constexpr size_t arena_size = 4096; + uint8_t arena[arena_size]; + tflite::MicroAllocator allocator(&context, model, arena, arena_size, + micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + + // Since all of the tensors are online planned and the model structure is + // identical to that in TestAllocationForModelsWithBranches, + // the offsets be should identical to that test. + uint8_t* start = context.tensors[0].data.uint8; + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[0].bytes); + TF_LITE_MICRO_EXPECT_EQ(96, context.tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[3].data.uint8 - start); +} + +TF_LITE_MICRO_TEST(OfflinePlannerBasic) { + int nbr_tensors = 4; + const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + + nbr_tensors] = {1, 0, nbr_tensors, + 0, // t0 + 48, // t1 + 0, // t2 + 48}; // t3 + + int t0 = 0; + int t1 = 1; + int t2 = 2; + int t3 = 3; + + std::vector node_list = {{ + {t0}, // input + {t1} // output + }, + { + {t1}, // input + {t2} // output + }, + { + {t2}, // input + {t3} // output + }}; + + const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( + nbr_tensors, metadata_buffer, node_list); + + TfLiteContext context; + constexpr size_t arena_size = 4096; + uint8_t arena[arena_size]; + tflite::MicroAllocator allocator(&context, model, arena, arena_size, + micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + + uint8_t* start = context.tensors[0].data.uint8; + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[3].data.uint8 - start); +} + +TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { + int nbr_tensors = 4; + const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + + nbr_tensors] = { + 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors + // memory offsets: + 0, // t0 + 0, // t1 + 48, // t2 + -1}; // t3 + + int t0 = 0; + int t1 = 1; + int t2 = 2; + int t3 = 3; + + std::vector node_list = { + { + {t0, t1}, // input, scratch + {t2} // output + }, + { + {t2}, // input + {t3} // output + }, + }; + + const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( + nbr_tensors, metadata_buffer, node_list); + + TfLiteContext context; + constexpr size_t arena_size = 4096; + uint8_t arena[arena_size]; + tflite::MicroAllocator allocator(&context, model, arena, arena_size, + micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + + uint8_t* start = context.tensors[0].data.uint8; + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[3].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[0].bytes); +} + +TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { + int nbr_tensors = 5; + const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + + nbr_tensors] = { + 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors + // memory offsets: + 0, // t0 + 48, // t1 + -1, // t2 + 0, // t3 + -1}; // t4 + + int t0 = 0; + int t1 = 1; + int t2 = 2; + int t3 = 3; + int t4 = 4; + + std::vector node_list = { + { + {t0, t1}, // input, scratch + {t2}, // output + }, + { + {t2}, // input + {t3, t4}, // output1, output2 + }, + }; + + const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( + nbr_tensors, metadata_buffer, node_list); + + TfLiteContext context; + constexpr size_t arena_size = 4096; + uint8_t arena[arena_size]; + tflite::MicroAllocator allocator(&context, model, arena, arena_size, + micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + + uint8_t* start = context.tensors[0].data.uint8; + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[1].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(96, context.tensors[2].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(48, context.tensors[4].data.uint8 - start); + TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[3].data.uint8 - start); +} + TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/lite/micro/test_helpers.cc b/tensorflow/lite/micro/test_helpers.cc index 77a1cc82f3b..b39d3b2916f 100644 --- a/tensorflow/lite/micro/test_helpers.cc +++ b/tensorflow/lite/micro/test_helpers.cc @@ -48,7 +48,7 @@ class StackAllocator : public flatbuffers::Allocator { return *inst; } - static constexpr size_t kStackAllocatorSize = 4096; + static constexpr size_t kStackAllocatorSize = 8192; private: uint8_t data_backing_[kStackAllocatorSize]; @@ -94,6 +94,10 @@ class ModelBuilder { Node AddNode(Operator op, std::initializer_list inputs, std::initializer_list outputs); + void AddMetadata(const char* description_string, + const int32_t* metadata_buffer_data, + size_t num_elements); + // Constructs the flatbuffer model using `builder_` and return a pointer to // it. The returned model has the same lifetime as `builder_`. const Model* BuildModel(std::initializer_list inputs, @@ -116,6 +120,16 @@ class ModelBuilder { static constexpr int kMaxTensors = 50; flatbuffers::Offset tensors_[kMaxTensors]; + + static constexpr int kMaxMetadataBuffers = 10; + + static constexpr int kMaxMetadatas = 10; + flatbuffers::Offset metadata_[kMaxMetadatas]; + + flatbuffers::Offset metadata_buffers_[kMaxMetadataBuffers]; + + int nbr_of_metadata_buffers_ = 0; + int next_tensor_id_ = 0; }; @@ -142,13 +156,34 @@ ModelBuilder::Node ModelBuilder::AddNode( return next_operator_id_ - 1; } +void ModelBuilder::AddMetadata(const char* description_string, + const int32_t* metadata_buffer_data, + size_t num_elements) { + metadata_[ModelBuilder::nbr_of_metadata_buffers_] = + CreateMetadata(*builder_, + builder_->CreateString(description_string), + 1 + ModelBuilder::nbr_of_metadata_buffers_); + + metadata_buffers_[nbr_of_metadata_buffers_] = tflite::CreateBuffer(*builder_, + builder_->CreateVector((uint8_t*)metadata_buffer_data, + sizeof(uint32_t) * num_elements)); + + ModelBuilder::nbr_of_metadata_buffers_++; +} + const Model* ModelBuilder::BuildModel( std::initializer_list inputs, std::initializer_list outputs) { // Model schema requires an empty buffer at idx 0. - constexpr size_t kBufferSize = 1; - const flatbuffers::Offset buffers[kBufferSize] = { - tflite::CreateBuffer(*builder_)}; + size_t kBufferSize = 1 + ModelBuilder::nbr_of_metadata_buffers_; + flatbuffers::Offset buffers[kBufferSize]; + buffers[0] = tflite::CreateBuffer(*builder_); + + // Place the metadata buffers first in the buffer since the indices for them + // have already been set in AddMetadata() + for (int i = 1; i < ModelBuilder::nbr_of_metadata_buffers_ + 1; ++i) { + buffers[i] = metadata_buffers_[i - 1]; + } // TFLM only supports single subgraph. constexpr size_t subgraphs_size = 1; @@ -159,12 +194,26 @@ const Model* ModelBuilder::BuildModel( builder_->CreateVector(outputs.begin(), outputs.size()), builder_->CreateVector(operators_, next_operator_id_), builder_->CreateString("test_subgraph"))}; - const flatbuffers::Offset model_offset = tflite::CreateModel( - *builder_, 0, - builder_->CreateVector(operator_codes_, next_operator_code_id_), - builder_->CreateVector(subgraphs, subgraphs_size), - builder_->CreateString("teset_model"), - builder_->CreateVector(buffers, kBufferSize)); + + flatbuffers::Offset model_offset; + if (ModelBuilder::nbr_of_metadata_buffers_ > 0) { + model_offset = tflite::CreateModel( + *builder_, 0, + builder_->CreateVector(operator_codes_, next_operator_code_id_), + builder_->CreateVector(subgraphs, subgraphs_size), + builder_->CreateString("teset_model"), + builder_->CreateVector(buffers, kBufferSize), + 0, + builder_->CreateVector(metadata_, ModelBuilder::nbr_of_metadata_buffers_)); + } else { + model_offset = tflite::CreateModel( + *builder_, 0, + builder_->CreateVector(operator_codes_, next_operator_code_id_), + builder_->CreateVector(subgraphs, subgraphs_size), + builder_->CreateString("teset_model"), + builder_->CreateVector(buffers, kBufferSize)); + } + tflite::FinishModelBuffer(*builder_, model_offset); void* model_pointer = builder_->GetBufferPointer(); const Model* model = flatbuffers::GetRoot(model_pointer); @@ -243,6 +292,35 @@ const Model* BuildSimpleModelWithBranch() { return model_builder.BuildModel({t0}, {t3}); } +const Model* BuildModelWithOfflinePlanning(int number_of_tensors, + const int32_t* metadata_buffer, + std::vector node_conn) { + using flatbuffers::Offset; + flatbuffers::FlatBufferBuilder* fb_builder = BuilderInstance(); + + ModelBuilder model_builder(fb_builder); + + const int op_id = + model_builder.RegisterOp(BuiltinOperator_CUSTOM, "mock_custom", + /* version= */ 0); + + int tensors[number_of_tensors]; + + for (int i = 0; i < number_of_tensors; ++i) { + tensors[i] = model_builder.AddTensor(TensorType_FLOAT32, {2, 2, 3}); + } + + for (int i = 0; i < node_conn.size(); i++) { + model_builder.AddNode(op_id, node_conn[i].input, node_conn[i].output); + } + + model_builder.AddMetadata("OfflineMemoryAllocation", + metadata_buffer, number_of_tensors + tflite::testing::kOfflinePlannerHeaderSize); + + return model_builder.BuildModel(node_conn[0].input, + node_conn[node_conn.size() - 1].output); +} + const Model* BuildSimpleMockModel() { using flatbuffers::Offset; flatbuffers::FlatBufferBuilder* builder = BuilderInstance(); @@ -496,6 +574,14 @@ const Model* GetSimpleModelWithBranch() { return model; } +const Model* GetModelWithOfflinePlanning(int num_tensors, + const int32_t* metadata_buffer, + std::vector node_conn) { + const Model* model = + BuildModelWithOfflinePlanning(num_tensors, metadata_buffer, node_conn); + return model; +} + const Model* GetSimpleStatefulModel() { static Model* model = nullptr; if (!model) { diff --git a/tensorflow/lite/micro/test_helpers.h b/tensorflow/lite/micro/test_helpers.h index f4e7fa8dfba..26aeeb086ef 100644 --- a/tensorflow/lite/micro/test_helpers.h +++ b/tensorflow/lite/micro/test_helpers.h @@ -27,6 +27,14 @@ limitations under the License. namespace tflite { namespace testing { +constexpr int kOfflinePlannerHeaderSize = 3; + +struct NodeConnection_ { + std::initializer_list input; + std::initializer_list output; +}; +typedef struct NodeConnection_ NodeConnection; + // Returns a simple example flatbuffer TensorFlow Lite model. Contains 1 input, // 1 layer of weights, 1 output Tensor, and 1 operator. const Model* GetSimpleMockModel(); @@ -38,6 +46,11 @@ const Model* GetComplexMockModel(); // Returns a simple flatbuffer model with two branches. const Model* GetSimpleModelWithBranch(); +// Returns a simple flatbuffer model with offline planned tensors +const Model* GetModelWithOfflinePlanning(int num_tensors, + const int32_t* metadata_buffer, + std::vector node_conn); + // Returns a flatbuffer model with `simple_stateful_op` const Model* GetSimpleStatefulModel(); From b8571e365d2e907f66693551b488ce2f72e3507b Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Thu, 2 Apr 2020 14:31:17 +0200 Subject: [PATCH 0030/1390] Fix compile error when building for ARM Cortex-M4 etc. --- tensorflow/lite/micro/micro_allocator.cc | 12 ++++++------ tensorflow/lite/micro/test_helpers.h | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index a9324eb35cd..fdfeb9d409c 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -37,7 +37,7 @@ struct AllocationInfo { int last_used; bool needs_allocating; void** output_ptr; - int offline_offset; + int32_t offline_offset; }; // We align tensor buffers to 16-byte boundaries, since this is a common @@ -206,10 +206,10 @@ class AllocationInfoBuilder { // - If there's metadata available, offline_planner_offsets will point to the // first offset in the metadata buffer list. TfLiteStatus GetOfflinePlannedOffsets(const Model* model, - int** offline_planner_offsets); + int32_t** offline_planner_offsets); // Add allocaiton information for the tensors. - TfLiteStatus AddTensors(const SubGraph* subgraph, int* offline_offsets, + TfLiteStatus AddTensors(const SubGraph* subgraph, int32_t* offline_offsets, TfLiteTensor* runtime_tensors); // Add allocation information for the scratch buffers. @@ -245,7 +245,7 @@ TfLiteStatus AllocationInfoBuilder::Allocate() { } TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, - int* offline_offsets, + int32_t* offline_offsets, TfLiteTensor* runtime_tensors) { // Set up allocation info for all tensors. for (size_t i = 0; i < tensor_count_; ++i) { @@ -337,7 +337,7 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, // | 4 | Arena byte offset of tensor #1 or -1 to allocate at runtime | // | 3+(n-1) | Arena byte offset of tensor #(n-1) or -1 to allocate at runtime | TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets( - const Model* model, int** offline_planner_offsets) { + const Model* model, int32_t** offline_planner_offsets) { if (model->metadata()) { for (int i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); @@ -702,7 +702,7 @@ TfLiteStatus MicroAllocator::FinishTensorAllocation() { AllocationInfoBuilder builder(error_reporter_, &tmp_allocator); TF_LITE_ENSURE_STATUS( builder.Init(tensors_->size(), scratch_buffer_count_)); - int* offline_planner_offsets = nullptr; + int32_t* offline_planner_offsets = nullptr; TF_LITE_ENSURE_STATUS( builder.GetOfflinePlannedOffsets(model_, &offline_planner_offsets)); TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph_, offline_planner_offsets, diff --git a/tensorflow/lite/micro/test_helpers.h b/tensorflow/lite/micro/test_helpers.h index 26aeeb086ef..81416e06dcc 100644 --- a/tensorflow/lite/micro/test_helpers.h +++ b/tensorflow/lite/micro/test_helpers.h @@ -30,8 +30,8 @@ namespace testing { constexpr int kOfflinePlannerHeaderSize = 3; struct NodeConnection_ { - std::initializer_list input; - std::initializer_list output; + std::initializer_list input; + std::initializer_list output; }; typedef struct NodeConnection_ NodeConnection; From 26e4ac1d76b43d9cf5288c914985e42fd0bfdbbf Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Mon, 6 Apr 2020 16:50:17 +0000 Subject: [PATCH 0031/1390] Update tensorflow/python/ops/array_ops.py --- tensorflow/python/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 9367374717e..fc29aafb7de 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -3507,7 +3507,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): For example: Given the following input, - * `hypothesis` is a `tf.SparseTensor` of shape `(2, 1, 1)` + * `hypothesis` is a `tf.SparseTensor` of shape `[2, 1, 1]` * `truth` is a `tf.SparseTensor` of shape `(2, 2, 2)` >>> hypothesis = tf.SparseTensor( From 05edf7e11bb3d91f0e2b75e67914c2379cf0cd89 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Mon, 6 Apr 2020 16:50:26 +0000 Subject: [PATCH 0032/1390] Update tensorflow/python/ops/array_ops.py --- tensorflow/python/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index fc29aafb7de..004bbeac7f6 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -3527,7 +3527,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): array([[inf, 1. ], [0.5, 1. ]], dtype=float32)> - The operaton returns a dense Tensor of shape `(2, 2)` with + The operaton returns a dense Tensor of shape `[2, 2]` with edit distances normalized by `truth` lengths. **Note**: It is possible to calculate edit distance between two From 28769b2b746570883ae8a57e042f50c4defba2b5 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Mon, 6 Apr 2020 16:50:35 +0000 Subject: [PATCH 0033/1390] Update tensorflow/python/ops/array_ops.py --- tensorflow/python/ops/array_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 004bbeac7f6..e508165a820 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -3508,7 +3508,7 @@ def edit_distance(hypothesis, truth, normalize=True, name="edit_distance"): Given the following input, * `hypothesis` is a `tf.SparseTensor` of shape `[2, 1, 1]` - * `truth` is a `tf.SparseTensor` of shape `(2, 2, 2)` + * `truth` is a `tf.SparseTensor` of shape `[2, 2, 2]` >>> hypothesis = tf.SparseTensor( ... [[0, 0, 0], From 2759cdca671dee4b7a2035710cf08725a97ce73c Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 7 Apr 2020 16:25:36 +0100 Subject: [PATCH 0034/1390] Fix after merging with master. --- tensorflow/lite/python/optimize/calibrator.py | 4 ++-- tensorflow/lite/python/optimize/calibrator_test.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorflow/lite/python/optimize/calibrator.py b/tensorflow/lite/python/optimize/calibrator.py index e31983b834e..90c43fcddfa 100644 --- a/tensorflow/lite/python/optimize/calibrator.py +++ b/tensorflow/lite/python/optimize/calibrator.py @@ -60,8 +60,8 @@ class Calibrator(object): input_type, output_type, allow_float, - resize_input=True, - activations_type = lite_constants.INT8): + activations_type=lite_constants.INT8, + resize_input=True): """Calibrates the model with specified generator and then quantizes it. The input shapes of the calibrator are resized with the calibration data if diff --git a/tensorflow/lite/python/optimize/calibrator_test.py b/tensorflow/lite/python/optimize/calibrator_test.py index f8a1171a629..f778c8a555d 100644 --- a/tensorflow/lite/python/optimize/calibrator_test.py +++ b/tensorflow/lite/python/optimize/calibrator_test.py @@ -148,7 +148,9 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): with self.assertRaisesRegex(ValueError, 'Size mismatch'): quantizer.calibrate_and_quantize(input_gen, constants.FLOAT, - constants.FLOAT, False, False) + constants.FLOAT, False, + constants.INT8, + False) def test_invalid_type_calibrator_gen(self): model_path = resource_loader.get_path_to_datafile( From cc08b5dff7346b2e2cd9fb80409c8f09bfebc089 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Tue, 7 Apr 2020 11:20:02 +0200 Subject: [PATCH 0035/1390] Update network_tester. New features: Multiple inputs Multiple outputs Output in json format Can call invoke() more than once Updated README --- .../examples/network_tester/Makefile.inc | 4 + .../micro/examples/network_tester/README.md | 34 ++++++- .../network_tester/expected_output_data.h | 2 +- .../examples/network_tester/input_data.h | 4 +- .../examples/network_tester/network_model.h | 5 +- .../network_tester/network_tester_test.cc | 88 +++++++++++++------ 6 files changed, 104 insertions(+), 33 deletions(-) diff --git a/tensorflow/lite/micro/examples/network_tester/Makefile.inc b/tensorflow/lite/micro/examples/network_tester/Makefile.inc index 27f54a66763..a5c911238c8 100644 --- a/tensorflow/lite/micro/examples/network_tester/Makefile.inc +++ b/tensorflow/lite/micro/examples/network_tester/Makefile.inc @@ -33,6 +33,10 @@ ifeq ($(COMPARE_OUTPUT_DATA),no) CXXFLAGS += -DNO_COMPARE_OUTPUT_DATA endif +ifdef NUM_INFERENCES + CXXFLAGS += -DNUM_INFERENCES=$(NUM_INFERENCES) +endif + # Builds a standalone object recognition binary. $(eval $(call microlite_test,network_tester_test,\ $(NETWORK_TESTER_TEST_SRCS),$(NETWORK_TESTER_TEST_HDRS))) diff --git a/tensorflow/lite/micro/examples/network_tester/README.md b/tensorflow/lite/micro/examples/network_tester/README.md index 7c4c48e4eb1..0cb709dce0a 100644 --- a/tensorflow/lite/micro/examples/network_tester/README.md +++ b/tensorflow/lite/micro/examples/network_tester/README.md @@ -34,8 +34,40 @@ make -f tensorflow/lite/micro/tools/make/Makefile network_tester_test \ `ARENA_SIZE`: The size of the memory to be allocated (in bytes) by the interpreter. \ `NUM_BYTES_TO_PRINT`: The number of bytes of the output data to print. \ -Defaults to 0 if not specified. \ +If set to 0, all bytes of the output are printed. \ `COMPARE_OUTPUT_DATA`: If set to "no" the output data is not compared to the expected output data. This could be useful e.g. if the execution time needs to be minimized, or there is no expected output data. If omitted, the output data is compared to the expected output. +`NUM_INFERENCES`: Define how many inferences that are made. Defaults to 1. \ + +The output is printed in JSON format using printf: +``` +num_of_outputs: 1 +output_begin +[ +{ +"dims": [4,1,2,2,1], +"data_address": "0x000000", +"data":"0x06,0x08,0x0e,0x10" +}] +output_end +``` + +If there are multiple output tensors, the output will look like this: +``` +num_of_outputs: 2 +output_begin +[ +{ +"dims": [4,1,2,2,1], +"data_address": "0x000000", +"data":"0x06,0x08,0x0e,0x10" +}, +{ +"dims": [4,1,2,2,1], +"data_address": "0x111111", +"data":"0x06,0x08,0x0e,0x10" +}] +output_end +``` diff --git a/tensorflow/lite/micro/examples/network_tester/expected_output_data.h b/tensorflow/lite/micro/examples/network_tester/expected_output_data.h index 03e21954b7f..934722bad94 100644 --- a/tensorflow/lite/micro/examples/network_tester/expected_output_data.h +++ b/tensorflow/lite/micro/examples/network_tester/expected_output_data.h @@ -17,6 +17,6 @@ limitations under the License. #define TENSORFLOW_LITE_MICRO_EXAMPLES_NETWORK_TESTER_EXPECTED_OUTPUT_DATA_H_ static unsigned int expected_output_data_len = 4; -static unsigned char expected_output_data[] = {6, 8, 14, 16}; +static unsigned char expected_output_data[1][4] = {6, 8, 14, 16}; #endif // TENSORFLOW_LITE_MICRO_EXAMPLES_NETWORK_TESTER_EXPECTED_OUTPUT_DATA_H_ diff --git a/tensorflow/lite/micro/examples/network_tester/input_data.h b/tensorflow/lite/micro/examples/network_tester/input_data.h index b47277cca93..b3710313dd2 100644 --- a/tensorflow/lite/micro/examples/network_tester/input_data.h +++ b/tensorflow/lite/micro/examples/network_tester/input_data.h @@ -17,7 +17,7 @@ limitations under the License. #define TENSORFLOW_LITE_MICRO_EXAMPLES_NETWORK_TESTER_INPUT_DATA_H_ static const int input_data_len = 16; -static const unsigned char input_data[] = {1, 2, 3, 4, 5, 6, 7, 8, - 9, 10, 11, 12, 13, 14, 15, 16}; +static const unsigned char input_data[1][16] = {{1, 2, 3, 4, 5, 6, 7, 8, + 9, 10, 11, 12, 13, 14, 15, 16}}; #endif // TENSORFLOW_LITE_MICRO_EXAMPLES_NETWORK_TESTER_INPUT_DATA_H_ diff --git a/tensorflow/lite/micro/examples/network_tester/network_model.h b/tensorflow/lite/micro/examples/network_tester/network_model.h index 4c275dbfbba..0431d7deee7 100644 --- a/tensorflow/lite/micro/examples/network_tester/network_model.h +++ b/tensorflow/lite/micro/examples/network_tester/network_model.h @@ -1,8 +1,11 @@ /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + 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 + http://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. @@ -64,4 +67,4 @@ const unsigned char network_model[] = { 0x08, 0x00, 0x07, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x11}; const unsigned int network_model_len = 576; -#endif // TENSORFLOW_LITE_MICRO_EXAMPLES_NETWORK_TESTER_NETWORK_MODEL_H_ +#endif diff --git a/tensorflow/lite/micro/examples/network_tester/network_tester_test.cc b/tensorflow/lite/micro/examples/network_tester/network_tester_test.cc index 0650222b970..5a307fb5c2a 100644 --- a/tensorflow/lite/micro/examples/network_tester/network_tester_test.cc +++ b/tensorflow/lite/micro/examples/network_tester/network_tester_test.cc @@ -1,8 +1,11 @@ /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + 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 + http://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. @@ -10,44 +13,54 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/lite/micro/examples/network_tester/expected_output_data.h" -#include "tensorflow/lite/micro/examples/network_tester/input_data.h" -#include "tensorflow/lite/micro/examples/network_tester/network_model.h" #include "tensorflow/lite/micro/kernels/all_ops_resolver.h" #include "tensorflow/lite/micro/micro_error_reporter.h" #include "tensorflow/lite/micro/micro_interpreter.h" -#include "tensorflow/lite/micro/testing/micro_test.h" -#include "tensorflow/lite/micro/testing/test_utils.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/version.h" +#include "tensorflow/lite/micro/examples/network_tester/expected_output_data.h" +#include "tensorflow/lite/micro/examples/network_tester/input_data.h" +#include "tensorflow/lite/micro/examples/network_tester/network_model.h" + +#include "tensorflow/lite/micro/testing/micro_test.h" +#include "tensorflow/lite/micro/testing/test_utils.h" + #ifndef TENSOR_ARENA_SIZE #define TENSOR_ARENA_SIZE (1024) #endif +#ifndef NUM_INFERENCES +#define NUM_INFERENCES 1 +#endif + uint8_t tensor_arena[TENSOR_ARENA_SIZE]; #ifdef NUM_BYTES_TO_PRINT inline void print_output_data(TfLiteTensor* output) { int num_bytes_to_print = - (output->bytes < NUM_BYTES_TO_PRINT) ? output->bytes : NUM_BYTES_TO_PRINT; + ((output->bytes < NUM_BYTES_TO_PRINT) || NUM_BYTES_TO_PRINT == 0) + ? output->bytes + : NUM_BYTES_TO_PRINT; int dims_size = output->dims->size; - printf("dims: {%d,", dims_size); + printf("{\n"); + printf("\"dims\": [%d,", dims_size); for (int i = 0; i < output->dims->size - 1; ++i) { printf("%d,", output->dims->data[i]); } - printf("%d}\n", output->dims->data[dims_size - 1]); + printf("%d],\n", output->dims->data[dims_size - 1]); - printf("data_address: %p\n", output->data.raw); - printf("data:\n{"); + printf("\"data_address\": \"%p\",\n", output->data.raw); + printf("\"data\":\""); for (int i = 0; i < num_bytes_to_print - 1; ++i) { - if (i % 16 == 0) { + if (i % 16 == 0 && i != 0) { printf("\n"); } printf("0x%02x,", output->data.uint8[i]); } - printf("0x%02x\n}\n", output->data.uint8[num_bytes_to_print - 1]); + printf("0x%02x\"\n", output->data.uint8[num_bytes_to_print - 1]); + printf("}"); } #endif @@ -63,7 +76,7 @@ TF_LITE_MICRO_TEST(TestInvoke) { "Model provided is schema version %d not equal " "to supported version %d.\n", model->version(), TFLITE_SCHEMA_VERSION); - return 1; + return kTfLiteError; } tflite::ops::micro::AllOpsResolver resolver; @@ -74,29 +87,48 @@ TF_LITE_MICRO_TEST(TestInvoke) { TfLiteStatus allocate_status = interpreter.AllocateTensors(); if (allocate_status != kTfLiteOk) { TF_LITE_REPORT_ERROR(error_reporter, "Tensor allocation failed\n"); + return kTfLiteError; } - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocate_status); - TfLiteTensor* input = interpreter.input(0); - memcpy(input->data.uint8, input_data, input->bytes); - - TfLiteStatus invoke_status = interpreter.Invoke(); - if (invoke_status != kTfLiteOk) { - TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed\n"); - } - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, invoke_status); - - TfLiteTensor* output = interpreter.output(0); + for (int n = 0; n < NUM_INFERENCES; n++) { + for (int i = 0; i < interpreter.inputs_size(); ++i) { + TfLiteTensor* input = interpreter.input(i); + memcpy(input->data.uint8, input_data[i], input->bytes); + } + TfLiteStatus invoke_status = interpreter.Invoke(); + if (invoke_status != kTfLiteOk) { + TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed\n"); + return kTfLiteError; + } + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, invoke_status); #ifdef NUM_BYTES_TO_PRINT - print_output_data(output); + // Print all of the output data, or the first NUM_BYTES_TO_PRINT bytes, + // whichever comes first as well as the output shape. + printf("num_of_outputs: %d\n", interpreter.outputs_size()); + printf("output_begin\n"); + printf("[\n"); + for (int i = 0; i < interpreter.outputs_size(); i++) { + TfLiteTensor* output = interpreter.output(i); + print_output_data(output); + if (i != interpreter.outputs_size() - 1) { + printf(",\n"); + } + } + printf("]\n"); + printf("output_end\n"); #endif #ifndef NO_COMPARE_OUTPUT_DATA - for (int i = 0; i < output->bytes; ++i) { - TF_LITE_MICRO_EXPECT_EQ(output->data.uint8[i], expected_output_data[i]); - } + for (int i = 0; i < interpreter.outputs_size(); i++) { + TfLiteTensor* output = interpreter.output(i); + for (int j = 0; j < output->bytes; ++j) { + TF_LITE_MICRO_EXPECT_EQ(output->data.uint8[j], + expected_output_data[i][j]); + } + } #endif + } TF_LITE_REPORT_ERROR(error_reporter, "Ran successfully\n"); } From 43036088fc3cfda3bd8d9c3a5df114d7d393618a Mon Sep 17 00:00:00 2001 From: Peng Meng Date: Mon, 13 Apr 2020 16:36:59 +0800 Subject: [PATCH 0036/1390] add unit test --- .../core/grappler/optimizers/remapper.cc | 6 +- .../core/grappler/optimizers/remapper_test.cc | 70 +++++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index fdfb8e379b2..239002c9da3 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -1476,8 +1476,8 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return true; }; - const auto is_conv2d_candidate = [&]() -> bool { - if (!IsConv2D(*node_def)) return false; + const auto is_relu_candidate = [&]() -> bool { + if (!IsRelu(*node_def)) return false; if (GetDataTypeFromAttr(*node_def, "T") != DT_FLOAT) return false; return true; @@ -1513,7 +1513,7 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return false; }; - return is_conv2d_candidate() || is_batch_norm_candidate() || + return is_relu_candidate() || is_batch_norm_candidate() || is_batch_norm_fusion_candidate(); } diff --git a/tensorflow/core/grappler/optimizers/remapper_test.cc b/tensorflow/core/grappler/optimizers/remapper_test.cc index 35e09b28205..831aa51ee72 100644 --- a/tensorflow/core/grappler/optimizers/remapper_test.cc +++ b/tensorflow/core/grappler/optimizers/remapper_test.cc @@ -449,6 +449,76 @@ TEST_F(RemapperTest, FuseMatMulWithBias) { test::ExpectTensorNear(tensors[0], tensors_expected[0], 1e-6); } +TEST_F(RemapperTest, FuseConv2DWithBiasAndActivationOnGPU) { + using ::tensorflow::ops::Placeholder; + + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + + auto input_shape = Placeholder::Shape({8, 32, 32, 3}); + auto filter_shape = Placeholder::Shape({3, 3, 3, 128}); + auto bias_shape = Placeholder::Shape({128}); + + auto input = Placeholder(s.WithOpName("input"), DT_FLOAT, input_shape); + auto filter = Placeholder(s.WithOpName("filter"), DT_FLOAT, filter_shape); + auto bias = Placeholder(s.WithOpName("bias"), DT_FLOAT, bias_shape); + + std::vector strides = {1, 1, 1, 1}; + auto conv = + ops::Conv2D(s.WithOpName("conv"), input, filter, strides, "SAME"); + auto bias_add = ops::BiasAdd(s.WithOpName("bias_add"), conv, bias); + + ops::Identity fetch = [&]() -> ops::Identity { + auto activate = s.WithOpName("activation"); + auto fetch = s.WithOpName("fetch"); + return ops::Identity(fetch, ops::Relu(activate, bias_add)); + }(); + + auto input_t = GenerateRandomTensor({8, 32, 32, 3}); + auto filter_t = GenerateRandomTensor({3, 3, 3, 128}); + auto bias_t = GenerateRandomTensor({128}); + + GrapplerItem item; + item.fetch = {"fetch"}; + item.feed = {{"input", input_t}, {"filter", filter_t}, {"bias", bias_t}}; + TF_ASSERT_OK(s.ToGraphDef(&item.graph)); + + // Place all nodes on GPU. + for (int i = 0; i < item.graph.node_size(); ++i) { + item.graph.mutable_node(i)->set_device("/device:GPU:0"); + } + + Remapper optimizer(RewriterConfig::AGGRESSIVE); // trust placeholders shape + //Remapper optimizer(RewriterConfig::ON); + GraphDef output; + TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output)); + + int found = 0; + for (const NodeDef& node : output.node()) { + if (node.name() == "activation") { + EXPECT_EQ(node.op(), "_FusedConv2D"); + ASSERT_GE(node.input_size(), 3); + EXPECT_EQ(node.input(0), "input"); + EXPECT_EQ(node.input(1), "filter"); + + EXPECT_EQ(node.attr().at("num_args").i(), 1); + EXPECT_EQ(node.input(2), "bias"); + + const auto fused_ops = node.attr().at("fused_ops").list().s(); + ASSERT_EQ(fused_ops.size(), 2); + EXPECT_EQ(fused_ops[0], "BiasAdd"); + EXPECT_EQ(fused_ops[1], "Relu"); + found++; + } + } + EXPECT_EQ(found, 1); + + auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed); + ASSERT_EQ(tensors_expected.size(), 1); + auto tensors = EvaluateNodes(output, item.fetch, item.feed); + ASSERT_EQ(tensors.size(), 1); + test::ExpectTensorNear(tensors[0], tensors_expected[0], 1e-6); +} + TEST_F(RemapperTest, FuseConv2DWithBiasAndActivation) { using ::tensorflow::ops::Placeholder; From 8b791cbf71aa68a91c69080b4313f115cb60fbf5 Mon Sep 17 00:00:00 2001 From: Peng Meng Date: Mon, 13 Apr 2020 17:36:16 +0800 Subject: [PATCH 0037/1390] optimize shape infer condition --- .../core/grappler/optimizers/remapper.cc | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index 239002c9da3..427a3f14aca 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -1461,6 +1461,7 @@ Status AddBatchNormNodes(RemapperContext* ctx, const FusedBatchNorm& matched) { // shapes: // (1) Splitting FusedBatchNorm into primitives. // (2) Fusing side input and/or activation into FusedBatchNorm. +// (3) Fusing Conv2D biasadd and relu on GPU bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { // Candidate for a FusedBatchNorm splitting. const auto* node_view = ctx.graph_view.GetNode(node_index); @@ -1476,10 +1477,30 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return true; }; - const auto is_relu_candidate = [&]() -> bool { + const auto is_relu_biasadd_conv2d_candidate = [&]() -> bool { if (!IsRelu(*node_def)) return false; if (GetDataTypeFromAttr(*node_def, "T") != DT_FLOAT) return false; + if (node_view->NumRegularFanins() < 1) return false; + const auto& relu_fanin_0 = node_view->GetRegularFanin(0); + const auto* relu_fanin_0_node_view = relu_fanin_0.node_view(); + const auto* relu_fanin_0_node_def = relu_fanin_0_node_view->node(); + + if (!IsBiasAdd(*relu_fanin_0_node_def)) return false; + if (GetDataTypeFromAttr(*relu_fanin_0_node_def, "T") != DT_FLOAT) + return false; + + if (relu_fanin_0_node_view->NumRegularFanins() < 1) return false; + + const auto& biasadd_fanin_0 = + relu_fanin_0_node_view->GetRegularFanin(0); + const auto* biasadd_fanin_0_node_def = + biasadd_fanin_0.node_view()->node(); + + if (!IsConv2D(*biasadd_fanin_0_node_def)) return false; + if (GetDataTypeFromAttr(*biasadd_fanin_0_node_def, "T") != DT_FLOAT) + return false; + return true; }; @@ -1513,7 +1534,7 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) { return false; }; - return is_relu_candidate() || is_batch_norm_candidate() || + return is_relu_biasadd_conv2d_candidate() || is_batch_norm_candidate() || is_batch_norm_fusion_candidate(); } From 9502277aa4a5e7219bc967a5016219fd970fff2f Mon Sep 17 00:00:00 2001 From: Biagio Montaruli Date: Tue, 21 Apr 2020 15:59:34 +0200 Subject: [PATCH 0038/1390] Update README.md of micro_speech example * Fix link to 'Animation on Arduino' * Substituite 'sample' with 'example' in the documentation --- .../lite/micro/examples/micro_speech/README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorflow/lite/micro/examples/micro_speech/README.md b/tensorflow/lite/micro/examples/micro_speech/README.md index 7ccaa806366..2a1898d25fb 100644 --- a/tensorflow/lite/micro/examples/micro_speech/README.md +++ b/tensorflow/lite/micro/examples/micro_speech/README.md @@ -7,7 +7,7 @@ The application listens to its surroundings with a microphone and indicates when it has detected a word by lighting an LED or displaying data on a screen, depending on the capabilities of the device. -![Animation on Arduino](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/hello_world/images/animation_on_arduino.gif) +![Animation on Arduino](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/images/animation_on_arduino.gif) The code has a small footprint (for example, around 22 kilobytes on a Cortex M3) and only uses about 10 kilobytes of RAM for working memory, so it's able to @@ -27,10 +27,10 @@ kilobytes of Flash. ## Deploy to Arduino -The following instructions will help you build and deploy this sample +The following instructions will help you build and deploy this example to [Arduino](https://www.arduino.cc/) devices. -The sample has been tested with the following devices: +The example has been tested with the following devices: - [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) @@ -84,11 +84,11 @@ If you don't see any output, repeat the process again. ## Deploy to ESP32 -The following instructions will help you build and deploy this sample to +The following instructions will help you build and deploy this example to [ESP32](https://www.espressif.com/en/products/hardware/esp32/overview) devices using the [ESP IDF](https://github.com/espressif/esp-idf). -The sample has been tested on ESP-IDF version 4.0 with the following devices: - +The example has been tested on ESP-IDF version 4.0 with the following devices: - [ESP32-DevKitC](http://esp-idf.readthedocs.io/en/latest/get-started/get-started-devkitc.html) - [ESP-EYE](https://github.com/espressif/esp-who/blob/master/docs/en/get-started/ESP-EYE_Getting_Started_Guide.md) @@ -139,7 +139,7 @@ monitor` ## Deploy to SparkFun Edge -The following instructions will help you build and deploy this sample on the +The following instructions will help you build and deploy this example on the [SparkFun Edge development board](https://sparkfun.com/products/15170). The program will toggle the blue LED on and off with each inference. It will @@ -288,7 +288,7 @@ followed by the `K` key, then hit the `Y` key. ## Deploy to STM32F746 -The following instructions will help you build and deploy the sample to the +The following instructions will help you build and deploy the example to the [STM32F7 discovery kit](https://os.mbed.com/platforms/ST-Discovery-F746NG/) using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). @@ -392,7 +392,7 @@ followed by the `K` key, then hit the `Y` key. ## Deploy to NXP FRDM K66F -The following instructions will help you build and deploy the sample to the +The following instructions will help you build and deploy the example to the [NXP FRDM K66F](https://www.nxp.com/design/development-boards/freedom-development-boards/mcu-boards/freedom-development-platform-for-kinetis-k66-k65-and-k26-mcus:FRDM-K66F) using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). From 8bcb69a2fc5e82c7b2b8028cfb1c3cc9a89e59f6 Mon Sep 17 00:00:00 2001 From: Biagio Montaruli Date: Tue, 21 Apr 2020 16:18:16 +0200 Subject: [PATCH 0039/1390] Update training documentation related to micro_speech example * Update link in the 'Overview' section * Update table in the 'Trained Models' section * Update spacing --- .../examples/micro_speech/train/README.md | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tensorflow/lite/micro/examples/micro_speech/train/README.md b/tensorflow/lite/micro/examples/micro_speech/train/README.md index 5793985a6e0..8228a71970f 100644 --- a/tensorflow/lite/micro/examples/micro_speech/train/README.md +++ b/tensorflow/lite/micro/examples/micro_speech/train/README.md @@ -34,14 +34,13 @@ go ## Overview -1. Training Jupyter Notebook: [`train_micro_speech_model.ipynb`] -(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train/train_micro_speech_model.ipynb) -. The training scripts used in this notebook is defined the +1. Training Jupyter Notebook: [`train_micro_speech_model.ipynb`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train/train_micro_speech_model.ipynb). +The training scripts used in this notebook is defined the [Simple Audio Recognition](https://www.tensorflow.org/tutorials/sequences/audio_recognition) tutorial. 2. Dataset Type: **Speech** -3. Dataset: Speech Commands, Version 2. ([Download Link](https://storage.cloud.google.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz) -, [Paper](https://arxiv.org/abs/1804.03209)) +3. Dataset: Speech Commands, Version 2. ([Download Link](https://storage.cloud.google.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz), +[Paper](https://arxiv.org/abs/1804.03209)) 4. Deep Learning Framework: **TensorFlow 1.5** 5. Language: **Python 3.7** 6. Model Size: **<20 kB** @@ -60,10 +59,8 @@ includes the following 3 model files: | Name | Format | Target Framework | Target Device | | :------------- |:-------------|:-------------|-----| | `model.pb` | Frozen GraphDef | TensorFlow | Large-Scale/Cloud/Servers | -| `model.tflite` *(<20 kB)* | Fully Quantized* TFLite Model | -TensorFlow Lite | Mobile Devices| -| `model.cc` | C Source File | TensorFlow Lite for Microcontrollers | -Microcontrollers | +| `model.tflite` *(<20 kB)* | Fully Quantized* TFLite Model | TensorFlow Lite | Mobile Devices| +| `model.cc` | C Source File | TensorFlow Lite for Microcontrollers | Microcontrollers | **Fully quantized implies that the model is **strictly int8** quantized including the input(s)and output(s).* @@ -154,8 +151,8 @@ simpler model for accurate results. ## Dataset -The Speech Commands Dataset. ([Download Link](https://storage.cloud.google.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz) -,[Paper](https://arxiv.org/abs/1804.03209)) consists of over 105,000 WAVE audio +The Speech Commands Dataset. ([Download Link](https://storage.cloud.google.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz), +[Paper](https://arxiv.org/abs/1804.03209)) consists of over 105,000 WAVE audio files of people saying thirty different words. This data was collected by Google and released under a CC BY license. You can help improve it by contributing five minutes of your own voice. The archive is over 2GB, so this From e9b965a29cda282bc5dd8b2bcbc7991bb623734e Mon Sep 17 00:00:00 2001 From: Fredrik Knutsson Date: Fri, 24 Apr 2020 15:19:39 +0200 Subject: [PATCH 0040/1390] Fix review comments, 24/4 --- .../person_detection/person_detection_test.cc | 2 + .../memory_planner/greedy_memory_planner.h | 6 ++- tensorflow/lite/micro/micro_allocator.cc | 52 +++++-------------- .../lite/micro/micro_optional_debug_tools.cc | 31 +++++++++++ .../lite/micro/micro_optional_debug_tools.h | 3 ++ 5 files changed, 54 insertions(+), 40 deletions(-) diff --git a/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc b/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc index 51a61881ead..f57e6ea88f1 100644 --- a/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc +++ b/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/lite/micro/micro_error_reporter.h" #include "tensorflow/lite/micro/micro_interpreter.h" #include "tensorflow/lite/micro/micro_mutable_op_resolver.h" +#include "tensorflow/lite/micro/micro_optional_debug_tools.h" #include "tensorflow/lite/micro/testing/micro_test.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/version.h" @@ -46,6 +47,7 @@ TF_LITE_MICRO_TEST(TestInvoke) { "to supported version %d.\n", model->version(), TFLITE_SCHEMA_VERSION); } + PrintModelData(model, error_reporter); // Pull in only the operation implementations we need. // This relies on a complete list of all the ops needed by this graph. diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h index c849e57645c..19a36f342fd 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.h @@ -144,8 +144,10 @@ class GreedyMemoryPlanner : public MemoryPlanner { int* buffer_sizes_sorted_; int* buffer_ids_sorted_; ListEntry* buffers_sorted_by_offset_; - int next_free_entry_; - int first_entry_index_; + int next_free_entry_; // Index of the next free entry of + // buffers_sorted_by_offset_ + int first_entry_index_; // Index of the first entry (smallest offset) of + // buffers_sorted_by_offset_ // Stores the outcome of the plan, the location of each buffer in the arena. int* buffer_offsets_; diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index daa7cbcb0c9..c57294f5745 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -45,6 +45,8 @@ struct AllocationInfo { // requirement for SIMD extensions. constexpr int kBufferAlignment = 16; +constexpr char kOfflineMemAllocMetadata[] = "OfflineMemoryAllocation"; + class MicroBuiltinDataAllocator : public BuiltinDataAllocator { public: explicit MicroBuiltinDataAllocator(SimpleMemoryAllocator* memory_allocator) @@ -81,33 +83,6 @@ TfLiteStatus AllocateVariables( return kTfLiteOk; } -// Helper function to print model flatbuffer data. This function is not called -// by default. Hence it's not linked in to the final binary code. -void PrintModelData(const Model* model, ErrorReporter* error_reporter) { - auto* subgraphs = model->subgraphs(); - const SubGraph* subgraph = (*subgraphs)[0]; - const flatbuffers::Vector>* tensors = - subgraph->tensors(); - const flatbuffers::Vector>* buffers = - model->buffers(); - TF_LITE_REPORT_ERROR(error_reporter, "==== Model info: ====="); - for (int i = 0; i < tensors->size(); ++i) { - const tflite::Tensor& flatbuffer_tensor = *tensors->Get(i); - auto* quantization = flatbuffer_tensor.quantization(); - size_t type_size, tensor_size; - auto* buffer = (*buffers)[flatbuffer_tensor.buffer()]; - auto* array = buffer->data(); - int array_size = 0; - if (array) { - array_size = array->size(); - } - BytesRequiredForTensor(flatbuffer_tensor, &tensor_size, &type_size, - error_reporter); - TF_LITE_REPORT_ERROR( - error_reporter, "Tensor index: %d arena tensor %d size %d", i, - !array_size && !flatbuffer_tensor.is_variable(), tensor_size); - } -} // Helper function to check flatbuffer metadata correctness. This function is // not called by default. Hence it's not linked in to the final binary code. @@ -116,8 +91,8 @@ TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, if (model->metadata()) { for (int i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); - if (strncmp(metadata->name()->c_str(), "OfflineMemoryAllocation", - strlen("OfflineMemoryAllocation")) == 0) { + if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata, + strlen(kOfflineMemAllocMetadata)) == 0) { auto* subgraphs = model->subgraphs(); const SubGraph* subgraph = (*subgraphs)[0]; const flatbuffers::Vector>* tensors = @@ -311,8 +286,9 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, // | name:string | “OfflineMemoryAllocation” | // | buffer:unit | Index of buffer containing memory allocation data | // -// The buffer contents for the memory allocation is a list of 32-bit integers of -// the following format: +// The buffer contents for the memory allocation is a list of 32-bit integers. +// The number of tensors, n, must be equal to the number of tensors defined in +// the model. The following encoding applies: // // | Offset | Value | // | 0 | Offline allocation format version – set to 0 | @@ -326,8 +302,8 @@ TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets( if (model->metadata()) { for (int i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); - if (strncmp(metadata->name()->c_str(), "OfflineMemoryAllocation", - strlen("OfflineMemoryAllocation")) == 0) { + if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata, + strlen(kOfflineMemAllocMetadata)) == 0) { const flatbuffers::Vector>* buffers = model->buffers(); auto* buffer = (*buffers)[metadata->buffer()]; @@ -365,7 +341,8 @@ TfLiteStatus AllocationInfoBuilder::AddScratchBuffers( return kTfLiteOk; } -TfLiteStatus CreatePlan(ErrorReporter* error_reporter, MemoryPlanner* planner, +TfLiteStatus CreatePlan(ErrorReporter* error_reporter, + GreedyMemoryPlanner* planner, const AllocationInfo* allocation_info, size_t allocation_info_size) { // Add the tensors to our allocation plan. @@ -380,10 +357,9 @@ TfLiteStatus CreatePlan(ErrorReporter* error_reporter, MemoryPlanner* planner, current->first_created, current->last_used)); } else { TF_LITE_ENSURE_STATUS( - (static_cast(planner)) - ->AddBuffer(error_reporter, aligned_bytes_required, - current->first_created, current->last_used, - current->offline_offset)); + planner->AddBuffer(error_reporter, aligned_bytes_required, + current->first_created, current->last_used, + current->offline_offset)); } } } diff --git a/tensorflow/lite/micro/micro_optional_debug_tools.cc b/tensorflow/lite/micro/micro_optional_debug_tools.cc index 70f16c78d79..10373d3c034 100644 --- a/tensorflow/lite/micro/micro_optional_debug_tools.cc +++ b/tensorflow/lite/micro/micro_optional_debug_tools.cc @@ -22,6 +22,8 @@ limitations under the License. #include #include "tensorflow/lite/schema/schema_generated.h" +#include "tensorflow/lite/micro/memory_helpers.h" + namespace tflite { namespace { @@ -100,6 +102,35 @@ const char* AllocTypeName(TfLiteAllocationType type) { } } // namespace +// Helper function to print model flatbuffer data. This function is not called +// by default. Hence it's not linked in to the final binary code. +void PrintModelData(const Model* model, ErrorReporter* error_reporter) { + auto* subgraphs = model->subgraphs(); + const SubGraph* subgraph = (*subgraphs)[0]; + const flatbuffers::Vector>* tensors = + subgraph->tensors(); + const flatbuffers::Vector>* buffers = + model->buffers(); + TF_LITE_REPORT_ERROR(error_reporter, "==== Model info: ====="); + for (int i = 0; i < tensors->size(); ++i) { + const tflite::Tensor& flatbuffer_tensor = *tensors->Get(i); + auto* quantization = flatbuffer_tensor.quantization(); + size_t type_size, tensor_size; + auto* buffer = (*buffers)[flatbuffer_tensor.buffer()]; + auto* array = buffer->data(); + int array_size = 0; + if (array) { + array_size = array->size(); + } + BytesRequiredForTensor(flatbuffer_tensor, &tensor_size, &type_size, + error_reporter); + TF_LITE_REPORT_ERROR( + error_reporter, + "Tensor index: %d arena tensor %d size %d ", + i, !array_size && !flatbuffer_tensor.is_variable(), tensor_size); + } +} + // Prints a dump of what tensors and what nodes are in the interpreter. void PrintInterpreterState(MicroInterpreter* interpreter) { printf("Interpreter has %zu tensors and %zu nodes\n", diff --git a/tensorflow/lite/micro/micro_optional_debug_tools.h b/tensorflow/lite/micro/micro_optional_debug_tools.h index ae96b62ab3c..cc9630e6f12 100644 --- a/tensorflow/lite/micro/micro_optional_debug_tools.h +++ b/tensorflow/lite/micro/micro_optional_debug_tools.h @@ -20,6 +20,9 @@ limitations under the License. #include "tensorflow/lite/micro/micro_interpreter.h" namespace tflite { +// Helper function to print model flatbuffer data. This function is not called +// by default. Hence it's not linked in to the final binary code. +void PrintModelData(const Model* model, ErrorReporter* error_reporter); // Prints a dump of what tensors and what nodes are in the interpreter. void PrintInterpreterState(MicroInterpreter* interpreter); } // namespace tflite From 7d3237ca0951e102dfcc04f5bd98e0bc1fa1e22c Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Mon, 27 Apr 2020 13:16:49 +0200 Subject: [PATCH 0041/1390] Address reviewer comments. --- tensorflow/lite/micro/micro_allocator_test.cc | 65 ++++++++++--------- tensorflow/lite/micro/test_helpers.cc | 56 ++++++++-------- tensorflow/lite/micro/test_helpers.h | 3 +- 3 files changed, 64 insertions(+), 60 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 45bc3b06b24..b5db8fdd626 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -247,22 +247,22 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { // The structure is identical to the one in // TestAllocationForModelsWithBranches - std::vector node_list = { - { - {0}, // input - {1} // output - }, - { - {0}, // input - {2} // output - }, - { - {1, 2}, // input1, input2 - {3} // output - }}; + int num_conns = 3; + tflite::testing::NodeConnection node_list[3] = {{ + {0}, // input + {1} // output + }, + { + {0}, // input + {2} // output + }, + { + {1, 2}, // input1, input2 + {3} // output + }}; const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( - nbr_tensors, metadata_buffer, node_list); + nbr_tensors, metadata_buffer, node_list, num_conns); TfLiteContext context; constexpr size_t arena_size = 4096; @@ -296,21 +296,22 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) { int t2 = 2; int t3 = 3; - std::vector node_list = {{ - {t0}, // input - {t1} // output - }, - { - {t1}, // input - {t2} // output - }, - { - {t2}, // input - {t3} // output - }}; + int num_conns = 3; + tflite::testing::NodeConnection node_list[3] = {{ + {t0}, // input + {t1} // output + }, + { + {t1}, // input + {t2} // output + }, + { + {t2}, // input + {t3} // output + }}; const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( - nbr_tensors, metadata_buffer, node_list); + nbr_tensors, metadata_buffer, node_list, num_conns); TfLiteContext context; constexpr size_t arena_size = 4096; @@ -342,7 +343,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { int t2 = 2; int t3 = 3; - std::vector node_list = { + int num_conns = 2; + tflite::testing::NodeConnection node_list[2] = { { {t0, t1}, // input, scratch {t2} // output @@ -354,7 +356,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { }; const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( - nbr_tensors, metadata_buffer, node_list); + nbr_tensors, metadata_buffer, node_list, num_conns); TfLiteContext context; constexpr size_t arena_size = 4096; @@ -389,7 +391,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { int t3 = 3; int t4 = 4; - std::vector node_list = { + int num_conns = 2; + tflite::testing::NodeConnection node_list[2] = { { {t0, t1}, // input, scratch {t2}, // output @@ -401,7 +404,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { }; const tflite::Model* model = tflite::testing::GetModelWithOfflinePlanning( - nbr_tensors, metadata_buffer, node_list); + nbr_tensors, metadata_buffer, node_list, num_conns); TfLiteContext context; constexpr size_t arena_size = 4096; diff --git a/tensorflow/lite/micro/test_helpers.cc b/tensorflow/lite/micro/test_helpers.cc index b39d3b2916f..f52ebdc4d45 100644 --- a/tensorflow/lite/micro/test_helpers.cc +++ b/tensorflow/lite/micro/test_helpers.cc @@ -95,8 +95,7 @@ class ModelBuilder { std::initializer_list outputs); void AddMetadata(const char* description_string, - const int32_t* metadata_buffer_data, - size_t num_elements); + const int32_t* metadata_buffer_data, size_t num_elements); // Constructs the flatbuffer model using `builder_` and return a pointer to // it. The returned model has the same lifetime as `builder_`. @@ -157,16 +156,15 @@ ModelBuilder::Node ModelBuilder::AddNode( } void ModelBuilder::AddMetadata(const char* description_string, - const int32_t* metadata_buffer_data, - size_t num_elements) { + const int32_t* metadata_buffer_data, + size_t num_elements) { metadata_[ModelBuilder::nbr_of_metadata_buffers_] = - CreateMetadata(*builder_, - builder_->CreateString(description_string), - 1 + ModelBuilder::nbr_of_metadata_buffers_); + CreateMetadata(*builder_, builder_->CreateString(description_string), + 1 + ModelBuilder::nbr_of_metadata_buffers_); - metadata_buffers_[nbr_of_metadata_buffers_] = tflite::CreateBuffer(*builder_, - builder_->CreateVector((uint8_t*)metadata_buffer_data, - sizeof(uint32_t) * num_elements)); + metadata_buffers_[nbr_of_metadata_buffers_] = tflite::CreateBuffer( + *builder_, builder_->CreateVector((uint8_t*)metadata_buffer_data, + sizeof(uint32_t) * num_elements)); ModelBuilder::nbr_of_metadata_buffers_++; } @@ -175,14 +173,14 @@ const Model* ModelBuilder::BuildModel( std::initializer_list inputs, std::initializer_list outputs) { // Model schema requires an empty buffer at idx 0. - size_t kBufferSize = 1 + ModelBuilder::nbr_of_metadata_buffers_; - flatbuffers::Offset buffers[kBufferSize]; + size_t buffer_size = 1 + ModelBuilder::nbr_of_metadata_buffers_; + flatbuffers::Offset buffers[kMaxMetadataBuffers]; buffers[0] = tflite::CreateBuffer(*builder_); // Place the metadata buffers first in the buffer since the indices for them // have already been set in AddMetadata() for (int i = 1; i < ModelBuilder::nbr_of_metadata_buffers_ + 1; ++i) { - buffers[i] = metadata_buffers_[i - 1]; + buffers[i] = metadata_buffers_[i - 1]; } // TFLM only supports single subgraph. @@ -202,16 +200,16 @@ const Model* ModelBuilder::BuildModel( builder_->CreateVector(operator_codes_, next_operator_code_id_), builder_->CreateVector(subgraphs, subgraphs_size), builder_->CreateString("teset_model"), - builder_->CreateVector(buffers, kBufferSize), - 0, - builder_->CreateVector(metadata_, ModelBuilder::nbr_of_metadata_buffers_)); + builder_->CreateVector(buffers, buffer_size), 0, + builder_->CreateVector(metadata_, + ModelBuilder::nbr_of_metadata_buffers_)); } else { model_offset = tflite::CreateModel( *builder_, 0, builder_->CreateVector(operator_codes_, next_operator_code_id_), builder_->CreateVector(subgraphs, subgraphs_size), builder_->CreateString("teset_model"), - builder_->CreateVector(buffers, kBufferSize)); + builder_->CreateVector(buffers, buffer_size)); } tflite::FinishModelBuffer(*builder_, model_offset); @@ -293,8 +291,9 @@ const Model* BuildSimpleModelWithBranch() { } const Model* BuildModelWithOfflinePlanning(int number_of_tensors, - const int32_t* metadata_buffer, - std::vector node_conn) { + const int32_t* metadata_buffer, + NodeConnection* node_conn, + int num_conns) { using flatbuffers::Offset; flatbuffers::FlatBufferBuilder* fb_builder = BuilderInstance(); @@ -310,15 +309,16 @@ const Model* BuildModelWithOfflinePlanning(int number_of_tensors, tensors[i] = model_builder.AddTensor(TensorType_FLOAT32, {2, 2, 3}); } - for (int i = 0; i < node_conn.size(); i++) { + for (int i = 0; i < num_conns; ++i) { model_builder.AddNode(op_id, node_conn[i].input, node_conn[i].output); } - model_builder.AddMetadata("OfflineMemoryAllocation", - metadata_buffer, number_of_tensors + tflite::testing::kOfflinePlannerHeaderSize); + model_builder.AddMetadata( + "OfflineMemoryAllocation", metadata_buffer, + number_of_tensors + tflite::testing::kOfflinePlannerHeaderSize); return model_builder.BuildModel(node_conn[0].input, - node_conn[node_conn.size() - 1].output); + node_conn[num_conns - 1].output); } const Model* BuildSimpleMockModel() { @@ -408,8 +408,7 @@ const Model* BuildComplexMockModel() { constexpr size_t buffers_size = 7; const Offset buffers[buffers_size] = { // Op 1 buffers: - CreateBuffer(*builder), - CreateBuffer(*builder), + CreateBuffer(*builder), CreateBuffer(*builder), CreateBuffer(*builder, builder->CreateVector(buffer_data_1, buffer_data_size)), // Op 2 buffers: @@ -576,9 +575,10 @@ const Model* GetSimpleModelWithBranch() { const Model* GetModelWithOfflinePlanning(int num_tensors, const int32_t* metadata_buffer, - std::vector node_conn) { - const Model* model = - BuildModelWithOfflinePlanning(num_tensors, metadata_buffer, node_conn); + NodeConnection* node_conn, + int num_conns) { + const Model* model = BuildModelWithOfflinePlanning( + num_tensors, metadata_buffer, node_conn, num_conns); return model; } diff --git a/tensorflow/lite/micro/test_helpers.h b/tensorflow/lite/micro/test_helpers.h index e31f5061de8..647ffb92cff 100644 --- a/tensorflow/lite/micro/test_helpers.h +++ b/tensorflow/lite/micro/test_helpers.h @@ -49,7 +49,8 @@ const Model* GetSimpleModelWithBranch(); // Returns a simple flatbuffer model with offline planned tensors const Model* GetModelWithOfflinePlanning(int num_tensors, const int32_t* metadata_buffer, - std::vector node_conn); + NodeConnection* node_conn, + int num_conns); // Returns a flatbuffer model with `simple_stateful_op` const Model* GetSimpleStatefulModel(); From d46acc6016666abfc27d28e35936ed70979c822e Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen" Date: Wed, 29 Apr 2020 15:05:06 +0800 Subject: [PATCH 0042/1390] debug_message --- tensorflow/lite/micro/we_i/debug_log.cc | 33 +++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 tensorflow/lite/micro/we_i/debug_log.cc diff --git a/tensorflow/lite/micro/we_i/debug_log.cc b/tensorflow/lite/micro/we_i/debug_log.cc new file mode 100644 index 00000000000..a115d476aff --- /dev/null +++ b/tensorflow/lite/micro/we_i/debug_log.cc @@ -0,0 +1,33 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +// Implementation for the DebugLog() function that prints to the UART on the +// SparkFun Edge microcontroller. The same should work for other targets using +// the Ambiq Apollo 3. + +#include "tensorflow/lite/micro/debug_log.h" +#include "xprintf.h" +#include "console_io.h" +#include + +extern "C" void DebugLog(const char* s) { + static bool is_initialized = false; + if (!is_initialized) { + xprintf_setup(); + is_initialized = true; + } + + xprintf("%s", s); +} From c8a2f59e2e2775808b82c877433ead9545fd7a84 Mon Sep 17 00:00:00 2001 From: Eugene Kuznetsov Date: Wed, 22 Jan 2020 12:53:00 -0800 Subject: [PATCH 0043/1390] Fixing and enabling //tensorflow/core/util:gpu_kernel_helper_test_gpu --- tensorflow/core/util/gpu_device_functions.h | 4 + .../core/util/gpu_kernel_helper_test.cu.cc | 127 +++++++++++------- 2 files changed, 84 insertions(+), 47 deletions(-) diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h index 61d1e3c9453..083d42b7de4 100644 --- a/tensorflow/core/util/gpu_device_functions.h +++ b/tensorflow/core/util/gpu_device_functions.h @@ -53,6 +53,8 @@ using gpuEvent_t = cudaEvent_t; #define gpuEventCreate cudaEventCreate #define gpuEventCreateWithFlags cudaEventCreateWithFlags #define gpuEventDisableTiming cudaEventDisableTiming +#define gpuDeviceSynchronize cudaDeviceSynchronize +#define gpuFree cudaFree #elif TENSORFLOW_USE_ROCM using gpuFloatComplex = hipFloatComplex; using gpuDoubleComplex = hipDoubleComplex; @@ -68,6 +70,8 @@ using cudaError_t = int; #define gpuEventCreate hipEventCreate #define gpuEventCreateWithFlags hipEventCreateWithFlags #define gpuEventDisableTiming hipEventDisableTiming +#define gpuDeviceSynchronize hipDeviceSynchronize +#define gpuFree hipFree static std::string cudaGetErrorString(int err) { return std::to_string(err); } #endif diff --git a/tensorflow/core/util/gpu_kernel_helper_test.cu.cc b/tensorflow/core/util/gpu_kernel_helper_test.cu.cc index c089511e964..3135d25d1b8 100644 --- a/tensorflow/core/util/gpu_kernel_helper_test.cu.cc +++ b/tensorflow/core/util/gpu_kernel_helper_test.cu.cc @@ -13,9 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#if GOOGLE_CUDA +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #define EIGEN_USE_GPU +#include #include #include "tensorflow/core/lib/core/status_test_util.h" @@ -25,14 +26,14 @@ limitations under the License. #define CUDA_EXPECT_SUCCESS \ { \ - cudaDeviceSynchronize(); \ + gpuDeviceSynchronize(); \ cudaError_t err = cudaGetLastError(); \ EXPECT_EQ(cudaSuccess, err) << cudaGetErrorString(err); \ } #define CUDA_ASSERT_SUCCESS \ { \ - cudaDeviceSynchronize(); \ + gpuDeviceSynchronize(); \ cudaError_t err = cudaGetLastError(); \ ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err); \ } @@ -94,8 +95,7 @@ __global__ void Count3D(Gpu3DLaunchConfig config, int bufsize, } } -__global__ void CudaShuffleGetSrcLaneTest( - unsigned* __restrict__ failure_count) { +__global__ void GpuShuffleGetSrcLaneTest(unsigned* __restrict__ failure_count) { unsigned lane_id = GpuLaneId(); for (int width = warpSize; width > 1; width /= 2) { auto check_result = [&](const char* op_name, int param, unsigned actual, @@ -103,31 +103,38 @@ __global__ void CudaShuffleGetSrcLaneTest( if (actual != expected) { printf("Cuda%sGetSrcLane(%d, %d) for lane %d returned %d, not %d\n", op_name, param, width, lane_id, actual, expected); - CudaAtomicAdd(failure_count, 1); + GpuAtomicAdd(failure_count, 1); } }; + for (int src_lane = -warpSize; src_lane <= warpSize; ++src_lane) { - unsigned actual_lane = detail::CudaShuffleGetSrcLane(src_lane, width); +#if TENSORFLOW_USE_ROCM + if (src_lane < 0 || src_lane >= width) continue; +#endif + unsigned actual_lane = detail::GpuShuffleGetSrcLane(src_lane, width); unsigned expect_lane = - CudaShuffleSync(kCudaWarpAll, lane_id, src_lane, width); + GpuShuffleSync(kCudaWarpAll, lane_id, src_lane, width); check_result("Shuffle", src_lane, actual_lane, expect_lane); } + for (unsigned delta = 0; delta <= warpSize; ++delta) { - unsigned actual_lane = detail::CudaShuffleUpGetSrcLane(delta, width); + unsigned actual_lane = detail::GpuShuffleUpGetSrcLane(delta, width); unsigned expect_lane = - CudaShuffleUpSync(kCudaWarpAll, lane_id, delta, width); + GpuShuffleUpSync(kCudaWarpAll, lane_id, delta, width); check_result("ShuffleUp", delta, actual_lane, expect_lane); } + for (unsigned delta = 0; delta <= warpSize; ++delta) { - unsigned actual_lane = detail::CudaShuffleDownGetSrcLane(delta, width); + unsigned actual_lane = detail::GpuShuffleDownGetSrcLane(delta, width); unsigned expect_lane = - CudaShuffleDownSync(kCudaWarpAll, lane_id, delta, width); + GpuShuffleDownSync(kCudaWarpAll, lane_id, delta, width); check_result("ShuffleDown", delta, actual_lane, expect_lane); } + for (int lane_lane = warpSize; lane_lane > 0; lane_lane /= 2) { - unsigned actual_lane = detail::CudaShuffleXorGetSrcLane(lane_lane, width); + unsigned actual_lane = detail::GpuShuffleXorGetSrcLane(lane_lane, width); unsigned expect_lane = - CudaShuffleXorSync(kCudaWarpAll, lane_id, lane_lane, width); + GpuShuffleXorSync(kCudaWarpAll, lane_id, lane_lane, width); check_result("ShuffleXor", lane_lane, actual_lane, expect_lane); } } @@ -137,19 +144,32 @@ __global__ void CudaShuffleGetSrcLaneTest( class GpuLaunchConfigTest : public ::testing::Test { protected: - const int bufsize = 1024; + static const int bufsize = 1024; int* outbuf = nullptr; + int* outbuf_host = nullptr; + int hostbuf[bufsize]; Eigen::GpuStreamDevice stream; Eigen::GpuDevice d = Eigen::GpuDevice(&stream); + void copyToHost() { +#if TENSORFLOW_USE_ROCM + hipMemcpy(hostbuf, outbuf, sizeof(int) * bufsize, hipMemcpyDeviceToHost); +#endif + } virtual void SetUp() { +#if GOOGLE_CUDA cudaError_t err = cudaMallocManaged(&outbuf, sizeof(int) * bufsize); + outbuf_host = outbuf; +#else + cudaError_t err = hipMalloc(&outbuf, sizeof(int) * bufsize); + outbuf_host = hostbuf; +#endif ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err); } virtual void TearDown() { - cudaDeviceSynchronize(); - cudaFree(outbuf); + gpuDeviceSynchronize(); + gpuFree(outbuf); outbuf = nullptr; } }; @@ -158,28 +178,32 @@ TEST_F(GpuLaunchConfigTest, GetGpuLaunchConfig) { GpuLaunchConfig cfg; // test valid inputs -#define TEST_LAUNCH_PARAMETER(work_element_count) \ - cfg = GetGpuLaunchConfig(bufsize, d); \ - TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count, \ - cfg.thread_per_block, 0, d.stream(), cfg, \ - outbuf)); \ - CUDA_ASSERT_SUCCESS \ - cfg = GetGpuLaunchConfig(work_element_count, d); \ - TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \ - 0, d.stream(), cfg, bufsize, outbuf)); \ - CUDA_EXPECT_SUCCESS \ - EXPECT_EQ(work_element_count, std::accumulate(outbuf, outbuf + bufsize, 0)); \ - \ - cfg = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0); \ - TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count, \ - cfg.thread_per_block, 0, d.stream(), cfg, \ - outbuf)); \ - CUDA_ASSERT_SUCCESS \ - cfg = GetGpuLaunchConfig(work_element_count, d, Count1D, 0, 0); \ - TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \ - 0, d.stream(), cfg, bufsize, outbuf)); \ - CUDA_EXPECT_SUCCESS \ - EXPECT_EQ(work_element_count, std::accumulate(outbuf, outbuf + bufsize, 0)) +#define TEST_LAUNCH_PARAMETER(work_element_count) \ + cfg = GetGpuLaunchConfig(bufsize, d); \ + TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count, \ + cfg.thread_per_block, 0, d.stream(), cfg, \ + outbuf)); \ + CUDA_ASSERT_SUCCESS \ + cfg = GetGpuLaunchConfig(work_element_count, d); \ + TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \ + 0, d.stream(), cfg, bufsize, outbuf)); \ + CUDA_EXPECT_SUCCESS \ + copyToHost(); \ + EXPECT_EQ(work_element_count, \ + std::accumulate(outbuf_host, outbuf_host + bufsize, 0)); \ + \ + cfg = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0); \ + TF_CHECK_OK(GpuLaunchKernel(SetOutbufZero, cfg.block_count, \ + cfg.thread_per_block, 0, d.stream(), cfg, \ + outbuf)); \ + CUDA_ASSERT_SUCCESS \ + cfg = GetGpuLaunchConfig(work_element_count, d, Count1D, 0, 0); \ + TF_CHECK_OK(GpuLaunchKernel(Count1D, cfg.block_count, cfg.thread_per_block, \ + 0, d.stream(), cfg, bufsize, outbuf)); \ + CUDA_EXPECT_SUCCESS \ + copyToHost(); \ + EXPECT_EQ(work_element_count, \ + std::accumulate(outbuf_host, outbuf_host + bufsize, 0)); TEST_LAUNCH_PARAMETER(128); TEST_LAUNCH_PARAMETER(129); @@ -221,7 +245,9 @@ TEST_F(GpuLaunchConfigTest, GetGpu2DLaunchConfig) { TF_EXPECT_OK(GpuLaunchKernel(Count2D, cfg.block_count, cfg.thread_per_block, \ 0, d.stream(), cfg, bufsize, outbuf)); \ CUDA_EXPECT_SUCCESS \ - EXPECT_EQ(dimx* dimy, std::accumulate(outbuf, outbuf + bufsize, 0)); \ + copyToHost(); \ + EXPECT_EQ(dimx* dimy, \ + std::accumulate(outbuf_host, outbuf_host + bufsize, 0)); \ \ cfg1d = GetGpuLaunchConfig(bufsize, d, SetOutbufZero, 0, 0); \ TF_EXPECT_OK(GpuLaunchKernel(SetOutbufZero, cfg1d.block_count, \ @@ -232,7 +258,8 @@ TEST_F(GpuLaunchConfigTest, GetGpu2DLaunchConfig) { TF_EXPECT_OK(GpuLaunchKernel(Count2D, cfg.block_count, cfg.thread_per_block, \ 0, d.stream(), cfg, bufsize, outbuf)); \ CUDA_EXPECT_SUCCESS \ - EXPECT_EQ(dimx* dimy, std::accumulate(outbuf, outbuf + bufsize, 0)) + copyToHost(); \ + EXPECT_EQ(dimx* dimy, std::accumulate(outbuf_host, outbuf_host + bufsize, 0)) TEST_LAUNCH_PARAMETER(128, 128); TEST_LAUNCH_PARAMETER(129, 64); @@ -263,7 +290,9 @@ TEST_F(GpuLaunchConfigTest, GetGpu3DLaunchConfig) { TF_EXPECT_OK(GpuLaunchKernel(Count3D, cfg.block_count, cfg.thread_per_block, \ 0, d.stream(), cfg, bufsize, outbuf)); \ CUDA_EXPECT_SUCCESS \ - EXPECT_EQ(dimx* dimy* dimz, std::accumulate(outbuf, outbuf + bufsize, 0)) + copyToHost(); \ + EXPECT_EQ(dimx* dimy* dimz, \ + std::accumulate(outbuf_host, outbuf_host + bufsize, 0)) TEST_LAUNCH_PARAMETER(128, 128, 128); TEST_LAUNCH_PARAMETER(129, 64, 1024); @@ -282,15 +311,19 @@ TEST_F(GpuLaunchConfigTest, GetGpu3DLaunchConfig) { TEST(CudaDeviceFunctionsTest, ShuffleGetSrcLane) { unsigned* failure_count; +#if GOOGLE_CUDA ASSERT_EQ(cudaMallocManaged(&failure_count, sizeof(unsigned)), cudaSuccess); +#else + ASSERT_EQ(hipHostMalloc(&failure_count, sizeof(unsigned), 0), cudaSuccess); +#endif *failure_count = 0; - TF_EXPECT_OK(GpuLaunchKernel(CudaShuffleGetSrcLaneTest, 1, 32, 0, nullptr, - failure_count)); - ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + TF_EXPECT_OK(GpuLaunchKernel(GpuShuffleGetSrcLaneTest, 1, TF_RED_WARPSIZE, 0, + nullptr, failure_count)); + ASSERT_EQ(gpuDeviceSynchronize(), cudaSuccess); ASSERT_EQ(*failure_count, 0); - cudaFree(failure_count); + gpuFree(failure_count); } } // namespace tensorflow -#endif // GOOGLE_CUDA +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM From e5d5522d827e1c60b3ac830000b4489206480f95 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Mon, 4 May 2020 14:12:51 +0100 Subject: [PATCH 0044/1390] Both inputs should be in int16 for MUL operator. Some networks have one of inputs as a constant. --- tensorflow/lite/tools/optimize/operator_property.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/lite/tools/optimize/operator_property.cc b/tensorflow/lite/tools/optimize/operator_property.cc index 38c34706fbe..2ffe9fa3671 100644 --- a/tensorflow/lite/tools/optimize/operator_property.cc +++ b/tensorflow/lite/tools/optimize/operator_property.cc @@ -807,6 +807,7 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index, case BuiltinOperator_MUL: property.inputs = {{0, {}}, {1, {}}}; property.outputs = {{0, {}}}; + property.quantize_input_as_activations = true; property.version = 2; break; case BuiltinOperator_PACK: From 0ee6b3a69da0f17d14d40e34b7008651012638da Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Wed, 6 May 2020 22:29:23 +0000 Subject: [PATCH 0045/1390] Return ValueError in case of empty list input for tf.map_fn This PR tries to address the issue raised in 39229 where empty lists input was not checked and throw out a non-obvious error: ```python >>> import numpy as np >>> import tensorflow as tf >>> fn = lambda x: x >>> tf.map_fn(fn, []) Traceback (most recent call last): File "", line 1, in File "/Library/Python/3.7/site-packages/tensorflow/python/util/deprecation.py", line 574, in new_func return func(*args, **kwargs) File "/Library/Python/3.7/site-packages/tensorflow/python/ops/map_fn.py", line 425, in map_fn_v2 name=name) File "/Library/Python/3.7/site-packages/tensorflow/python/ops/map_fn.py", line 213, in map_fn static_shape = elems_flat[0].shape IndexError: list index out of range >>> ``` In case of empty list the behavior is undefined as we even don't know the output dtype. This PR update to perform a check and thrown out `ValueError("elems must not be empty")` to help clarify. This PR fixes 39229. Signed-off-by: Yong Tang --- tensorflow/python/kernel_tests/map_fn_test.py | 6 ++++++ tensorflow/python/ops/map_fn.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/tensorflow/python/kernel_tests/map_fn_test.py b/tensorflow/python/kernel_tests/map_fn_test.py index 1e10d689886..9825939933e 100644 --- a/tensorflow/python/kernel_tests/map_fn_test.py +++ b/tensorflow/python/kernel_tests/map_fn_test.py @@ -217,6 +217,12 @@ class MapFnTest(test.TestCase): self.assertAllEqual([0, 3, 2], map_return.get_shape().dims) self.assertAllEqual([0, 3, 2], self.evaluate(map_return).shape) + @test_util.run_in_graph_and_eager_modes + def testMapEmptyList(self): + x = [] + with self.assertRaisesRegexp( + ValueError, r"elems must be a Tensor or"): + _ = map_fn.map_fn(lambda e: e, x) if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/map_fn.py b/tensorflow/python/ops/map_fn.py index 2c9c678336e..6f59bcf5599 100644 --- a/tensorflow/python/ops/map_fn.py +++ b/tensorflow/python/ops/map_fn.py @@ -375,6 +375,13 @@ def map_fn(fn, # Flatten the input tensors, and get the TypeSpec for each one. elems_flat = nest.flatten(elems) + + # Check in case this is an empty list + if len(elems_flat) == 0: + raise ValueError( + "elems must be a Tensor or (possibly nested) sequence of Tensors. " + "Got {}, which does not contain any Tensors.".format(elems)) + elems_flat_signature = [type_spec.type_spec_from_value(e) for e in elems_flat] elems_unflatten = lambda x: nest.pack_sequence_as(elems, x) From 28f2af10ecdde4ab8e24247a728032ea1891d730 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Tue, 12 May 2020 17:41:28 +0200 Subject: [PATCH 0046/1390] Realign AllocationInfo struct. After adding offline_offset, sizeof(AllocationInfo) = 40, which caused hello_world_test to crash. After realigning it's back to its original size (32). --- tensorflow/lite/micro/micro_allocator.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index f1c1d65f1cc..5ffda9209d9 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -34,11 +34,11 @@ namespace { // Used to hold information used during allocation calculations. struct AllocationInfo { size_t bytes; + void** output_ptr; int first_created; int last_used; - bool needs_allocating; - void** output_ptr; int32_t offline_offset; + bool needs_allocating; }; // We align tensor buffers to 16-byte boundaries, since this is a common From 489926629dea271e28417a3c427bf698b7d21d64 Mon Sep 17 00:00:00 2001 From: Niranjan Hasabnis Date: Tue, 12 May 2020 15:26:43 -0700 Subject: [PATCH 0047/1390] [Intel MKL] Adding MklTanh op --- .../core/common_runtime/mkl_layout_pass.cc | 14 +- .../common_runtime/mkl_layout_pass_test.cc | 60 ++++++ tensorflow/core/kernels/BUILD | 21 ++ tensorflow/core/kernels/mkl_relu_op.cc | 65 ++++-- tensorflow/core/kernels/mkl_relu_op_test.cc | 193 ++++++++++++++++++ 5 files changed, 325 insertions(+), 28 deletions(-) create mode 100644 tensorflow/core/kernels/mkl_relu_op_test.cc diff --git a/tensorflow/core/common_runtime/mkl_layout_pass.cc b/tensorflow/core/common_runtime/mkl_layout_pass.cc index 2941845a604..f4de923cb2d 100644 --- a/tensorflow/core/common_runtime/mkl_layout_pass.cc +++ b/tensorflow/core/common_runtime/mkl_layout_pass.cc @@ -675,18 +675,12 @@ class MklLayoutRewritePass : public GraphOptimizationPass { rinfo_.push_back( {csinfo_.requantize, mkl_op_registry::GetMklOpName(csinfo_.requantize), CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation}); - // Disable these two MKL operators for now due to some test failures caused - // by these two ops - /* - rinfo_.push_back({csinfo_.tanh, - mkl_op_registry::GetMklOpName(csinfo_.tanh), + rinfo_.push_back({csinfo_.tanh, mkl_op_registry::GetMklOpName(csinfo_.tanh), CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation}); - rinfo_.push_back({csinfo_.tanh_grad, - mkl_op_registry::GetMklOpName(csinfo_.tanh_grad), - CopyAttrsAll, AlwaysRewrite, - kRewriteForLayoutPropagation}); - */ + rinfo_.push_back( + {csinfo_.tanh_grad, mkl_op_registry::GetMklOpName(csinfo_.tanh_grad), + CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation}); rinfo_.push_back( {csinfo_.reshape, mkl_op_registry::GetMklOpName(csinfo_.reshape), CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation}); diff --git a/tensorflow/core/common_runtime/mkl_layout_pass_test.cc b/tensorflow/core/common_runtime/mkl_layout_pass_test.cc index c6d5331852e..daa7f42620c 100644 --- a/tensorflow/core/common_runtime/mkl_layout_pass_test.cc +++ b/tensorflow/core/common_runtime/mkl_layout_pass_test.cc @@ -2949,6 +2949,66 @@ TEST_F(MklLayoutPassTest, NodeRewrite_LeakyReluLeakyReluGrad_Positive) { "DMT/_1->C:2"); } +// clang-format off +#define REGISTER_TEST(NAME, T, INPUT) \ + TEST_F(MklLayoutPassTest, NAME##_##T) { \ + DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS); \ + InitGraph( \ + "node { name: 'A' op: '" #INPUT "'}" \ + "node { name: 'B' op: 'Tanh'" \ + " attr { key: 'T' value { type: " #T " } }" \ + " input: ['A'] }" \ + "node { name: 'C' op: 'Zeta' attr { key: 'T' value { type: " #T " } }" \ + " input: ['A', 'B'] }"); \ + EXPECT_EQ(DoMklLayoutOptimizationPass(), \ + "A(" #INPUT ");B(_MklTanh);C(Zeta);DMT/_0(Const)|A->B;A->C;" \ + "A:control->DMT/_0:control;B->C:1;DMT/_0->B:1"); \ +} +REGISTER_TEST_ALL_TYPES(NodeRewrite_Tanh_Positive); +#undef REGISTER_TEST + +#define REGISTER_TEST(NAME, T, INPUT) \ + TEST_F(MklLayoutPassTest, NAME##_##T) { \ + DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS); \ + InitGraph( \ + "node { name: 'A' op: '" #INPUT "'}" \ + "node { name: 'B' op: '" #INPUT "'}" \ + "node { name: 'C' op: 'TanhGrad'" \ + " attr { key: 'T' value { type: " #T " } }" \ + " input: ['A', 'B'] }" \ + "node { name: 'D' op: 'Zeta' attr { key: 'T' value { type: " #T " } }" \ + " input: ['A', 'C'] }"); \ + EXPECT_EQ(DoMklLayoutOptimizationPass(), \ + "A(" #INPUT ");B(" #INPUT ");C(_MklTanhGrad);D(Zeta);DMT/_0(Const);" \ + "DMT/_1(Const)|A->C;A->D;A:control->DMT/_0:control;" \ + "A:control->DMT/_1:control;B->C:1;C->D:1;DMT/_0->C:2;DMT/_1->C:3"); \ +} +REGISTER_TEST_ALL_TYPES(NodeRewrite_TanhGrad_Positive); +#undef REGISTER_TEST + +#define REGISTER_TEST(NAME, T, INPUT) \ + TEST_F(MklLayoutPassTest, NAME##_##T) { \ + DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS); \ + InitGraph( \ + "node { name: 'A' op: '" #INPUT "'}" \ + "node { name: 'B' op: 'Tanh'" \ + " attr { key: 'T' value { type: " #T " } }" \ + " input: ['A'] }" \ + "node { name: 'C' op: 'TanhGrad'" \ + " attr { key: 'T' value { type: " #T " } }" \ + " input: ['B', 'A'] }" \ + "node { name: 'D' op: 'Zeta' attr { key: 'T' value { type: " #T " } }" \ + " input: ['A', 'C'] }"); \ + EXPECT_EQ(DoMklLayoutOptimizationPass(), \ + "A(" #INPUT ");B(_MklTanh);C(_MklTanhGrad);D(Zeta);DMT/_0(Const);" \ + "DMT/_1(Const)|A->B;A->C:1;A->D;A:control->DMT/_0:control;" \ + "B->C;B:1->C:2;B:control->DMT/_1:control;C->D:1;DMT/_0->B:1;" \ + "DMT/_1->C:3"); \ +} +REGISTER_TEST_ALL_TYPES(NodeRewrite_TanhTanhGrad_Positive); +#undef REGISTER_TEST +// clang-format on + TEST_F(MklLayoutPassTest, NodeRewrite_AvgPool_Positive) { InitGraph( "node { name: 'A' op: 'Input'}" diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index e47c681bb61..3f6c2fbfb04 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -8084,6 +8084,27 @@ tf_cc_test_mkl( ], ) +tf_cc_test_mkl( + name = "mkl_relu_op_test", + size = "small", + srcs = ["mkl_relu_op_test.cc"], + linkstatic = 1, # Fixes dyld error on MacOS. + deps = [ + ":ops_testutil", + ":ops_util", + "//tensorflow/cc:cc_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensorflow", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ], +) + tf_mkl_kernel_library( name = "mkl_tfconv_op", prefix = "mkl_tfconv", diff --git a/tensorflow/core/kernels/mkl_relu_op.cc b/tensorflow/core/kernels/mkl_relu_op.cc index ffbc1e28355..7f885bfdc01 100644 --- a/tensorflow/core/kernels/mkl_relu_op.cc +++ b/tensorflow/core/kernels/mkl_relu_op.cc @@ -19,7 +19,6 @@ limitations under the License. #include #include "mkldnn.hpp" -#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" @@ -27,6 +26,7 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/util/mkl_types.h" #include "tensorflow/core/util/mkl_util.h" +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" using mkldnn::algorithm; using mkldnn::eltwise_forward; @@ -266,15 +266,19 @@ class MklEltwiseBwdParams { algorithm alg_kind; float alpha; float beta; + // Whether the input that grad op gets from forward op is SRC + // of forward op or DST of forward op. + int forward_input_type; MklEltwiseBwdParams(const memory::dims& src_dims, const memory::desc& common_md, algorithm alg_kind, - float alpha, float beta) + float alpha, float beta, int forward_input_type) : src_dims(src_dims), common_md(common_md), alg_kind(alg_kind), alpha(alpha), - beta(beta) {} + beta(beta), + forward_input_type(forward_input_type) {} }; template @@ -430,7 +434,7 @@ class MklEltwiseBwdPrimitive : public MklPrimitive { // Create eltwise primitive and add it to net. context_.eltwise_bwd.reset(new mkldnn::eltwise_backward(*context_.bwd_pd)); context_.bwd_primitives_args.push_back( - {{MKLDNN_ARG_SRC, *context_.src_mem}, + {{bwdParams.forward_input_type, *context_.src_mem}, {MKLDNN_ARG_DIFF_DST, *context_.diff_dst_mem}, { MKLDNN_ARG_DIFF_SRC, *context_.diff_src_mem }}); @@ -631,14 +635,30 @@ class MklReluGradOpBase : public OpKernel { virtual void Compute_Scalar(OpKernelContext* context) = 0; + // All activation functions that are part of NN ops, such as Relu, Elu, + // LeakyRelu, Relu6, etc have dy at index 0 and y at index 1. + // + // if forward op is defined as: y = f(x), + // {Relu,Elu,Relu6,LeakyRelu}Grad is: z = f_grad(dy,x) + // TanhGrad is: z = tanh_grad(y,dy) + // + // Src below refers to a tensor that gradient op receives from forward + // operator. From Relu-family ops, it is 'x'; while for TanhGrad, it is 'y'. + virtual int GetDiffDstIndex() const { return 0; } + virtual int GetSrcIndex() const { return 1; } + virtual int GetDiffSrcIndex() const { return 0; } + // What is the type of input tensor that grad op receives from forward op -- + // is it 'x' (SRC) or 'y' (DST). For Relu-family, it is 'x', so fwd op SRC. + virtual int GetTypeOfInputTensorFromFwdOp() const { return MKLDNN_ARG_SRC; } + void Compute(OpKernelContext* context) { try { MklDnnData src(&cpu_engine); MklDnnData diff_dst(&cpu_engine); - const size_t diff_dst_index = 0; // index of diff_dst input tensor - const size_t src_index = 1; // index of src input tensor - const size_t diff_src_index = 0; // index of diff_src output tensor + size_t diff_dst_index = GetDiffDstIndex(); + size_t src_index = GetSrcIndex(); + const size_t diff_src_index = GetDiffSrcIndex(); const Tensor& src_tensor = MklGetInput(context, src_index); const Tensor& diff_dst_tensor = MklGetInput(context, diff_dst_index); @@ -722,7 +742,7 @@ class MklReluGradOpBase : public OpKernel { } MklEltwiseBwdParams bwdParams(src_dims, common_md, alg_kind, alpha_, - beta_); + beta_, GetTypeOfInputTensorFromFwdOp()); MklEltwiseBwdPrimitive* eltwise_bwd = MklEltwiseBwdPrimitiveFactory::Get(bwdParams); @@ -976,18 +996,28 @@ class MklTanhOp : public MklReluOpBase { template class MklTanhGradOp - : public MklReluGradOpBase { + : public MklReluGradOpBase { public: ~MklTanhGradOp() {} explicit MklTanhGradOp(OpKernelConstruction* context) - : MklReluGradOpBase(context, 0.0f, - 0.0f) {} + : MklReluGradOpBase( + context, 0.0f, 0.0f) {} + + virtual int GetDiffDstIndex() const { return 1; } + virtual int GetSrcIndex() const { return 0; } + virtual int GetDiffSrcIndex() const { return 0; } + + // TanhGrad gets 'y' from Tanh, where 'y' is output of Tanh(x). + virtual int GetTypeOfInputTensorFromFwdOp() const { return MKLDNN_ARG_DST; } virtual void Compute_Scalar(OpKernelContext* context) { - const size_t diff_dst_index = 0; // index of diff_dst input tensor - const size_t src_index = 1; // index of src input tensor - const size_t diff_src_index = 0; // index of diff_src output tensor + // NOTE: Order of y and dy for Tanh is reverse of that for Relu/Elu/other + // element-wise ops. Tanh is math op in Tensorflow; others are NN ops. + const size_t diff_dst_index = GetDiffDstIndex(); + const size_t src_index = GetSrcIndex(); + const size_t diff_src_index = GetDiffSrcIndex(); const Tensor& src_tensor = MklGetInput(context, src_index); const Tensor& diff_dst_tensor = MklGetInput(context, diff_dst_index); Tensor* diff_src_tensor = nullptr; @@ -1003,10 +1033,9 @@ class MklTanhGradOp void* user_i = static_cast(const_cast(src_tensor.flat().data())); // gradient of tanh(x) = 1 - tanh(x)^2 - T feature = (static_cast(user_i))[0]; - T e1 = std::exp(feature); - T e2 = std::exp(-feature); - T tanh = (e1 - e2) / (e1 + e2); + // Input to TanhGrad is output of Tanh. So we do not need to compute + // Tanh again. + T tanh = (static_cast(user_i))[0]; void* user_g = static_cast(const_cast(diff_dst_tensor.flat().data())); (static_cast(out_o))[0] = diff --git a/tensorflow/core/kernels/mkl_relu_op_test.cc b/tensorflow/core/kernels/mkl_relu_op_test.cc new file mode 100644 index 00000000000..7a3dffef0de --- /dev/null +++ b/tensorflow/core/kernels/mkl_relu_op_test.cc @@ -0,0 +1,193 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#undef INTEL_MKL + +#ifdef INTEL_MKL + +#include "tensorflow/cc/ops/const_op.h" +#include "tensorflow/cc/ops/nn_ops.h" +#include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h" +#include "tensorflow/core/framework/fake_input.h" +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/kernels/ops_util.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/test_benchmark.h" +#include "tensorflow/core/public/session.h" + +#include "mkldnn.hpp" +#include "tensorflow/core/util/mkl_util.h" + +// Compare performance of default Tensorflow convolution kernels (Eigen) with +// MKL kernels on CPU. + +// Before running these benchmarks configure OpenMP environment variables: +// export KMP_BLOCKTIME=0 +// export OMP_NUM_THREADS=${num_threads} + +namespace tensorflow { +static Tensor NonMklTensor() { + MklDnnShape non_mkl_shape; + non_mkl_shape.SetMklTensor(false); + + auto size = static_cast(non_mkl_shape.GetSerializeBufferSize()); + Tensor tensor(DT_UINT8, {size}); + + non_mkl_shape.SerializeMklDnnShape(tensor.flat().data(), + size * sizeof(uint8)); + return tensor; +} + +static Tensor GetRandomTensor(const TensorShape& shape) { + Tensor tensor(DT_FLOAT, TensorShape(shape)); + tensor.flat() = tensor.flat().setRandom(); + return tensor; +} + +#define CREATE_DEFAULT_FWD_OP(NODE_NAME, OP_NAME) \ + static Graph* NODE_NAME(const TensorShape& shape) { \ + auto* graph = new Graph(OpRegistry::Global()); \ + Tensor input_t = GetRandomTensor(shape); \ + Node* input = test::graph::Constant(graph, input_t, "input"); \ + Node* op; \ + TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ + .Input(input) \ + .Attr("T", DT_FLOAT) \ + .Finalize(graph, &op)); \ + return graph; \ + } +CREATE_DEFAULT_FWD_OP(Default_Tanh, Tanh) +CREATE_DEFAULT_FWD_OP(Default_Elu, Elu) +CREATE_DEFAULT_FWD_OP(Default_Relu, Relu) +CREATE_DEFAULT_FWD_OP(Default_Relu6, Relu6) +CREATE_DEFAULT_FWD_OP(Default_LeakyRelu, LeakyRelu) + +#define CREATE_DEFAULT_BWD_OP(NODE_NAME, OP_NAME) \ + static Graph* NODE_NAME(const TensorShape& shape) { \ + auto* graph = new Graph(OpRegistry::Global()); \ + Tensor input_t = GetRandomTensor(shape); \ + Node* input = test::graph::Constant(graph, input_t, "input"); \ + Tensor grad_t = GetRandomTensor(shape); \ + Node* grad = test::graph::Constant(graph, grad_t, "grad"); \ + Node* op; \ + TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ + .Input(grad) \ + .Input(input) \ + .Attr("T", DT_FLOAT) \ + .Finalize(graph, &op)); \ + return graph; \ + } +CREATE_DEFAULT_BWD_OP(Default_TanhGrad, TanhGrad) +CREATE_DEFAULT_BWD_OP(Default_EluGrad, EluGrad) +CREATE_DEFAULT_BWD_OP(Default_ReluGrad, ReluGrad) +CREATE_DEFAULT_BWD_OP(Default_Relu6Grad, Relu6Grad) +CREATE_DEFAULT_BWD_OP(Default_LeakyReluGrad, LeakyReluGrad) + +#define CREATE_MKL_FWD_OP(NODE_NAME, OP_NAME) \ + static Graph* NODE_NAME(const TensorShape& shape) { \ + auto* graph = new Graph(OpRegistry::Global()); \ + \ + Tensor input_t = GetRandomTensor(shape); \ + Node* input = test::graph::Constant(graph, input_t, "input"); \ + \ + Node* not_mkl_shape = \ + test::graph::Constant(graph, NonMklTensor(), "not_mkl"); \ + \ + Node* op; \ + TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ + .Input(input) \ + .Input(not_mkl_shape) \ + .Attr("T", DT_FLOAT) \ + .Attr("_kernel", "MklLayoutDependentOp") \ + .Finalize(graph, &op)); \ + \ + return graph; \ + } + +CREATE_MKL_FWD_OP(Mkl_Tanh, _MklTanh) +CREATE_MKL_FWD_OP(Mkl_Elu, _MklElu) +CREATE_MKL_FWD_OP(Mkl_Relu, _MklRelu) +CREATE_MKL_FWD_OP(Mkl_Relu6, _MklRelu6) +CREATE_MKL_FWD_OP(Mkl_LeakyRelu, _MklLeakyRelu) + +#define CREATE_MKL_BWD_OP(NODE_NAME, OP_NAME) \ + static Graph* NODE_NAME(const TensorShape& shape) { \ + auto* graph = new Graph(OpRegistry::Global()); \ + \ + Tensor input_t = GetRandomTensor(shape); \ + Node* input = test::graph::Constant(graph, input_t, "input"); \ + Tensor grad_t = GetRandomTensor(shape); \ + Node* grad = test::graph::Constant(graph, grad_t, "grad"); \ + \ + Node* not_mkl_shape = \ + test::graph::Constant(graph, NonMklTensor(), "not_mkl"); \ + \ + Node* op; \ + TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ + .Input(grad) \ + .Input(input) \ + .Input(not_mkl_shape) \ + .Input(not_mkl_shape) \ + .Attr("T", DT_FLOAT) \ + .Attr("_kernel", "MklLayoutDependentOp") \ + .Finalize(graph, &op)); \ + \ + return graph; \ + } + +CREATE_MKL_BWD_OP(Mkl_TanhGrad, _MklTanhGrad) +CREATE_MKL_BWD_OP(Mkl_EluGrad, _MklEluGrad) +CREATE_MKL_BWD_OP(Mkl_ReluGrad, _MklReluGrad) +CREATE_MKL_BWD_OP(Mkl_Relu6Grad, _MklRelu6Grad) +CREATE_MKL_BWD_OP(Mkl_LeakyReluGrad, _MklLeakyReluGrad) + +#define BM_Activation(op, kind, A, B, C, D, type) \ + static void BM_##op##_##kind##_##type##_##A##_##B##_##C##_##D(int iters) { \ + int64 num_computed_elements = (A) * (B) * (C) * (D); \ + int64 flops_per_iter = num_computed_elements; \ + testing::ItemsProcessed(static_cast(iters) * flops_per_iter); \ + \ + test::Benchmark(#type, kind##_##op({A, B, C, D})).Run(iters); \ + } \ + BENCHMARK(BM_##op##_##kind##_##type##_##A##_##B##_##C##_##D) + +#define BM(op, A, B, C, D, type) \ + BM_Activation(op, Default, A, B, C, D, type); \ + BM_Activation(op, Mkl, A, B, C, D, type); + +#define TEST_ALL_SIZES(OP) \ + BM(OP, 2, 4, 8, 16, cpu); \ + BM(OP, 3, 5, 9, 17, cpu); \ + BM(OP, 32, 64, 128, 256, cpu); \ + BM(OP, 33, 65, 129, 257, cpu); + +TEST_ALL_SIZES(Tanh) +TEST_ALL_SIZES(TanhGrad) +TEST_ALL_SIZES(Relu) +TEST_ALL_SIZES(ReluGrad) +TEST_ALL_SIZES(Elu) +TEST_ALL_SIZES(EluGrad) +TEST_ALL_SIZES(Relu6) +TEST_ALL_SIZES(Relu6Grad) +TEST_ALL_SIZES(LeakyRelu) +TEST_ALL_SIZES(LeakyReluGrad) + +} // namespace tensorflow + +#endif // INTEL_MKL From f87cf9c19b393657fb4971dd9513f44d73589b16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Kubov=C4=8D=C3=ADk?= Date: Mon, 18 May 2020 00:49:25 +0200 Subject: [PATCH 0048/1390] add lock file --- tensorflow/python/keras/saving/hdf5_format.py | 31 ++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/keras/saving/hdf5_format.py b/tensorflow/python/keras/saving/hdf5_format.py index f3adb2d0695..93f06c018d1 100644 --- a/tensorflow/python/keras/saving/hdf5_format.py +++ b/tensorflow/python/keras/saving/hdf5_format.py @@ -53,7 +53,24 @@ sequential_lib = LazyLoader( # pylint:enable=g-inconsistent-quotes -def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True): +# create lock file +def create_lockfile(filepath): + lockfile_path = f"{filepath}.lock" + + f = open(lockfile_path, 'w') + f.write(f"{os.getpid()}") + f.close() + + return lockfile_path + +def check_lockfile(filepath): + lockfile_path = f"{filepath}.lock" + if os.path.exists(lockfile_path): + # use PID? + return True + return False + +def save_model_to_hdf5(model, filepath, overwrite=True, lockFile=True, include_optimizer=True): """Saves a model to a HDF5 file. The saved model contains: @@ -99,6 +116,10 @@ def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True): if not proceed: return + # create lock file + if (lockFile == True): + lockfile_path = create_lockfile(filepath) + f = h5py.File(filepath, mode='w') opened_new_file = True else: @@ -129,6 +150,10 @@ def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True): if opened_new_file: f.close() + # remove lock file + if (lockFile == True): + os.remove(lockfile_path) + def load_model_from_hdf5(filepath, custom_objects=None, compile=True): # pylint: disable=redefined-builtin """Loads a model saved via `save_model_to_hdf5`. @@ -163,6 +188,10 @@ def load_model_from_hdf5(filepath, custom_objects=None, compile=True): # pylint opened_new_file = not isinstance(filepath, h5py.File) if opened_new_file: + # check if lock file exist + if check_lockfile(filepath) == True: + raise ValueError('Cannot read from file at this time.') + f = h5py.File(filepath, mode='r') else: f = filepath From 93cc43bef97f4371379c2ea6e87b260a2a2cf7af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Kubov=C4=8D=C3=ADk?= Date: Mon, 18 May 2020 00:54:34 +0200 Subject: [PATCH 0049/1390] add lockFile argument to save_model() --- tensorflow/python/keras/saving/save.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/keras/saving/save.py b/tensorflow/python/keras/saving/save.py index 43c09a62ea9..100fb05943a 100644 --- a/tensorflow/python/keras/saving/save.py +++ b/tensorflow/python/keras/saving/save.py @@ -48,6 +48,7 @@ _KERAS_SAVED_MODEL_STILL_EXPERIMENTAL = True def save_model(model, filepath, overwrite=True, + lockFile=True, include_optimizer=True, save_format=None, signatures=None, @@ -95,6 +96,7 @@ def save_model(model, overwrite: Whether we should overwrite any existing model at the target location, or instead ask the user with a manual prompt. include_optimizer: If True, save optimizer's state together. + lockFile: If True, protect model file while saving model. save_format: Either 'tf' or 'h5', indicating whether to save the model to Tensorflow SavedModel or HDF5. Defaults to 'tf' in TF 2.X, and 'h5' in TF 1.X. @@ -128,7 +130,7 @@ def save_model(model, 'to the Tensorflow SavedModel format (by setting save_format="tf") ' 'or using `save_weights`.') hdf5_format.save_model_to_hdf5( - model, filepath, overwrite, include_optimizer) + model, filepath, overwrite, lockFile, include_optimizer) else: saved_model_save.save(model, filepath, overwrite, include_optimizer, signatures, options) From 80fc2a1dd49afe2e22d7ea6621cba6aefb38818c Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 19 May 2020 16:09:25 -0700 Subject: [PATCH 0050/1390] Added aarch64 dockerfiles --- tensorflow/tools/dockerfiles/spec.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml index 436ef41c15a..f68f627e1d0 100644 --- a/tensorflow/tools/dockerfiles/spec.yml +++ b/tensorflow/tools/dockerfiles/spec.yml @@ -50,6 +50,7 @@ releases: - "{ubuntu-devel-ppc64le}{jupyter}" - "{ubuntu-horovod}{jupyter}" - "{ubuntu-devel-horovod}{jupyter}" + - "{ubuntu-devel-aarch64}{jupyter}" slice_sets: @@ -109,6 +110,22 @@ slice_sets: args: - CHECKOUT_TF_SRC=1 + ubuntu-devel-aarch64: + - add_to_name: "devel-aarch64" + dockerfile_exclusive_name: "devel-cpu-aarch64" + dockerfile_subdirectory: "aarch64" + partials: + - ubuntu/version + - ubuntu/devel-cpu + - ubuntu/python + - ubuntu/bazelbuild + - shell + tests: + - build-cpu.sh + args: + - UBUNTU_VERSION=18.04 + - CHECKOUT_TF_SRC=1 + ubuntu-horovod: - add_to_name: "-horovod" dockerfile_exclusive_name: "horovod" From 9df2c846ee87174dc35aabd301f9f9eb1df89c7f Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 19 May 2020 16:15:50 -0700 Subject: [PATCH 0051/1390] Update bazelbuild.partial.Dockerfile --- .../ubuntu/bazelbuild.partial.Dockerfile | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile index 2b4761abc39..d8f88672a84 100644 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile @@ -7,17 +7,20 @@ RUN apt-get update && apt-get install -y \ virtualenv \ swig +RUN apt-get update && apt-get install -y \ + python3-pil \ + python3-h5py \ + python3-keras-preprocessing \ + python3-matplotlib \ + python3-mock \ + python3-numpy \ + python3-scipy \ + python3-sklearn \ + python3-pandas \ + python3-future \ + python3-portpicker + RUN python3 -m pip --no-cache-dir install \ - Pillow \ - h5py \ - keras_preprocessing \ - matplotlib \ - mock \ - numpy \ - scipy \ - sklearn \ - pandas \ - portpicker \ enum34 # Build and install bazel From 3d2e5f8c1dd13f9d066621867a4316ca190fcd06 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 19 May 2020 16:39:44 -0700 Subject: [PATCH 0052/1390] Update bazelbuild.partial.Dockerfile --- .../dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile index d8f88672a84..6c050496324 100644 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile @@ -17,7 +17,6 @@ RUN apt-get update && apt-get install -y \ python3-scipy \ python3-sklearn \ python3-pandas \ - python3-future \ python3-portpicker RUN python3 -m pip --no-cache-dir install \ From fffea34c27c8b177c5ca5e2de00b3d8628c669f7 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Thu, 21 May 2020 10:55:42 -0700 Subject: [PATCH 0053/1390] Update bazelbuild.partial.Dockerfile --- .../dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile index 6c050496324..b833657aa69 100644 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile @@ -22,7 +22,7 @@ RUN apt-get update && apt-get install -y \ RUN python3 -m pip --no-cache-dir install \ enum34 - # Build and install bazel +# Build and install bazel ENV BAZEL_VERSION 3.0.0 WORKDIR / RUN mkdir /bazel && \ From ab809024a4a5b0887c360b3e5542c149f4a5f14d Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Tue, 19 May 2020 15:53:07 -0700 Subject: [PATCH 0054/1390] Enable wider vector for reluGrad --- tensorflow/core/kernels/relu_op_gpu.cu.cc | 87 +++++++++++++++++++++-- 1 file changed, 80 insertions(+), 7 deletions(-) diff --git a/tensorflow/core/kernels/relu_op_gpu.cu.cc b/tensorflow/core/kernels/relu_op_gpu.cu.cc index 27fd5f64249..ca1a4235f3a 100644 --- a/tensorflow/core/kernels/relu_op_gpu.cu.cc +++ b/tensorflow/core/kernels/relu_op_gpu.cu.cc @@ -35,6 +35,7 @@ namespace tensorflow { typedef Eigen::GpuDevice GPUDevice; +static constexpr int VectorSize = 8; namespace functor { // This kernel computes ReluGrad by processing one half2, two fp16, at a time. @@ -93,6 +94,64 @@ __global__ void ReluGradHalfKernel(const Eigen::half* __restrict__ gradient, } } +__global__ void ReluGradHalfKernelVector( + const Eigen::half* __restrict__ gradient, + const Eigen::half* __restrict__ feature, + Eigen::half* __restrict__ backprop, int32 count) { + int32 half8_count = count / VectorSize; + int32 index = blockIdx.x * blockDim.x + threadIdx.x; + + if (index < half8_count) { + float4 gradient_h8 = reinterpret_cast(gradient)[index]; + float4 feature_h8 = reinterpret_cast(feature)[index]; + float4* p_backprop_h8 = reinterpret_cast(backprop) + index; + + half2 *gradient_h2 = reinterpret_cast(&gradient_h8); + half2 *feature_h2 = reinterpret_cast(&feature_h8); + float4 backprop_h8; + half2* p_backprop_h2 = reinterpret_cast(&backprop_h8); + + // Fast path, when half2 primitives are available. +#if __CUDA_ARCH__ >= 530 + const half2 kZeroH2 = __float2half2_rn(0.f); +#endif + for (int i = 0; i < VectorSize / 2; i++) { +#if __CUDA_ARCH__ >= 530 + // mask = (feature > 0) + half2 mask_h2 = __hgt2(feature_h2[i], kZeroH2); + // backprop = mask * gradient + half2 backprop_h2 = __hmul2(mask_h2, gradient_h2[i]); +#else + // Fall back: convert half2 to float2 for processing. + float2 feature_f2 = __half22float2(feature_h2[i]); + float2 gradient_f2 = __half22float2(gradient_h2[i]); + float2 backprop_f2 = make_float2((feature_f2.x > 0) ? gradient_f2.x : 0, + (feature_f2.y > 0) ? gradient_f2.y : 0); + // Convert back to half2. + half2 backprop_h2 = __float22half2_rn(backprop_f2); +#endif + p_backprop_h2[i] = backprop_h2; + } + // Write back the result. + *p_backprop_h8 = backprop_h8; + } + + int remaining_count = (count % VectorSize); + + if (index < remaining_count) { + // Use first threads to process the remaining elements. + Eigen::half grad_h = gradient[half8_count * VectorSize + index]; + Eigen::half feature_h = feature[half8_count * VectorSize + index]; + + float grad_f = static_cast(grad_h); + float feature_f = static_cast(feature_h); + float backprop_f = (feature_f > 0) ? grad_f : 0; + + Eigen::half backprop_h(backprop_f); + backprop[half8_count * VectorSize + index] = backprop_h; + } +} + template struct ReluGrad { // Computes ReluGrad backprop. @@ -108,15 +167,29 @@ struct ReluGrad { // NOTE: When the activation is exactly zero, we do not propagate the // associated gradient value. This allows the output of the Relu to be used, // as well as its input. + auto gradient_ptr = reinterpret_cast(gradient.data()); + auto feature_ptr = reinterpret_cast(feature.data()); + auto backprop_ptr = reinterpret_cast(backprop.data()); + bool aligned = gradient_ptr % 16 == 0 && feature_ptr % 16 == 0 && + backprop_ptr % 16 == 0; int32 count = gradient.size(); - if (count == 0) return; - int32 half2_count = Eigen::divup(count, 2); constexpr int32 kThreadInBlock = 512; - GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize( - half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock); - TF_CHECK_OK(GpuLaunchKernel( - ReluGradHalfKernel, config.block_count, config.thread_per_block, 0, - d.stream(), gradient.data(), feature.data(), backprop.data(), count)); + if (count == 0) return; + if (aligned) { + int32 half8_count = Eigen::divup(count, VectorSize); + int32 kBlock = Eigen::divup(half8_count, kThreadInBlock); + TF_CHECK_OK(GpuLaunchKernel( + ReluGradHalfKernelVector, kBlock, kThreadInBlock, + 0, d.stream(), gradient.data(), feature.data(), backprop.data(), + count)); + } else { + int32 half2_count = Eigen::divup(count, 2); + GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize( + half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock); + TF_CHECK_OK(GpuLaunchKernel( + ReluGradHalfKernel, config.block_count, config.thread_per_block, 0, + d.stream(), gradient.data(), feature.data(), backprop.data(), count)); + } } }; From 695acfc91defb5e55110abd64bcc6b7f6ac65e41 Mon Sep 17 00:00:00 2001 From: Steenu Johnson Date: Sat, 23 May 2020 08:09:25 +0530 Subject: [PATCH 0055/1390] Exclude col property in CSVDataset. Signed-off-by: Steenu Johnson --- .../data/experimental/csv_dataset_op.cc | 73 ++++++++++++++----- .../core/ops/experimental_dataset_ops.cc | 10 ++- .../kernel_tests/csv_dataset_test.py | 10 +++ .../python/data/experimental/ops/readers.py | 13 +++- 4 files changed, 85 insertions(+), 21 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc index 8d1bd7acfd9..62d27294f04 100644 --- a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc @@ -62,6 +62,11 @@ class CSVDatasetOp : public DatasetOpKernel { OP_REQUIRES(ctx, select_cols_tensor->dims() == 1, errors::InvalidArgument("`select_cols` must be a vector.")); + const Tensor* exclude_cols_tensor; + OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); + OP_REQUIRES(ctx, exclude_cols_tensor->dims() == 1, + errors::InvalidArgument("`exclude_cols` must be a vector")); + int64 buffer_size = 0; OP_REQUIRES_OK( ctx, ParseScalarArgument(ctx, "buffer_size", &buffer_size)); @@ -126,11 +131,29 @@ class CSVDatasetOp : public DatasetOpKernel { ctx, select_cols.empty() || select_cols.front() >= 0, errors::InvalidArgument("select_cols should be non-negative indices")); - *output = new Dataset(ctx, std::move(filenames), header, - std::move(compression_type), zlib_compression_options, - output_types_, output_shapes_, - std::move(record_defaults), std::move(select_cols), - use_quote_delim, delim[0], std::move(na_value)); + std::vector exclude_cols; + exclude_cols.reserve(exclude_cols_tensor->NumElements()); + for (int i = 0; i < exclude_cols_tensor->NumElements(); ++i) { + exclude_cols.push_back(exclude_cols_tensor->flat()(i)); + } + OP_REQUIRES(ctx, select_cols.empty() || exclude_cols.empty(), + errors::InvalidArgument( + "Either select_cols or exlcude_cols should be empty")); + for (int i = 1; i < exclude_cols.size(); i++) { + OP_REQUIRES(ctx, exclude_cols[i - 1] < exclude_cols[i], + errors::InvalidArgument( + "exclude_cols should be strictly increasing indices")); + } + OP_REQUIRES( + ctx, exclude_cols.empty() || exclude_cols.front() >= 0, + errors::InvalidArgument("exclude_cols should be non-negative indices")); + + *output = + new Dataset(ctx, std::move(filenames), header, + std::move(compression_type), zlib_compression_options, + output_types_, output_shapes_, std::move(record_defaults), + std::move(select_cols), std::move(exclude_cols), + use_quote_delim, delim[0], std::move(na_value)); } private: @@ -141,7 +164,8 @@ class CSVDatasetOp : public DatasetOpKernel { const DataTypeVector& output_types, const std::vector& output_shapes, std::vector record_defaults, std::vector select_cols, - bool use_quote_delim, char delim, string na_value) + std::vector exclude_cols, bool use_quote_delim, char delim, + string na_value) : DatasetBase(DatasetContext(ctx)), filenames_(std::move(filenames)), header_(header), @@ -149,6 +173,7 @@ class CSVDatasetOp : public DatasetOpKernel { output_shapes_(output_shapes), record_defaults_(std::move(record_defaults)), select_cols_(std::move(select_cols)), + exclude_cols_(std::move(exclude_cols)), use_quote_delim_(use_quote_delim), delim_(delim), na_value_(std::move(na_value)), @@ -184,6 +209,7 @@ class CSVDatasetOp : public DatasetOpKernel { Node* use_quote_delim = nullptr; Node* na_value = nullptr; Node* select_cols = nullptr; + Node* exclude_cols = nullptr; std::vector record_defaults; record_defaults.reserve(record_defaults_.size()); @@ -204,16 +230,18 @@ class CSVDatasetOp : public DatasetOpKernel { TF_RETURN_IF_ERROR(b->AddScalar(use_quote_delim_, &use_quote_delim)); TF_RETURN_IF_ERROR(b->AddScalar(na_value_, &na_value)); TF_RETURN_IF_ERROR(b->AddVector(select_cols_, &select_cols)); + TF_RETURN_IF_ERROR(b->AddVector(exclude_cols_, &exclude_cols)); TF_RETURN_IF_ERROR(b->AddDataset( this, {std::make_pair(0, filenames), std::make_pair(1, compression_type), std::make_pair(2, buffer_size), std::make_pair(3, header), std::make_pair(4, delim), std::make_pair(5, use_quote_delim), - std::make_pair(6, na_value), - std::make_pair(7, select_cols)}, // Single tensor inputs - {std::make_pair(8, record_defaults)}, // Tensor list inputs - {}, output)); + std::make_pair(6, na_value), std::make_pair(7, select_cols), + std::make_pair(8, exclude_cols)}, // Single tensor inputs + {std::make_pair(9, record_defaults)}, // Tensor list inputs + {}, + output)); return Status::OK(); } @@ -227,12 +255,14 @@ class CSVDatasetOp : public DatasetOpKernel { std::vector* out_tensors, bool* end_of_sequence) override { mutex_lock l(mu_); - bool select_all = dataset()->select_cols_.empty(); + bool select_all = + dataset()->select_cols_.empty() && dataset()->exclude_cols_.empty(); do { // We are currently processing a file, so try to read the next record if (input_stream_) { - Status s = ReadRecord(ctx, out_tensors, select_all, - dataset()->select_cols_); + Status s = + ReadRecord(ctx, out_tensors, select_all, + dataset()->select_cols_, dataset()->exclude_cols_); if (s.ok()) { // Validate output if (out_tensors->size() != dataset()->out_type_.size()) { @@ -336,7 +366,8 @@ class CSVDatasetOp : public DatasetOpKernel { // Note: ctx and out_tensors are only used in this function // when fields are included in the record. Status ReadRecord(IteratorContext* ctx, std::vector* out_tensors, - bool select_all, const std::vector& selected) + bool select_all, const std::vector& selected, + const std::vector& excluded) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (pos_ >= buffer_.size()) { // At the end of the file, this will return errors::OutOfRange @@ -350,13 +381,17 @@ class CSVDatasetOp : public DatasetOpKernel { bool end_of_record = false; // Keep track of when we find \n, \r or EOF size_t num_parsed = 0; size_t num_selected_parsed = 0; + size_t num_excluded_parsed = 0; Status result; while (!end_of_record) { // Read till we reach \n, \r or EOF - bool include = - select_all || (num_selected_parsed < selected.size() && - selected[num_selected_parsed] == num_parsed); + bool exclude = num_excluded_parsed < excluded.size() && + excluded[num_excluded_parsed] == num_parsed; + bool include = select_all || + (num_selected_parsed < selected.size() && + selected[num_selected_parsed] == num_parsed) || + (!excluded.empty() && !exclude); // Don't fail fast, so that the next call to GetNext may still return // a valid record @@ -365,6 +400,7 @@ class CSVDatasetOp : public DatasetOpKernel { num_parsed++; if (include) num_selected_parsed++; + if (exclude) num_excluded_parsed++; } return result; @@ -815,7 +851,7 @@ class CSVDatasetOp : public DatasetOpKernel { // the first newline because it might contain quoted fields with // newlines in the header as well std::vector empty; - Status s = ReadRecord(nullptr, nullptr, false, empty); + Status s = ReadRecord(nullptr, nullptr, false, empty, empty); if (!s.ok()) { return errors::InvalidArgument("Can't read header of file"); } @@ -849,6 +885,7 @@ class CSVDatasetOp : public DatasetOpKernel { const std::vector output_shapes_; const std::vector record_defaults_; const std::vector select_cols_; + const std::vector exclude_cols_; const bool use_quote_delim_; const char delim_; const tstring na_value_; diff --git a/tensorflow/core/ops/experimental_dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc index aa4bd64270a..9910472a2c6 100644 --- a/tensorflow/core/ops/experimental_dataset_ops.cc +++ b/tensorflow/core/ops/experimental_dataset_ops.cc @@ -154,6 +154,7 @@ REGISTER_OP("CSVDataset") .Input("use_quote_delim: bool") .Input("na_value: string") .Input("select_cols: int64") + .Input("exclude_cols: int64") .Input("record_defaults: output_types") .Output("handle: variant") .Attr("output_types: list({float,double,int32,int64,string}) >= 1") @@ -174,8 +175,10 @@ REGISTER_OP("CSVDataset") TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); // `select_cols` must be a vector TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); + //`exclude_cols` must be a vecotr + TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused)); // `record_defaults` must be lists of scalars - for (size_t i = 8; i < c->num_inputs(); ++i) { + for (size_t i = 9; i < c->num_inputs(); ++i) { shape_inference::ShapeHandle v; TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { @@ -196,6 +199,7 @@ REGISTER_OP("ExperimentalCSVDataset") .Input("use_quote_delim: bool") .Input("na_value: string") .Input("select_cols: int64") + .Input("exclude_cols: int64") .Input("record_defaults: output_types") .Output("handle: variant") .Attr("output_types: list({float,double,int32,int64,string}) >= 1") @@ -216,8 +220,10 @@ REGISTER_OP("ExperimentalCSVDataset") TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); // `select_cols` must be a vector TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); + // `exclude_cols` must be a vector + TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused)); // `record_defaults` must be lists of scalars - for (size_t i = 8; i < c->num_inputs(); ++i) { + for (size_t i = 9; i < c->num_inputs(); ++i) { shape_inference::ShapeHandle v; TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { diff --git a/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py index 13948305aea..42a96812bb4 100644 --- a/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py @@ -414,6 +414,16 @@ class CsvDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): record_defaults=record_defaults, select_cols=[0]) + @combinations.generate(test_base.default_test_combinations()) + def testCsvDataset_withExcludeCol(self): + record_defaults = [['']] + inputs = [['1,2,3', '5,6,7']] + self._test_dataset( + inputs, + expected_output=[['1'], ['5']], + record_defaults=record_defaults, + exclude_cols=[1, 2]) + @combinations.generate(test_base.default_test_combinations()) def testCsvDataset_withMultipleNewLines(self): # In this case, we expect it to behave differently from diff --git a/tensorflow/python/data/experimental/ops/readers.py b/tensorflow/python/data/experimental/ops/readers.py index b8f4c34f40e..fcf92f5aaf9 100644 --- a/tensorflow/python/data/experimental/ops/readers.py +++ b/tensorflow/python/data/experimental/ops/readers.py @@ -612,7 +612,8 @@ class CsvDatasetV2(dataset_ops.DatasetSource): field_delim=",", use_quote_delim=True, na_value="", - select_cols=None): + select_cols=None, + exclude_cols=None): """Creates a `CsvDataset` by reading and decoding CSV files. The elements of this dataset correspond to records from the file(s). @@ -679,6 +680,9 @@ class CsvDatasetV2(dataset_ops.DatasetSource): select_cols: (Optional.) A sorted list of column indices to select from the input data. If specified, only this subset of columns will be parsed. Defaults to parsing all columns. + exclude_cols:(Optional.) A sorted list of column indices to exclude from + the input data. If specified, only the complement of this set of column + will be parsed. Defaults to parsing all columns. """ self._filenames = ops.convert_to_tensor( filenames, dtype=dtypes.string, name="filenames") @@ -710,6 +714,12 @@ class CsvDatasetV2(dataset_ops.DatasetSource): argument_default=[], argument_dtype=dtypes.int64, ) + self._exclude_cols = convert.optional_param_to_tensor( + "exclude_cols", + exclude_cols, + argument_default=[], + argument_dtype=dtypes.int64, + ) self._element_spec = tuple( tensor_spec.TensorSpec([], d.dtype) for d in self._record_defaults) variant_tensor = gen_experimental_dataset_ops.csv_dataset( @@ -722,6 +732,7 @@ class CsvDatasetV2(dataset_ops.DatasetSource): use_quote_delim=self._use_quote_delim, na_value=self._na_value, select_cols=self._select_cols, + exclude_cols=self._exclude_cols, compression_type=self._compression_type) super(CsvDatasetV2, self).__init__(variant_tensor) From fb190dcf851b96c3175417f253865d94f9eb4785 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 17:56:14 -0700 Subject: [PATCH 0056/1390] Update spec.yml --- tensorflow/tools/dockerfiles/spec.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml index f68f627e1d0..9a7325c7441 100644 --- a/tensorflow/tools/dockerfiles/spec.yml +++ b/tensorflow/tools/dockerfiles/spec.yml @@ -50,7 +50,7 @@ releases: - "{ubuntu-devel-ppc64le}{jupyter}" - "{ubuntu-horovod}{jupyter}" - "{ubuntu-devel-horovod}{jupyter}" - - "{ubuntu-devel-aarch64}{jupyter}" + - "{ubuntu-devel-arm64v8}{jupyter}" slice_sets: @@ -111,9 +111,9 @@ slice_sets: - CHECKOUT_TF_SRC=1 ubuntu-devel-aarch64: - - add_to_name: "devel-aarch64" - dockerfile_exclusive_name: "devel-cpu-aarch64" - dockerfile_subdirectory: "aarch64" + - add_to_name: "devel-arm64v8" + dockerfile_exclusive_name: "devel-cpu-arm64v8" + dockerfile_subdirectory: "arm64v8" partials: - ubuntu/version - ubuntu/devel-cpu From 3dc3f491c4c53a87e6a40b4f9d00415d0c8634e8 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 17:57:47 -0700 Subject: [PATCH 0057/1390] Update spec.yml --- tensorflow/tools/dockerfiles/spec.yml | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml index 9a7325c7441..6928cca8bce 100644 --- a/tensorflow/tools/dockerfiles/spec.yml +++ b/tensorflow/tools/dockerfiles/spec.yml @@ -50,6 +50,7 @@ releases: - "{ubuntu-devel-ppc64le}{jupyter}" - "{ubuntu-horovod}{jupyter}" - "{ubuntu-devel-horovod}{jupyter}" + - "{ubuntu-devel-arm32v7}{jupyter}" - "{ubuntu-devel-arm64v8}{jupyter}" slice_sets: @@ -110,7 +111,23 @@ slice_sets: args: - CHECKOUT_TF_SRC=1 - ubuntu-devel-aarch64: + ubuntu-devel-arm32v7: + - add_to_name: "devel-arm32v7" + dockerfile_exclusive_name: "devel-cpu-arm32v7" + dockerfile_subdirectory: "arm32v7" + partials: + - ubuntu/version + - ubuntu/devel-cpu + - ubuntu/python + - ubuntu/bazelbuild + - shell + tests: + - build-cpu.sh + args: + - UBUNTU_VERSION=18.04 + - CHECKOUT_TF_SRC=1 + + ubuntu-devel-arm64v8: - add_to_name: "devel-arm64v8" dockerfile_exclusive_name: "devel-cpu-arm64v8" dockerfile_subdirectory: "arm64v8" From 57032278576a43f654c45230ae055f5e0ea06cd6 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:00:52 -0700 Subject: [PATCH 0058/1390] Create bazelbuild-arm32v7.partial.Dockerfile --- .../bazelbuild-arm32v7.partial.Dockerfile | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile new file mode 100644 index 00000000000..b833657aa69 --- /dev/null +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile @@ -0,0 +1,35 @@ +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + openjdk-8-jdk \ + python3-dev \ + virtualenv \ + swig + +RUN apt-get update && apt-get install -y \ + python3-pil \ + python3-h5py \ + python3-keras-preprocessing \ + python3-matplotlib \ + python3-mock \ + python3-numpy \ + python3-scipy \ + python3-sklearn \ + python3-pandas \ + python3-portpicker + +RUN python3 -m pip --no-cache-dir install \ + enum34 + +# Build and install bazel +ENV BAZEL_VERSION 3.0.0 +WORKDIR / +RUN mkdir /bazel && \ + cd /bazel && \ + curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip && \ + unzip bazel-$BAZEL_VERSION-dist.zip && \ + bash ./compile.sh && \ + cp output/bazel /usr/local/bin/ && \ + rm -rf /bazel && \ + cd - From a1e2ea51681b06086d2ce041fc976951e3faa29f Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:01:39 -0700 Subject: [PATCH 0059/1390] Create bazelbuild-arm64v8.partial.Dockerfile --- .../bazelbuild-arm64v8.partial.Dockerfile | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile new file mode 100644 index 00000000000..b833657aa69 --- /dev/null +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile @@ -0,0 +1,35 @@ +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + openjdk-8-jdk \ + python3-dev \ + virtualenv \ + swig + +RUN apt-get update && apt-get install -y \ + python3-pil \ + python3-h5py \ + python3-keras-preprocessing \ + python3-matplotlib \ + python3-mock \ + python3-numpy \ + python3-scipy \ + python3-sklearn \ + python3-pandas \ + python3-portpicker + +RUN python3 -m pip --no-cache-dir install \ + enum34 + +# Build and install bazel +ENV BAZEL_VERSION 3.0.0 +WORKDIR / +RUN mkdir /bazel && \ + cd /bazel && \ + curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip && \ + unzip bazel-$BAZEL_VERSION-dist.zip && \ + bash ./compile.sh && \ + cp output/bazel /usr/local/bin/ && \ + rm -rf /bazel && \ + cd - From 40094b18f499d59bfcd21fde380ae6abb24a959e Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:02:49 -0700 Subject: [PATCH 0060/1390] Update spec.yml --- tensorflow/tools/dockerfiles/spec.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml index 6928cca8bce..2556e740e1d 100644 --- a/tensorflow/tools/dockerfiles/spec.yml +++ b/tensorflow/tools/dockerfiles/spec.yml @@ -119,7 +119,7 @@ slice_sets: - ubuntu/version - ubuntu/devel-cpu - ubuntu/python - - ubuntu/bazelbuild + - ubuntu/bazelbuild-arm32v7 - shell tests: - build-cpu.sh @@ -135,7 +135,7 @@ slice_sets: - ubuntu/version - ubuntu/devel-cpu - ubuntu/python - - ubuntu/bazelbuild + - ubuntu/bazelbuild-arm64v8 - shell tests: - build-cpu.sh From e71c33a671f38c77f48cd30e89972ef47aa56adb Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:08:37 -0700 Subject: [PATCH 0061/1390] Update bazelbuild.partial.Dockerfile --- .../ubuntu/bazelbuild.partial.Dockerfile | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile index b833657aa69..628e56890a8 100644 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild.partial.Dockerfile @@ -7,19 +7,17 @@ RUN apt-get update && apt-get install -y \ virtualenv \ swig -RUN apt-get update && apt-get install -y \ - python3-pil \ - python3-h5py \ - python3-keras-preprocessing \ - python3-matplotlib \ - python3-mock \ - python3-numpy \ - python3-scipy \ - python3-sklearn \ - python3-pandas \ - python3-portpicker - RUN python3 -m pip --no-cache-dir install \ + Pillow \ + h5py \ + keras_preprocessing \ + matplotlib \ + mock \ + numpy \ + scipy \ + sklearn \ + pandas \ + portpicker \ enum34 # Build and install bazel From d3775b48823caeff6f9393bfe3e3d6ddca670d84 Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:14:00 -0700 Subject: [PATCH 0062/1390] Delete bazelbuild-arm32v7.partial.Dockerfile --- .../bazelbuild-arm32v7.partial.Dockerfile | 35 ------------------- 1 file changed, 35 deletions(-) delete mode 100644 tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile deleted file mode 100644 index b833657aa69..00000000000 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm32v7.partial.Dockerfile +++ /dev/null @@ -1,35 +0,0 @@ -RUN apt-get update && apt-get install -y \ - build-essential \ - curl \ - git \ - openjdk-8-jdk \ - python3-dev \ - virtualenv \ - swig - -RUN apt-get update && apt-get install -y \ - python3-pil \ - python3-h5py \ - python3-keras-preprocessing \ - python3-matplotlib \ - python3-mock \ - python3-numpy \ - python3-scipy \ - python3-sklearn \ - python3-pandas \ - python3-portpicker - -RUN python3 -m pip --no-cache-dir install \ - enum34 - -# Build and install bazel -ENV BAZEL_VERSION 3.0.0 -WORKDIR / -RUN mkdir /bazel && \ - cd /bazel && \ - curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip && \ - unzip bazel-$BAZEL_VERSION-dist.zip && \ - bash ./compile.sh && \ - cp output/bazel /usr/local/bin/ && \ - rm -rf /bazel && \ - cd - From 8882e4b89a707d68ab6f71704d388a43ae7de76f Mon Sep 17 00:00:00 2001 From: settle <31239886+settle@users.noreply.github.com> Date: Tue, 26 May 2020 18:14:41 -0700 Subject: [PATCH 0063/1390] Update spec.yml --- tensorflow/tools/dockerfiles/spec.yml | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/tensorflow/tools/dockerfiles/spec.yml b/tensorflow/tools/dockerfiles/spec.yml index 2556e740e1d..ea05d77d001 100644 --- a/tensorflow/tools/dockerfiles/spec.yml +++ b/tensorflow/tools/dockerfiles/spec.yml @@ -50,7 +50,6 @@ releases: - "{ubuntu-devel-ppc64le}{jupyter}" - "{ubuntu-horovod}{jupyter}" - "{ubuntu-devel-horovod}{jupyter}" - - "{ubuntu-devel-arm32v7}{jupyter}" - "{ubuntu-devel-arm64v8}{jupyter}" slice_sets: @@ -111,22 +110,6 @@ slice_sets: args: - CHECKOUT_TF_SRC=1 - ubuntu-devel-arm32v7: - - add_to_name: "devel-arm32v7" - dockerfile_exclusive_name: "devel-cpu-arm32v7" - dockerfile_subdirectory: "arm32v7" - partials: - - ubuntu/version - - ubuntu/devel-cpu - - ubuntu/python - - ubuntu/bazelbuild-arm32v7 - - shell - tests: - - build-cpu.sh - args: - - UBUNTU_VERSION=18.04 - - CHECKOUT_TF_SRC=1 - ubuntu-devel-arm64v8: - add_to_name: "devel-arm64v8" dockerfile_exclusive_name: "devel-cpu-arm64v8" From b4062c71c659bde06e5e15c53424a6853aa1ee3c Mon Sep 17 00:00:00 2001 From: rangjiaheng Date: Wed, 27 May 2020 13:27:18 +0800 Subject: [PATCH 0064/1390] add TF_LOCKS_EXCLUDED to MutableDenseHashTable::MemoryUsed --- tensorflow/core/kernels/lookup_table_op.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/lookup_table_op.cc b/tensorflow/core/kernels/lookup_table_op.cc index 9807247ed4f..f269aa65b4e 100644 --- a/tensorflow/core/kernels/lookup_table_op.cc +++ b/tensorflow/core/kernels/lookup_table_op.cc @@ -557,7 +557,7 @@ class MutableDenseHashTable final : public LookupInterface { TensorShape value_shape() const override { return value_shape_; } - int64 MemoryUsed() const override { + int64 MemoryUsed() const override TF_LOCKS_EXCLUDED(mu_) { tf_shared_lock l(mu_); return sizeof(MutableDenseHashTable) + key_buckets_.AllocatedBytes() + value_buckets_.AllocatedBytes() + empty_key_.AllocatedBytes(); From 18aa35c75fcf64dfd99e2a3e7cdcd62bafbc030f Mon Sep 17 00:00:00 2001 From: Fredrik Knutsson Date: Wed, 27 May 2020 19:18:06 +0200 Subject: [PATCH 0065/1390] Fixed review comments 27/5 --- .../micro/memory_planner/greedy_memory_planner.cc | 2 +- tensorflow/lite/micro/micro_allocator.cc | 13 +++++++------ tensorflow/lite/micro/micro_optional_debug_tools.cc | 3 +-- tensorflow/lite/micro/test_helpers.cc | 4 +--- 4 files changed, 10 insertions(+), 12 deletions(-) diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc index 47bb7cfb8c0..8f21a167f67 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc @@ -292,7 +292,7 @@ size_t GreedyMemoryPlanner::GetMaximumMemorySize() { if (buffer_count_ == 0) { return 0; } - ListEntry* entry = &buffers_sorted_by_offset_[0]; + ListEntry* entry = &buffers_sorted_by_offset_[first_entry_index_]; size_t max_size = 0; while (entry) { BufferRequirements* requirements = diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 9c4de6e3035..870e466b4e4 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -91,7 +91,7 @@ TfLiteStatus AllocateVariables( TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, ErrorReporter* error_reporter) { if (model->metadata()) { - for (int i = 0; i < model->metadata()->size(); ++i) { + for (size_t i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata, strlen(kOfflineMemAllocMetadata)) == 0) { @@ -115,11 +115,11 @@ TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, "Offline planner metadata found, version %d, " "subgraph %d, nbr offline offsets %d", version, subgraph_idx, nbr_offline_offsets); - for (int i = 0; i < nbr_offline_offsets; ++i) { + for (int j = 0; j < nbr_offline_offsets; ++j) { TF_LITE_REPORT_ERROR( error_reporter, - "Offline planner tensor index %d, offline offset: %d", i, - offline_planner_offsets[i]); + "Offline planner tensor index %d, offline offset: %d", j, + offline_planner_offsets[j]); } if (version != 1) { @@ -302,7 +302,7 @@ TfLiteStatus AllocationInfoBuilder::AddTensors(const SubGraph* subgraph, TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets( const Model* model, int32_t** offline_planner_offsets) { if (model->metadata()) { - for (int i = 0; i < model->metadata()->size(); ++i) { + for (size_t i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); if (strncmp(metadata->name()->c_str(), kOfflineMemAllocMetadata, strlen(kOfflineMemAllocMetadata)) == 0) { @@ -311,7 +311,7 @@ TfLiteStatus AllocationInfoBuilder::GetOfflinePlannedOffsets( auto* buffer = (*buffers)[metadata->buffer()]; auto* array = buffer->data(); const uint32_t* metadata_buffer = (uint32_t*)array->data(); - const int32_t nbr_tensors = metadata_buffer[2]; + const size_t nbr_tensors = (size_t)metadata_buffer[2]; *offline_planner_offsets = (int32_t*)&metadata_buffer[3]; if (tensor_count_ != nbr_tensors) { @@ -339,6 +339,7 @@ TfLiteStatus AllocationInfoBuilder::AddScratchBuffers( current->first_created = handle->node_idx; current->last_used = handle->node_idx; current->needs_allocating = true; + current->offline_offset = kOnlinePlannedBuffer; } return kTfLiteOk; } diff --git a/tensorflow/lite/micro/micro_optional_debug_tools.cc b/tensorflow/lite/micro/micro_optional_debug_tools.cc index 418347a5b25..22b170094d5 100644 --- a/tensorflow/lite/micro/micro_optional_debug_tools.cc +++ b/tensorflow/lite/micro/micro_optional_debug_tools.cc @@ -122,9 +122,8 @@ void PrintModelData(const Model* model, ErrorReporter* error_reporter) { const flatbuffers::Vector>* buffers = model->buffers(); TF_LITE_REPORT_ERROR(error_reporter, "==== Model info: ====="); - for (int i = 0; i < tensors->size(); ++i) { + for (size_t i = 0; i < tensors->size(); ++i) { const tflite::Tensor& flatbuffer_tensor = *tensors->Get(i); - auto* quantization = flatbuffer_tensor.quantization(); size_t type_size, tensor_size; auto* buffer = (*buffers)[flatbuffer_tensor.buffer()]; auto* array = buffer->data(); diff --git a/tensorflow/lite/micro/test_helpers.cc b/tensorflow/lite/micro/test_helpers.cc index b60d3065020..96e000b1b6d 100644 --- a/tensorflow/lite/micro/test_helpers.cc +++ b/tensorflow/lite/micro/test_helpers.cc @@ -308,10 +308,8 @@ const Model* BuildModelWithOfflinePlanning(int number_of_tensors, model_builder.RegisterOp(BuiltinOperator_CUSTOM, "mock_custom", /* version= */ 0); - int tensors[number_of_tensors]; - for (int i = 0; i < number_of_tensors; ++i) { - tensors[i] = model_builder.AddTensor(TensorType_FLOAT32, {2, 2, 3}); + model_builder.AddTensor(TensorType_FLOAT32, {2, 2, 3}); } for (int i = 0; i < num_conns; ++i) { From 0c7160babc2213269798ea64a292801123edd0d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Tue, 21 Apr 2020 15:25:11 +0200 Subject: [PATCH 0066/1390] TFLu: Add MVE flag to cmsis-nn glue for clarity __ARM_FEATURE_MVE is now autodetected so ARM_MATH_MVEI is no longer needed. --- tensorflow/lite/micro/kernels/cmsis-nn/conv.cc | 4 ++-- tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc | 6 +++--- tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc | 7 ++++--- tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc | 4 ++-- tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc | 5 ----- 5 files changed, 11 insertions(+), 15 deletions(-) diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/conv.cc b/tensorflow/lite/micro/kernels/cmsis-nn/conv.cc index 6e8272b221a..c505670d747 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/conv.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/conv.cc @@ -115,7 +115,7 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) { } TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) OpData data; int32_t buf_size = 0; @@ -240,7 +240,7 @@ TfLiteStatus EvalQuantizedPerChannel( quant_params.multiplier = data->per_channel_output_multiplier; quant_params.shift = data->per_channel_output_shift; -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) RuntimeShape filter_shape = GetTensorShape(filter); RuntimeShape input_shape = GetTensorShape(input); RuntimeShape output_shape = GetTensorShape(output); diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc b/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc index 7ba03d3890d..f18c3170174 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc @@ -104,7 +104,7 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) { } TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) auto* params = reinterpret_cast(node->builtin_data); @@ -186,7 +186,7 @@ TfLiteStatus EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node, op_params.quantized_activation_min = std::numeric_limits::min(); op_params.quantized_activation_max = std::numeric_limits::max(); -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) RuntimeShape filter_shape = GetTensorShape(filter); const int filter_height = filter_shape.Dims(1); const int filter_width = filter_shape.Dims(2); @@ -284,7 +284,7 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node, // Legacy ops used mixed left and right shifts. Now all are +ve-means-left. op_params.output_shift = -data->output_shift; -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) // optimizations utilize loop unrolling which requires the following power // of two kernel dimensions RuntimeShape filter_shape = GetTensorShape(filter); diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc index 78787ea2547..0ab32ecc3c3 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc @@ -84,11 +84,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, input->type, output->type); TF_LITE_ENSURE_MSG(context, input->type == filter->type, "Hybrid models are not supported on TFLite Micro."); -#if defined(__ARM_FEATURE_DSP) + +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) RuntimeShape filter_shape = GetTensorShape(filter); const int filter_dim_count = filter_shape.DimensionsCount(); const int accum_depth = filter_shape.Dims(filter_dim_count - 1); - const int32_t buf_size = arm_fully_connected_s8_get_buffer_size(accum_depth); int* buffer_idx = reinterpret_cast(node->user_data); @@ -101,6 +101,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { *buffer_idx = -1; } #endif + return kTfLiteOk; } @@ -116,7 +117,7 @@ TfLiteStatus EvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node, const int filter_dim_count = filter_shape.DimensionsCount(); const int accum_depth = filter_shape.Dims(filter_dim_count - 1); -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) int16_t* buf = nullptr; auto* buffer_idx = reinterpret_cast(node->user_data); diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc index a12f628e721..8447ad041cc 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc @@ -106,7 +106,7 @@ TfLiteStatus AverageEvalInt8(TfLiteContext* context, const TfLiteNode* node, TFLITE_DCHECK_LE(activation_min, activation_max); -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) RuntimeShape input_shape = GetTensorShape(input); TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); @@ -283,7 +283,7 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) { } TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { -#if defined(__ARM_FEATURE_DSP) +#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE) const TfLiteTensor* input = GetInput(context, node, kInputTensor); const TfLiteTensor* output = GetOutput(context, node, kOutputTensor); diff --git a/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc b/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc index 657d4fa87cf..cfd87089a84 100644 --- a/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc +++ b/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc @@ -8,11 +8,6 @@ ifneq ($(filter cmsis-nn,$(ALL_TAGS)),) THIRD_PARTY_DOWNLOADS += \ $(eval $(call add_third_party_download,$(CMSIS_URL),$(CMSIS_MD5),cmsis,)) - ifneq (,$(filter $(TARGET_ARCH), cortex-m55)) - CCFLAGS += -DARM_MATH_MVEI - CXXFLAGS += -DARM_MATH_MVEI - endif - CMSIS_PATH = $(MAKEFILE_DIR)/downloads/cmsis/ # Include CMSIS-NN files From 1faa442c02c6bb479b958a46eb2f858c1985966a Mon Sep 17 00:00:00 2001 From: storypku Date: Fri, 29 May 2020 23:00:06 +0800 Subject: [PATCH 0067/1390] Make third_party/gpus/... build pass. fix issue #39759 --- third_party/gpus/cuda/BUILD | 10 +--------- third_party/gpus/rocm/BUILD | 7 +------ 2 files changed, 2 insertions(+), 15 deletions(-) diff --git a/third_party/gpus/cuda/BUILD b/third_party/gpus/cuda/BUILD index 413d28a2723..9c988668ebc 100644 --- a/third_party/gpus/cuda/BUILD +++ b/third_party/gpus/cuda/BUILD @@ -1,9 +1 @@ -load("@bazel_skylib//:bzl_library.bzl", "bzl_library") - -bzl_library( - name = "build_defs_bzl", - srcs = ["build_defs.bzl"], - deps = [ - "@bazel_skylib//lib:selects", - ], -) +# Dummy BUILD for Bazel diff --git a/third_party/gpus/rocm/BUILD b/third_party/gpus/rocm/BUILD index bc2dd419259..dd64bfb0f36 100644 --- a/third_party/gpus/rocm/BUILD +++ b/third_party/gpus/rocm/BUILD @@ -1,6 +1 @@ -load("@bazel_skylib//:bzl_library.bzl", "bzl_library") - -bzl_library( - name = "build_defs_bzl", - srcs = ["build_defs.bzl"], -) +# placeholder for bazel From cbf60b5223997eceb6c4221ef4868fd6c792622c Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 3 Jun 2020 12:38:44 +0100 Subject: [PATCH 0068/1390] Addressed reviewer's comments. Change-Id: I3b7842c42b8c905ed44e0cd556134210cb45479c --- tensorflow/lite/python/lite.py | 37 ++++++++++++++++------------------ 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 92c0d5a95d9..010952820b9 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -258,7 +258,7 @@ class QuantizationMode(object): if self.training_time_int8_allow_float(): return { - "inference_type": inference_ty if inference_ty else constants.INT8, + "inference_type": inference_ty if inference_ty else self.activations_type(), "inference_input_type": inference_input_ty if inference_input_ty else constants.FLOAT, "post_training_quantize": False, # disable dynamic range quantization @@ -297,12 +297,28 @@ class QuantizationMode(object): return True, { "inference_input_type": inference_input_type, "inference_output_type": inference_output_type, + "activations_type": constants.INT8, "allow_float": False } elif self.post_training_int8_allow_float(): return True, { "inference_input_type": inference_input_type, "inference_output_type": inference_output_type, + "activations_type": constants.INT8, + "allow_float": True + } + elif self.post_training_int16x8_no_float(): + return True, { + "inference_input_type": inference_input_type, + "inference_output_type": inference_output_type, + "activations_type": constants.INT16, + "allow_float": False + } + elif self.post_training_int16x8_allow_float(): + return True, { + "inference_input_type": inference_input_type, + "inference_output_type": inference_output_type, + "activations_type": constants.INT16, "allow_float": True } else: @@ -573,25 +589,6 @@ class TFLiteConverterBaseV2(TFLiteConverterBase): output_tensors=output_tensors, **converter_kwargs) - activations_type = quant_mode.activations_type() - - if quant_mode.post_training_int8_no_float(): - result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, activations_type, - False) - elif quant_mode.post_training_int8_allow_float(): - result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, activations_type, - True) - elif quant_mode.post_training_int16x8_no_float(): - result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, activations_type, - False) - elif quant_mode.post_training_int16x8_allow_float(): - result = self._calibrate_quantize_model(result, constants.FLOAT, - constants.FLOAT, activations_type, - True) - calibrate_and_quantize, flags = quant_mode.quantizer_flags() if calibrate_and_quantize: result = self._calibrate_quantize_model(result, **flags) From 67ea57b15bb223e72a60265c24082ae5a31d0f0e Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 3 Jun 2020 15:57:18 +0100 Subject: [PATCH 0069/1390] Small fix for inference. Change-Id: Ifd8670ccb9604ecced3d013f529ddbe16fcd75cf --- tensorflow/lite/python/util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/lite/python/util.py b/tensorflow/lite/python/util.py index 32a2d596629..b56b9f49b7a 100644 --- a/tensorflow/lite/python/util.py +++ b/tensorflow/lite/python/util.py @@ -49,6 +49,7 @@ _MAP_TF_TO_TFLITE_TYPES = { dtypes.string: _types_pb2.STRING, dtypes.uint8: _types_pb2.QUANTIZED_UINT8, dtypes.int8: _types_pb2.INT8, + dtypes.int16: _types_pb2.QUANTIZED_INT16, dtypes.complex64: _types_pb2.COMPLEX64, dtypes.bool: _types_pb2.BOOL, } From 29fdee8e85e750d04f6e9d378e85443ba5c7a239 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 3 Jun 2020 17:50:28 +0100 Subject: [PATCH 0070/1390] Fix for error_reporter. Change-Id: I58745cc97872af74b1ad5b0af3ad778b39f01555 --- .../lite/tools/optimize/quantization_utils.cc | 3 ++- tensorflow/lite/tools/optimize/quantize_model.cc | 16 +++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/tensorflow/lite/tools/optimize/quantization_utils.cc b/tensorflow/lite/tools/optimize/quantization_utils.cc index abbcb642287..cdf2743585e 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils.cc +++ b/tensorflow/lite/tools/optimize/quantization_utils.cc @@ -113,7 +113,8 @@ TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type, tensor->quantization->max[0], quantized_range, quantization_params); } else { - error_reporter->Report( + TF_LITE_REPORT_ERROR( + error_reporter, "Unsupported activation type for quantize-activation: %s", activations_type); return kTfLiteError; diff --git a/tensorflow/lite/tools/optimize/quantize_model.cc b/tensorflow/lite/tools/optimize/quantize_model.cc index 6dd8ddd2d8c..0cf69eee3b4 100644 --- a/tensorflow/lite/tools/optimize/quantize_model.cc +++ b/tensorflow/lite/tools/optimize/quantize_model.cc @@ -370,9 +370,9 @@ TfLiteStatus ApplyConstraints(ModelT* model, std::unique_ptr additional_tensor; const string requant_tensor_name = input_tensor->name + "_requantized"; utils::MakeTensorWithQuantParam( - requant_tensor_name, input_tensor->shape, - input_tensor->shape_signature, activations_type, - output_scale, output_zp, &additional_tensor); + requant_tensor_name, input_tensor->shape, + input_tensor->shape_signature, activations_type, output_scale, + output_zp, &additional_tensor); const int32_t additional_tensor_idx = subgraph->tensors.size(); subgraph->tensors.push_back(std::move(additional_tensor)); @@ -869,13 +869,15 @@ TfLiteStatus QuantizeWeightsInputOutput( if (activations_type == TensorType_INT16 && !property.quantizable && !allow_float) { - error_reporter->Report( - "Quantization to 16x8-bit not yet supported for op: %s", + TF_LITE_REPORT_ERROR( + error_reporter, + "Quantization to 16x8-bit not yet supported for op: %", EnumNameBuiltinOperator(op_code)); return kTfLiteError; } else if (!property.quantizable && !allow_float) { - error_reporter->Report("Quantization not yet supported for op: %s", - EnumNameBuiltinOperator(op_code)); + TF_LITE_REPORT_ERROR(error_reporter, + "Quantization not yet supported for op: %", + EnumNameBuiltinOperator(op_code)); return kTfLiteError; } From 761d850ac6456aed93ab250ff49af3f0a6a62960 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 3 Jun 2020 17:56:39 +0100 Subject: [PATCH 0071/1390] Renamed option with the prefix EXPERIMENTAL_. Change-Id: Idb84736507d5c07ebdf182b8a15d55906d0d7fc0 --- tensorflow/lite/python/convert.py | 2 +- tensorflow/lite/python/lite.py | 2 +- tensorflow/lite/python/lite_test.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index c30987a5898..939de61c608 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -98,7 +98,7 @@ class OpsSet(enum.Enum): # and int16 activations. # Specifying this will throw an error for operations that do not yet have # quantized implementations. - TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = "TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" + EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" def __str__(self): return self.value diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 010952820b9..781007241b4 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -356,7 +356,7 @@ class QuantizationMode(object): def _is_int16x8_target_required(self): return bool( set(self._target_spec.supported_ops).intersection([ - OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 + OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 ])) def _is_allow_float(self): diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index 4075a887943..1d052b88c10 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -885,7 +885,7 @@ class FromSessionTest(TestModels, parameterized.TestCase): # Quantize model to Int8: with disable mlir ('UseTfliteBuiltinsIntDisableMLIR', [lite.OpsSet.TFLITE_BUILTINS_INT8], False), # Quantize model to Int16: with disable mlir - ('UseTfliteBuiltinsInt16DisableMLIR', [lite.OpsSet.TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], False)) + ('UseTfliteBuiltinsInt16DisableMLIR', [lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], False)) def testCalibrateAndQuantizeBuiltinInt(self, supported_ops, enable_mlir): with ops.Graph().as_default(): inp, output, calibration_gen = self._getCalibrationQuantizeModel() From ab6b2ffde37bf4443c7dadc39312f0429f417db6 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Fri, 15 May 2020 11:40:22 -0500 Subject: [PATCH 0072/1390] Removed TENSOR_OP disable env vars. * TF_DISABLE_CUBLAS_TENSOR_OP_MATH * TF_DISABLE_CUDNN_TENSOR_OP_MATH * TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH --- tensorflow/stream_executor/cuda/cuda_blas.cc | 21 ++------ tensorflow/stream_executor/cuda/cuda_dnn.cc | 55 +++++--------------- 2 files changed, 16 insertions(+), 60 deletions(-) diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc index c9f0fc462c9..65c07e72154 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.cc +++ b/tensorflow/stream_executor/cuda/cuda_blas.cc @@ -101,18 +101,6 @@ static std::string ToString(cublasStatus_t status) { } } -// Decide whether to enable TENSOR_OP_MATH -static bool TensorOpMathEnabled() { - static bool is_enabled = [] { - bool is_disabled; - TF_CHECK_OK( - tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUBLAS_TENSOR_OP_MATH", - /*default_val=*/false, &is_disabled)); - return !is_disabled; - }(); - return is_enabled; -} - // cuBLAS has interfaces that permit pointers to be passed from either the host // memory space or the device memory space; however, you must instruct it as to // which address space those pointers are in with cublasSetPointerMode. @@ -1640,7 +1628,7 @@ bool CUDABlas::DoBlasGemm( &cc_minor); // GPUs < sm_70 don't support tensor ops. - if (cc_major >= 7 && TensorOpMathEnabled()) { + if (cc_major >= 7) { use_tensor_ops = true; } #endif @@ -1921,8 +1909,7 @@ static bool TensorOpsAvailable(int cc_major) { // strictly correct. We can't simply enable it, though, as that would change // clients' behavior significantly: Using tensor ops on fp32 inputs cause them // to be rounded to fp16. - if (cc_major >= 7 && TensorOpMathEnabled() && - std::is_same::value) { + if (cc_major >= 7 && std::is_same::value) { return true; } #endif @@ -2270,7 +2257,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal( if (stream->parent()->GetDeviceDescription().cuda_compute_capability( &cc_major, &cc_minor) && cc_major >= 5) { - bool use_tensor_ops = TensorOpMathEnabled() && data_type == CUDA_R_16F; + bool use_tensor_ops = data_type == CUDA_R_16F; cublasGemmAlgo_t algo = (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT); cudaDataType_t compute_type = @@ -2425,7 +2412,7 @@ bool CUDABlas::DoBlasGemmStridedBatched( if (stream->parent()->GetDeviceDescription().cuda_compute_capability( &cc_major, &cc_minor)) { // GPUs < sm_70 don't support tensor ops. - if (cc_major >= 7 && TensorOpMathEnabled()) { + if (cc_major >= 7) { use_tensor_ops = true; } #if CUDA_VERSION >= 9010 diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index 6122877f91f..780f1475c2c 100755 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -603,31 +603,6 @@ class CudnnFilterDescriptor { SE_DISALLOW_COPY_AND_ASSIGN(CudnnFilterDescriptor); }; -// A helper function to decide whether to enable the TENSOR_OP_MATH math type -bool TensorOpMathEnabled() { - static bool is_enabled = [] { - bool is_disabled = false; - TF_CHECK_OK( - tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_TENSOR_OP_MATH", - /*default_val=*/false, &is_disabled)); - return !is_disabled; - }(); - return is_enabled; -} - -// A helper function to decide whether to enable the TENSOR_OP_MATH math type -// for RNNs. -bool RnnTensorOpMathEnabled() { - static bool is_enabled = [] { - bool is_disabled = false; - TF_CHECK_OK( - tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH", - /*default_val=*/false, &is_disabled)); - return !is_disabled; - }(); - return is_enabled; -} - // A helper function to decide whether to use // CUDNN_BATCHNORM_SPATIAL_PERSISTENT in batchnorm. This mode can be faster in // some tasks because an optimized path may be selected for CUDNN_DATA_FLOAT @@ -751,9 +726,7 @@ class CudnnConvolutionDescriptor { #if CUDNN_VERSION >= 7000 cudnnMathType_t math_type = (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH); - if (TensorOpMathEnabled()) { - CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type)); - } + CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type)); #endif } @@ -1157,21 +1130,19 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor { // in profile mode, which is run with algorithms returned from // GetRnnAlgorithms() (which are non-default and explicitly set whether to // use tensor ops). CuDNN 7.2.1 fixed this issue - if (RnnTensorOpMathEnabled()) { - cudnnMathType_t math_type; - if (algorithm_config.algorithm().has_value()) { - math_type = algorithm_config.algorithm()->tensor_ops_enabled() - ? CUDNN_TENSOR_OP_MATH - : CUDNN_DEFAULT_MATH; - } else { + cudnnMathType_t math_type; + if (algorithm_config.algorithm().has_value()) { + math_type = algorithm_config.algorithm()->tensor_ops_enabled() + ? CUDNN_TENSOR_OP_MATH + : CUDNN_DEFAULT_MATH; + } else { #if CUDNN_VERSION >= 7201 - math_type = CUDNN_TENSOR_OP_MATH; + math_type = CUDNN_TENSOR_OP_MATH; #else - math_type = CUDNN_DEFAULT_MATH; + math_type = CUDNN_DEFAULT_MATH; #endif // CUDNN_VERSION >= 7201 - } - CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type)); } + CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type)); #endif // CUDNN_VERSION >= 7000 return CudnnRnnDescriptor(cudnn, std::move(rnn_desc), std::move(rnn_plan), @@ -2605,7 +2576,7 @@ AllocateCudnnConvolutionBackwardFilterWorkspace( } static bool TensorOpMathAvailable(int cc_major) { - return cc_major >= 7 && CUDNN_VERSION >= 7000 && TensorOpMathEnabled(); + return cc_major >= 7 && CUDNN_VERSION >= 7000; } port::StatusOr GetCudnnConvolutionForwardAlgorithm( @@ -3399,9 +3370,7 @@ bool CudnnSupport::GetRnnAlgorithms( for (auto i : algo_types) { out_algorithms->push_back({i, /*use_tensor_ops=*/false}); #if CUDNN_VERSION >= 7100 - if (RnnTensorOpMathEnabled()) { - out_algorithms->push_back({i, /*use_tensor_ops=*/true}); - } + out_algorithms->push_back({i, /*use_tensor_ops=*/true}); #endif } return true; From 3cad62e356b8c72d03af13ba29f7ace29a6f0772 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Fri, 15 May 2020 11:46:41 -0500 Subject: [PATCH 0073/1390] Add global setting control TF32 execution --- tensorflow/core/platform/BUILD | 7 +++++++ tensorflow/core/platform/tf32_utils.cc | 27 ++++++++++++++++++++++++++ tensorflow/core/platform/tf32_utils.h | 27 ++++++++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tensorflow/core/platform/tf32_utils.cc create mode 100644 tensorflow/core/platform/tf32_utils.h diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD index c7ff378d2ac..f27d2f09208 100644 --- a/tensorflow/core/platform/BUILD +++ b/tensorflow/core/platform/BUILD @@ -937,6 +937,13 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "tf32_utils", + srcs = ["tf32_utils.cc"], + hdrs = ["tf32_utils.h"], + copts = tf_copts(), +) + tf_cc_tests( name = "low_level_library_tests", size = "small", diff --git a/tensorflow/core/platform/tf32_utils.cc b/tensorflow/core/platform/tf32_utils.cc new file mode 100644 index 00000000000..715b5996dc3 --- /dev/null +++ b/tensorflow/core/platform/tf32_utils.cc @@ -0,0 +1,27 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/core/platform/tf32_utils.h" + +namespace tensorflow { + +// TODO(nluehr): enable tf32 execution by default after TF32 Ampere testing. +static bool tf32_enabled = false; + +void allow_tf32_execution(bool allow) { tf32_enabled = allow; } + +bool tf32_execution_allowed() { return tf32_enabled; } + +} // namespace tensorflow diff --git a/tensorflow/core/platform/tf32_utils.h b/tensorflow/core/platform/tf32_utils.h new file mode 100644 index 00000000000..a0ce58f9bbd --- /dev/null +++ b/tensorflow/core/platform/tf32_utils.h @@ -0,0 +1,27 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_ +#define TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_ + +namespace tensorflow { + +void allow_tf32_execution(bool allow); + +bool tf32_execution_allowed(); + +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_ From 8bfee17f5880eccdb759fb47ab11b782f201cf0f Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Fri, 15 May 2020 13:33:02 -0500 Subject: [PATCH 0074/1390] Python tf.config tf32 interface --- tensorflow/python/BUILD | 11 +++++++++++ tensorflow/python/framework/config.py | 26 ++++++++++++++++++++++++++ tensorflow/python/util/tf32.cc | 22 ++++++++++++++++++++++ 3 files changed, 59 insertions(+) create mode 100644 tensorflow/python/util/tf32.cc diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index a49e4b74def..997ec6c924f 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -746,6 +746,16 @@ tf_python_pybind_extension( ], ) +tf_python_pybind_extension( + name = "_pywrap_tf32_execution", + srcs = ["util/tf32.cc"], + module_name = "_pywrap_tf32_execution", + deps = [ + "//tensorflow/core/platform:tf32_utils", + "@pybind11", + ], +) + tf_python_pybind_extension( name = "_pywrap_util_port", srcs = ["util/port_wrapper.cc"], @@ -5573,6 +5583,7 @@ py_library( "//tensorflow:composite_tensor_whitelist", ], deps = [ + ":_pywrap_tf32_execution", ":tf_decorator", ":tf_export", ":tf_stack", diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py index 5361d7290e8..042af4d1023 100644 --- a/tensorflow/python/framework/config.py +++ b/tensorflow/python/framework/config.py @@ -18,10 +18,36 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python import _pywrap_tf32_execution from tensorflow.python.eager import context from tensorflow.python.util import deprecation from tensorflow.python.util.tf_export import tf_export +def tensor_float32_execution_allowed(): + """Get if TensorFloat-32 operations are enabled on supported hardware. + + Returns: + True if TensorFloat-32 execution is enabled and False otherwise. + """ + return _pywrap_tf32_execution.is_allowed() + +def allow_tensor_float_32_execution(allow): + """Allow use of TensorFloat-32 with float32 ops on supported hardware. + + TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture. + TensorFloat-32 kernels take float32 inputs and produce float32 outputs. + Internally, the inputs are cast to a custom representation with 10-bit + mantissa (similar to float16) and 8-bit exponent (similar to float32) and are + executed using TensorCores with float32 accumulation. For more information, + see https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/. + + TensorFloat-32 execution is disabled by default, but this may change in a + future version. + + Args: + allow: whether to allow TensorFloat-32 execution + """ + _pywrap_tf32_execution.allow(allow) @tf_export('config.threading.get_intra_op_parallelism_threads') def get_intra_op_parallelism_threads(): diff --git a/tensorflow/python/util/tf32.cc b/tensorflow/python/util/tf32.cc new file mode 100644 index 00000000000..7dece6ccdae --- /dev/null +++ b/tensorflow/python/util/tf32.cc @@ -0,0 +1,22 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "pybind11/pybind11.h" +#include "tensorflow/core/platform/tf32_utils.h" + +PYBIND11_MODULE(_pywrap_tf32_execution, m) { + m.def("allow", &tensorflow::allow_tf32_execution); + m.def("is_allowed", &tensorflow::tf32_execution_allowed); +} From dedb51aec2a766bdeb8b4c2ab1700bfcf7687966 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Tue, 19 May 2020 14:58:30 -0500 Subject: [PATCH 0075/1390] Convolution TF32 Plumbing --- tensorflow/stream_executor/cuda/BUILD | 1 + tensorflow/stream_executor/cuda/cuda_dnn.cc | 200 +++++++++++++------- 2 files changed, 135 insertions(+), 66 deletions(-) mode change 100755 => 100644 tensorflow/stream_executor/cuda/cuda_dnn.cc diff --git a/tensorflow/stream_executor/cuda/BUILD b/tensorflow/stream_executor/cuda/BUILD index 1457a36beaf..2749281335e 100644 --- a/tensorflow/stream_executor/cuda/BUILD +++ b/tensorflow/stream_executor/cuda/BUILD @@ -353,6 +353,7 @@ cc_library( "@local_config_cuda//cuda:cudnn_header", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core/platform:tf32_utils", "//tensorflow/stream_executor:dnn", "//tensorflow/stream_executor:event", "//tensorflow/stream_executor:plugin_registry", diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc old mode 100755 new mode 100644 index 780f1475c2c..53296f4eea5 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -20,8 +20,8 @@ limitations under the License. #include #include "absl/strings/str_cat.h" -#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/tf32_utils.h" #include "tensorflow/core/util/env_var.h" #include "tensorflow/stream_executor/cuda/cuda_activation.h" #include "tensorflow/stream_executor/cuda/cuda_diagnostics.h" @@ -42,6 +42,7 @@ limitations under the License. #include "tensorflow/stream_executor/scratch_allocator.h" #include "tensorflow/stream_executor/stream.h" #include "tensorflow/stream_executor/stream_executor_pimpl.h" +#include "third_party/eigen3/Eigen/Core" // clang-format off #include "third_party/gpus/cudnn/cudnn.h" #include "absl/strings/string_view.h" @@ -707,10 +708,6 @@ class CudnnConvolutionDescriptor { : CUDNN_CROSS_CORRELATION, data_type)); - // NOTE(benbarsdell): This only applies if tensor op math is enabled - // and algo selection is set to Default. - this->set_use_tensor_op_math(true); - #if CUDNN_MAJOR >= 7 VLOG(2) << "Requesting grouped convolution: " << convolution_descriptor.group_count(); @@ -722,10 +719,14 @@ class CudnnConvolutionDescriptor { #endif } - void set_use_tensor_op_math(bool use_tensor_op_math) const { + void set_use_tensor_op_math(bool use_tensor_op_math) { #if CUDNN_VERSION >= 7000 cudnnMathType_t math_type = +#if CUDNN_VERSION >= 8000 + (use_tensor_op_math ? CUDNN_DEFAULT_MATH : CUDNN_FMA_MATH); +#else (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH); +#endif CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type)); #endif } @@ -738,6 +739,38 @@ class CudnnConvolutionDescriptor { SE_DISALLOW_COPY_AND_ASSIGN(CudnnConvolutionDescriptor); }; +// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math +// set +static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) { + cudnnMathType_t math_type; + CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type)); +#if CUDNN_VERSION >= 8000 + return math_type != CUDNN_FMA_MATH; +#else + return math_type == CUDNN_TENSOR_OP_MATH; +#endif +} + +static bool TensorOpMathAvailable(int cc_major) { + return cc_major >= 7 && CUDNN_VERSION >= 7000; +} + +static bool IsTensorMathAllowed(Stream* stream, dnn::DataType input_type) { + int cc_major, cc_minor; + std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream); + if (!TensorOpMathAvailable(cc_major)) { + return false; + } + if (input_type == dnn::DataType::kFloat) { + if (CUDNN_VERSION < 8000) { + return false; + } else if (!tensorflow::tf32_execution_allowed()) { + return false; + } + } + return true; +} + // Turns a PoolingDescriptor structure into a cudnn pooling descriptor handle // within a scope. class CudnnPoolingDescriptor { @@ -2450,10 +2483,11 @@ port::StatusOr> AllocateCudnnConvolutionForwardWorkspace( const CudnnTensorDescriptor& output_nd, const dnn::AlgorithmDesc& algorithm_desc, ScratchAllocator* scratch_allocator) { - // TODO(csigg): This has side effects on the convolution descriptor. It is - // functionally correct because the convolution is run with the algorithm of - // the last call to this function, but should be fixed anyway. - conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled()); + if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) { + return port::Status( + port::error::INTERNAL, + "Mismatch between cudnn conv and algorithm descriptors."); + } // Query the size of the workspace and allocate it. size_t size_in_bytes; @@ -2493,10 +2527,11 @@ AllocateCudnnConvolutionBackwardDataWorkspace( const CudnnTensorDescriptor& output_nd, const dnn::AlgorithmDesc& algorithm_desc, ScratchAllocator* scratch_allocator) { - // TODO(csigg): This has side effects on the convolution descriptor. It is - // functionally correct because the convolution is run with the algorithm of - // the last call to this function, but should be fixed anyway. - conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled()); + if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) { + return port::Status( + port::error::INTERNAL, + "Mismatch between cudnn conv and algorithm descriptors."); + } // Query the size of the workspace and allocate it. size_t size_in_bytes; @@ -2538,10 +2573,11 @@ AllocateCudnnConvolutionBackwardFilterWorkspace( const CudnnTensorDescriptor& output_nd, const dnn::AlgorithmDesc& algorithm_desc, ScratchAllocator* scratch_allocator) { - // TODO(csigg): This has side effects on the convolution descriptor. It is - // functionally correct because the convolution is run with the algorithm of - // the last call to this function, but should be fixed anyway. - conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled()); + if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) { + return port::Status( + port::error::INTERNAL, + "Mismatch between cudnn conv and algorithm descriptors."); + } // Query the size of the workspace and allocate it. size_t size_in_bytes; @@ -2575,18 +2611,39 @@ AllocateCudnnConvolutionBackwardFilterWorkspace( return scratch_allocator->AllocateBytes(size_in_bytes); } -static bool TensorOpMathAvailable(int cc_major) { - return cc_major >= 7 && CUDNN_VERSION >= 7000; +port::StatusOr UseTensorOps(Stream* stream, dnn::DataType type, + absl::optional desc) { + bool use_tensor_ops; + if (desc.has_value()) { + use_tensor_ops = desc->tensor_ops_enabled(); + if (use_tensor_ops && !IsTensorMathAllowed(stream, type)) { + return port::Status(port::error::INVALID_ARGUMENT, + "Algo requests disallowed tensor op evaluation."); + } + } else { + use_tensor_ops = IsTensorMathAllowed(stream, type); + } + return use_tensor_ops; } port::StatusOr GetCudnnConvolutionForwardAlgorithm( Stream* stream, const CudnnHandle& cudnn, const dnn::AlgorithmConfig& algorithm_config, const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter, - const CudnnConvolutionDescriptor& conv, + dnn::DataType element_type, + const dnn::ConvolutionDescriptor& convolution_descriptor, const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, DeviceMemory* scratch) { absl::optional algo_desc = algorithm_config.algorithm(); + + CudnnConvolutionDescriptor conv( + convolution_descriptor, + ToCudnnDataType(GetConvAccumulatorType(element_type))); + bool use_tensor_ops; + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); + if (!algo_desc.has_value()) { // Pick fastest algorithm within memory limit according to cuDNN's // heuristics. @@ -2599,10 +2656,7 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm( GetCudnnConvolutionForwardAlgo( cudnn, input_nd, filter, conv, output_nd, specify_workspace_limit, memory_limit_bytes)); - int cc_major, cc_minor; - std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream); - algo_desc = dnn::AlgorithmDesc( - algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major)); + algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops); } const auto scratch_or = AllocateCudnnConvolutionForwardWorkspace( @@ -2626,6 +2680,9 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm( "Returned status: ", scratch_or.status().ToString())); } + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionForwardWorkspace( stream, cudnn, input_nd, filter, conv, output_nd, *algo_desc, scratch_allocator)); @@ -2636,10 +2693,19 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm( Stream* stream, const CudnnHandle& cudnn, const dnn::AlgorithmConfig& algorithm_config, const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter, - const CudnnConvolutionDescriptor& conv, + dnn::DataType element_type, + const dnn::ConvolutionDescriptor& convolution_descriptor, const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, DeviceMemory* scratch) { absl::optional algo_desc = algorithm_config.algorithm(); + CudnnConvolutionDescriptor conv( + convolution_descriptor, + ToCudnnDataType(GetConvAccumulatorType(element_type))); + bool use_tensor_ops; + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); + if (!algo_desc.has_value()) { // Pick fastest algorithm within memory limit according to cuDNN's // heuristics. @@ -2652,10 +2718,7 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm( GetCudnnConvolutionBackwardDataAlgo( cudnn, input_nd, filter, conv, output_nd, specify_workspace_limit, memory_limit_bytes)); - int cc_major, cc_minor; - std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream); - algo_desc = dnn::AlgorithmDesc( - algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major)); + algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops); } const auto scratch_or = AllocateCudnnConvolutionBackwardDataWorkspace( @@ -2678,6 +2741,9 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm( "while a secondary algorithm is not provided."); } + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardDataWorkspace( stream, cudnn, input_nd, filter, conv, output_nd, *algo_desc, scratch_allocator)); @@ -2688,10 +2754,19 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm( Stream* stream, const CudnnHandle& cudnn, const dnn::AlgorithmConfig& algorithm_config, const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter, - const CudnnConvolutionDescriptor& conv, + dnn::DataType element_type, + const dnn::ConvolutionDescriptor& convolution_descriptor, const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator, DeviceMemory* scratch) { absl::optional algo_desc = algorithm_config.algorithm(); + CudnnConvolutionDescriptor conv( + convolution_descriptor, + ToCudnnDataType(GetConvAccumulatorType(element_type))); + bool use_tensor_ops; + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); + if (!algo_desc.has_value()) { // Pick fastest algorithm within memory limit according to cuDNN's // heuristics. @@ -2704,10 +2779,7 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm( GetCudnnConvolutionBackwardFilterAlgo( cudnn, input_nd, filter, conv, output_nd, specify_workspace_limit, memory_limit_bytes)); - int cc_major, cc_minor; - std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream); - algo_desc = dnn::AlgorithmDesc( - algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major)); + algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops); } auto scratch_or = AllocateCudnnConvolutionBackwardFilterWorkspace( @@ -2730,6 +2802,9 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm( "while a secondary algorithm is not provided."); } + SE_ASSIGN_OR_RETURN(use_tensor_ops, + UseTensorOps(stream, element_type, algo_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardFilterWorkspace( stream, cudnn, input_nd, filter, conv, output_nd, *algo_desc, scratch_allocator)); @@ -2894,35 +2969,32 @@ port::Status CudnnSupport::DoPrepareForConvolution( CudnnTensorDescriptor output_nd( output_descriptor, ToCudnnDataType(element_type, output_descriptor.layout())); - CudnnConvolutionDescriptor conv( - convolution_descriptor, - ToCudnnDataType(GetConvAccumulatorType(element_type))); auto cudnn = cudnn_->GetHandle(parent_, stream); switch (kind) { case dnn::ConvolutionKind::FORWARD: { - SE_ASSIGN_OR_RETURN( - *algorithm_desc, - GetCudnnConvolutionForwardAlgorithm( - stream, cudnn, algorithm_config, input_nd, filter_nd, conv, - output_nd, scratch_allocator, scratch_memory)); + SE_ASSIGN_OR_RETURN(*algorithm_desc, + GetCudnnConvolutionForwardAlgorithm( + stream, cudnn, algorithm_config, input_nd, + filter_nd, element_type, convolution_descriptor, + output_nd, scratch_allocator, scratch_memory)); break; } case dnn::ConvolutionKind::BACKWARD_DATA: { - SE_ASSIGN_OR_RETURN( - *algorithm_desc, - GetCudnnConvolutionBackwardDataAlgorithm( - stream, cudnn, algorithm_config, input_nd, filter_nd, conv, - output_nd, scratch_allocator, scratch_memory)); + SE_ASSIGN_OR_RETURN(*algorithm_desc, + GetCudnnConvolutionBackwardDataAlgorithm( + stream, cudnn, algorithm_config, input_nd, + filter_nd, element_type, convolution_descriptor, + output_nd, scratch_allocator, scratch_memory)); break; } case dnn::ConvolutionKind::BACKWARD_FILTER: { - SE_ASSIGN_OR_RETURN( - *algorithm_desc, - GetCudnnConvolutionBackwardFilterAlgorithm( - stream, cudnn, algorithm_config, input_nd, filter_nd, conv, - output_nd, scratch_allocator, scratch_memory)); + SE_ASSIGN_OR_RETURN(*algorithm_desc, + GetCudnnConvolutionBackwardFilterAlgorithm( + stream, cudnn, algorithm_config, input_nd, + filter_nd, element_type, convolution_descriptor, + output_nd, scratch_allocator, scratch_memory)); break; } default: @@ -2951,8 +3023,9 @@ port::Status CudnnSupport::DoConvolve( auto accumulator_type = GetConvAccumulatorType(element_type); CudnnConvolutionDescriptor conv(convolution_descriptor, ToCudnnDataType(accumulator_type)); - // Set use_tensor_math param to correct value - conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled()); + SE_ASSIGN_OR_RETURN(bool use_tensor_ops, + UseTensorOps(stream, element_type, algorithm_desc)); + conv.set_use_tensor_op_math(use_tensor_ops); auto cudnn = cudnn_->GetHandle(parent_, stream); // Alpha is the scaling factor for input. @@ -3185,14 +3258,6 @@ port::Status CudnnSupport::DoConvolve( return port::Status::OK(); } -// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math -// set -static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) { - cudnnMathType_t math_type; - CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type)); - return math_type == CUDNN_TENSOR_OP_MATH; -} - template port::Status CudnnSupport::DoFusedConvolveImpl( @@ -3226,8 +3291,6 @@ port::Status CudnnSupport::DoFusedConvolveImpl( filter_descriptor, GetCudnnDataType(conv_input_descriptor.layout())); CudnnTensorDescriptor bias_nd(bias_descriptor, GetCudnnDataType()); - CudnnConvolutionDescriptor conv(convolution_descriptor, - ToCudnnDataType(accumulator_type)); auto cudnn = cudnn_->GetHandle(parent_, stream); @@ -3237,9 +3300,14 @@ port::Status CudnnSupport::DoFusedConvolveImpl( SE_ASSIGN_OR_RETURN( dnn::AlgorithmDesc algo_desc, GetCudnnConvolutionForwardAlgorithm( - stream, cudnn, algorithm_config, conv_input_nd, filter, conv, + stream, cudnn, algorithm_config, conv_input_nd, filter, + dnn::ToDataType::value, convolution_descriptor, output_nd, scratch_allocator, &scratch)); + CudnnConvolutionDescriptor conv(convolution_descriptor, + ToCudnnDataType(accumulator_type)); + conv.set_use_tensor_op_math(algo_desc.tensor_ops_enabled()); + std::unique_ptr timer; if (is_profiling) { timer.reset(new GpuTimer(parent_)); // NOLINT From 0f58bb63090222cef0eebe74630b2d4d9d886a2f Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Tue, 19 May 2020 15:54:10 -0500 Subject: [PATCH 0076/1390] Plumb TF32 for RNN --- tensorflow/stream_executor/cuda/cuda_dnn.cc | 30 ++++++++++++++------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index 53296f4eea5..fa06d410323 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -1163,17 +1163,26 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor { // in profile mode, which is run with algorithms returned from // GetRnnAlgorithms() (which are non-default and explicitly set whether to // use tensor ops). CuDNN 7.2.1 fixed this issue - cudnnMathType_t math_type; + bool allow_tensor_ops = + data_type != CUDNN_DATA_FLOAT || tensorflow::tf32_execution_allowed(); + bool use_tensor_ops; if (algorithm_config.algorithm().has_value()) { - math_type = algorithm_config.algorithm()->tensor_ops_enabled() - ? CUDNN_TENSOR_OP_MATH - : CUDNN_DEFAULT_MATH; + use_tensor_ops = algorithm_config.algorithm()->tensor_ops_enabled(); } else { -#if CUDNN_VERSION >= 7201 - math_type = CUDNN_TENSOR_OP_MATH; -#else - math_type = CUDNN_DEFAULT_MATH; -#endif // CUDNN_VERSION >= 7201 + use_tensor_ops = CUDNN_VERSION >= 7201 && allow_tensor_ops; + } + + if (use_tensor_ops && !allow_tensor_ops) { + return port::Status(port::error::INVALID_ARGUMENT, + "Algo requests disallowed tensor op evaluation."); + } + + cudnnMathType_t math_type; + if (use_tensor_ops) { + math_type = + CUDNN_VERSION >= 8000 ? CUDNN_DEFAULT_MATH : CUDNN_TENSOR_OP_MATH; + } else { + math_type = CUDNN_VERSION >= 8000 ? CUDNN_FMA_MATH : CUDNN_DEFAULT_MATH; } CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type)); #endif // CUDNN_VERSION >= 7000 @@ -2626,6 +2635,9 @@ port::StatusOr UseTensorOps(Stream* stream, dnn::DataType type, return use_tensor_ops; } +cudnnDataType_t GetRnnComputeType(dnn::DataType data_type); +dnn::DataType GetConvAccumulatorType(dnn::DataType data_type); + port::StatusOr GetCudnnConvolutionForwardAlgorithm( Stream* stream, const CudnnHandle& cudnn, const dnn::AlgorithmConfig& algorithm_config, From b67608e66c54224fa52200095fba09df0f2b3c71 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Wed, 20 May 2020 10:06:35 -0500 Subject: [PATCH 0077/1390] Plumb TF32 for cublas gemm --- tensorflow/stream_executor/cuda/BUILD | 1 + tensorflow/stream_executor/cuda/cuda_blas.cc | 84 +++++++++----------- tensorflow/stream_executor/cuda/cuda_blas.h | 8 +- 3 files changed, 43 insertions(+), 50 deletions(-) diff --git a/tensorflow/stream_executor/cuda/BUILD b/tensorflow/stream_executor/cuda/BUILD index 2749281335e..519033a62d8 100644 --- a/tensorflow/stream_executor/cuda/BUILD +++ b/tensorflow/stream_executor/cuda/BUILD @@ -251,6 +251,7 @@ cc_library( "@local_config_cuda//cuda:cuda_headers", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core/platform:tf32_utils", "//tensorflow/stream_executor", "//tensorflow/stream_executor:event", "//tensorflow/stream_executor:host_or_device_scalar", diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc index 65c07e72154..e2cbb0b75df 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.cc +++ b/tensorflow/stream_executor/cuda/cuda_blas.cc @@ -48,7 +48,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" -#include "third_party/eigen3/Eigen/Core" +#include "tensorflow/core/platform/tf32_utils.h" #include "tensorflow/core/util/env_var.h" #include "tensorflow/stream_executor/cuda/cuda_activation.h" #include "tensorflow/stream_executor/cuda/cuda_gpu_executor.h" @@ -66,6 +66,7 @@ limitations under the License. #include "tensorflow/stream_executor/plugin_registry.h" #include "tensorflow/stream_executor/scratch_allocator.h" #include "tensorflow/stream_executor/stream_executor.h" +#include "third_party/eigen3/Eigen/Core" namespace stream_executor { namespace gpu { @@ -225,6 +226,18 @@ bool CUDABlas::Init() { return false; } +#if CUDA_VERSION >= 9000 +#if CUBLAS_VER_MAJOR >= 11 + ret = cublasSetMathMode(blas_, CUBLAS_TF32_TENSOR_OP_MATH); +#else + ret = cublasSetMathMode(blas_, CUBLAS_TENSOR_OP_MATH); +#endif + if (ret != CUBLAS_STATUS_SUCCESS) { + LOG(ERROR) << "failed to set cublas default math mode: " << ToString(ret); + return false; + } +#endif + return true; } @@ -387,7 +400,7 @@ cudaDataType_t CUDAComputationType(blas::ComputationType ty) { template bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream, bool pointer_mode_host, bool err_on_failure, - bool use_tensor_op_math, Args... args) { + Args... args) { absl::MutexLock lock(&mu_); CHECK(blas_ != nullptr); @@ -401,10 +414,10 @@ bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream, : CUBLAS_POINTER_MODE_DEVICE)) { return false; } -#if CUDA_VERSION >= 9000 +#if CUBLAS_VER_MAJOR >= 11 ScopedCublasMathMode math_mode{blas_}; - if (use_tensor_op_math) { - if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) { + if (!tensorflow::tf32_execution_allowed()) { + if (!math_mode.Init(CUBLAS_DEFAULT_MATH)) { return false; } } @@ -1621,21 +1634,9 @@ bool CUDABlas::DoBlasGemm( } } - bool use_tensor_ops = false; -#if CUDA_VERSION >= 9000 - int cc_major, cc_minor; - stream->parent()->GetDeviceDescription().cuda_compute_capability(&cc_major, - &cc_minor); - - // GPUs < sm_70 don't support tensor ops. - if (cc_major >= 7) { - use_tensor_ops = true; - } -#endif - return DoBlasInternalImpl( cublasSgemmEx, stream, true /* = pointer_mode_host */, - true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa), + true /* = err_on_failure= */, CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha, GpuMemory(a), SE_CUDA_DATA_HALF, lda, GpuMemory(b), SE_CUDA_DATA_HALF, ldb, &beta, GpuMemoryMutable(c), SE_CUDA_DATA_HALF, ldc); @@ -2257,7 +2258,8 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal( if (stream->parent()->GetDeviceDescription().cuda_compute_capability( &cc_major, &cc_minor) && cc_major >= 5) { - bool use_tensor_ops = data_type == CUDA_R_16F; + bool use_tensor_ops = + data_type == CUDA_R_16F || tensorflow::tf32_execution_allowed(); cublasGemmAlgo_t algo = (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT); cudaDataType_t compute_type = @@ -2271,7 +2273,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal( bool ok; ok = DoBlasInternalImpl( AS_LAMBDA(cublasGemmBatchedEx), stream, true /* = pointer_mode_host */, - true /* = err_on_failure */, use_tensor_ops, CUDABlasTranspose(transa), + true /* = err_on_failure */, CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha, a_void_ptrs, data_type, lda, b_void_ptrs, data_type, ldb, &beta, c_void_ptrs, data_type, ldc, batch_count, compute_type, algo); @@ -2406,33 +2408,25 @@ bool CUDABlas::DoBlasGemmStridedBatched( int lda, int64 stride_a, const DeviceMemory &b, int ldb, int64 stride_b, float beta, DeviceMemory *c, int ldc, int64 stride_c, int batch_count) { - bool use_tensor_ops = false; -#if CUDA_VERSION >= 9000 +#if CUDA_VERSION >= 9010 int cc_major, cc_minor; if (stream->parent()->GetDeviceDescription().cuda_compute_capability( - &cc_major, &cc_minor)) { - // GPUs < sm_70 don't support tensor ops. - if (cc_major >= 7) { - use_tensor_ops = true; + &cc_major, &cc_minor) && + cc_major >= 5) { + cublasGemmAlgo_t algo = + (cc_major >= 7 ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT); + bool ok = DoBlasInternalImpl( + AS_LAMBDA(cublasGemmStridedBatchedEx), stream, + true /* = pointer_mode_host */, true /* = err_on_failure */, + CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha, + GpuMemory(a), CUDA_R_16F, lda, stride_a, GpuMemory(b), CUDA_R_16F, ldb, + stride_b, &beta, GpuMemoryMutable(c), CUDA_R_16F, ldc, stride_c, + batch_count, CUDA_R_32F, algo); + if (ok) { + return true; } -#if CUDA_VERSION >= 9010 - if (cc_major >= 5) { - cublasGemmAlgo_t algo = - (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT); - bool ok = DoBlasInternalImpl( - AS_LAMBDA(cublasGemmStridedBatchedEx), stream, - true /* = pointer_mode_host */, true /* = err_on_failure */, - use_tensor_ops, CUDABlasTranspose(transa), CUDABlasTranspose(transb), - m, n, k, &alpha, GpuMemory(a), CUDA_R_16F, lda, stride_a, - GpuMemory(b), CUDA_R_16F, ldb, stride_b, &beta, GpuMemoryMutable(c), - CUDA_R_16F, ldc, stride_c, batch_count, CUDA_R_32F, algo); - if (ok) { - return true; - } - LOG(ERROR) << "failed BLAS call, see log for details"; - return false; - } -#endif + LOG(ERROR) << "failed BLAS call, see log for details"; + return false; } #endif // Either CUDA_VERSION < 9.1 or SM < 5.0. Fall back to a loop. @@ -2445,7 +2439,7 @@ bool CUDABlas::DoBlasGemmStridedBatched( reinterpret_cast<__half *>(GpuMemoryMutable(c) + batch * stride_c); bool ok = DoBlasInternalImpl( cublasSgemmEx, stream, true /* = pointer_mode_host */, - true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa), + true /* = err_on_failure= */, CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha, a_matrix, SE_CUDA_DATA_HALF, lda, b_matrix, SE_CUDA_DATA_HALF, ldb, &beta, c_matrix, SE_CUDA_DATA_HALF, ldc); diff --git a/tensorflow/stream_executor/cuda/cuda_blas.h b/tensorflow/stream_executor/cuda/cuda_blas.h index 817bdb72777..556456c83db 100644 --- a/tensorflow/stream_executor/cuda/cuda_blas.h +++ b/tensorflow/stream_executor/cuda/cuda_blas.h @@ -83,7 +83,7 @@ class CUDABlas : public blas::BlasSupport { template bool DoBlasInternalImpl(FuncT cublas_func, Stream *stream, bool pointer_mode_host, bool err_on_failure, - bool use_tensor_op_math, Args... args); + Args... args); // Convenience functions that call DoBlasInternalImpl with different values // for err_on_failure. @@ -91,8 +91,7 @@ class CUDABlas : public blas::BlasSupport { bool DoBlasInternal(FuncT cublas_func, Stream *stream, bool pointer_mode_host, Args... args) { return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host, - /*err_on_failure=*/true, /*use_tensor_ops=*/false, - args...); + /*err_on_failure=*/true, args...); } template bool DoBlasInternalFailureOK(FuncT cublas_func, Stream *stream, @@ -100,8 +99,7 @@ class CUDABlas : public blas::BlasSupport { // Tensor ops are hard-coded off in this path, but can still be enabled with // a specific algorithm choice as in DoBlasGemmWithAlgorithmImpl(). return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host, - /*err_on_failure=*/false, - /*use_tensor_ops=*/false, args...); + /*err_on_failure=*/false, args...); } // A helper function to implement DoBlasGemmBatched interfaces for generic From b1e74b227c681588a62768042816702a9518f642 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Thu, 4 Jun 2020 09:22:37 +0200 Subject: [PATCH 0078/1390] Fix compile errors. --- tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc | 2 +- tensorflow/lite/micro/micro_allocator.cc | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc index 8f21a167f67..39991ab758b 100644 --- a/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc +++ b/tensorflow/lite/micro/memory_planner/greedy_memory_planner.cc @@ -191,7 +191,7 @@ void GreedyMemoryPlanner::CalculateOffsetsIfNeeded() { // Work through the rest of the buffers to find a good gap to place each one. for (int i = 1; i < buffer_count_; ++i) { // The id is the order the buffer was originally added by the client. - const int buffer_id = buffer_ids_sorted_[i]; + buffer_id = buffer_ids_sorted_[i]; // Look at what size and time range the buffer needs to be active. BufferRequirements* wanted_requirements = &requirements_[buffer_id]; const int wanted_size = wanted_requirements->size; diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index c204f4460b4..8fac421750d 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -99,6 +99,9 @@ TfLiteStatus AllocateVariables( // not called by default. Hence it's not linked in to the final binary code. TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, ErrorReporter* error_reporter) { + // Suppress compile warning for unused function + (void)CheckOfflinePlannedOffsets; + if (model->metadata()) { for (size_t i = 0; i < model->metadata()->size(); ++i) { auto metadata = model->metadata()->Get(i); From 521b7595b7adba5627a2687befc3fb41bc5c2bec Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Thu, 4 Jun 2020 17:19:31 +0800 Subject: [PATCH 0079/1390] TFLM: add HIMAX WE1 EVB to support TFLM example(hello word and person detection INT8) --- .../lite/micro/examples/hello_world/README.md | 132 +++- .../micro/examples/hello_world/README.md~ | 595 ++++++++++++++++++ .../hello_world/himax_we1_evb/constants.cc | 19 + .../himax_we1_evb/output_handler.cc | 35 ++ .../himax_we1_evb/output_handler.cc~ | 53 ++ .../person_detection_experimental/README.md | 129 +++- .../himax_we1_evb/detection_responder.cc | 34 + .../himax_we1_evb/image_provider.cc | 44 ++ .../himax_we1_evb/image_provider.cc~ | 44 ++ .../himax_we1_evb/main_functions.cc | 127 ++++ .../{we_i => himax_we1_evb}/debug_log.cc | 9 +- .../make/targets/himax_we1_evb_makefile.inc | 91 +++ .../make/targets/himax_we1_evb_makefile.inc~ | 93 +++ .../tools/make/third_party_downloads.inc | 4 + .../tools/make/third_party_downloads.inc~ | 86 +++ 15 files changed, 1476 insertions(+), 19 deletions(-) create mode 100644 tensorflow/lite/micro/examples/hello_world/README.md~ create mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc create mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc create mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc rename tensorflow/lite/micro/{we_i => himax_we1_evb}/debug_log.cc (90%) create mode 100644 tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc create mode 100644 tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ create mode 100644 tensorflow/lite/micro/tools/make/third_party_downloads.inc~ diff --git a/tensorflow/lite/micro/examples/hello_world/README.md b/tensorflow/lite/micro/examples/hello_world/README.md index 3b633890306..9c0a5e2306a 100644 --- a/tensorflow/lite/micro/examples/hello_world/README.md +++ b/tensorflow/lite/micro/examples/hello_world/README.md @@ -14,13 +14,34 @@ of the device. ## Table of contents -- [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) -- [Deploy to Arduino](#deploy-to-arduino) -- [Deploy to ESP32](#deploy-to-esp32) -- [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) -- [Deploy to STM32F746](#deploy-to-STM32F746) -- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) -- [Train your own model](#train-your-own-model) +- [Hello World Example](#hello-world-example) + - [Table of contents](#table-of-contents) + - [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) + - [Initial Setup](#initial-setup) + - [Generate Example Project](#generate-example-project) + - [Build and Run Example](#build-and-run-example) + - [Deploy to Arduino](#deploy-to-arduino) + - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) + - [Load and run the example](#load-and-run-the-example) + - [Deploy to ESP32](#deploy-to-esp32) + - [Install the ESP IDF](#install-the-esp-idf) + - [Generate the examples](#generate-the-examples) + - [Building the example](#building-the-example) + - [Load and run the example](#load-and-run-the-example-1) + - [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) + - [Initial Setup](#initial-setup-1) + - [MetaWare Development Toolkit](#metaware-development-toolkit) + - [Make Tool version](#make-tool-version) + - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) + - [Generate Example Project](#generate-example-project-1) + - [Build and Burn Example](#build-and-burn-example) + - [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) + - [Compile the binary](#compile-the-binary) + - [Sign the binary](#sign-the-binary) + - [Flash the binary](#flash-the-binary) + - [Deploy to STM32F746](#deploy-to-stm32f746) + - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) + - [Train your own model](#train-your-own-model) ## Deploy to ARC EM SDP @@ -191,6 +212,103 @@ The previous two commands can be combined: idf.py --port /dev/ttyUSB0 flash monitor ``` +## Deploy to himax WE1 EVB + +The following instructions will help you build and deploy this example to +[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) +board. To undstand more about using this board, please check +[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). + +### Initial Setup + +To use the HIMAX WE1 EVB, please make sure following software are installed: + +#### MetaWare Development Toolkit + +See +[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) +section for instructions on toolchain installation. + +#### Make Tool version + +A `'make'` tool is required for deploying Tensorflow Lite Micro +applications on HIMAX WE1 EVB, See +[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) +section for proper environment. + +#### Serial Terminal Emulation Application + +There are 2 main purposes for HIMAX WE1 EVB Debug UART port + +- print application output +- burn application to flash by using xmodem send application binary + +You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). + + +### Generate Example Project + +The example project for HIMAX WE1 EVB platform can be generated with the following +command: + +Download related third party data + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads +``` + +Generate hello world project + +``` +make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb TAGS=no_arc_mli +``` + +### Build and Burn Example + +Following the Steps to run hello world example at HIMAX WE1 EVB platform. + +1. Go to the generated example project directory. + + ``` + cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/hello_world/make + ``` + +2. Build the example using + + ``` + make app + ``` + +3. After example build finish, copy ELF file and map file to image generate tool directory. + image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` + + ``` + cp hello_world.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +4. Go to flash image generate tool directory. + + ``` + cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +5. run image generate tool, generate flash image file. + + * Before running image generate tool, by typing `sudo chmod +x image_gen` + and `sudo chmod +x sign_tool` to make sure it is executable. + + ``` + image_gen -e hello_world.elf -m himax_we1_evb.map -o out.img + ``` + + +6. Download flash image file to HIMAX WE1 EVB by UART: + + * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) + +After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial +terminal. + ## Deploy to SparkFun Edge The following instructions will help you build and deploy this sample on the diff --git a/tensorflow/lite/micro/examples/hello_world/README.md~ b/tensorflow/lite/micro/examples/hello_world/README.md~ new file mode 100644 index 00000000000..011711493d5 --- /dev/null +++ b/tensorflow/lite/micro/examples/hello_world/README.md~ @@ -0,0 +1,595 @@ +# Hello World Example + +This example is designed to demonstrate the absolute basics of using [TensorFlow +Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers). +It includes the full end-to-end workflow of training a model, converting it for +use with TensorFlow Lite for Microcontrollers for running inference on a +microcontroller. + +The model is trained to replicate a `sine` function and generates a pattern of +data to either blink LEDs or control an animation, depending on the capabilities +of the device. + +![Animation on STM32F746](images/animation_on_STM32F746.gif) + +## Table of contents + +- [Hello World Example](#hello-world-example) + - [Table of contents](#table-of-contents) + - [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) + - [Initial Setup](#initial-setup) + - [Generate Example Project](#generate-example-project) + - [Build and Run Example](#build-and-run-example) + - [Deploy to Arduino](#deploy-to-arduino) + - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) + - [Load and run the example](#load-and-run-the-example) + - [Deploy to ESP32](#deploy-to-esp32) + - [Install the ESP IDF](#install-the-esp-idf) + - [Generate the examples](#generate-the-examples) + - [Building the example](#building-the-example) + - [Load and run the example](#load-and-run-the-example-1) + - [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) + - [Initial Setup](#initial-setup-1) + - [MetaWare Development Toolkit](#metaware-development-toolkit) + - [Make Tool version](#make-tool-version) + - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) + - [Generate Example Project](#generate-example-project-1) + - [Build and Burn Example](#build-and-burn-example) + - [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) + - [Compile the binary](#compile-the-binary) + - [Sign the binary](#sign-the-binary) + - [Flash the binary](#flash-the-binary) + - [Deploy to STM32F746](#deploy-to-stm32f746) + - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) + - [Train your own model](#train-your-own-model) + +## Deploy to ARC EM SDP + +The following instructions will help you to build and deploy this example to +[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) +board. General information and instructions on using the board with TensorFlow +Lite Micro can be found in the common +[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). + +### Initial Setup + +Follow the instructions on the +[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) +to get and install all required tools for work with ARC EM SDP. + +### Generate Example Project + +The example project for ARC EM SDP platform can be generated with the following +command: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp TAGS=no_arc_mli generate_hello_world_make_project +``` + +### Build and Run Example + +For more detailed information on building and running examples see the +appropriate sections of general descriptions of the +[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). +In the directory with generated project you can also find a +*README_ARC_EMSDP.md* file with instructions and options on building and +running. Here we only briefly mention main steps which are typically enough to +get it started. + +1. You need to + [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) + and open an serial connection. + +2. Go to the generated example project director + + ``` + cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/hello_world/make + ``` + +3. Build the example using + + ``` + make app + ``` + +4. To generate artefacts for self-boot of example from the board use + + ``` + make flash + ``` + +5. To run application from the board using microSD card: + + * Copy the content of the created /bin folder into the root of microSD + card. Note that the card must be formatted as FAT32 with default cluster + size (but less than 32 Kbytes) + * Plug in the microSD card into the J11 connector. + * Push the RST button. If a red LED is lit beside RST button, push the CFG + button. + +6. If you have the MetaWare Debugger installed in your environment: + + * To run application from the console using it type `make run`. + * To stop the execution type `Ctrl+C` in the console several times. + +In both cases (step 5 and 6) you will see the application output in the serial +terminal. + +## Deploy to Arduino + +The following instructions will help you build and deploy this sample +to [Arduino](https://www.arduino.cc/) devices. + +![Animation on Arduino MKRZERO](images/animation_on_arduino_mkrzero.gif) + +The sample has been tested with the following devices: + +- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) +- [Arduino MKRZERO](https://store.arduino.cc/usa/arduino-mkrzero) + +The sample will use PWM to fade an LED on and off according to the model's +output. In the code, the `LED_BUILTIN` constant is used to specify the board's +built-in LED as the one being controlled. However, on some boards, this built-in +LED is not attached to a pin with PWM capabilities. In this case, the LED will +blink instead of fading. + +### Install the Arduino_TensorFlowLite library + +This example application is included as part of the official TensorFlow Lite +Arduino library. To install it, open the Arduino library manager in +`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. + +### Load and run the example + +Once the library has been added, go to `File -> Examples`. You should see an +example near the bottom of the list named `TensorFlowLite:hello_world`. Select +it and click `hello_world` to load the example. + +Use the Arduino IDE to build and upload the example. Once it is running, +you should see the built-in LED on your device flashing. + +The Arduino Desktop IDE includes a plotter that we can use to display the sine +wave graphically. To view it, go to `Tools -> Serial Plotter`. You will see one +datapoint being logged for each inference cycle, expressed as a number between 0 +and 255. + +## Deploy to ESP32 + +The following instructions will help you build and deploy this sample +to [ESP32](https://www.espressif.com/en/products/hardware/esp32/overview) +devices using the [ESP IDF](https://github.com/espressif/esp-idf). + +The sample has been tested on ESP-IDF version 4.0 with the following devices: +- [ESP32-DevKitC](http://esp-idf.readthedocs.io/en/latest/get-started/get-started-devkitc.html) +- [ESP-EYE](https://github.com/espressif/esp-who/blob/master/docs/en/get-started/ESP-EYE_Getting_Started_Guide.md) + +### Install the ESP IDF + +Follow the instructions of the +[ESP-IDF get started guide](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html) +to setup the toolchain and the ESP-IDF itself. + +The next steps assume that the +[IDF environment variables are set](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html#step-4-set-up-the-environment-variables) : + + * The `IDF_PATH` environment variable is set + * `idf.py` and Xtensa-esp32 tools (e.g. `xtensa-esp32-elf-gcc`) are in `$PATH` + +### Generate the examples +The example project can be generated with the following command: +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=esp generate_hello_world_esp_project +``` + +### Building the example + +Go the the example project directory +``` +cd tensorflow/lite/micro/tools/make/gen/esp_xtensa-esp32/prj/hello_world/esp-idf +``` + +Then build with `idf.py` +``` +idf.py build +``` + +### Load and run the example + +To flash (replace `/dev/ttyUSB0` with the device serial port): +``` +idf.py --port /dev/ttyUSB0 flash +``` + +Monitor the serial output: +``` +idf.py --port /dev/ttyUSB0 monitor +``` + +Use `Ctrl+]` to exit. + +The previous two commands can be combined: +``` +idf.py --port /dev/ttyUSB0 flash monitor +``` + +## Deploy to himax WE1 EVB + +The following instructions will help you build and deploy this example to +[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) +board. To undstand more about using this board, please check +[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). + +### Initial Setup + +To use the HIMAX WE1 EVB, please make sure following software are installed: + +#### MetaWare Development Toolkit + +See +[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) +section for instructions on toolchain installation. + +#### Make Tool version + +A `'make'` tool is required for deploying Tensorflow Lite Micro +applications on HIMAX WE1 EVB, See +[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) +section for proper environment. + +#### Serial Terminal Emulation Application + +There are 2 main purposes for HIMAX WE1 EVB Debug UART port + +- print application output +- burn application to flash by using xmodem send application binary + +You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). + + +### Generate Example Project + +The example project for HIMAX WE1 EVB platform can be generated with the following +command: + +Download related third party data + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads +``` + +Generate hello world project + +``` +make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb +``` + +### Build and Burn Example + +Following the Steps to run hello world example at HIMAX WE1 EVB platform. + +1. Go to the generated example project directory. + + ``` + cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/hello_world/make + ``` + +2. Build the example using + + ``` + make app + ``` + +3. After example build finish, copy ELF file and map file to image generate tool directory. + image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` + + ``` + cp hello_world.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +4. Go to flash image generate tool directory. + + ``` + cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +5. run image generate tool, generate flash image file. + + * Before running image generate tool, by typing `sudo chmod +x image_gen` + and `sudo chmod +x sign_tool` to make sure it is executable. + + ``` + image_gen -e hello_world.elf -m himax_we1_evb.map -o out.img + ``` + + +6. Download flash image file to HIMAX WE1 EVB by UART: + + * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) + +After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial +terminal. + +## Deploy to SparkFun Edge + +The following instructions will help you build and deploy this sample on the +[SparkFun Edge development board](https://sparkfun.com/products/15170). + +![Animation on SparkFun Edge](images/animation_on_sparkfun_edge.gif) + +If you're new to using this board, we recommend walking through the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab to get an understanding of the workflow. + +### Compile the binary + +The following command will download the required dependencies and then compile a +binary for the SparkFun Edge: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge hello_world_bin +``` + +The binary will be created in the following location: + +``` +tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin +``` + +### Sign the binary + +The binary must be signed with cryptographic keys to be deployed to the device. +We'll now run some commands that will sign our binary so it can be flashed to +the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is +downloaded when the `Makefile` is run. + +Enter the following command to set up some dummy cryptographic keys we can use +for development: + +``` +cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info0.py \ +tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info.py +``` + +Next, run the following command to create a signed binary: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_image_blob.py \ +--bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin \ +--load-address 0xC000 \ +--magic-num 0xCB \ +-o main_nonsecure_ota \ +--version 0x0 +``` + +This will create the file `main_nonsecure_ota.bin`. We'll now run another +command to create a final version of the file that can be used to flash our +device with the bootloader script we will use in the next step: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ +--load-address 0x20000 \ +--bin main_nonsecure_ota.bin \ +-i 6 \ +-o main_nonsecure_wire \ +--options 0x1 +``` + +You should now have a file called `main_nonsecure_wire.bin` in the directory +where you ran the commands. This is the file we'll be flashing to the device. + +### Flash the binary + +Next, attach the board to your computer via a USB-to-serial adapter. + +**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), +you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) +before you continue. + +Once connected, assign the USB device name to an environment variable: + +``` +export DEVICENAME=put your device name here +``` + +Set another variable with the baud rate: + +``` +export BAUD_RATE=921600 +``` + +Now, hold the button marked `14` on the device. While still holding the button, +hit the button marked `RST`. Continue holding the button marked `14` while +running the following command: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/uart_wired_update.py \ +-b ${BAUD_RATE} ${DEVICENAME} \ +-r 1 \ +-f main_nonsecure_wire.bin \ +-i 6 +``` + +You should see a long stream of output as the binary is flashed to the device. +Once you see the following lines, flashing is complete: + +``` +Sending Reset Command. +Done. +``` + +If you don't see these lines, flashing may have failed. Try running through the +steps in [Flash the binary](#flash-the-binary) again (you can skip over setting +the environment variables). If you continue to run into problems, follow the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab, which includes more comprehensive instructions for the flashing +process. + +The binary should now be deployed to the device. Hit the button marked `RST` to +reboot the board. You should see the device's four LEDs flashing in sequence. + +Debug information is logged by the board while the program is running. To view +it, establish a serial connection to the board using a baud rate of `115200`. +On OSX and Linux, the following command should work: + +``` +screen ${DEVICENAME} 115200 +``` + +You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, +immediately followed by `Esc`. You can then use the arrow keys to explore the +output, which will contain the results of running inference on various `x` +values: + +``` +x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 +``` + +To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately +followed by the `K` key, then hit the `Y` key. + + +## Deploy to STM32F746 + +The following instructions will help you build and deploy the sample to the +[STM32F7 discovery kit](https://os.mbed.com/platforms/ST-Discovery-F746NG/) +using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). + +![Animation on STM32F746](images/animation_on_STM32F746.gif) + +Before we begin, you'll need the following: + +- STM32F7 discovery kit board +- Mini-USB cable +- ARM Mbed CLI ([installation instructions](https://os.mbed.com/docs/mbed-os/v5.12/tools/installation-and-setup.html)) +- Python 2.7 and pip + +Since Mbed requires a special folder structure for projects, we'll first run a +command to generate a subfolder containing the required source files in this +structure: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=mbed TAGS="CMSIS disco_f746ng" generate_hello_world_mbed_project +``` + +This will result in the creation of a new folder: + +``` +tensorflow/lite/micro/tools/make/gen/mbed_cortex-m4/prj/hello_world/mbed +``` + +This folder contains all of the example's dependencies structured in the correct +way for Mbed to be able to build it. + +Change into the directory and run the following commands, making sure you are +using Python 2.7.15. + +First, tell Mbed that the current directory is the root of an Mbed project: + +``` +mbed config root . +``` + +Next, tell Mbed to download the dependencies and prepare to build: + +``` +mbed deploy +``` + +By default, Mbed will build the project using C++98. However, TensorFlow Lite +requires C++11. Run the following Python snippet to modify the Mbed +configuration files so that it uses C++11: + +``` +python -c 'import fileinput, glob; +for filename in glob.glob("mbed-os/tools/profiles/*.json"): + for line in fileinput.input(filename, inplace=True): + print line.replace("\"-std=gnu++98\"","\"-std=c++11\", \"-fpermissive\"")' + +``` + +Finally, run the following command to compile: + +``` +mbed compile -m DISCO_F746NG -t GCC_ARM +``` + +This should result in a binary at the following path: + +``` +./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin +``` + +To deploy, plug in your STM board and copy the file to it. On MacOS, you can do +this with the following command: + +``` +cp ./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin /Volumes/DIS_F746NG/ +``` + +Copying the file will initiate the flashing process. Once this is complete, you +should see an animation on the device's screen. + + +``` +screen /dev/tty.usbmodem14403 9600 +``` + +In addition to this animation, debug information is logged by the board while +the program is running. To view it, establish a serial connection to the board +using a baud rate of `9600`. On OSX and Linux, the following command should +work, replacing `/dev/tty.devicename` with the name of your device as it appears +in `/dev`: + +``` +screen /dev/tty.devicename 9600 +``` + +You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, +immediately followed by `Esc`. You can then use the arrow keys to explore the +output, which will contain the results of running inference on various `x` +values: + +``` +x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 +``` + +To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately +followed by the `K` key, then hit the `Y` key. + +### Run the tests on a development machine + +To compile and test this example on a desktop Linux or macOS machine, first +clone the TensorFlow repository from GitHub to a convenient place: + +```bash +git clone --depth 1 https://github.com/tensorflow/tensorflow.git +``` + +Next, `cd` into the source directory from a terminal, and then run the following +command: + +```bash +make -f tensorflow/lite/micro/tools/make/Makefile test_hello_world_test +``` + +This will take a few minutes, and downloads frameworks the code uses. Once the +process has finished, you should see a series of files get compiled, followed by +some logging output from a test, which should conclude with +`~~~ALL TESTS PASSED~~~`. + +If you see this, it means that a small program has been built and run that loads +the trained TensorFlow model, runs some example inputs through it, and got the +expected outputs. + +To understand how TensorFlow Lite does this, you can look at the source in +[hello_world_test.cc](hello_world_test.cc). +It's a fairly small amount of code that creates an interpreter, gets a handle to +a model that's been compiled into the program, and then invokes the interpreter +with the model and sample inputs. + +### Train your own model + +So far you have used an existing trained model to run inference on +microcontrollers. If you wish to train your own model, follow the instructions +given in the [train/](train/) directory. + diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc new file mode 100644 index 00000000000..1816a2f3207 --- /dev/null +++ b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc @@ -0,0 +1,19 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/hello_world/constants.h" + +// This is tuned so that a full cycle takes ~4 seconds on a SparkFun Edge. +const int kInferencesPerCycle = 1000; diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc new file mode 100644 index 00000000000..8ca028acc55 --- /dev/null +++ b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc @@ -0,0 +1,35 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/hello_world/output_handler.h" + + +/* +This function trigger different device's LEDthrough y value. +y value range -1 <= y <= 1. +| Range is from -1~1 | LEDs | +| 0 <= y <= 1 | [ 0 1 ] | +| -1 <= y < 0 | [ 1 0 ] | + +*/ +void HandleOutput(tflite::ErrorReporter* error_reporter, float x_value, + float y_value) { + // The first time this method runs, set up our LEDs correctly + + // Log the current X and Y values + TF_LITE_REPORT_ERROR(error_reporter, "x_value: %f, y_value: %f\n", + static_cast(x_value), + static_cast(y_value)); +} diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ new file mode 100644 index 00000000000..b59242d0b6f --- /dev/null +++ b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ @@ -0,0 +1,53 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/hello_world/output_handler.h" + + +/* +This function trigger different device's LEDthrough y value. +y value range -1 <= y <= 1. +| Range is from -1~1 | LEDs | +| 0 <= y <= 1 | [ 0 1 ] | +| -1 <= y < 0 | [ 1 0 ] | + +*/ +void HandleOutput(tflite::ErrorReporter* error_reporter, float x_value, + float y_value) { + // The first time this method runs, set up our LEDs correctly +/* static bool is_initialized = false; + if (!is_initialized) { + // TODO Setup LED's as outputs + + // end of setup + is_initialized = true; + } + + // Set the LEDs to represent negative values + if (y_value < 0) { + //enable LED1 + + //enable LED0 + } else if (y_value > 0) { + //enable LED0 + + //enable LED1 + } + */ + // Log the current X and Y values + TF_LITE_REPORT_ERROR(error_reporter, "x_value: %f, y_value: %f\n", + static_cast(x_value), + static_cast(y_value)); +} diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/README.md b/tensorflow/lite/micro/examples/person_detection_experimental/README.md index bf99b40d776..4d53e551431 100644 --- a/tensorflow/lite/micro/examples/person_detection_experimental/README.md +++ b/tensorflow/lite/micro/examples/person_detection_experimental/README.md @@ -7,13 +7,31 @@ This uses the experimental int8 quantized version of the person detection model. ## Table of contents -- [Getting started](#getting-started) -- [Running on ARC EM SDP](#running-on-arc-em-sdp) -- [Running on Arduino](#running-on-arduino) -- [Running on SparkFun Edge](#running-on-sparkfun-edge) -- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) -- [Debugging image capture](#debugging-image-capture) -- [Training your own model](#training-your-own-model) +- [Person detection example](#person-detection-example) + - [Table of contents](#table-of-contents) + - [Running on ARC EM SDP](#running-on-arc-em-sdp) + - [Initial setup](#initial-setup) + - [Generate Example Project](#generate-example-project) + - [Build and Run Example](#build-and-run-example) + - [Running on Arduino](#running-on-arduino) + - [Hardware](#hardware) + - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) + - [Install other libraries](#install-other-libraries) + - [Load and run the example](#load-and-run-the-example) + - [Running on HIMAX WE1 EVB](#running-on-himax-we1-evb) + - [Initial Setup](#initial-setup) + - [MetaWare Development Toolkit](#metaware-development-toolkit) + - [Make Tool version](#make-tool-version) + - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) + - [Generate Example Project](#generate-example-project-1) + - [Build and Burn Example](#build-and-burn-example) + - [Running on SparkFun Edge](#running-on-sparkfun-edge) + - [Compile the binary](#compile-the-binary) + - [Sign the binary](#sign-the-binary) + - [Flash the binary](#flash-the-binary) + - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) + - [Debugging image capture](#debugging-image-capture) + - [Training your own model](#training-your-own-model) ## Running on ARC EM SDP @@ -260,6 +278,103 @@ From the log, we can see that it took around 170 ms to capture and read the image data from the camera module, 180 ms to decode the JPEG and convert it to greyscale, and 18.6 seconds to run inference. +## Running on HIMAX WE1 EVB + +The following instructions will help you build and deploy this example to +[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) +board. To undstand more about using this board, please check +[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). + +### Initial Setup + +To use the HIMAX WE1 EVB, please make sure following software are installed: + +#### MetaWare Development Toolkit + +See +[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) +section for instructions on toolchain installation. + +#### Make Tool version + +A `'make'` tool is required for deploying Tensorflow Lite Micro +applications on HIMAX WE1 EVB, See +[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) +section for proper environment. + +#### Serial Terminal Emulation Application + +There are 2 main purposes for HIMAX WE1 EVB Debug UART port + +- print application output +- burn application to flash by using xmodem send application binary + +You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). + + +### Generate Example Project + +The example project for HIMAX WE1 EVB platform can be generated with the following +command: + +Download related third party data + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads +``` + +Generate person detection project + +``` +make -f tensorflow/lite/micro/tools/make/Makefile generate_person_detection_int8_make_project TARGET=himax_we1_evb +``` + +### Build and Burn Example + +Following the Steps to run person detection example at HIMAX WE1 EVB platform. + +1. Go to the generated example project directory. + + ``` + cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/person_detection_int8/make + ``` + +2. Build the example using + + ``` + make app + ``` + +3. After example build finish, copy ELF file and map file to image generate tool directory. + image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` + + ``` + cp person_detection_int8.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +4. Go to flash image generate tool directory. + + ``` + cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +5. run image generate tool, generate flash image file. + + * Before running image generate tool, by typing `sudo chmod +x image_gen` + and `sudo chmod +x sign_tool` to make sure it is executable. + + ``` + image_gen -e person_detection_int8.elf -m himax_we1_evb.map -o out.img + ``` + + +6. Download flash image file to HIMAX WE1 EVB by UART: + + * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) + +After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial +terminal. + ## Running on SparkFun Edge The following instructions will help you build and deploy this sample on the diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc new file mode 100644 index 00000000000..a353dc8a9b8 --- /dev/null +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc @@ -0,0 +1,34 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/person_detection_experimental/detection_responder.h" + +#include "hx_drv_tflm.h" + +// This dummy implementation writes person and no person scores to the error +// console. Real applications will want to take some custom action instead, and +// should implement their own versions of this function. +void RespondToDetection(tflite::ErrorReporter* error_reporter, + int8_t person_score, int8_t no_person_score) { + + if (person_score > no_person_score) { + hx_drv_led_on(HX_DRV_LED_GREEN); + } else { + hx_drv_led_off(HX_DRV_LED_GREEN); + } + + TF_LITE_REPORT_ERROR(error_reporter, "person score:%d no person score %d", + person_score, no_person_score); +} diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc new file mode 100644 index 00000000000..727d93c61d1 --- /dev/null +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc @@ -0,0 +1,44 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/person_detection_experimental/image_provider.h" + +#include "tensorflow/lite/micro/examples/person_detection_experimental/model_settings.h" + +#include "hx_drv_tflm.h" + +hx_drv_sensor_image_config_t g_pimg_config; + + +TfLiteStatus GetImage(tflite::ErrorReporter* error_reporter, int image_width, + int image_height, int channels, int8_t* image_data) { + static bool is_initialized = false; + + if (!is_initialized) { + if(hx_drv_sensor_initial(&g_pimg_config)!= HX_DRV_LIB_PASS) + { + return kTfLiteError; + } + is_initialized = true; + } + + hx_drv_sensor_capture(&g_pimg_config); + + hx_drv_image_rescale((uint8_t*)g_pimg_config.raw_address, g_pimg_config.img_width, g_pimg_config.img_height, + image_data, image_width, image_height); + + + return kTfLiteOk; +} diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ new file mode 100644 index 00000000000..d5b4d136642 --- /dev/null +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ @@ -0,0 +1,44 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/person_detection_experimental/image_provider.h" + +#include "tensorflow/lite/micro/examples/person_detection_experimental/model_settings.h" + +#include "hx_drv_tflm.h" + +hx_drv_sensor_image_config_t g_pimg_config; + + +TfLiteStatus GetImage(tflite::ErrorReporter* error_reporter, int image_width, + int image_height, int channels, int8_t* image_data) { + static bool is_initialized = false; + + if (!is_initialized) { + if(hx_drv_sensor_initial(&g_pimg_config)!= HX_DRV_LIB_PASS) + { + return kTfLiteError; + } + is_initialized = true; + } + + hx_drv_sensor_capture(&g_pimg_config); + + hx_drv_image_rescale((uint8_t*)g_pimg_config.raw_address, g_pimg_config.img_width, g_pimg_config.img_height, + image_data, image_data, image_height); + + + return kTfLiteOk; +} diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc new file mode 100644 index 00000000000..552b52c9c51 --- /dev/null +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc @@ -0,0 +1,127 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#include "tensorflow/lite/micro/examples/person_detection_experimental/main_functions.h" + +#include "tensorflow/lite/micro/examples/person_detection_experimental/detection_responder.h" +#include "tensorflow/lite/micro/examples/person_detection_experimental/image_provider.h" +#include "tensorflow/lite/micro/examples/person_detection_experimental/model_settings.h" +#include "tensorflow/lite/micro/examples/person_detection_experimental/person_detect_model_data.h" +#include "tensorflow/lite/micro/kernels/micro_ops.h" +#include "tensorflow/lite/micro/micro_error_reporter.h" +#include "tensorflow/lite/micro/micro_interpreter.h" +#include "tensorflow/lite/micro/micro_mutable_op_resolver.h" +#include "tensorflow/lite/schema/schema_generated.h" +#include "tensorflow/lite/version.h" + +// Globals, used for compatibility with Arduino-style sketches. +namespace { +tflite::ErrorReporter* error_reporter = nullptr; +const tflite::Model* model = nullptr; +tflite::MicroInterpreter* interpreter = nullptr; +TfLiteTensor* input = nullptr; + +// In order to use optimized tensorflow lite kernels, a signed int8 quantized +// model is preferred over the legacy unsigned model format. This means that +// throughout this project, input images must be converted from unisgned to +// signed format. The easiest and quickest way to convert from unsigned to +// signed 8-bit integers is to subtract 128 from the unsigned value to get a +// signed value. + +// An area of memory to use for input, output, and intermediate arrays. +constexpr int kTensorArenaSize = 125 * 1024; +#pragma Bss(".tensor_arena") +static uint8_t tensor_arena[kTensorArenaSize]; +#pragma Bss() +} // namespace + +// The name of this function is important for Arduino compatibility. +void setup() { + // Set up logging. Google style is to avoid globals or statics because of + // lifetime uncertainty, but since this has a trivial destructor it's okay. + // NOLINTNEXTLINE(runtime-global-variables) + static tflite::MicroErrorReporter micro_error_reporter; + error_reporter = µ_error_reporter; + + // Map the model into a usable data structure. This doesn't involve any + // copying or parsing, it's a very lightweight operation. + model = tflite::GetModel(g_person_detect_model_data); + if (model->version() != TFLITE_SCHEMA_VERSION) { + TF_LITE_REPORT_ERROR(error_reporter, + "Model provided is schema version %d not equal " + "to supported version %d.", + model->version(), TFLITE_SCHEMA_VERSION); + return; + } + + // Pull in only the operation implementations we need. + // This relies on a complete list of all the ops needed by this graph. + // An easier approach is to just use the AllOpsResolver, but this will + // incur some penalty in code space for op implementations that are not + // needed by this graph. + // + // tflite::ops::micro::AllOpsResolver resolver; + // NOLINTNEXTLINE(runtime-global-variables) + static tflite::MicroOpResolver<12> micro_op_resolver; + micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_DEPTHWISE_CONV_2D, + tflite::ops::micro::Register_DEPTHWISE_CONV_2D(), + 1, 3); + micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D, + tflite::ops::micro::Register_CONV_2D(), 1, 3); + micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D, + tflite::ops::micro::Register_AVERAGE_POOL_2D(), + 1, 2); + micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE, + tflite::ops::micro::Register_RESHAPE()); + micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX, + tflite::ops::micro::Register_SOFTMAX(), 1, 3); + + // Build an interpreter to run the model with. + // NOLINTNEXTLINE(runtime-global-variables) + static tflite::MicroInterpreter static_interpreter( + model, micro_op_resolver, tensor_arena, kTensorArenaSize, error_reporter); + interpreter = &static_interpreter; + + // Allocate memory from the tensor_arena for the model's tensors. + TfLiteStatus allocate_status = interpreter->AllocateTensors(); + if (allocate_status != kTfLiteOk) { + TF_LITE_REPORT_ERROR(error_reporter, "AllocateTensors() failed"); + return; + } + + // Get information about the memory area to use for the model's input. + input = interpreter->input(0); +} + +// The name of this function is important for Arduino compatibility. +void loop() { + // Get image from provider. + if (kTfLiteOk != GetImage(error_reporter, kNumCols, kNumRows, kNumChannels, + input->data.int8)) { + TF_LITE_REPORT_ERROR(error_reporter, "Image capture failed."); + } + + // Run the model on this input and make sure it succeeds. + if (kTfLiteOk != interpreter->Invoke()) { + TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed."); + } + + TfLiteTensor* output = interpreter->output(0); + + // Process the inference results. + int8_t person_score = output->data.uint8[kPersonIndex]; + int8_t no_person_score = output->data.uint8[kNotAPersonIndex]; + RespondToDetection(error_reporter, person_score, no_person_score); +} diff --git a/tensorflow/lite/micro/we_i/debug_log.cc b/tensorflow/lite/micro/himax_we1_evb/debug_log.cc similarity index 90% rename from tensorflow/lite/micro/we_i/debug_log.cc rename to tensorflow/lite/micro/himax_we1_evb/debug_log.cc index a115d476aff..32af2625630 100644 --- a/tensorflow/lite/micro/we_i/debug_log.cc +++ b/tensorflow/lite/micro/himax_we1_evb/debug_log.cc @@ -18,16 +18,15 @@ limitations under the License. // the Ambiq Apollo 3. #include "tensorflow/lite/micro/debug_log.h" -#include "xprintf.h" -#include "console_io.h" -#include +#include "hx_drv_tflm.h" + extern "C" void DebugLog(const char* s) { static bool is_initialized = false; if (!is_initialized) { - xprintf_setup(); + hx_drv_uart_initial(); is_initialized = true; } - xprintf("%s", s); + hx_drv_uart_print("%s", s); } diff --git a/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc b/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc new file mode 100644 index 00000000000..60fc2e7cca1 --- /dev/null +++ b/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc @@ -0,0 +1,91 @@ +# Settings for himax WE_1 evb. +ifeq ($(TARGET), himax_we1_evb) + + CC_TOOL = ccac + AR_TOOL = arac + CXX_TOOL = ccac + LD_TOOL := ccac + TARGET_ARCH := arc + #ARC_TOOLCHAIN := mwdt + + BUILD_ARC_MLI := false + ARC_MLI_PRE_COMPILED_TARGET := himax_arcem9d_r16 + + include $(MAKEFILE_DIR)/targets/arc/arc_common.inc + + #download SDK & MLI + HIMAX_WE1_SDK_NAME := himax_we1_sdk + $(eval $(call add_third_party_download,$(HIMAX_WE1_SDK_URL),$(HIMAX_WE1_SDK_MD5),$(HIMAX_WE1_SDK_NAME),)) + + #export path of toolchain + #export PATH := $(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/image_gen_linux_v3/:$(PATH) + + TCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/arcem9d_wei_r16.tcf + LCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/memory.lcf + ARCLIB_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/libembarc.a + LIB_HEADER_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/hx_drv_tflm.h + + + DEFAULT_HEAPSZ := 8192 + DEFAULT_STACKSZ := 8192 + + TCF_FILE_NAME = $(notdir $(TCF_FILE)) + ARC_TARGET_FILES_DIRS = $(dir $(TCF_FILE_NAME)) + MAKE_PROJECT_FILES += $(TCF_FILE_NAME) + + LCF_FILE_NAME = $(notdir $(LCF_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(LCF_FILE)) + MAKE_PROJECT_FILES += $(LCF_FILE_NAME) + + ARCLIB_FILE_NAME = $(notdir $(ARCLIB_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(ARCLIB_FILE)) + MAKE_PROJECT_FILES += $(ARCLIB_FILE_NAME) + + LIB_HEADER_FILE_NAME = $(notdir $(LIB_HEADER_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(LIB_HEADER_FILE)) + MAKE_PROJECT_FILES += $(LIB_HEADER_FILE_NAME) + + + + # Need a pointer to the TCF and lcf file + + PLATFORM_FLAGS = \ + -DNDEBUG \ + -g \ + -DCPU_ARC \ + -Hnosdata \ + -DTF_LITE_STATIC_MEMORY \ + -tcf=$(TCF_FILE_NAME) \ + -Hnocopyr \ + -Hpurge \ + -Hcl \ + -fslp-vectorize-aggressive \ + -ffunction-sections \ + -fdata-sections \ + -tcf_core_config \ + + CXXFLAGS += -fno-rtti -DSCRATCH_MEM_Z_SIZE=0x10000 $(PLATFORM_FLAGS) + CCFLAGS += $(PLATFORM_FLAGS) + + INCLUDES+= \ + -I $(MAKEFILE_DIR)/downloads/$(WEI_SDK_NAME) \ + -I $(MAKEFILE_DIR)/downloads/kissfft + + GENERATED_PROJECT_INCLUDES += \ + -I. \ + -I./third_party/kissfft + + LDFLAGS += \ + -Hheap=8192 \ + -tcf=$(TCF_FILE_NAME) \ + -Hnocopyr \ + -m \ + -Hldopt=-Coutput=$(TARGET).map \ + $(LCF_FILE_NAME) \ + -Hldopt=-Bgrouplib $(ARCLIB_FILE_NAME) + + CXXFLAGS := $(filter-out -std=c++11,$(CXXFLAGS)) + CCFLAGS := $(filter-out -std=c11,$(CCFLAGS)) + MICROLITE_LIBS := $(filter-out -lm,$(MICROLITE_LIBS)) + +endif diff --git a/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ b/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ new file mode 100644 index 00000000000..733f258fbbb --- /dev/null +++ b/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ @@ -0,0 +1,93 @@ +# Settings for himax WE_1 evb. +ifeq ($(TARGET), himax_we1_evb) + + CC_TOOL = ccac + AR_TOOL = arac + CXX_TOOL = ccac + LD_TOOL := ccac + TARGET_ARCH := arc + #ARC_TOOLCHAIN := mwdt + + BUILD_ARC_MLI := false + ARC_MLI_PRE_COMPILED_TARGET := himax_arcem9d_r16 + +include $(MAKEFILE_DIR)/targets/arc/arc_common.inc + #download SDK & MLI + HIMAX_WE1_SDK_NAME := himax_we1_sdk + #MLI_LIB_DIR = arc_mli_package + #MLI_LIB_DIR = arc_mli_package + #$(eval $(call add_third_party_download,$(EMBARC_MLI_PRE_COMPILED_URL),$(EMBARC_MLI_PRE_COMPILED_MD5),$(MLI_LIB_DIR),)) + $(eval $(call add_third_party_download,$(HIMAX_WE1_SDK_URL),$(HIMAX_WE1_SDK_MD5),$(HIMAX_WE1_SDK_NAME),)) + + #export path of toolchain + #export PATH := $(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/image_gen_linux_v3/:$(PATH) + + TCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/arcem9d_wei_r16.tcf + LCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/memory.lcf + ARCLIB_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/libembarc.a + LIB_HEADER_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/hx_drv_tflm.h + + + DEFAULT_HEAPSZ := 8192 + DEFAULT_STACKSZ := 8192 + + TCF_FILE_NAME = $(notdir $(TCF_FILE)) + ARC_TARGET_FILES_DIRS = $(dir $(TCF_FILE_NAME)) + MAKE_PROJECT_FILES += $(TCF_FILE_NAME) + + LCF_FILE_NAME = $(notdir $(LCF_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(LCF_FILE)) + MAKE_PROJECT_FILES += $(LCF_FILE_NAME) + + ARCLIB_FILE_NAME = $(notdir $(ARCLIB_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(ARCLIB_FILE)) + MAKE_PROJECT_FILES += $(ARCLIB_FILE_NAME) + + LIB_HEADER_FILE_NAME = $(notdir $(LIB_HEADER_FILE)) + ARC_TARGET_FILES_DIRS += $(dir $(LIB_HEADER_FILE)) + MAKE_PROJECT_FILES += $(LIB_HEADER_FILE_NAME) + + + + # Need a pointer to the TCF and lcf file + + PLATFORM_FLAGS = \ + -DNDEBUG \ + -g \ + -DCPU_ARC \ + -Hnosdata \ + -DTF_LITE_STATIC_MEMORY \ + -tcf=$(TCF_FILE_NAME) \ + -Hnocopyr \ + -Hpurge \ + -Hcl \ + -fslp-vectorize-aggressive \ + -ffunction-sections \ + -fdata-sections \ + -tcf_core_config \ + + CXXFLAGS += -fno-rtti -DSCRATCH_MEM_Z_SIZE=0x10000 $(PLATFORM_FLAGS) + CCFLAGS += $(PLATFORM_FLAGS) + + INCLUDES+= \ + -I $(MAKEFILE_DIR)/downloads/$(WEI_SDK_NAME) \ + -I $(MAKEFILE_DIR)/downloads/kissfft + + GENERATED_PROJECT_INCLUDES += \ + -I. \ + -I./third_party/kissfft + + LDFLAGS += \ + -Hheap=8192 \ + -tcf=$(TCF_FILE_NAME) \ + -Hnocopyr \ + -m \ + -Hldopt=-Coutput=$(TARGET).map \ + $(LCF_FILE_NAME) \ + -Hldopt=-Bgrouplib $(ARCLIB_FILE_NAME) + + CXXFLAGS := $(filter-out -std=c++11,$(CXXFLAGS)) + CCFLAGS := $(filter-out -std=c11,$(CCFLAGS)) + MICROLITE_LIBS := $(filter-out -lm,$(MICROLITE_LIBS)) + +endif diff --git a/tensorflow/lite/micro/tools/make/third_party_downloads.inc b/tensorflow/lite/micro/tools/make/third_party_downloads.inc index 806501a004a..75a51e0df10 100644 --- a/tensorflow/lite/micro/tools/make/third_party_downloads.inc +++ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc @@ -80,3 +80,7 @@ EMBARC_MLI_PRE_COMPILED_MD5 := "a95ff9e0370434484f14e7e4114327f6" XTENSA_HIFI4_URL :="https://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_04_07.zip" XTENSA_HIFI4_MD5 :="f234764928f9a42901df33a27e118c8b" +HIMAX_WE1_SDK_URL ="https://www.himax.com.tw/we-i/himax_we1_sdk_v02.zip" +HIMAX_WE1_SDK_MD5 ="9a4b2f29b16052764e437b64bdcba816" + + diff --git a/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ new file mode 100644 index 00000000000..3c7ee1b64d2 --- /dev/null +++ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ @@ -0,0 +1,86 @@ +# Add URLs and MD5 checksums for third-party libraries here. + +GEMMLOWP_URL := "https://github.com/google/gemmlowp/archive/719139ce755a0f31cbf1c37f7f98adcc7fc9f425.zip" +GEMMLOWP_MD5 := "7e8191b24853d75de2af87622ad293ba" + +ifeq ($(HOST_OS),windows) + FLATBUFFERS_URL := "https://github.com/google/flatbuffers/archive/v1.12.0.zip" + FLATBUFFERS_MD5 := "a1afdbf114dec01a861c1b8c917d0fc7" +else + FLATBUFFERS_URL := "https://github.com/google/flatbuffers/archive/v1.12.0.tar.gz" + FLATBUFFERS_MD5 := "c62ffefb3d4548b127cca14ce047f16c" +endif + +ifeq ($(HOST_OS),osx) + GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-mac.tar.bz2" + GCC_EMBEDDED_MD5 := "a66be9828cf3c57d7d21178e07cd8904" +else ifeq ($(HOST_OS),windows) + GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-win32.zip" + GCC_EMBEDDED_MD5 := "bc8ae26d7c429f30d583a605a4bcf9bc" +else + GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-linux.tar.bz2" + GCC_EMBEDDED_MD5 := "299ebd3f1c2c90930d28ab82e5d8d6c0" +endif + +LEON_BCC2_URL := "https://www.gaisler.com/anonftp/bcc2/bin/bcc-2.0.7-gcc-linux64.tar.xz" +LEON_BCC2_MD5 := "cdf78082be4882da2a92c9baa82fe765" + +TSIM_URL := "https://www.gaisler.com/anonftp/tsim/tsim-eval-2.0.63.tar.gz" +TSIM_MD5 := "afa0095d3ed989a949e1467f94e41d2f" + +CMSIS_URL := "https://github.com/ARM-software/CMSIS_5/archive/1150e71e07c79b538efd842aba5b210a31827ae5.zip" +CMSIS_MD5 := "e05f4222ef58825193910b41a0871dcb" + +AM_SDK_URL := "http://s3.asia.ambiqmicro.com/downloads/AmbiqSuite-Rel2.2.0.zip" +AM_SDK_MD5 := "7605fa2d4d97e6bb7a1190c92b66b597" +AM_SDK_DEST := AmbiqSuite-Rel2.2.0 + +SF_BSPS_URL := "https://github.com/sparkfun/SparkFun_Apollo3_AmbiqSuite_BSPs/archive/v0.0.7.zip" +SF_BSPS_MD5 := "34199f7e754735661d1c8a70a40ca7a3" +SF_BSPS_DEST := boards_sfe + +STM32_BARE_LIB_URL := "https://github.com/google/stm32_bare_lib/archive/c07d611fb0af58450c5a3e0ab4d52b47f99bc82d.zip" +STM32_BARE_LIB_MD5 := "282bff40d4d0b92278fd123a3b6e3123" + +ifeq ($(HOST_OS),osx) + RISCV_TOOLCHAIN_URL := "https://static.dev.sifive.com/dev-tools/riscv64-unknown-elf-gcc-8.1.0-2019.01.0-x86_64-apple-darwin.tar.gz" + RISCV_TOOLCHAIN_MD5 := "2ac2fa00618b9ab7fa0c7d0ec173de94" +else + RISCV_TOOLCHAIN_URL := "https://static.dev.sifive.com/dev-tools/riscv64-unknown-elf-gcc-20181030-x86_64-linux-ubuntu14.tar.gz" + RISCV_TOOLCHAIN_MD5="2366b7afe36a54dc94fb0ff8a0830934" +endif + +SIFIVE_FE310_LIB_URL := "https://github.com/sifive/freedom-e-sdk/archive/baeeb8fd497a99b3c141d7494309ec2e64f19bdf.zip" +SIFIVE_FE310_LIB_MD5 := "06ee24c4956f8e21670ab3395861fe64" + +KISSFFT_URL="https://github.com/mborgerding/kissfft/archive/v130.zip" +KISSFFT_MD5="438ba1fef5783cc5f5f201395cc477ca" + +RUY_URL="https://github.com/google/ruy/archive/1b313682ef8b8fc8ed08719c610d1c3503b016bf.zip" +RUY_MD5="2d54f058f8f7120dfc1ecee79dbf259e" + +CIFAR10_DATASET_URL="https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz" +CIFAR10_DATASET_MD5="c32a1d4ab5d03f1284b67883e8d87530" + +IMAGE_RECOGNITION_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/models/tflite/cifar_image_recognition_model_2020_4_14.zip" +IMAGE_RECOGNITION_MODEL_MD5 := "2b886156e7ef4d6e53d0f1a4bc800e56" + +PERSON_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_grayscale_2019_11_21.zip" +PERSON_MODEL_MD5 := "fe2934bd0788f1dcc7af3f0a954542ab" + +PERSON_MODEL_INT8_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_01_13.zip" +PERSON_MODEL_INT8_MD5 := "8a7d2c70325f53136faea6dde517b8cc" + +EMBARC_MLI_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/archive/58284867ca52d1f43b25045e8601999d7359d986.zip" +EMBARC_MLI_MD5 := "2bf4982a327fdaa9d475803ce014d1ef" + +EMBARC_MLI_PRE_COMPILED_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/releases/download/Release_1.1_RC2/embARC_MLI_package.zip" +EMBARC_MLI_PRE_COMPILED_MD5 := "a95ff9e0370434484f14e7e4114327f6" + +XTENSA_HIFI4_URL :="https://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_04_07.zip" +XTENSA_HIFI4_MD5 :="f234764928f9a42901df33a27e118c8b" + +HIMAX_WE1_SDK_URL ="https://www.himax.com.tw/we-i/himax_we1_sdk_v02.zip" +HIMAX_WE1_SDK_MD5 ="5063c24d298fbcfe118163f3ccc43079" + + From 96eb311826d68179bf85a228e294ae55f39ef2e4 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Thu, 4 Jun 2020 18:08:43 +0800 Subject: [PATCH 0080/1390] TFLM: update example readme --- .../lite/micro/examples/hello_world/README.md | 36 +- .../micro/examples/hello_world/README.md~ | 2 +- .../hello_world/himax_we1_evb/constants.cc | 19 - .../himax_we1_evb/output_handler.cc | 35 -- .../himax_we1_evb/output_handler.cc~ | 53 -- .../person_detection_experimental/README.md | 33 +- .../person_detection_experimental/README.md~ | 568 ++++++++++++++++++ .../tools/make/third_party_downloads.inc~ | 86 --- 8 files changed, 585 insertions(+), 247 deletions(-) delete mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc delete mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc delete mode 100644 tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/README.md~ delete mode 100644 tensorflow/lite/micro/tools/make/third_party_downloads.inc~ diff --git a/tensorflow/lite/micro/examples/hello_world/README.md b/tensorflow/lite/micro/examples/hello_world/README.md index 9c0a5e2306a..d3762ada790 100644 --- a/tensorflow/lite/micro/examples/hello_world/README.md +++ b/tensorflow/lite/micro/examples/hello_world/README.md @@ -14,34 +14,14 @@ of the device. ## Table of contents -- [Hello World Example](#hello-world-example) - - [Table of contents](#table-of-contents) - - [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) - - [Initial Setup](#initial-setup) - - [Generate Example Project](#generate-example-project) - - [Build and Run Example](#build-and-run-example) - - [Deploy to Arduino](#deploy-to-arduino) - - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) - - [Load and run the example](#load-and-run-the-example) - - [Deploy to ESP32](#deploy-to-esp32) - - [Install the ESP IDF](#install-the-esp-idf) - - [Generate the examples](#generate-the-examples) - - [Building the example](#building-the-example) - - [Load and run the example](#load-and-run-the-example-1) - - [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) - - [Initial Setup](#initial-setup-1) - - [MetaWare Development Toolkit](#metaware-development-toolkit) - - [Make Tool version](#make-tool-version) - - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) - - [Generate Example Project](#generate-example-project-1) - - [Build and Burn Example](#build-and-burn-example) - - [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) - - [Compile the binary](#compile-the-binary) - - [Sign the binary](#sign-the-binary) - - [Flash the binary](#flash-the-binary) - - [Deploy to STM32F746](#deploy-to-stm32f746) - - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) - - [Train your own model](#train-your-own-model) +- [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) +- [Deploy to Arduino](#deploy-to-arduino) +- [Deploy to ESP32](#deploy-to-esp32) +- [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) +- [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) +- [Deploy to STM32F746](#deploy-to-STM32F746) +- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) +- [Train your own model](#train-your-own-model) ## Deploy to ARC EM SDP diff --git a/tensorflow/lite/micro/examples/hello_world/README.md~ b/tensorflow/lite/micro/examples/hello_world/README.md~ index 011711493d5..9c0a5e2306a 100644 --- a/tensorflow/lite/micro/examples/hello_world/README.md~ +++ b/tensorflow/lite/micro/examples/hello_world/README.md~ @@ -260,7 +260,7 @@ make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_par Generate hello world project ``` -make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb +make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb TAGS=no_arc_mli ``` ### Build and Burn Example diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc deleted file mode 100644 index 1816a2f3207..00000000000 --- a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/constants.cc +++ /dev/null @@ -1,19 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -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 - - http://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. -==============================================================================*/ - -#include "tensorflow/lite/micro/examples/hello_world/constants.h" - -// This is tuned so that a full cycle takes ~4 seconds on a SparkFun Edge. -const int kInferencesPerCycle = 1000; diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc deleted file mode 100644 index 8ca028acc55..00000000000 --- a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc +++ /dev/null @@ -1,35 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -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 - - http://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. -==============================================================================*/ - -#include "tensorflow/lite/micro/examples/hello_world/output_handler.h" - - -/* -This function trigger different device's LEDthrough y value. -y value range -1 <= y <= 1. -| Range is from -1~1 | LEDs | -| 0 <= y <= 1 | [ 0 1 ] | -| -1 <= y < 0 | [ 1 0 ] | - -*/ -void HandleOutput(tflite::ErrorReporter* error_reporter, float x_value, - float y_value) { - // The first time this method runs, set up our LEDs correctly - - // Log the current X and Y values - TF_LITE_REPORT_ERROR(error_reporter, "x_value: %f, y_value: %f\n", - static_cast(x_value), - static_cast(y_value)); -} diff --git a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ b/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ deleted file mode 100644 index b59242d0b6f..00000000000 --- a/tensorflow/lite/micro/examples/hello_world/himax_we1_evb/output_handler.cc~ +++ /dev/null @@ -1,53 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -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 - - http://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. -==============================================================================*/ - -#include "tensorflow/lite/micro/examples/hello_world/output_handler.h" - - -/* -This function trigger different device's LEDthrough y value. -y value range -1 <= y <= 1. -| Range is from -1~1 | LEDs | -| 0 <= y <= 1 | [ 0 1 ] | -| -1 <= y < 0 | [ 1 0 ] | - -*/ -void HandleOutput(tflite::ErrorReporter* error_reporter, float x_value, - float y_value) { - // The first time this method runs, set up our LEDs correctly -/* static bool is_initialized = false; - if (!is_initialized) { - // TODO Setup LED's as outputs - - // end of setup - is_initialized = true; - } - - // Set the LEDs to represent negative values - if (y_value < 0) { - //enable LED1 - - //enable LED0 - } else if (y_value > 0) { - //enable LED0 - - //enable LED1 - } - */ - // Log the current X and Y values - TF_LITE_REPORT_ERROR(error_reporter, "x_value: %f, y_value: %f\n", - static_cast(x_value), - static_cast(y_value)); -} diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/README.md b/tensorflow/lite/micro/examples/person_detection_experimental/README.md index 4d53e551431..06f5640986f 100644 --- a/tensorflow/lite/micro/examples/person_detection_experimental/README.md +++ b/tensorflow/lite/micro/examples/person_detection_experimental/README.md @@ -7,31 +7,14 @@ This uses the experimental int8 quantized version of the person detection model. ## Table of contents -- [Person detection example](#person-detection-example) - - [Table of contents](#table-of-contents) - - [Running on ARC EM SDP](#running-on-arc-em-sdp) - - [Initial setup](#initial-setup) - - [Generate Example Project](#generate-example-project) - - [Build and Run Example](#build-and-run-example) - - [Running on Arduino](#running-on-arduino) - - [Hardware](#hardware) - - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) - - [Install other libraries](#install-other-libraries) - - [Load and run the example](#load-and-run-the-example) - - [Running on HIMAX WE1 EVB](#running-on-himax-we1-evb) - - [Initial Setup](#initial-setup) - - [MetaWare Development Toolkit](#metaware-development-toolkit) - - [Make Tool version](#make-tool-version) - - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) - - [Generate Example Project](#generate-example-project-1) - - [Build and Burn Example](#build-and-burn-example) - - [Running on SparkFun Edge](#running-on-sparkfun-edge) - - [Compile the binary](#compile-the-binary) - - [Sign the binary](#sign-the-binary) - - [Flash the binary](#flash-the-binary) - - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) - - [Debugging image capture](#debugging-image-capture) - - [Training your own model](#training-your-own-model) +- [Getting started](#getting-started) +- [Running on ARC EM SDP](#running-on-arc-em-sdp) +- [Running on Arduino](#running-on-arduino) +- [Running on HIMAX WE1 EVB](#running-on-himax-we1-evb) +- [Running on SparkFun Edge](#running-on-sparkfun-edge) +- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) +- [Debugging image capture](#debugging-image-capture) +- [Training your own model](#training-your-own-model) ## Running on ARC EM SDP diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ b/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ new file mode 100644 index 00000000000..4d53e551431 --- /dev/null +++ b/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ @@ -0,0 +1,568 @@ +# Person detection example + +This example shows how you can use Tensorflow Lite to run a 250 kilobyte neural +network to recognize people in images captured by a camera. It is designed to +run on systems with small amounts of memory such as microcontrollers and DSPs. +This uses the experimental int8 quantized version of the person detection model. + +## Table of contents + +- [Person detection example](#person-detection-example) + - [Table of contents](#table-of-contents) + - [Running on ARC EM SDP](#running-on-arc-em-sdp) + - [Initial setup](#initial-setup) + - [Generate Example Project](#generate-example-project) + - [Build and Run Example](#build-and-run-example) + - [Running on Arduino](#running-on-arduino) + - [Hardware](#hardware) + - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) + - [Install other libraries](#install-other-libraries) + - [Load and run the example](#load-and-run-the-example) + - [Running on HIMAX WE1 EVB](#running-on-himax-we1-evb) + - [Initial Setup](#initial-setup) + - [MetaWare Development Toolkit](#metaware-development-toolkit) + - [Make Tool version](#make-tool-version) + - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) + - [Generate Example Project](#generate-example-project-1) + - [Build and Burn Example](#build-and-burn-example) + - [Running on SparkFun Edge](#running-on-sparkfun-edge) + - [Compile the binary](#compile-the-binary) + - [Sign the binary](#sign-the-binary) + - [Flash the binary](#flash-the-binary) + - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) + - [Debugging image capture](#debugging-image-capture) + - [Training your own model](#training-your-own-model) + +## Running on ARC EM SDP + +The following instructions will help you to build and deploy this example to +[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) +board. General information and instructions on using the board with TensorFlow +Lite Micro can be found in the common +[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). + +This example uses asymmetric int8 quantization and can therefore leverage +optimized int8 kernels from the embARC MLI library + +The ARC EM SDP board contains a rich set of extension interfaces. You can choose +any compatible camera and modify +[image_provider.cc](/tensorflow/lite/micro/examples/person_detection_experimental/image_provider.cc) +file accordingly to use input from your specific camera. By default, results of +running this example are printed to the console. If you would like to instead +implement some target-specific actions, you need to modify +[detection_responder.cc](/tensorflow/lite/micro/examples/person_detection_experimental/detection_responder.cc) +accordingly. + +The reference implementations of these files are used by default on the EM SDP. + +### Initial setup + +Follow the instructions on the +[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) +to get and install all required tools for work with ARC EM SDP. + +### Generate Example Project + +The example project for ARC EM SDP platform can be generated with the following +command: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp generate_person_detection_int8_make_project +``` + +### Build and Run Example + +For more detailed information on building and running examples see the +appropriate sections of general descriptions of the +[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). +In the directory with generated project you can also find a +*README_ARC_EMSDP.md* file with instructions and options on building and +running. Here we only briefly mention main steps which are typically enough to +get it started. + +1. You need to + [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) + and open an serial connection. + +2. Go to the generated example project director + + ``` + cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/person_detection_int8/make + ``` + +3. Build the example using + + ``` + make app + ``` + +4. To generate artefacts for self-boot of example from the board use + + ``` + make flash + ``` + +5. To run application from the board using microSD card: + + * Copy the content of the created /bin folder into the root of microSD + card. Note that the card must be formatted as FAT32 with default cluster + size (but less than 32 Kbytes) + * Plug in the microSD card into the J11 connector. + * Push the RST button. If a red LED is lit beside RST button, push the CFG + button. + +6. If you have the MetaWare Debugger installed in your environment: + + * To run application from the console using it type `make run`. + * To stop the execution type `Ctrl+C` in the console several times. + +In both cases (step 5 and 6) you will see the application output in the serial +terminal. + +## Running on Arduino + +The following instructions will help you build and deploy this sample +to [Arduino](https://www.arduino.cc/) devices. + +The sample has been tested with the following device: + +- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) + +You will also need the following camera module: + +- [Arducam Mini 2MP Plus](https://www.amazon.com/Arducam-Module-Megapixels-Arduino-Mega2560/dp/B012UXNDOY) + +### Hardware + +Connect the Arducam pins as follows: + +|Arducam pin name|Arduino pin name| +|----------------|----------------| +|CS|D7 (unlabelled, immediately to the right of D6)| +|MOSI|D11| +|MISO|D12| +|SCK|D13| +|GND|GND (either pin marked GND is fine)| +|VCC|3.3 V| +|SDA|A4| +|SCL|A5| + +### Install the Arduino_TensorFlowLite library + +Download the current nightly build of the library: +[person_detection.zip](https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_01_13.zip) + +This example application is included as part of the official TensorFlow Lite +Arduino library. To install it, open the Arduino library manager in +`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. + +### Install other libraries + +In addition to the TensorFlow library, you'll also need to install two +libraries: + +* The Arducam library, so our code can interface with the hardware +* The JPEGDecoder library, so we can decode JPEG-encoded images + +The Arducam Arduino library is available from GitHub at +[https://github.com/ArduCAM/Arduino](https://github.com/ArduCAM/Arduino). +To install it, download or clone the repository. Next, copy its `ArduCAM` +subdirectory into your `Arduino/libraries` directory. To find this directory on +your machine, check the *Sketchbook location* in the Arduino IDE's +*Preferences* window. + +After downloading the library, you'll need to edit one of its files to make sure +it is configured for the Arducam Mini 2MP Plus. To do so, open the following +file: + +``` +Arduino/libraries/ArduCAM/memorysaver.h +``` + +You'll see a bunch of `#define` statements listed. Make sure that they are all +commented out, except for `#define OV2640_MINI_2MP_PLUS`, as so: + +``` +//Step 1: select the hardware platform, only one at a time +//#define OV2640_MINI_2MP +//#define OV3640_MINI_3MP +//#define OV5642_MINI_5MP +//#define OV5642_MINI_5MP_BIT_ROTATION_FIXED +#define OV2640_MINI_2MP_PLUS +//#define OV5642_MINI_5MP_PLUS +//#define OV5640_MINI_5MP_PLUS +``` + +Once you save the file, we're done configuring the Arducam library. + +Our next step is to install the JPEGDecoder library. We can do this from within +the Arduino IDE. First, go to the *Manage Libraries...* option in the *Tools* +menu and search for `JPEGDecoder`. You should install version _1.8.0_ of the +library. + +Once the library has installed, we'll need to configure it to disable some +optional components that are not compatible with the Arduino Nano 33 BLE Sense. +Open the following file: + +``` +Arduino/libraries/JPEGDecoder/src/User_Config.h +``` + +Make sure that both `#define LOAD_SD_LIBRARY` and `#define LOAD_SDFAT_LIBRARY` +are commented out, as shown in this excerpt from the file: + +```c++ +// Comment out the next #defines if you are not using an SD Card to store the JPEGs +// Commenting out the line is NOT essential but will save some FLASH space if +// SD Card access is not needed. Note: use of SdFat is currently untested! + +//#define LOAD_SD_LIBRARY // Default SD Card library +//#define LOAD_SDFAT_LIBRARY // Use SdFat library instead, so SD Card SPI can be bit bashed +``` + +Once you've saved the file, you are done installing libraries. + +### Load and run the example + +Go to `File -> Examples`. You should see an +example near the bottom of the list named `TensorFlowLite`. Select +it and click `person_detection` to load the example. Connect your device, then +build and upload the example. + +To test the camera, start by pointing the device's camera at something that is +definitely not a person, or just covering it up. The next time the blue LED +flashes, the device will capture a frame from the camera and begin to run +inference. Since the vision model we are using for person detection is +relatively large, it takes a long time to run inference—around 19 seconds at the +time of writing, though it's possible TensorFlow Lite has gotten faster since +then. + +After 19 seconds or so, the inference result will be translated into another LED +being lit. Since you pointed the camera at something that isn't a person, the +red LED should light up. + +Now, try pointing the device's camera at yourself! The next time the blue LED +flashes, the device will capture another image and begin to run inference. After +19 seconds, the green LED should light up! + +Remember, image data is captured as a snapshot before each inference, whenever +the blue LED flashes. Whatever the camera is pointed at during that moment is +what will be fed into the model. It doesn't matter where the camera is pointed +until the next time an image is captured, when the blue LED will flash again. + +If you're getting seemingly incorrect results, make sure you are in an +environment with good lighting. You should also make sure that the camera is +oriented correctly, with the pins pointing downwards, so that the images it +captures are the right way up—the model was not trained to recognize upside-down +people! In addition, it's good to remember that this is a tiny model, which +trades accuracy for small size. It works very well, but it isn't accurate 100% +of the time. + +We can also see the results of inference via the Arduino Serial Monitor. To do +this, open the *Serial Monitor* from the *Tools* menu. You'll see a detailed +log of what is happening while our application runs. It's also interesting to +check the *Show timestamp* box, so you can see how long each part of the process +takes: + +``` +14:17:50.714 -> Starting capture +14:17:50.714 -> Image captured +14:17:50.784 -> Reading 3080 bytes from ArduCAM +14:17:50.887 -> Finished reading +14:17:50.887 -> Decoding JPEG and converting to greyscale +14:17:51.074 -> Image decoded and processed +14:18:09.710 -> Person score: 246 No person score: 66 +``` + +From the log, we can see that it took around 170 ms to capture and read the +image data from the camera module, 180 ms to decode the JPEG and convert it to +greyscale, and 18.6 seconds to run inference. + +## Running on HIMAX WE1 EVB + +The following instructions will help you build and deploy this example to +[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) +board. To undstand more about using this board, please check +[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). + +### Initial Setup + +To use the HIMAX WE1 EVB, please make sure following software are installed: + +#### MetaWare Development Toolkit + +See +[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) +section for instructions on toolchain installation. + +#### Make Tool version + +A `'make'` tool is required for deploying Tensorflow Lite Micro +applications on HIMAX WE1 EVB, See +[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) +section for proper environment. + +#### Serial Terminal Emulation Application + +There are 2 main purposes for HIMAX WE1 EVB Debug UART port + +- print application output +- burn application to flash by using xmodem send application binary + +You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). + + +### Generate Example Project + +The example project for HIMAX WE1 EVB platform can be generated with the following +command: + +Download related third party data + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads +``` + +Generate person detection project + +``` +make -f tensorflow/lite/micro/tools/make/Makefile generate_person_detection_int8_make_project TARGET=himax_we1_evb +``` + +### Build and Burn Example + +Following the Steps to run person detection example at HIMAX WE1 EVB platform. + +1. Go to the generated example project directory. + + ``` + cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/person_detection_int8/make + ``` + +2. Build the example using + + ``` + make app + ``` + +3. After example build finish, copy ELF file and map file to image generate tool directory. + image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` + + ``` + cp person_detection_int8.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +4. Go to flash image generate tool directory. + + ``` + cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +5. run image generate tool, generate flash image file. + + * Before running image generate tool, by typing `sudo chmod +x image_gen` + and `sudo chmod +x sign_tool` to make sure it is executable. + + ``` + image_gen -e person_detection_int8.elf -m himax_we1_evb.map -o out.img + ``` + + +6. Download flash image file to HIMAX WE1 EVB by UART: + + * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) + +After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial +terminal. + +## Running on SparkFun Edge + +The following instructions will help you build and deploy this sample on the +[SparkFun Edge development board](https://sparkfun.com/products/15170). This +sample requires the Sparkfun Himax camera for the Sparkfun Edge board. It is +not available for purchase yet. + +If you're new to using this board, we recommend walking through the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab to get an understanding of the workflow. + +### Compile the binary + +The following command will download the required dependencies and then compile a +binary for the SparkFun Edge: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge person_detection_bin +``` + +The binary will be created in the following location: + +``` +tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/person_detection.bin +``` + +### Sign the binary + +The binary must be signed with cryptographic keys to be deployed to the device. +We'll now run some commands that will sign our binary so it can be flashed to +the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is +downloaded when the `Makefile` is run. + +Enter the following command to set up some dummy cryptographic keys we can use +for development: + +``` +cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/keys_info0.py \ +tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/keys_info.py +``` + +Next, run the following command to create a signed binary: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/create_cust_image_blob.py \ +--bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/person_detection.bin \ +--load-address 0xC000 \ +--magic-num 0xCB \ +-o main_nonsecure_ota \ +--version 0x0 +``` + +This will create the file `main_nonsecure_ota.bin`. We'll now run another +command to create a final version of the file that can be used to flash our +device with the bootloader script we will use in the next step: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ +--load-address 0x20000 \ +--bin main_nonsecure_ota.bin \ +-i 6 \ +-o main_nonsecure_wire \ +--options 0x1 +``` + +You should now have a file called `main_nonsecure_wire.bin` in the directory +where you ran the commands. This is the file we'll be flashing to the device. + +### Flash the binary + +Next, attach the board to your computer via a USB-to-serial adapter. + +**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), +you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) +before you continue. + +Once connected, assign the USB device name to an environment variable: + +``` +export DEVICENAME=put your device name here +``` + +Set another variable with the baud rate: + +``` +export BAUD_RATE=921600 +``` + +Now, hold the button marked `14` on the device. While still holding the button, +hit the button marked `RST`. Continue holding the button marked `14` while +running the following command: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/uart_wired_update.py \ +-b ${BAUD_RATE} ${DEVICENAME} \ +-r 1 \ +-f main_nonsecure_wire.bin \ +-i 6 +``` + +You should see a long stream of output as the binary is flashed to the device. +Once you see the following lines, flashing is complete: + +``` +Sending Reset Command. +Done. +``` + +If you don't see these lines, flashing may have failed. Try running through the +steps in [Flash the binary](#flash-the-binary) again (you can skip over setting +the environment variables). If you continue to run into problems, follow the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab, which includes more comprehensive instructions for the flashing +process. + +The binary should now be deployed to the device. Hit the button marked `RST` to +reboot the board. You should see the device's four LEDs flashing in sequence. + +Debug information is logged by the board while the program is running. To view +it, establish a serial connection to the board using a baud rate of `115200`. +On OSX and Linux, the following command should work: + +``` +screen ${DEVICENAME} 115200 +``` + +To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately +followed by the `K` key, then hit the `Y` key. + +## Run the tests on a development machine + +To compile and test this example on a desktop Linux or MacOS machine, download +[the TensorFlow source code](https://github.com/tensorflow/tensorflow), `cd` +into the source directory from a terminal, and then run the following command: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile +``` + +This will take a few minutes, and downloads frameworks the code uses like +[CMSIS](https://developer.arm.com/embedded/cmsis) and +[flatbuffers](https://google.github.io/flatbuffers/). Once that process has +finished, run: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile test_person_detection_test +``` + +You should see a series of files get compiled, followed by some logging output +from a test, which should conclude with `~~~ALL TESTS PASSED~~~`. If you see +this, it means that a small program has been built and run that loads a trained +TensorFlow model, runs some example images through it, and got the expected +outputs. This particular test runs images with a and without a person in them, +and checks that the network correctly identifies them. + +To understand how TensorFlow Lite does this, you can look at the `TestInvoke()` +function in +[person_detection_test.cc](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc). +It's a fairly small amount of code, creating an interpreter, getting a handle to +a model that's been compiled into the program, and then invoking the interpreter +with the model and sample inputs. + +## Debugging image capture +When the sample is running, check the LEDs to determine whether the inference is +running correctly. If the red light is stuck on, it means there was an error +communicating with the camera. This is likely due to an incorrectly connected +or broken camera. + +During inference, the blue LED will toggle every time inference is complete. The +orange LED indicates that no person was found, and the green LED indicates a +person was found. The red LED should never turn on, since it indicates an error. + +In order to view the captured image, set the DUMP_IMAGE define in main.cc.  This +causes the board to log raw image info to the console. After the board has been +flashed and reset, dump the log to a text file: + + +``` +screen -L -Logfile ${DEVICENAME} 115200 +``` + +Next, run the raw to bitmap converter to view captured images: + +``` +python3 raw_to_bitmap.py -r GRAY -i +``` + +## Training your own model + +You can train your own model with some easy-to-use scripts. See +[training_a_model.md](training_a_model.md) for instructions. diff --git a/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ deleted file mode 100644 index 3c7ee1b64d2..00000000000 --- a/tensorflow/lite/micro/tools/make/third_party_downloads.inc~ +++ /dev/null @@ -1,86 +0,0 @@ -# Add URLs and MD5 checksums for third-party libraries here. - -GEMMLOWP_URL := "https://github.com/google/gemmlowp/archive/719139ce755a0f31cbf1c37f7f98adcc7fc9f425.zip" -GEMMLOWP_MD5 := "7e8191b24853d75de2af87622ad293ba" - -ifeq ($(HOST_OS),windows) - FLATBUFFERS_URL := "https://github.com/google/flatbuffers/archive/v1.12.0.zip" - FLATBUFFERS_MD5 := "a1afdbf114dec01a861c1b8c917d0fc7" -else - FLATBUFFERS_URL := "https://github.com/google/flatbuffers/archive/v1.12.0.tar.gz" - FLATBUFFERS_MD5 := "c62ffefb3d4548b127cca14ce047f16c" -endif - -ifeq ($(HOST_OS),osx) - GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-mac.tar.bz2" - GCC_EMBEDDED_MD5 := "a66be9828cf3c57d7d21178e07cd8904" -else ifeq ($(HOST_OS),windows) - GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-win32.zip" - GCC_EMBEDDED_MD5 := "bc8ae26d7c429f30d583a605a4bcf9bc" -else - GCC_EMBEDDED_URL := "https://developer.arm.com/-/media/Files/downloads/gnu-rm/7-2018q2/gcc-arm-none-eabi-7-2018-q2-update-linux.tar.bz2" - GCC_EMBEDDED_MD5 := "299ebd3f1c2c90930d28ab82e5d8d6c0" -endif - -LEON_BCC2_URL := "https://www.gaisler.com/anonftp/bcc2/bin/bcc-2.0.7-gcc-linux64.tar.xz" -LEON_BCC2_MD5 := "cdf78082be4882da2a92c9baa82fe765" - -TSIM_URL := "https://www.gaisler.com/anonftp/tsim/tsim-eval-2.0.63.tar.gz" -TSIM_MD5 := "afa0095d3ed989a949e1467f94e41d2f" - -CMSIS_URL := "https://github.com/ARM-software/CMSIS_5/archive/1150e71e07c79b538efd842aba5b210a31827ae5.zip" -CMSIS_MD5 := "e05f4222ef58825193910b41a0871dcb" - -AM_SDK_URL := "http://s3.asia.ambiqmicro.com/downloads/AmbiqSuite-Rel2.2.0.zip" -AM_SDK_MD5 := "7605fa2d4d97e6bb7a1190c92b66b597" -AM_SDK_DEST := AmbiqSuite-Rel2.2.0 - -SF_BSPS_URL := "https://github.com/sparkfun/SparkFun_Apollo3_AmbiqSuite_BSPs/archive/v0.0.7.zip" -SF_BSPS_MD5 := "34199f7e754735661d1c8a70a40ca7a3" -SF_BSPS_DEST := boards_sfe - -STM32_BARE_LIB_URL := "https://github.com/google/stm32_bare_lib/archive/c07d611fb0af58450c5a3e0ab4d52b47f99bc82d.zip" -STM32_BARE_LIB_MD5 := "282bff40d4d0b92278fd123a3b6e3123" - -ifeq ($(HOST_OS),osx) - RISCV_TOOLCHAIN_URL := "https://static.dev.sifive.com/dev-tools/riscv64-unknown-elf-gcc-8.1.0-2019.01.0-x86_64-apple-darwin.tar.gz" - RISCV_TOOLCHAIN_MD5 := "2ac2fa00618b9ab7fa0c7d0ec173de94" -else - RISCV_TOOLCHAIN_URL := "https://static.dev.sifive.com/dev-tools/riscv64-unknown-elf-gcc-20181030-x86_64-linux-ubuntu14.tar.gz" - RISCV_TOOLCHAIN_MD5="2366b7afe36a54dc94fb0ff8a0830934" -endif - -SIFIVE_FE310_LIB_URL := "https://github.com/sifive/freedom-e-sdk/archive/baeeb8fd497a99b3c141d7494309ec2e64f19bdf.zip" -SIFIVE_FE310_LIB_MD5 := "06ee24c4956f8e21670ab3395861fe64" - -KISSFFT_URL="https://github.com/mborgerding/kissfft/archive/v130.zip" -KISSFFT_MD5="438ba1fef5783cc5f5f201395cc477ca" - -RUY_URL="https://github.com/google/ruy/archive/1b313682ef8b8fc8ed08719c610d1c3503b016bf.zip" -RUY_MD5="2d54f058f8f7120dfc1ecee79dbf259e" - -CIFAR10_DATASET_URL="https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz" -CIFAR10_DATASET_MD5="c32a1d4ab5d03f1284b67883e8d87530" - -IMAGE_RECOGNITION_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/models/tflite/cifar_image_recognition_model_2020_4_14.zip" -IMAGE_RECOGNITION_MODEL_MD5 := "2b886156e7ef4d6e53d0f1a4bc800e56" - -PERSON_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_grayscale_2019_11_21.zip" -PERSON_MODEL_MD5 := "fe2934bd0788f1dcc7af3f0a954542ab" - -PERSON_MODEL_INT8_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_01_13.zip" -PERSON_MODEL_INT8_MD5 := "8a7d2c70325f53136faea6dde517b8cc" - -EMBARC_MLI_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/archive/58284867ca52d1f43b25045e8601999d7359d986.zip" -EMBARC_MLI_MD5 := "2bf4982a327fdaa9d475803ce014d1ef" - -EMBARC_MLI_PRE_COMPILED_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/releases/download/Release_1.1_RC2/embARC_MLI_package.zip" -EMBARC_MLI_PRE_COMPILED_MD5 := "a95ff9e0370434484f14e7e4114327f6" - -XTENSA_HIFI4_URL :="https://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_04_07.zip" -XTENSA_HIFI4_MD5 :="f234764928f9a42901df33a27e118c8b" - -HIMAX_WE1_SDK_URL ="https://www.himax.com.tw/we-i/himax_we1_sdk_v02.zip" -HIMAX_WE1_SDK_MD5 ="5063c24d298fbcfe118163f3ccc43079" - - From 233f1d53f829f157788ed687c07698b12cc8e091 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Thu, 4 Jun 2020 18:16:32 +0800 Subject: [PATCH 0081/1390] remove temp readme in example directory --- .../micro/examples/hello_world/README.md~ | 595 ------------------ .../person_detection_experimental/README.md~ | 568 ----------------- 2 files changed, 1163 deletions(-) delete mode 100644 tensorflow/lite/micro/examples/hello_world/README.md~ delete mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/README.md~ diff --git a/tensorflow/lite/micro/examples/hello_world/README.md~ b/tensorflow/lite/micro/examples/hello_world/README.md~ deleted file mode 100644 index 9c0a5e2306a..00000000000 --- a/tensorflow/lite/micro/examples/hello_world/README.md~ +++ /dev/null @@ -1,595 +0,0 @@ -# Hello World Example - -This example is designed to demonstrate the absolute basics of using [TensorFlow -Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers). -It includes the full end-to-end workflow of training a model, converting it for -use with TensorFlow Lite for Microcontrollers for running inference on a -microcontroller. - -The model is trained to replicate a `sine` function and generates a pattern of -data to either blink LEDs or control an animation, depending on the capabilities -of the device. - -![Animation on STM32F746](images/animation_on_STM32F746.gif) - -## Table of contents - -- [Hello World Example](#hello-world-example) - - [Table of contents](#table-of-contents) - - [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) - - [Initial Setup](#initial-setup) - - [Generate Example Project](#generate-example-project) - - [Build and Run Example](#build-and-run-example) - - [Deploy to Arduino](#deploy-to-arduino) - - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) - - [Load and run the example](#load-and-run-the-example) - - [Deploy to ESP32](#deploy-to-esp32) - - [Install the ESP IDF](#install-the-esp-idf) - - [Generate the examples](#generate-the-examples) - - [Building the example](#building-the-example) - - [Load and run the example](#load-and-run-the-example-1) - - [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) - - [Initial Setup](#initial-setup-1) - - [MetaWare Development Toolkit](#metaware-development-toolkit) - - [Make Tool version](#make-tool-version) - - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) - - [Generate Example Project](#generate-example-project-1) - - [Build and Burn Example](#build-and-burn-example) - - [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) - - [Compile the binary](#compile-the-binary) - - [Sign the binary](#sign-the-binary) - - [Flash the binary](#flash-the-binary) - - [Deploy to STM32F746](#deploy-to-stm32f746) - - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) - - [Train your own model](#train-your-own-model) - -## Deploy to ARC EM SDP - -The following instructions will help you to build and deploy this example to -[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) -board. General information and instructions on using the board with TensorFlow -Lite Micro can be found in the common -[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). - -### Initial Setup - -Follow the instructions on the -[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) -to get and install all required tools for work with ARC EM SDP. - -### Generate Example Project - -The example project for ARC EM SDP platform can be generated with the following -command: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp TAGS=no_arc_mli generate_hello_world_make_project -``` - -### Build and Run Example - -For more detailed information on building and running examples see the -appropriate sections of general descriptions of the -[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). -In the directory with generated project you can also find a -*README_ARC_EMSDP.md* file with instructions and options on building and -running. Here we only briefly mention main steps which are typically enough to -get it started. - -1. You need to - [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) - and open an serial connection. - -2. Go to the generated example project director - - ``` - cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/hello_world/make - ``` - -3. Build the example using - - ``` - make app - ``` - -4. To generate artefacts for self-boot of example from the board use - - ``` - make flash - ``` - -5. To run application from the board using microSD card: - - * Copy the content of the created /bin folder into the root of microSD - card. Note that the card must be formatted as FAT32 with default cluster - size (but less than 32 Kbytes) - * Plug in the microSD card into the J11 connector. - * Push the RST button. If a red LED is lit beside RST button, push the CFG - button. - -6. If you have the MetaWare Debugger installed in your environment: - - * To run application from the console using it type `make run`. - * To stop the execution type `Ctrl+C` in the console several times. - -In both cases (step 5 and 6) you will see the application output in the serial -terminal. - -## Deploy to Arduino - -The following instructions will help you build and deploy this sample -to [Arduino](https://www.arduino.cc/) devices. - -![Animation on Arduino MKRZERO](images/animation_on_arduino_mkrzero.gif) - -The sample has been tested with the following devices: - -- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) -- [Arduino MKRZERO](https://store.arduino.cc/usa/arduino-mkrzero) - -The sample will use PWM to fade an LED on and off according to the model's -output. In the code, the `LED_BUILTIN` constant is used to specify the board's -built-in LED as the one being controlled. However, on some boards, this built-in -LED is not attached to a pin with PWM capabilities. In this case, the LED will -blink instead of fading. - -### Install the Arduino_TensorFlowLite library - -This example application is included as part of the official TensorFlow Lite -Arduino library. To install it, open the Arduino library manager in -`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. - -### Load and run the example - -Once the library has been added, go to `File -> Examples`. You should see an -example near the bottom of the list named `TensorFlowLite:hello_world`. Select -it and click `hello_world` to load the example. - -Use the Arduino IDE to build and upload the example. Once it is running, -you should see the built-in LED on your device flashing. - -The Arduino Desktop IDE includes a plotter that we can use to display the sine -wave graphically. To view it, go to `Tools -> Serial Plotter`. You will see one -datapoint being logged for each inference cycle, expressed as a number between 0 -and 255. - -## Deploy to ESP32 - -The following instructions will help you build and deploy this sample -to [ESP32](https://www.espressif.com/en/products/hardware/esp32/overview) -devices using the [ESP IDF](https://github.com/espressif/esp-idf). - -The sample has been tested on ESP-IDF version 4.0 with the following devices: -- [ESP32-DevKitC](http://esp-idf.readthedocs.io/en/latest/get-started/get-started-devkitc.html) -- [ESP-EYE](https://github.com/espressif/esp-who/blob/master/docs/en/get-started/ESP-EYE_Getting_Started_Guide.md) - -### Install the ESP IDF - -Follow the instructions of the -[ESP-IDF get started guide](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html) -to setup the toolchain and the ESP-IDF itself. - -The next steps assume that the -[IDF environment variables are set](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html#step-4-set-up-the-environment-variables) : - - * The `IDF_PATH` environment variable is set - * `idf.py` and Xtensa-esp32 tools (e.g. `xtensa-esp32-elf-gcc`) are in `$PATH` - -### Generate the examples -The example project can be generated with the following command: -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=esp generate_hello_world_esp_project -``` - -### Building the example - -Go the the example project directory -``` -cd tensorflow/lite/micro/tools/make/gen/esp_xtensa-esp32/prj/hello_world/esp-idf -``` - -Then build with `idf.py` -``` -idf.py build -``` - -### Load and run the example - -To flash (replace `/dev/ttyUSB0` with the device serial port): -``` -idf.py --port /dev/ttyUSB0 flash -``` - -Monitor the serial output: -``` -idf.py --port /dev/ttyUSB0 monitor -``` - -Use `Ctrl+]` to exit. - -The previous two commands can be combined: -``` -idf.py --port /dev/ttyUSB0 flash monitor -``` - -## Deploy to himax WE1 EVB - -The following instructions will help you build and deploy this example to -[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) -board. To undstand more about using this board, please check -[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). - -### Initial Setup - -To use the HIMAX WE1 EVB, please make sure following software are installed: - -#### MetaWare Development Toolkit - -See -[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) -section for instructions on toolchain installation. - -#### Make Tool version - -A `'make'` tool is required for deploying Tensorflow Lite Micro -applications on HIMAX WE1 EVB, See -[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) -section for proper environment. - -#### Serial Terminal Emulation Application - -There are 2 main purposes for HIMAX WE1 EVB Debug UART port - -- print application output -- burn application to flash by using xmodem send application binary - -You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). - - -### Generate Example Project - -The example project for HIMAX WE1 EVB platform can be generated with the following -command: - -Download related third party data - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads -``` - -Generate hello world project - -``` -make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb TAGS=no_arc_mli -``` - -### Build and Burn Example - -Following the Steps to run hello world example at HIMAX WE1 EVB platform. - -1. Go to the generated example project directory. - - ``` - cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/hello_world/make - ``` - -2. Build the example using - - ``` - make app - ``` - -3. After example build finish, copy ELF file and map file to image generate tool directory. - image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` - - ``` - cp hello_world.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -4. Go to flash image generate tool directory. - - ``` - cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -5. run image generate tool, generate flash image file. - - * Before running image generate tool, by typing `sudo chmod +x image_gen` - and `sudo chmod +x sign_tool` to make sure it is executable. - - ``` - image_gen -e hello_world.elf -m himax_we1_evb.map -o out.img - ``` - - -6. Download flash image file to HIMAX WE1 EVB by UART: - - * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) - -After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial -terminal. - -## Deploy to SparkFun Edge - -The following instructions will help you build and deploy this sample on the -[SparkFun Edge development board](https://sparkfun.com/products/15170). - -![Animation on SparkFun Edge](images/animation_on_sparkfun_edge.gif) - -If you're new to using this board, we recommend walking through the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab to get an understanding of the workflow. - -### Compile the binary - -The following command will download the required dependencies and then compile a -binary for the SparkFun Edge: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge hello_world_bin -``` - -The binary will be created in the following location: - -``` -tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin -``` - -### Sign the binary - -The binary must be signed with cryptographic keys to be deployed to the device. -We'll now run some commands that will sign our binary so it can be flashed to -the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is -downloaded when the `Makefile` is run. - -Enter the following command to set up some dummy cryptographic keys we can use -for development: - -``` -cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info0.py \ -tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info.py -``` - -Next, run the following command to create a signed binary: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_image_blob.py \ ---bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin \ ---load-address 0xC000 \ ---magic-num 0xCB \ --o main_nonsecure_ota \ ---version 0x0 -``` - -This will create the file `main_nonsecure_ota.bin`. We'll now run another -command to create a final version of the file that can be used to flash our -device with the bootloader script we will use in the next step: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ ---load-address 0x20000 \ ---bin main_nonsecure_ota.bin \ --i 6 \ --o main_nonsecure_wire \ ---options 0x1 -``` - -You should now have a file called `main_nonsecure_wire.bin` in the directory -where you ran the commands. This is the file we'll be flashing to the device. - -### Flash the binary - -Next, attach the board to your computer via a USB-to-serial adapter. - -**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), -you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) -before you continue. - -Once connected, assign the USB device name to an environment variable: - -``` -export DEVICENAME=put your device name here -``` - -Set another variable with the baud rate: - -``` -export BAUD_RATE=921600 -``` - -Now, hold the button marked `14` on the device. While still holding the button, -hit the button marked `RST`. Continue holding the button marked `14` while -running the following command: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/uart_wired_update.py \ --b ${BAUD_RATE} ${DEVICENAME} \ --r 1 \ --f main_nonsecure_wire.bin \ --i 6 -``` - -You should see a long stream of output as the binary is flashed to the device. -Once you see the following lines, flashing is complete: - -``` -Sending Reset Command. -Done. -``` - -If you don't see these lines, flashing may have failed. Try running through the -steps in [Flash the binary](#flash-the-binary) again (you can skip over setting -the environment variables). If you continue to run into problems, follow the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab, which includes more comprehensive instructions for the flashing -process. - -The binary should now be deployed to the device. Hit the button marked `RST` to -reboot the board. You should see the device's four LEDs flashing in sequence. - -Debug information is logged by the board while the program is running. To view -it, establish a serial connection to the board using a baud rate of `115200`. -On OSX and Linux, the following command should work: - -``` -screen ${DEVICENAME} 115200 -``` - -You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, -immediately followed by `Esc`. You can then use the arrow keys to explore the -output, which will contain the results of running inference on various `x` -values: - -``` -x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 -``` - -To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately -followed by the `K` key, then hit the `Y` key. - - -## Deploy to STM32F746 - -The following instructions will help you build and deploy the sample to the -[STM32F7 discovery kit](https://os.mbed.com/platforms/ST-Discovery-F746NG/) -using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). - -![Animation on STM32F746](images/animation_on_STM32F746.gif) - -Before we begin, you'll need the following: - -- STM32F7 discovery kit board -- Mini-USB cable -- ARM Mbed CLI ([installation instructions](https://os.mbed.com/docs/mbed-os/v5.12/tools/installation-and-setup.html)) -- Python 2.7 and pip - -Since Mbed requires a special folder structure for projects, we'll first run a -command to generate a subfolder containing the required source files in this -structure: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=mbed TAGS="CMSIS disco_f746ng" generate_hello_world_mbed_project -``` - -This will result in the creation of a new folder: - -``` -tensorflow/lite/micro/tools/make/gen/mbed_cortex-m4/prj/hello_world/mbed -``` - -This folder contains all of the example's dependencies structured in the correct -way for Mbed to be able to build it. - -Change into the directory and run the following commands, making sure you are -using Python 2.7.15. - -First, tell Mbed that the current directory is the root of an Mbed project: - -``` -mbed config root . -``` - -Next, tell Mbed to download the dependencies and prepare to build: - -``` -mbed deploy -``` - -By default, Mbed will build the project using C++98. However, TensorFlow Lite -requires C++11. Run the following Python snippet to modify the Mbed -configuration files so that it uses C++11: - -``` -python -c 'import fileinput, glob; -for filename in glob.glob("mbed-os/tools/profiles/*.json"): - for line in fileinput.input(filename, inplace=True): - print line.replace("\"-std=gnu++98\"","\"-std=c++11\", \"-fpermissive\"")' - -``` - -Finally, run the following command to compile: - -``` -mbed compile -m DISCO_F746NG -t GCC_ARM -``` - -This should result in a binary at the following path: - -``` -./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin -``` - -To deploy, plug in your STM board and copy the file to it. On MacOS, you can do -this with the following command: - -``` -cp ./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin /Volumes/DIS_F746NG/ -``` - -Copying the file will initiate the flashing process. Once this is complete, you -should see an animation on the device's screen. - - -``` -screen /dev/tty.usbmodem14403 9600 -``` - -In addition to this animation, debug information is logged by the board while -the program is running. To view it, establish a serial connection to the board -using a baud rate of `9600`. On OSX and Linux, the following command should -work, replacing `/dev/tty.devicename` with the name of your device as it appears -in `/dev`: - -``` -screen /dev/tty.devicename 9600 -``` - -You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, -immediately followed by `Esc`. You can then use the arrow keys to explore the -output, which will contain the results of running inference on various `x` -values: - -``` -x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 -``` - -To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately -followed by the `K` key, then hit the `Y` key. - -### Run the tests on a development machine - -To compile and test this example on a desktop Linux or macOS machine, first -clone the TensorFlow repository from GitHub to a convenient place: - -```bash -git clone --depth 1 https://github.com/tensorflow/tensorflow.git -``` - -Next, `cd` into the source directory from a terminal, and then run the following -command: - -```bash -make -f tensorflow/lite/micro/tools/make/Makefile test_hello_world_test -``` - -This will take a few minutes, and downloads frameworks the code uses. Once the -process has finished, you should see a series of files get compiled, followed by -some logging output from a test, which should conclude with -`~~~ALL TESTS PASSED~~~`. - -If you see this, it means that a small program has been built and run that loads -the trained TensorFlow model, runs some example inputs through it, and got the -expected outputs. - -To understand how TensorFlow Lite does this, you can look at the source in -[hello_world_test.cc](hello_world_test.cc). -It's a fairly small amount of code that creates an interpreter, gets a handle to -a model that's been compiled into the program, and then invokes the interpreter -with the model and sample inputs. - -### Train your own model - -So far you have used an existing trained model to run inference on -microcontrollers. If you wish to train your own model, follow the instructions -given in the [train/](train/) directory. - diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ b/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ deleted file mode 100644 index 4d53e551431..00000000000 --- a/tensorflow/lite/micro/examples/person_detection_experimental/README.md~ +++ /dev/null @@ -1,568 +0,0 @@ -# Person detection example - -This example shows how you can use Tensorflow Lite to run a 250 kilobyte neural -network to recognize people in images captured by a camera. It is designed to -run on systems with small amounts of memory such as microcontrollers and DSPs. -This uses the experimental int8 quantized version of the person detection model. - -## Table of contents - -- [Person detection example](#person-detection-example) - - [Table of contents](#table-of-contents) - - [Running on ARC EM SDP](#running-on-arc-em-sdp) - - [Initial setup](#initial-setup) - - [Generate Example Project](#generate-example-project) - - [Build and Run Example](#build-and-run-example) - - [Running on Arduino](#running-on-arduino) - - [Hardware](#hardware) - - [Install the Arduino_TensorFlowLite library](#install-the-arduinotensorflowlite-library) - - [Install other libraries](#install-other-libraries) - - [Load and run the example](#load-and-run-the-example) - - [Running on HIMAX WE1 EVB](#running-on-himax-we1-evb) - - [Initial Setup](#initial-setup) - - [MetaWare Development Toolkit](#metaware-development-toolkit) - - [Make Tool version](#make-tool-version) - - [Serial Terminal Emulation Application](#serial-terminal-emulation-application) - - [Generate Example Project](#generate-example-project-1) - - [Build and Burn Example](#build-and-burn-example) - - [Running on SparkFun Edge](#running-on-sparkfun-edge) - - [Compile the binary](#compile-the-binary) - - [Sign the binary](#sign-the-binary) - - [Flash the binary](#flash-the-binary) - - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) - - [Debugging image capture](#debugging-image-capture) - - [Training your own model](#training-your-own-model) - -## Running on ARC EM SDP - -The following instructions will help you to build and deploy this example to -[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) -board. General information and instructions on using the board with TensorFlow -Lite Micro can be found in the common -[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). - -This example uses asymmetric int8 quantization and can therefore leverage -optimized int8 kernels from the embARC MLI library - -The ARC EM SDP board contains a rich set of extension interfaces. You can choose -any compatible camera and modify -[image_provider.cc](/tensorflow/lite/micro/examples/person_detection_experimental/image_provider.cc) -file accordingly to use input from your specific camera. By default, results of -running this example are printed to the console. If you would like to instead -implement some target-specific actions, you need to modify -[detection_responder.cc](/tensorflow/lite/micro/examples/person_detection_experimental/detection_responder.cc) -accordingly. - -The reference implementations of these files are used by default on the EM SDP. - -### Initial setup - -Follow the instructions on the -[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) -to get and install all required tools for work with ARC EM SDP. - -### Generate Example Project - -The example project for ARC EM SDP platform can be generated with the following -command: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp generate_person_detection_int8_make_project -``` - -### Build and Run Example - -For more detailed information on building and running examples see the -appropriate sections of general descriptions of the -[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). -In the directory with generated project you can also find a -*README_ARC_EMSDP.md* file with instructions and options on building and -running. Here we only briefly mention main steps which are typically enough to -get it started. - -1. You need to - [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) - and open an serial connection. - -2. Go to the generated example project director - - ``` - cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/person_detection_int8/make - ``` - -3. Build the example using - - ``` - make app - ``` - -4. To generate artefacts for self-boot of example from the board use - - ``` - make flash - ``` - -5. To run application from the board using microSD card: - - * Copy the content of the created /bin folder into the root of microSD - card. Note that the card must be formatted as FAT32 with default cluster - size (but less than 32 Kbytes) - * Plug in the microSD card into the J11 connector. - * Push the RST button. If a red LED is lit beside RST button, push the CFG - button. - -6. If you have the MetaWare Debugger installed in your environment: - - * To run application from the console using it type `make run`. - * To stop the execution type `Ctrl+C` in the console several times. - -In both cases (step 5 and 6) you will see the application output in the serial -terminal. - -## Running on Arduino - -The following instructions will help you build and deploy this sample -to [Arduino](https://www.arduino.cc/) devices. - -The sample has been tested with the following device: - -- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) - -You will also need the following camera module: - -- [Arducam Mini 2MP Plus](https://www.amazon.com/Arducam-Module-Megapixels-Arduino-Mega2560/dp/B012UXNDOY) - -### Hardware - -Connect the Arducam pins as follows: - -|Arducam pin name|Arduino pin name| -|----------------|----------------| -|CS|D7 (unlabelled, immediately to the right of D6)| -|MOSI|D11| -|MISO|D12| -|SCK|D13| -|GND|GND (either pin marked GND is fine)| -|VCC|3.3 V| -|SDA|A4| -|SCL|A5| - -### Install the Arduino_TensorFlowLite library - -Download the current nightly build of the library: -[person_detection.zip](https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_01_13.zip) - -This example application is included as part of the official TensorFlow Lite -Arduino library. To install it, open the Arduino library manager in -`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. - -### Install other libraries - -In addition to the TensorFlow library, you'll also need to install two -libraries: - -* The Arducam library, so our code can interface with the hardware -* The JPEGDecoder library, so we can decode JPEG-encoded images - -The Arducam Arduino library is available from GitHub at -[https://github.com/ArduCAM/Arduino](https://github.com/ArduCAM/Arduino). -To install it, download or clone the repository. Next, copy its `ArduCAM` -subdirectory into your `Arduino/libraries` directory. To find this directory on -your machine, check the *Sketchbook location* in the Arduino IDE's -*Preferences* window. - -After downloading the library, you'll need to edit one of its files to make sure -it is configured for the Arducam Mini 2MP Plus. To do so, open the following -file: - -``` -Arduino/libraries/ArduCAM/memorysaver.h -``` - -You'll see a bunch of `#define` statements listed. Make sure that they are all -commented out, except for `#define OV2640_MINI_2MP_PLUS`, as so: - -``` -//Step 1: select the hardware platform, only one at a time -//#define OV2640_MINI_2MP -//#define OV3640_MINI_3MP -//#define OV5642_MINI_5MP -//#define OV5642_MINI_5MP_BIT_ROTATION_FIXED -#define OV2640_MINI_2MP_PLUS -//#define OV5642_MINI_5MP_PLUS -//#define OV5640_MINI_5MP_PLUS -``` - -Once you save the file, we're done configuring the Arducam library. - -Our next step is to install the JPEGDecoder library. We can do this from within -the Arduino IDE. First, go to the *Manage Libraries...* option in the *Tools* -menu and search for `JPEGDecoder`. You should install version _1.8.0_ of the -library. - -Once the library has installed, we'll need to configure it to disable some -optional components that are not compatible with the Arduino Nano 33 BLE Sense. -Open the following file: - -``` -Arduino/libraries/JPEGDecoder/src/User_Config.h -``` - -Make sure that both `#define LOAD_SD_LIBRARY` and `#define LOAD_SDFAT_LIBRARY` -are commented out, as shown in this excerpt from the file: - -```c++ -// Comment out the next #defines if you are not using an SD Card to store the JPEGs -// Commenting out the line is NOT essential but will save some FLASH space if -// SD Card access is not needed. Note: use of SdFat is currently untested! - -//#define LOAD_SD_LIBRARY // Default SD Card library -//#define LOAD_SDFAT_LIBRARY // Use SdFat library instead, so SD Card SPI can be bit bashed -``` - -Once you've saved the file, you are done installing libraries. - -### Load and run the example - -Go to `File -> Examples`. You should see an -example near the bottom of the list named `TensorFlowLite`. Select -it and click `person_detection` to load the example. Connect your device, then -build and upload the example. - -To test the camera, start by pointing the device's camera at something that is -definitely not a person, or just covering it up. The next time the blue LED -flashes, the device will capture a frame from the camera and begin to run -inference. Since the vision model we are using for person detection is -relatively large, it takes a long time to run inference—around 19 seconds at the -time of writing, though it's possible TensorFlow Lite has gotten faster since -then. - -After 19 seconds or so, the inference result will be translated into another LED -being lit. Since you pointed the camera at something that isn't a person, the -red LED should light up. - -Now, try pointing the device's camera at yourself! The next time the blue LED -flashes, the device will capture another image and begin to run inference. After -19 seconds, the green LED should light up! - -Remember, image data is captured as a snapshot before each inference, whenever -the blue LED flashes. Whatever the camera is pointed at during that moment is -what will be fed into the model. It doesn't matter where the camera is pointed -until the next time an image is captured, when the blue LED will flash again. - -If you're getting seemingly incorrect results, make sure you are in an -environment with good lighting. You should also make sure that the camera is -oriented correctly, with the pins pointing downwards, so that the images it -captures are the right way up—the model was not trained to recognize upside-down -people! In addition, it's good to remember that this is a tiny model, which -trades accuracy for small size. It works very well, but it isn't accurate 100% -of the time. - -We can also see the results of inference via the Arduino Serial Monitor. To do -this, open the *Serial Monitor* from the *Tools* menu. You'll see a detailed -log of what is happening while our application runs. It's also interesting to -check the *Show timestamp* box, so you can see how long each part of the process -takes: - -``` -14:17:50.714 -> Starting capture -14:17:50.714 -> Image captured -14:17:50.784 -> Reading 3080 bytes from ArduCAM -14:17:50.887 -> Finished reading -14:17:50.887 -> Decoding JPEG and converting to greyscale -14:17:51.074 -> Image decoded and processed -14:18:09.710 -> Person score: 246 No person score: 66 -``` - -From the log, we can see that it took around 170 ms to capture and read the -image data from the camera module, 180 ms to decode the JPEG and convert it to -greyscale, and 18.6 seconds to run inference. - -## Running on HIMAX WE1 EVB - -The following instructions will help you build and deploy this example to -[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) -board. To undstand more about using this board, please check -[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). - -### Initial Setup - -To use the HIMAX WE1 EVB, please make sure following software are installed: - -#### MetaWare Development Toolkit - -See -[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) -section for instructions on toolchain installation. - -#### Make Tool version - -A `'make'` tool is required for deploying Tensorflow Lite Micro -applications on HIMAX WE1 EVB, See -[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) -section for proper environment. - -#### Serial Terminal Emulation Application - -There are 2 main purposes for HIMAX WE1 EVB Debug UART port - -- print application output -- burn application to flash by using xmodem send application binary - -You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). - - -### Generate Example Project - -The example project for HIMAX WE1 EVB platform can be generated with the following -command: - -Download related third party data - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads -``` - -Generate person detection project - -``` -make -f tensorflow/lite/micro/tools/make/Makefile generate_person_detection_int8_make_project TARGET=himax_we1_evb -``` - -### Build and Burn Example - -Following the Steps to run person detection example at HIMAX WE1 EVB platform. - -1. Go to the generated example project directory. - - ``` - cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/person_detection_int8/make - ``` - -2. Build the example using - - ``` - make app - ``` - -3. After example build finish, copy ELF file and map file to image generate tool directory. - image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` - - ``` - cp person_detection_int8.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -4. Go to flash image generate tool directory. - - ``` - cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -5. run image generate tool, generate flash image file. - - * Before running image generate tool, by typing `sudo chmod +x image_gen` - and `sudo chmod +x sign_tool` to make sure it is executable. - - ``` - image_gen -e person_detection_int8.elf -m himax_we1_evb.map -o out.img - ``` - - -6. Download flash image file to HIMAX WE1 EVB by UART: - - * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) - -After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial -terminal. - -## Running on SparkFun Edge - -The following instructions will help you build and deploy this sample on the -[SparkFun Edge development board](https://sparkfun.com/products/15170). This -sample requires the Sparkfun Himax camera for the Sparkfun Edge board. It is -not available for purchase yet. - -If you're new to using this board, we recommend walking through the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab to get an understanding of the workflow. - -### Compile the binary - -The following command will download the required dependencies and then compile a -binary for the SparkFun Edge: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge person_detection_bin -``` - -The binary will be created in the following location: - -``` -tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/person_detection.bin -``` - -### Sign the binary - -The binary must be signed with cryptographic keys to be deployed to the device. -We'll now run some commands that will sign our binary so it can be flashed to -the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is -downloaded when the `Makefile` is run. - -Enter the following command to set up some dummy cryptographic keys we can use -for development: - -``` -cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/keys_info0.py \ -tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/keys_info.py -``` - -Next, run the following command to create a signed binary: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/create_cust_image_blob.py \ ---bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/person_detection.bin \ ---load-address 0xC000 \ ---magic-num 0xCB \ --o main_nonsecure_ota \ ---version 0x0 -``` - -This will create the file `main_nonsecure_ota.bin`. We'll now run another -command to create a final version of the file that can be used to flash our -device with the bootloader script we will use in the next step: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ ---load-address 0x20000 \ ---bin main_nonsecure_ota.bin \ --i 6 \ --o main_nonsecure_wire \ ---options 0x1 -``` - -You should now have a file called `main_nonsecure_wire.bin` in the directory -where you ran the commands. This is the file we'll be flashing to the device. - -### Flash the binary - -Next, attach the board to your computer via a USB-to-serial adapter. - -**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), -you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) -before you continue. - -Once connected, assign the USB device name to an environment variable: - -``` -export DEVICENAME=put your device name here -``` - -Set another variable with the baud rate: - -``` -export BAUD_RATE=921600 -``` - -Now, hold the button marked `14` on the device. While still holding the button, -hit the button marked `RST`. Continue holding the button marked `14` while -running the following command: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.0.0/tools/apollo3_scripts/uart_wired_update.py \ --b ${BAUD_RATE} ${DEVICENAME} \ --r 1 \ --f main_nonsecure_wire.bin \ --i 6 -``` - -You should see a long stream of output as the binary is flashed to the device. -Once you see the following lines, flashing is complete: - -``` -Sending Reset Command. -Done. -``` - -If you don't see these lines, flashing may have failed. Try running through the -steps in [Flash the binary](#flash-the-binary) again (you can skip over setting -the environment variables). If you continue to run into problems, follow the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab, which includes more comprehensive instructions for the flashing -process. - -The binary should now be deployed to the device. Hit the button marked `RST` to -reboot the board. You should see the device's four LEDs flashing in sequence. - -Debug information is logged by the board while the program is running. To view -it, establish a serial connection to the board using a baud rate of `115200`. -On OSX and Linux, the following command should work: - -``` -screen ${DEVICENAME} 115200 -``` - -To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately -followed by the `K` key, then hit the `Y` key. - -## Run the tests on a development machine - -To compile and test this example on a desktop Linux or MacOS machine, download -[the TensorFlow source code](https://github.com/tensorflow/tensorflow), `cd` -into the source directory from a terminal, and then run the following command: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile -``` - -This will take a few minutes, and downloads frameworks the code uses like -[CMSIS](https://developer.arm.com/embedded/cmsis) and -[flatbuffers](https://google.github.io/flatbuffers/). Once that process has -finished, run: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile test_person_detection_test -``` - -You should see a series of files get compiled, followed by some logging output -from a test, which should conclude with `~~~ALL TESTS PASSED~~~`. If you see -this, it means that a small program has been built and run that loads a trained -TensorFlow model, runs some example images through it, and got the expected -outputs. This particular test runs images with a and without a person in them, -and checks that the network correctly identifies them. - -To understand how TensorFlow Lite does this, you can look at the `TestInvoke()` -function in -[person_detection_test.cc](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc). -It's a fairly small amount of code, creating an interpreter, getting a handle to -a model that's been compiled into the program, and then invoking the interpreter -with the model and sample inputs. - -## Debugging image capture -When the sample is running, check the LEDs to determine whether the inference is -running correctly. If the red light is stuck on, it means there was an error -communicating with the camera. This is likely due to an incorrectly connected -or broken camera. - -During inference, the blue LED will toggle every time inference is complete. The -orange LED indicates that no person was found, and the green LED indicates a -person was found. The red LED should never turn on, since it indicates an error. - -In order to view the captured image, set the DUMP_IMAGE define in main.cc.  This -causes the board to log raw image info to the console. After the board has been -flashed and reset, dump the log to a text file: - - -``` -screen -L -Logfile ${DEVICENAME} 115200 -``` - -Next, run the raw to bitmap converter to view captured images: - -``` -python3 raw_to_bitmap.py -r GRAY -i -``` - -## Training your own model - -You can train your own model with some easy-to-use scripts. See -[training_a_model.md](training_a_model.md) for instructions. From 39e65d52e400f8c343195e2f8ac34f286648a415 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Thu, 4 Jun 2020 18:19:58 +0800 Subject: [PATCH 0082/1390] remove temp makefile in target --- .../make/targets/himax_we1_evb_makefile.inc~ | 93 ------------------- 1 file changed, 93 deletions(-) delete mode 100644 tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ diff --git a/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ b/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ deleted file mode 100644 index 733f258fbbb..00000000000 --- a/tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc~ +++ /dev/null @@ -1,93 +0,0 @@ -# Settings for himax WE_1 evb. -ifeq ($(TARGET), himax_we1_evb) - - CC_TOOL = ccac - AR_TOOL = arac - CXX_TOOL = ccac - LD_TOOL := ccac - TARGET_ARCH := arc - #ARC_TOOLCHAIN := mwdt - - BUILD_ARC_MLI := false - ARC_MLI_PRE_COMPILED_TARGET := himax_arcem9d_r16 - -include $(MAKEFILE_DIR)/targets/arc/arc_common.inc - #download SDK & MLI - HIMAX_WE1_SDK_NAME := himax_we1_sdk - #MLI_LIB_DIR = arc_mli_package - #MLI_LIB_DIR = arc_mli_package - #$(eval $(call add_third_party_download,$(EMBARC_MLI_PRE_COMPILED_URL),$(EMBARC_MLI_PRE_COMPILED_MD5),$(MLI_LIB_DIR),)) - $(eval $(call add_third_party_download,$(HIMAX_WE1_SDK_URL),$(HIMAX_WE1_SDK_MD5),$(HIMAX_WE1_SDK_NAME),)) - - #export path of toolchain - #export PATH := $(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/image_gen_linux_v3/:$(PATH) - - TCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/arcem9d_wei_r16.tcf - LCF_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/memory.lcf - ARCLIB_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/libembarc.a - LIB_HEADER_FILE := $(PWD)/$(MAKEFILE_DIR)/downloads/$(HIMAX_WE1_SDK_NAME)/hx_drv_tflm.h - - - DEFAULT_HEAPSZ := 8192 - DEFAULT_STACKSZ := 8192 - - TCF_FILE_NAME = $(notdir $(TCF_FILE)) - ARC_TARGET_FILES_DIRS = $(dir $(TCF_FILE_NAME)) - MAKE_PROJECT_FILES += $(TCF_FILE_NAME) - - LCF_FILE_NAME = $(notdir $(LCF_FILE)) - ARC_TARGET_FILES_DIRS += $(dir $(LCF_FILE)) - MAKE_PROJECT_FILES += $(LCF_FILE_NAME) - - ARCLIB_FILE_NAME = $(notdir $(ARCLIB_FILE)) - ARC_TARGET_FILES_DIRS += $(dir $(ARCLIB_FILE)) - MAKE_PROJECT_FILES += $(ARCLIB_FILE_NAME) - - LIB_HEADER_FILE_NAME = $(notdir $(LIB_HEADER_FILE)) - ARC_TARGET_FILES_DIRS += $(dir $(LIB_HEADER_FILE)) - MAKE_PROJECT_FILES += $(LIB_HEADER_FILE_NAME) - - - - # Need a pointer to the TCF and lcf file - - PLATFORM_FLAGS = \ - -DNDEBUG \ - -g \ - -DCPU_ARC \ - -Hnosdata \ - -DTF_LITE_STATIC_MEMORY \ - -tcf=$(TCF_FILE_NAME) \ - -Hnocopyr \ - -Hpurge \ - -Hcl \ - -fslp-vectorize-aggressive \ - -ffunction-sections \ - -fdata-sections \ - -tcf_core_config \ - - CXXFLAGS += -fno-rtti -DSCRATCH_MEM_Z_SIZE=0x10000 $(PLATFORM_FLAGS) - CCFLAGS += $(PLATFORM_FLAGS) - - INCLUDES+= \ - -I $(MAKEFILE_DIR)/downloads/$(WEI_SDK_NAME) \ - -I $(MAKEFILE_DIR)/downloads/kissfft - - GENERATED_PROJECT_INCLUDES += \ - -I. \ - -I./third_party/kissfft - - LDFLAGS += \ - -Hheap=8192 \ - -tcf=$(TCF_FILE_NAME) \ - -Hnocopyr \ - -m \ - -Hldopt=-Coutput=$(TARGET).map \ - $(LCF_FILE_NAME) \ - -Hldopt=-Bgrouplib $(ARCLIB_FILE_NAME) - - CXXFLAGS := $(filter-out -std=c++11,$(CXXFLAGS)) - CCFLAGS := $(filter-out -std=c11,$(CCFLAGS)) - MICROLITE_LIBS := $(filter-out -lm,$(MICROLITE_LIBS)) - -endif From bdf6adfc0e78b5d1e21df527200b9cfaad5830c2 Mon Sep 17 00:00:00 2001 From: Nishidha Panpaliya Date: Thu, 4 Jun 2020 14:52:03 +0000 Subject: [PATCH 0083/1390] Review comments addressed --- tensorflow/compiler/xla/service/cpu/BUILD | 10 +++++++ .../service/cpu/test_target_triple_helper.h | 27 +++++++++++++++++++ .../compiler/xla/service/cpu/tests/BUILD | 4 +++ .../service/cpu/tests/cpu_dyn_shape_test.cc | 3 ++- .../cpu/tests/cpu_eigen_dot_operation_test.cc | 3 ++- .../cpu/tests/cpu_key_value_sort_test.cc | 3 ++- .../cpu/tests/cpu_literal_caching_test.cc | 5 ++-- .../xla/service/cpu/tests/cpu_outfeed_test.cc | 5 ++-- ...ed_reduce_with_no_vector_registers_test.cc | 3 ++- .../xla/tests/local_client_aot_test_helper.cc | 2 ++ 10 files changed, 57 insertions(+), 8 deletions(-) create mode 100644 tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index 3460e65b0a2..7be4d3e724a 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -30,6 +30,15 @@ filegroup( ]), ) +cc_library( + name = "test_header_helper", + testonly = True, + hdrs = ["test_target_triple_helper.h"], + deps = [ + "//tensorflow/core:test", + ], +) + filegroup( name = "single_threaded_runtime_srcs", srcs = [ @@ -1071,6 +1080,7 @@ tf_cc_test( deps = [ ":cpu_compiler", ":cpu_transfer_manager", + ":test_header_helper", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", diff --git a/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h new file mode 100644 index 00000000000..e248f6de8bd --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h @@ -0,0 +1,27 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +#ifndef TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_ +#define TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_ + +#if (defined(__powerpc__) || defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__)) +static const std::string kTargetCpuForHost="ppc"; +static const std::string kTargetTripleForHost="ppc64le-ibm-linux-gnu"; +#else +static const std::string kTargetCpuForHost=""; +static const std::string kTargetTripleForHost="x86_64-pc-linux"; +#endif + +#endif diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD index 1ac8509cdb1..18624330a26 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/BUILD +++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD @@ -42,6 +42,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", + "//tensorflow/compiler/xla/service/cpu:test_header_helper", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -216,6 +217,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", + "//tensorflow/compiler/xla/service/cpu:test_header_helper", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -229,6 +231,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", + "//tensorflow/compiler/xla/service/cpu:test_header_helper", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -242,6 +245,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", + "//tensorflow/compiler/xla/service/cpu:test_header_helper", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_dyn_shape_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_dyn_shape_test.cc index 46249caa0c7..7f9fab5dab0 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_dyn_shape_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_dyn_shape_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" namespace xla { namespace cpu { @@ -46,7 +47,7 @@ TEST_F(CpuDynamicShapeTest, DynamicShapeR2) { )"; CpuAotCompilationOptions options{ - /*triple=*/"x86_64", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc index 8b7f843582b..40314ae5158 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/tests/test_utils.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" #include "tensorflow/core/platform/test.h" namespace xla { @@ -45,7 +46,7 @@ class CpuEigenDotOperationTest void CompileAndCheck(std::unique_ptr entry_computation, const string& filecheck_lines) { CpuAotCompilationOptions options{ - /*triple=*/"x86_64", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc index f3b7b91b2b5..8bb8acb557d 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" namespace xla { namespace cpu { @@ -48,7 +49,7 @@ CHECK: call void @__xla_cpu_runtime_KeyValueSort TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ - /*triple=*/"x86_64", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc index fc670201125..b86e23fc7b6 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" namespace xla { namespace cpu { @@ -64,7 +65,7 @@ CHECK-NOT: private unnamed_addr constant [48 x i8] ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ - /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; @@ -112,7 +113,7 @@ CHECK-NOT: private unnamed_addr constant [8 x i8] ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ - /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc index ad83c485998..ea5f282fb2b 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" namespace xla { namespace cpu { @@ -46,7 +47,7 @@ CHECK: private unnamed_addr constant [48 x i8] TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ - /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; @@ -73,7 +74,7 @@ CHECK: Outfeed TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ - /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", + /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost, /*features=*/"", /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; diff --git a/tensorflow/compiler/xla/service/cpu/vectorized_reduce_with_no_vector_registers_test.cc b/tensorflow/compiler/xla/service/cpu/vectorized_reduce_with_no_vector_registers_test.cc index 754885d8744..e59a531b114 100644 --- a/tensorflow/compiler/xla/service/cpu/vectorized_reduce_with_no_vector_registers_test.cc +++ b/tensorflow/compiler/xla/service/cpu/vectorized_reduce_with_no_vector_registers_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h" namespace xla { namespace { @@ -75,7 +76,7 @@ ENTRY main { // Check that the GetTargetVectorRegisterByteSize is itself working. TF_ASSERT_OK_AND_ASSIGN(unsigned vector_register_byte_size_for_x86_64, - GetTargetVectorRegisterByteSize("x86_64-pc-linux")); + GetTargetVectorRegisterByteSize(kTargetTripleForHost)); ASSERT_EQ(vector_register_byte_size_for_x86_64, 16); std::string triple = "i686-none-android"; diff --git a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc index 53c0d84854e..3e9a3ec2314 100644 --- a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc +++ b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc @@ -71,6 +71,8 @@ int main(int argc, char** argv) { triple_string = "aarch64-none-linux-gnu"; } else if (target_cpu == "x64_windows") { triple_string = "x86_64-pc-windows-msvc19"; + } else if (target_cpu == "ppc") { + triple_string = "ppc64le-ibm-linux-gnu"; } else if (target_cpu == "local") { triple_string = llvm::sys::getDefaultTargetTriple(); } else { From 17b7e169135127e0e866b50577ad8b213abc1d97 Mon Sep 17 00:00:00 2001 From: Dominic Jack Date: Fri, 5 Jun 2020 08:54:38 +1000 Subject: [PATCH 0084/1390] ensure model initialized on ANY trackable attr set --- tensorflow/python/keras/engine/training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py index 29ff31d56db..2b4ba8af3f5 100644 --- a/tensorflow/python/keras/engine/training.py +++ b/tensorflow/python/keras/engine/training.py @@ -319,7 +319,7 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector): super(Model, self).__setattr__(name, value) return - if all( + if any( isinstance(v, (base_layer.Layer, data_structures.TrackableDataStructure)) or trackable_layer_utils.has_weights(v) for v in nest.flatten(value)): From ac123654efc63ffa17240479a5b926ca6357c766 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Fri, 5 Jun 2020 10:03:47 +0800 Subject: [PATCH 0085/1390] TFLM: update hello world example readme --- .../lite/micro/examples/hello_world/README.md | 4 +- .../micro/examples/hello_world/README.md~ | 575 ++++++++++++++++++ 2 files changed, 577 insertions(+), 2 deletions(-) create mode 100644 tensorflow/lite/micro/examples/hello_world/README.md~ diff --git a/tensorflow/lite/micro/examples/hello_world/README.md b/tensorflow/lite/micro/examples/hello_world/README.md index d3762ada790..26b0f12c83a 100644 --- a/tensorflow/lite/micro/examples/hello_world/README.md +++ b/tensorflow/lite/micro/examples/hello_world/README.md @@ -17,7 +17,7 @@ of the device. - [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) - [Deploy to Arduino](#deploy-to-arduino) - [Deploy to ESP32](#deploy-to-esp32) -- [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) +- [Deploy to Himax WE1 EVB](#deploy-to-himax-we1-evb) - [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) - [Deploy to STM32F746](#deploy-to-STM32F746) - [Run the tests on a development machine](#run-the-tests-on-a-development-machine) @@ -192,7 +192,7 @@ The previous two commands can be combined: idf.py --port /dev/ttyUSB0 flash monitor ``` -## Deploy to himax WE1 EVB +## Deploy to Himax WE1 EVB The following instructions will help you build and deploy this example to [HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) diff --git a/tensorflow/lite/micro/examples/hello_world/README.md~ b/tensorflow/lite/micro/examples/hello_world/README.md~ new file mode 100644 index 00000000000..d3762ada790 --- /dev/null +++ b/tensorflow/lite/micro/examples/hello_world/README.md~ @@ -0,0 +1,575 @@ +# Hello World Example + +This example is designed to demonstrate the absolute basics of using [TensorFlow +Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers). +It includes the full end-to-end workflow of training a model, converting it for +use with TensorFlow Lite for Microcontrollers for running inference on a +microcontroller. + +The model is trained to replicate a `sine` function and generates a pattern of +data to either blink LEDs or control an animation, depending on the capabilities +of the device. + +![Animation on STM32F746](images/animation_on_STM32F746.gif) + +## Table of contents + +- [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) +- [Deploy to Arduino](#deploy-to-arduino) +- [Deploy to ESP32](#deploy-to-esp32) +- [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) +- [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) +- [Deploy to STM32F746](#deploy-to-STM32F746) +- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) +- [Train your own model](#train-your-own-model) + +## Deploy to ARC EM SDP + +The following instructions will help you to build and deploy this example to +[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) +board. General information and instructions on using the board with TensorFlow +Lite Micro can be found in the common +[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). + +### Initial Setup + +Follow the instructions on the +[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) +to get and install all required tools for work with ARC EM SDP. + +### Generate Example Project + +The example project for ARC EM SDP platform can be generated with the following +command: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp TAGS=no_arc_mli generate_hello_world_make_project +``` + +### Build and Run Example + +For more detailed information on building and running examples see the +appropriate sections of general descriptions of the +[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). +In the directory with generated project you can also find a +*README_ARC_EMSDP.md* file with instructions and options on building and +running. Here we only briefly mention main steps which are typically enough to +get it started. + +1. You need to + [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) + and open an serial connection. + +2. Go to the generated example project director + + ``` + cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/hello_world/make + ``` + +3. Build the example using + + ``` + make app + ``` + +4. To generate artefacts for self-boot of example from the board use + + ``` + make flash + ``` + +5. To run application from the board using microSD card: + + * Copy the content of the created /bin folder into the root of microSD + card. Note that the card must be formatted as FAT32 with default cluster + size (but less than 32 Kbytes) + * Plug in the microSD card into the J11 connector. + * Push the RST button. If a red LED is lit beside RST button, push the CFG + button. + +6. If you have the MetaWare Debugger installed in your environment: + + * To run application from the console using it type `make run`. + * To stop the execution type `Ctrl+C` in the console several times. + +In both cases (step 5 and 6) you will see the application output in the serial +terminal. + +## Deploy to Arduino + +The following instructions will help you build and deploy this sample +to [Arduino](https://www.arduino.cc/) devices. + +![Animation on Arduino MKRZERO](images/animation_on_arduino_mkrzero.gif) + +The sample has been tested with the following devices: + +- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) +- [Arduino MKRZERO](https://store.arduino.cc/usa/arduino-mkrzero) + +The sample will use PWM to fade an LED on and off according to the model's +output. In the code, the `LED_BUILTIN` constant is used to specify the board's +built-in LED as the one being controlled. However, on some boards, this built-in +LED is not attached to a pin with PWM capabilities. In this case, the LED will +blink instead of fading. + +### Install the Arduino_TensorFlowLite library + +This example application is included as part of the official TensorFlow Lite +Arduino library. To install it, open the Arduino library manager in +`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. + +### Load and run the example + +Once the library has been added, go to `File -> Examples`. You should see an +example near the bottom of the list named `TensorFlowLite:hello_world`. Select +it and click `hello_world` to load the example. + +Use the Arduino IDE to build and upload the example. Once it is running, +you should see the built-in LED on your device flashing. + +The Arduino Desktop IDE includes a plotter that we can use to display the sine +wave graphically. To view it, go to `Tools -> Serial Plotter`. You will see one +datapoint being logged for each inference cycle, expressed as a number between 0 +and 255. + +## Deploy to ESP32 + +The following instructions will help you build and deploy this sample +to [ESP32](https://www.espressif.com/en/products/hardware/esp32/overview) +devices using the [ESP IDF](https://github.com/espressif/esp-idf). + +The sample has been tested on ESP-IDF version 4.0 with the following devices: +- [ESP32-DevKitC](http://esp-idf.readthedocs.io/en/latest/get-started/get-started-devkitc.html) +- [ESP-EYE](https://github.com/espressif/esp-who/blob/master/docs/en/get-started/ESP-EYE_Getting_Started_Guide.md) + +### Install the ESP IDF + +Follow the instructions of the +[ESP-IDF get started guide](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html) +to setup the toolchain and the ESP-IDF itself. + +The next steps assume that the +[IDF environment variables are set](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html#step-4-set-up-the-environment-variables) : + + * The `IDF_PATH` environment variable is set + * `idf.py` and Xtensa-esp32 tools (e.g. `xtensa-esp32-elf-gcc`) are in `$PATH` + +### Generate the examples +The example project can be generated with the following command: +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=esp generate_hello_world_esp_project +``` + +### Building the example + +Go the the example project directory +``` +cd tensorflow/lite/micro/tools/make/gen/esp_xtensa-esp32/prj/hello_world/esp-idf +``` + +Then build with `idf.py` +``` +idf.py build +``` + +### Load and run the example + +To flash (replace `/dev/ttyUSB0` with the device serial port): +``` +idf.py --port /dev/ttyUSB0 flash +``` + +Monitor the serial output: +``` +idf.py --port /dev/ttyUSB0 monitor +``` + +Use `Ctrl+]` to exit. + +The previous two commands can be combined: +``` +idf.py --port /dev/ttyUSB0 flash monitor +``` + +## Deploy to himax WE1 EVB + +The following instructions will help you build and deploy this example to +[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) +board. To undstand more about using this board, please check +[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). + +### Initial Setup + +To use the HIMAX WE1 EVB, please make sure following software are installed: + +#### MetaWare Development Toolkit + +See +[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) +section for instructions on toolchain installation. + +#### Make Tool version + +A `'make'` tool is required for deploying Tensorflow Lite Micro +applications on HIMAX WE1 EVB, See +[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) +section for proper environment. + +#### Serial Terminal Emulation Application + +There are 2 main purposes for HIMAX WE1 EVB Debug UART port + +- print application output +- burn application to flash by using xmodem send application binary + +You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). + + +### Generate Example Project + +The example project for HIMAX WE1 EVB platform can be generated with the following +command: + +Download related third party data + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads +``` + +Generate hello world project + +``` +make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb TAGS=no_arc_mli +``` + +### Build and Burn Example + +Following the Steps to run hello world example at HIMAX WE1 EVB platform. + +1. Go to the generated example project directory. + + ``` + cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/hello_world/make + ``` + +2. Build the example using + + ``` + make app + ``` + +3. After example build finish, copy ELF file and map file to image generate tool directory. + image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` + + ``` + cp hello_world.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +4. Go to flash image generate tool directory. + + ``` + cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ + ``` + +5. run image generate tool, generate flash image file. + + * Before running image generate tool, by typing `sudo chmod +x image_gen` + and `sudo chmod +x sign_tool` to make sure it is executable. + + ``` + image_gen -e hello_world.elf -m himax_we1_evb.map -o out.img + ``` + + +6. Download flash image file to HIMAX WE1 EVB by UART: + + * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) + +After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial +terminal. + +## Deploy to SparkFun Edge + +The following instructions will help you build and deploy this sample on the +[SparkFun Edge development board](https://sparkfun.com/products/15170). + +![Animation on SparkFun Edge](images/animation_on_sparkfun_edge.gif) + +If you're new to using this board, we recommend walking through the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab to get an understanding of the workflow. + +### Compile the binary + +The following command will download the required dependencies and then compile a +binary for the SparkFun Edge: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge hello_world_bin +``` + +The binary will be created in the following location: + +``` +tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin +``` + +### Sign the binary + +The binary must be signed with cryptographic keys to be deployed to the device. +We'll now run some commands that will sign our binary so it can be flashed to +the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is +downloaded when the `Makefile` is run. + +Enter the following command to set up some dummy cryptographic keys we can use +for development: + +``` +cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info0.py \ +tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info.py +``` + +Next, run the following command to create a signed binary: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_image_blob.py \ +--bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin \ +--load-address 0xC000 \ +--magic-num 0xCB \ +-o main_nonsecure_ota \ +--version 0x0 +``` + +This will create the file `main_nonsecure_ota.bin`. We'll now run another +command to create a final version of the file that can be used to flash our +device with the bootloader script we will use in the next step: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ +--load-address 0x20000 \ +--bin main_nonsecure_ota.bin \ +-i 6 \ +-o main_nonsecure_wire \ +--options 0x1 +``` + +You should now have a file called `main_nonsecure_wire.bin` in the directory +where you ran the commands. This is the file we'll be flashing to the device. + +### Flash the binary + +Next, attach the board to your computer via a USB-to-serial adapter. + +**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), +you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) +before you continue. + +Once connected, assign the USB device name to an environment variable: + +``` +export DEVICENAME=put your device name here +``` + +Set another variable with the baud rate: + +``` +export BAUD_RATE=921600 +``` + +Now, hold the button marked `14` on the device. While still holding the button, +hit the button marked `RST`. Continue holding the button marked `14` while +running the following command: + +``` +python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/uart_wired_update.py \ +-b ${BAUD_RATE} ${DEVICENAME} \ +-r 1 \ +-f main_nonsecure_wire.bin \ +-i 6 +``` + +You should see a long stream of output as the binary is flashed to the device. +Once you see the following lines, flashing is complete: + +``` +Sending Reset Command. +Done. +``` + +If you don't see these lines, flashing may have failed. Try running through the +steps in [Flash the binary](#flash-the-binary) again (you can skip over setting +the environment variables). If you continue to run into problems, follow the +[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) +codelab, which includes more comprehensive instructions for the flashing +process. + +The binary should now be deployed to the device. Hit the button marked `RST` to +reboot the board. You should see the device's four LEDs flashing in sequence. + +Debug information is logged by the board while the program is running. To view +it, establish a serial connection to the board using a baud rate of `115200`. +On OSX and Linux, the following command should work: + +``` +screen ${DEVICENAME} 115200 +``` + +You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, +immediately followed by `Esc`. You can then use the arrow keys to explore the +output, which will contain the results of running inference on various `x` +values: + +``` +x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 +``` + +To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately +followed by the `K` key, then hit the `Y` key. + + +## Deploy to STM32F746 + +The following instructions will help you build and deploy the sample to the +[STM32F7 discovery kit](https://os.mbed.com/platforms/ST-Discovery-F746NG/) +using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). + +![Animation on STM32F746](images/animation_on_STM32F746.gif) + +Before we begin, you'll need the following: + +- STM32F7 discovery kit board +- Mini-USB cable +- ARM Mbed CLI ([installation instructions](https://os.mbed.com/docs/mbed-os/v5.12/tools/installation-and-setup.html)) +- Python 2.7 and pip + +Since Mbed requires a special folder structure for projects, we'll first run a +command to generate a subfolder containing the required source files in this +structure: + +``` +make -f tensorflow/lite/micro/tools/make/Makefile TARGET=mbed TAGS="CMSIS disco_f746ng" generate_hello_world_mbed_project +``` + +This will result in the creation of a new folder: + +``` +tensorflow/lite/micro/tools/make/gen/mbed_cortex-m4/prj/hello_world/mbed +``` + +This folder contains all of the example's dependencies structured in the correct +way for Mbed to be able to build it. + +Change into the directory and run the following commands, making sure you are +using Python 2.7.15. + +First, tell Mbed that the current directory is the root of an Mbed project: + +``` +mbed config root . +``` + +Next, tell Mbed to download the dependencies and prepare to build: + +``` +mbed deploy +``` + +By default, Mbed will build the project using C++98. However, TensorFlow Lite +requires C++11. Run the following Python snippet to modify the Mbed +configuration files so that it uses C++11: + +``` +python -c 'import fileinput, glob; +for filename in glob.glob("mbed-os/tools/profiles/*.json"): + for line in fileinput.input(filename, inplace=True): + print line.replace("\"-std=gnu++98\"","\"-std=c++11\", \"-fpermissive\"")' + +``` + +Finally, run the following command to compile: + +``` +mbed compile -m DISCO_F746NG -t GCC_ARM +``` + +This should result in a binary at the following path: + +``` +./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin +``` + +To deploy, plug in your STM board and copy the file to it. On MacOS, you can do +this with the following command: + +``` +cp ./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin /Volumes/DIS_F746NG/ +``` + +Copying the file will initiate the flashing process. Once this is complete, you +should see an animation on the device's screen. + + +``` +screen /dev/tty.usbmodem14403 9600 +``` + +In addition to this animation, debug information is logged by the board while +the program is running. To view it, establish a serial connection to the board +using a baud rate of `9600`. On OSX and Linux, the following command should +work, replacing `/dev/tty.devicename` with the name of your device as it appears +in `/dev`: + +``` +screen /dev/tty.devicename 9600 +``` + +You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, +immediately followed by `Esc`. You can then use the arrow keys to explore the +output, which will contain the results of running inference on various `x` +values: + +``` +x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 +``` + +To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately +followed by the `K` key, then hit the `Y` key. + +### Run the tests on a development machine + +To compile and test this example on a desktop Linux or macOS machine, first +clone the TensorFlow repository from GitHub to a convenient place: + +```bash +git clone --depth 1 https://github.com/tensorflow/tensorflow.git +``` + +Next, `cd` into the source directory from a terminal, and then run the following +command: + +```bash +make -f tensorflow/lite/micro/tools/make/Makefile test_hello_world_test +``` + +This will take a few minutes, and downloads frameworks the code uses. Once the +process has finished, you should see a series of files get compiled, followed by +some logging output from a test, which should conclude with +`~~~ALL TESTS PASSED~~~`. + +If you see this, it means that a small program has been built and run that loads +the trained TensorFlow model, runs some example inputs through it, and got the +expected outputs. + +To understand how TensorFlow Lite does this, you can look at the source in +[hello_world_test.cc](hello_world_test.cc). +It's a fairly small amount of code that creates an interpreter, gets a handle to +a model that's been compiled into the program, and then invokes the interpreter +with the model and sample inputs. + +### Train your own model + +So far you have used an existing trained model to run inference on +microcontrollers. If you wish to train your own model, follow the instructions +given in the [train/](train/) directory. + From 1c26e6abd76fe700ecf87d892ceed1dc5bfa90d3 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Fri, 5 Jun 2020 10:12:46 +0800 Subject: [PATCH 0086/1390] TFLM: delete temp readme file in hello world example --- .../micro/examples/hello_world/README.md~ | 575 ------------------ 1 file changed, 575 deletions(-) delete mode 100644 tensorflow/lite/micro/examples/hello_world/README.md~ diff --git a/tensorflow/lite/micro/examples/hello_world/README.md~ b/tensorflow/lite/micro/examples/hello_world/README.md~ deleted file mode 100644 index d3762ada790..00000000000 --- a/tensorflow/lite/micro/examples/hello_world/README.md~ +++ /dev/null @@ -1,575 +0,0 @@ -# Hello World Example - -This example is designed to demonstrate the absolute basics of using [TensorFlow -Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers). -It includes the full end-to-end workflow of training a model, converting it for -use with TensorFlow Lite for Microcontrollers for running inference on a -microcontroller. - -The model is trained to replicate a `sine` function and generates a pattern of -data to either blink LEDs or control an animation, depending on the capabilities -of the device. - -![Animation on STM32F746](images/animation_on_STM32F746.gif) - -## Table of contents - -- [Deploy to ARC EM SDP](#deploy-to-arc-em-sdp) -- [Deploy to Arduino](#deploy-to-arduino) -- [Deploy to ESP32](#deploy-to-esp32) -- [Deploy to himax WE1 EVB](#deploy-to-himax-we1-evb) -- [Deploy to SparkFun Edge](#deploy-to-sparkfun-edge) -- [Deploy to STM32F746](#deploy-to-STM32F746) -- [Run the tests on a development machine](#run-the-tests-on-a-development-machine) -- [Train your own model](#train-your-own-model) - -## Deploy to ARC EM SDP - -The following instructions will help you to build and deploy this example to -[ARC EM SDP](https://www.synopsys.com/dw/ipdir.php?ds=arc-em-software-development-platform) -board. General information and instructions on using the board with TensorFlow -Lite Micro can be found in the common -[ARC targets description](/tensorflow/lite/micro/tools/make/targets/arc/README.md). - -### Initial Setup - -Follow the instructions on the -[ARC EM SDP Initial Setup](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP) -to get and install all required tools for work with ARC EM SDP. - -### Generate Example Project - -The example project for ARC EM SDP platform can be generated with the following -command: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=arc_emsdp TAGS=no_arc_mli generate_hello_world_make_project -``` - -### Build and Run Example - -For more detailed information on building and running examples see the -appropriate sections of general descriptions of the -[ARC EM SDP usage with TFLM](/tensorflow/lite/micro/tools/make/targets/arc/README.md#ARC-EM-Software-Development-Platform-ARC-EM-SDP). -In the directory with generated project you can also find a -*README_ARC_EMSDP.md* file with instructions and options on building and -running. Here we only briefly mention main steps which are typically enough to -get it started. - -1. You need to - [connect the board](/tensorflow/lite/micro/tools/make/targets/arc/README.md#connect-the-board) - and open an serial connection. - -2. Go to the generated example project director - - ``` - cd tensorflow/lite/micro/tools/make/gen/arc_emsdp_arc/prj/hello_world/make - ``` - -3. Build the example using - - ``` - make app - ``` - -4. To generate artefacts for self-boot of example from the board use - - ``` - make flash - ``` - -5. To run application from the board using microSD card: - - * Copy the content of the created /bin folder into the root of microSD - card. Note that the card must be formatted as FAT32 with default cluster - size (but less than 32 Kbytes) - * Plug in the microSD card into the J11 connector. - * Push the RST button. If a red LED is lit beside RST button, push the CFG - button. - -6. If you have the MetaWare Debugger installed in your environment: - - * To run application from the console using it type `make run`. - * To stop the execution type `Ctrl+C` in the console several times. - -In both cases (step 5 and 6) you will see the application output in the serial -terminal. - -## Deploy to Arduino - -The following instructions will help you build and deploy this sample -to [Arduino](https://www.arduino.cc/) devices. - -![Animation on Arduino MKRZERO](images/animation_on_arduino_mkrzero.gif) - -The sample has been tested with the following devices: - -- [Arduino Nano 33 BLE Sense](https://store.arduino.cc/usa/nano-33-ble-sense-with-headers) -- [Arduino MKRZERO](https://store.arduino.cc/usa/arduino-mkrzero) - -The sample will use PWM to fade an LED on and off according to the model's -output. In the code, the `LED_BUILTIN` constant is used to specify the board's -built-in LED as the one being controlled. However, on some boards, this built-in -LED is not attached to a pin with PWM capabilities. In this case, the LED will -blink instead of fading. - -### Install the Arduino_TensorFlowLite library - -This example application is included as part of the official TensorFlow Lite -Arduino library. To install it, open the Arduino library manager in -`Tools -> Manage Libraries...` and search for `Arduino_TensorFlowLite`. - -### Load and run the example - -Once the library has been added, go to `File -> Examples`. You should see an -example near the bottom of the list named `TensorFlowLite:hello_world`. Select -it and click `hello_world` to load the example. - -Use the Arduino IDE to build and upload the example. Once it is running, -you should see the built-in LED on your device flashing. - -The Arduino Desktop IDE includes a plotter that we can use to display the sine -wave graphically. To view it, go to `Tools -> Serial Plotter`. You will see one -datapoint being logged for each inference cycle, expressed as a number between 0 -and 255. - -## Deploy to ESP32 - -The following instructions will help you build and deploy this sample -to [ESP32](https://www.espressif.com/en/products/hardware/esp32/overview) -devices using the [ESP IDF](https://github.com/espressif/esp-idf). - -The sample has been tested on ESP-IDF version 4.0 with the following devices: -- [ESP32-DevKitC](http://esp-idf.readthedocs.io/en/latest/get-started/get-started-devkitc.html) -- [ESP-EYE](https://github.com/espressif/esp-who/blob/master/docs/en/get-started/ESP-EYE_Getting_Started_Guide.md) - -### Install the ESP IDF - -Follow the instructions of the -[ESP-IDF get started guide](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html) -to setup the toolchain and the ESP-IDF itself. - -The next steps assume that the -[IDF environment variables are set](https://docs.espressif.com/projects/esp-idf/en/latest/get-started/index.html#step-4-set-up-the-environment-variables) : - - * The `IDF_PATH` environment variable is set - * `idf.py` and Xtensa-esp32 tools (e.g. `xtensa-esp32-elf-gcc`) are in `$PATH` - -### Generate the examples -The example project can be generated with the following command: -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=esp generate_hello_world_esp_project -``` - -### Building the example - -Go the the example project directory -``` -cd tensorflow/lite/micro/tools/make/gen/esp_xtensa-esp32/prj/hello_world/esp-idf -``` - -Then build with `idf.py` -``` -idf.py build -``` - -### Load and run the example - -To flash (replace `/dev/ttyUSB0` with the device serial port): -``` -idf.py --port /dev/ttyUSB0 flash -``` - -Monitor the serial output: -``` -idf.py --port /dev/ttyUSB0 monitor -``` - -Use `Ctrl+]` to exit. - -The previous two commands can be combined: -``` -idf.py --port /dev/ttyUSB0 flash monitor -``` - -## Deploy to himax WE1 EVB - -The following instructions will help you build and deploy this example to -[HIMAX WE1 EVB](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_board_brief) -board. To undstand more about using this board, please check -[HIMAX WE1 EVB user guide](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide). - -### Initial Setup - -To use the HIMAX WE1 EVB, please make sure following software are installed: - -#### MetaWare Development Toolkit - -See -[Install the Synopsys DesignWare ARC MetaWare Development Toolkit](/tensorflow/lite/micro/tools/make/targets/arc/README.md#install-the-synopsys-designware-arc-metaware-development-toolkit) -section for instructions on toolchain installation. - -#### Make Tool version - -A `'make'` tool is required for deploying Tensorflow Lite Micro -applications on HIMAX WE1 EVB, See -[Check make tool version](/tensorflow/lite/micro/tools/make/targets/arc/README.md#make-tool) -section for proper environment. - -#### Serial Terminal Emulation Application - -There are 2 main purposes for HIMAX WE1 EVB Debug UART port - -- print application output -- burn application to flash by using xmodem send application binary - -You can use any terminal emulation program (like [PuTTY](https://www.putty.org/) or [minicom](https://linux.die.net/man/1/minicom)). - - -### Generate Example Project - -The example project for HIMAX WE1 EVB platform can be generated with the following -command: - -Download related third party data - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=himax_we1_evb third_party_downloads -``` - -Generate hello world project - -``` -make -f tensorflow/lite/micro/tools/make/Makefile generate_hello_world_make_project TARGET=himax_we1_evb TAGS=no_arc_mli -``` - -### Build and Burn Example - -Following the Steps to run hello world example at HIMAX WE1 EVB platform. - -1. Go to the generated example project directory. - - ``` - cd tensorflow/lite/micro/tools/make/gen/himax_we1_evb_arc/prj/hello_world/make - ``` - -2. Build the example using - - ``` - make app - ``` - -3. After example build finish, copy ELF file and map file to image generate tool directory. - image generate tool directory located at `'tensorflow/lite/micro/tools/make/downloads/himax_we1_sdk/image_gen_linux_v3/'` - - ``` - cp hello_world.elf himax_we1_evb.map ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -4. Go to flash image generate tool directory. - - ``` - cd ../../../../../downloads/himax_we1_sdk/image_gen_linux_v3/ - ``` - -5. run image generate tool, generate flash image file. - - * Before running image generate tool, by typing `sudo chmod +x image_gen` - and `sudo chmod +x sign_tool` to make sure it is executable. - - ``` - image_gen -e hello_world.elf -m himax_we1_evb.map -o out.img - ``` - - -6. Download flash image file to HIMAX WE1 EVB by UART: - - * more detail about download image through UART can be found at [HIMAX WE1 EVB update Flash image](https://github.com/HimaxWiseEyePlus/bsp_tflu/tree/master/HIMAX_WE1_EVB_user_guide#flash-image-update) - -After these steps, press reset button on the HIMAX WE1 EVB, you will see application output in the serial -terminal. - -## Deploy to SparkFun Edge - -The following instructions will help you build and deploy this sample on the -[SparkFun Edge development board](https://sparkfun.com/products/15170). - -![Animation on SparkFun Edge](images/animation_on_sparkfun_edge.gif) - -If you're new to using this board, we recommend walking through the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab to get an understanding of the workflow. - -### Compile the binary - -The following command will download the required dependencies and then compile a -binary for the SparkFun Edge: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=sparkfun_edge hello_world_bin -``` - -The binary will be created in the following location: - -``` -tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin -``` - -### Sign the binary - -The binary must be signed with cryptographic keys to be deployed to the device. -We'll now run some commands that will sign our binary so it can be flashed to -the SparkFun Edge. The scripts we are using come from the Ambiq SDK, which is -downloaded when the `Makefile` is run. - -Enter the following command to set up some dummy cryptographic keys we can use -for development: - -``` -cp tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info0.py \ -tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/keys_info.py -``` - -Next, run the following command to create a signed binary: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_image_blob.py \ ---bin tensorflow/lite/micro/tools/make/gen/sparkfun_edge_cortex-m4/bin/hello_world.bin \ ---load-address 0xC000 \ ---magic-num 0xCB \ --o main_nonsecure_ota \ ---version 0x0 -``` - -This will create the file `main_nonsecure_ota.bin`. We'll now run another -command to create a final version of the file that can be used to flash our -device with the bootloader script we will use in the next step: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/create_cust_wireupdate_blob.py \ ---load-address 0x20000 \ ---bin main_nonsecure_ota.bin \ --i 6 \ --o main_nonsecure_wire \ ---options 0x1 -``` - -You should now have a file called `main_nonsecure_wire.bin` in the directory -where you ran the commands. This is the file we'll be flashing to the device. - -### Flash the binary - -Next, attach the board to your computer via a USB-to-serial adapter. - -**Note:** If you're using the [SparkFun Serial Basic Breakout](https://www.sparkfun.com/products/15096), -you should [install the latest drivers](https://learn.sparkfun.com/tutorials/sparkfun-serial-basic-ch340c-hookup-guide#drivers-if-you-need-them) -before you continue. - -Once connected, assign the USB device name to an environment variable: - -``` -export DEVICENAME=put your device name here -``` - -Set another variable with the baud rate: - -``` -export BAUD_RATE=921600 -``` - -Now, hold the button marked `14` on the device. While still holding the button, -hit the button marked `RST`. Continue holding the button marked `14` while -running the following command: - -``` -python3 tensorflow/lite/micro/tools/make/downloads/AmbiqSuite-Rel2.2.0/tools/apollo3_scripts/uart_wired_update.py \ --b ${BAUD_RATE} ${DEVICENAME} \ --r 1 \ --f main_nonsecure_wire.bin \ --i 6 -``` - -You should see a long stream of output as the binary is flashed to the device. -Once you see the following lines, flashing is complete: - -``` -Sending Reset Command. -Done. -``` - -If you don't see these lines, flashing may have failed. Try running through the -steps in [Flash the binary](#flash-the-binary) again (you can skip over setting -the environment variables). If you continue to run into problems, follow the -[AI on a microcontroller with TensorFlow Lite and SparkFun Edge](https://codelabs.developers.google.com/codelabs/sparkfun-tensorflow) -codelab, which includes more comprehensive instructions for the flashing -process. - -The binary should now be deployed to the device. Hit the button marked `RST` to -reboot the board. You should see the device's four LEDs flashing in sequence. - -Debug information is logged by the board while the program is running. To view -it, establish a serial connection to the board using a baud rate of `115200`. -On OSX and Linux, the following command should work: - -``` -screen ${DEVICENAME} 115200 -``` - -You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, -immediately followed by `Esc`. You can then use the arrow keys to explore the -output, which will contain the results of running inference on various `x` -values: - -``` -x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 -``` - -To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately -followed by the `K` key, then hit the `Y` key. - - -## Deploy to STM32F746 - -The following instructions will help you build and deploy the sample to the -[STM32F7 discovery kit](https://os.mbed.com/platforms/ST-Discovery-F746NG/) -using [ARM Mbed](https://github.com/ARMmbed/mbed-cli). - -![Animation on STM32F746](images/animation_on_STM32F746.gif) - -Before we begin, you'll need the following: - -- STM32F7 discovery kit board -- Mini-USB cable -- ARM Mbed CLI ([installation instructions](https://os.mbed.com/docs/mbed-os/v5.12/tools/installation-and-setup.html)) -- Python 2.7 and pip - -Since Mbed requires a special folder structure for projects, we'll first run a -command to generate a subfolder containing the required source files in this -structure: - -``` -make -f tensorflow/lite/micro/tools/make/Makefile TARGET=mbed TAGS="CMSIS disco_f746ng" generate_hello_world_mbed_project -``` - -This will result in the creation of a new folder: - -``` -tensorflow/lite/micro/tools/make/gen/mbed_cortex-m4/prj/hello_world/mbed -``` - -This folder contains all of the example's dependencies structured in the correct -way for Mbed to be able to build it. - -Change into the directory and run the following commands, making sure you are -using Python 2.7.15. - -First, tell Mbed that the current directory is the root of an Mbed project: - -``` -mbed config root . -``` - -Next, tell Mbed to download the dependencies and prepare to build: - -``` -mbed deploy -``` - -By default, Mbed will build the project using C++98. However, TensorFlow Lite -requires C++11. Run the following Python snippet to modify the Mbed -configuration files so that it uses C++11: - -``` -python -c 'import fileinput, glob; -for filename in glob.glob("mbed-os/tools/profiles/*.json"): - for line in fileinput.input(filename, inplace=True): - print line.replace("\"-std=gnu++98\"","\"-std=c++11\", \"-fpermissive\"")' - -``` - -Finally, run the following command to compile: - -``` -mbed compile -m DISCO_F746NG -t GCC_ARM -``` - -This should result in a binary at the following path: - -``` -./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin -``` - -To deploy, plug in your STM board and copy the file to it. On MacOS, you can do -this with the following command: - -``` -cp ./BUILD/DISCO_F746NG/GCC_ARM/mbed.bin /Volumes/DIS_F746NG/ -``` - -Copying the file will initiate the flashing process. Once this is complete, you -should see an animation on the device's screen. - - -``` -screen /dev/tty.usbmodem14403 9600 -``` - -In addition to this animation, debug information is logged by the board while -the program is running. To view it, establish a serial connection to the board -using a baud rate of `9600`. On OSX and Linux, the following command should -work, replacing `/dev/tty.devicename` with the name of your device as it appears -in `/dev`: - -``` -screen /dev/tty.devicename 9600 -``` - -You will see a lot of output flying past! To stop the scrolling, hit `Ctrl+A`, -immediately followed by `Esc`. You can then use the arrow keys to explore the -output, which will contain the results of running inference on various `x` -values: - -``` -x_value: 1.1843798*2^2, y_value: -1.9542645*2^-1 -``` - -To stop viewing the debug output with `screen`, hit `Ctrl+A`, immediately -followed by the `K` key, then hit the `Y` key. - -### Run the tests on a development machine - -To compile and test this example on a desktop Linux or macOS machine, first -clone the TensorFlow repository from GitHub to a convenient place: - -```bash -git clone --depth 1 https://github.com/tensorflow/tensorflow.git -``` - -Next, `cd` into the source directory from a terminal, and then run the following -command: - -```bash -make -f tensorflow/lite/micro/tools/make/Makefile test_hello_world_test -``` - -This will take a few minutes, and downloads frameworks the code uses. Once the -process has finished, you should see a series of files get compiled, followed by -some logging output from a test, which should conclude with -`~~~ALL TESTS PASSED~~~`. - -If you see this, it means that a small program has been built and run that loads -the trained TensorFlow model, runs some example inputs through it, and got the -expected outputs. - -To understand how TensorFlow Lite does this, you can look at the source in -[hello_world_test.cc](hello_world_test.cc). -It's a fairly small amount of code that creates an interpreter, gets a handle to -a model that's been compiled into the program, and then invokes the interpreter -with the model and sample inputs. - -### Train your own model - -So far you have used an existing trained model to run inference on -microcontrollers. If you wish to train your own model, follow the instructions -given in the [train/](train/) directory. - From d6676205f20e6a9476f6e0eca8f5b00367f9c623 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Fri, 5 Jun 2020 10:16:04 +0800 Subject: [PATCH 0087/1390] TFLM: remove temp file in person detection example --- .../himax_we1_evb/image_provider.cc~ | 44 ------------------- 1 file changed, 44 deletions(-) delete mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ deleted file mode 100644 index d5b4d136642..00000000000 --- a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc~ +++ /dev/null @@ -1,44 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -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 - - http://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. -==============================================================================*/ - -#include "tensorflow/lite/micro/examples/person_detection_experimental/image_provider.h" - -#include "tensorflow/lite/micro/examples/person_detection_experimental/model_settings.h" - -#include "hx_drv_tflm.h" - -hx_drv_sensor_image_config_t g_pimg_config; - - -TfLiteStatus GetImage(tflite::ErrorReporter* error_reporter, int image_width, - int image_height, int channels, int8_t* image_data) { - static bool is_initialized = false; - - if (!is_initialized) { - if(hx_drv_sensor_initial(&g_pimg_config)!= HX_DRV_LIB_PASS) - { - return kTfLiteError; - } - is_initialized = true; - } - - hx_drv_sensor_capture(&g_pimg_config); - - hx_drv_image_rescale((uint8_t*)g_pimg_config.raw_address, g_pimg_config.img_width, g_pimg_config.img_height, - image_data, image_data, image_height); - - - return kTfLiteOk; -} From 4f222bed159eeb2743fb8c97fa5e36d5bea4e5da Mon Sep 17 00:00:00 2001 From: Nishidha Panpaliya Date: Fri, 5 Jun 2020 07:49:19 +0000 Subject: [PATCH 0088/1390] Fixed copyright --- tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h index e248f6de8bd..a28d48b7eb0 100644 --- a/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h +++ b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h @@ -1,4 +1,4 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. From 66eb52eaee54576c7c1b5fe887d5d0400b557f86 Mon Sep 17 00:00:00 2001 From: Nishidha Panpaliya Date: Fri, 5 Jun 2020 08:19:55 +0000 Subject: [PATCH 0089/1390] Fixed build error in one of the xla tests --- tensorflow/compiler/xla/service/cpu/tests/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD index 18624330a26..9036c5c9024 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/BUILD +++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD @@ -138,6 +138,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test", "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/service/cpu:test_header_helper", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", From 4d7aadeda14cdc41606faaf2bf397a3904847d7d Mon Sep 17 00:00:00 2001 From: nihui Date: Fri, 5 Jun 2020 16:49:23 +0800 Subject: [PATCH 0090/1390] Update tf_generated_ops.td --- tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td index 1c9297485b4..467a119c174 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td @@ -8651,11 +8651,11 @@ particular, begin = [1, 2, x, x, 0, x] # x denotes don't care (usually 0) end = [2, 4, x, x, -3, x] strides = [1, 1, x, x, -1, 1] -begin_mask = 1<<4 | 1 << 5 = 48 +begin_mask = 1<<4 | 1<<5 = 48 end_mask = 1<<5 = 32 ellipsis_mask = 1<<3 = 8 -new_axis_mask = 1<<2 4 -shrink_axis_mask = 1<<0 +new_axis_mask = 1<<2 = 4 +shrink_axis_mask = 1<<0 = 1 ``` In this case if `foo.shape` is (5, 5, 5, 5, 5, 5) the final shape of From d8bb6569bc6c820b098fc0b56c3c1a7a318422a9 Mon Sep 17 00:00:00 2001 From: Yixing Fu Date: Fri, 5 Jun 2020 18:55:32 -0400 Subject: [PATCH 0091/1390] correct summing total blocks --- tensorflow/python/keras/applications/efficientnet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/keras/applications/efficientnet.py b/tensorflow/python/keras/applications/efficientnet.py index e1413b08533..e3c1a261e80 100644 --- a/tensorflow/python/keras/applications/efficientnet.py +++ b/tensorflow/python/keras/applications/efficientnet.py @@ -334,7 +334,7 @@ def EfficientNet( blocks_args = copy.deepcopy(blocks_args) b = 0 - blocks = float(sum(args['repeats'] for args in blocks_args)) + blocks = float(sum(round_repeats(args['repeats']) for args in blocks_args)) for (i, args) in enumerate(blocks_args): assert args['repeats'] > 0 # Update block input and output filters based on depth multiplier. From 8b81960fcb60222b03f04751e5f99a36b24d27ca Mon Sep 17 00:00:00 2001 From: nihui Date: Sat, 6 Jun 2020 12:00:15 +0800 Subject: [PATCH 0092/1390] Update api_def_StridedSlice.pbtxt --- tensorflow/core/api_def/base_api/api_def_StridedSlice.pbtxt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/api_def/base_api/api_def_StridedSlice.pbtxt b/tensorflow/core/api_def/base_api/api_def_StridedSlice.pbtxt index 9a89a4e8e75..2714e31ac28 100644 --- a/tensorflow/core/api_def/base_api/api_def_StridedSlice.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_StridedSlice.pbtxt @@ -123,11 +123,11 @@ particular, begin = [1, 2, x, x, 0, x] # x denotes don't care (usually 0) end = [2, 4, x, x, -3, x] strides = [1, 1, x, x, -1, 1] -begin_mask = 1<<4 | 1 << 5 = 48 +begin_mask = 1<<4 | 1<<5 = 48 end_mask = 1<<5 = 32 ellipsis_mask = 1<<3 = 8 -new_axis_mask = 1<<2 4 -shrink_axis_mask = 1<<0 +new_axis_mask = 1<<2 = 4 +shrink_axis_mask = 1<<0 = 1 ``` In this case if `foo.shape` is (5, 5, 5, 5, 5, 5) the final shape of From b8f57874ba13cf24a54878c3f4f6cd3f387a0c44 Mon Sep 17 00:00:00 2001 From: cclauss Date: Sun, 7 Jun 2020 15:56:59 +0200 Subject: [PATCH 0093/1390] Fix SyntaxWarnings on Python >= 3.8 --- tensorflow/python/kernel_tests/matrix_band_part_op_test.py | 2 +- tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py | 2 +- .../python/ops/ragged/ragged_batch_gather_with_default_op.py | 3 +-- tensorflow/python/ops/random_ops.py | 4 ++-- tensorflow/tools/docs/doc_generator_visitor.py | 5 ++--- 5 files changed, 7 insertions(+), 9 deletions(-) diff --git a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py index fdb7e4a1a4e..25b502cf814 100644 --- a/tensorflow/python/kernel_tests/matrix_band_part_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_band_part_op_test.py @@ -56,7 +56,7 @@ def _GetMatrixBandPartTest(dtype_, batch_shape_, shape_): band_np = np.triu(band_np, -lower) if upper >= 0: band_np = np.tril(band_np, upper) - if batch_shape_ is not (): + if batch_shape_ != (): band_np = np.tile(band_np, batch_shape_ + (1, 1)) for index_dtype in [dtypes_lib.int32, dtypes_lib.int64]: with self.cached_session(use_gpu=False): diff --git a/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py b/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py index b7a159e2eff..889ea0dbd6c 100644 --- a/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py +++ b/tensorflow/python/kernel_tests/matrix_solve_ls_op_test.py @@ -107,7 +107,7 @@ class MatrixSolveLsOpTest(test_lib.TestCase): np_ans = _SolveWithNumpy(x, y, l2_regularizer=l2_regularizer) np_r = np.dot(np.conj(a.T), b - np.dot(a, np_ans)) np_r_norm = np.sqrt(np.sum(np.conj(np_r) * np_r)) - if batch_shape is not (): + if batch_shape != (): a = np.tile(a, batch_shape + (1, 1)) b = np.tile(b, batch_shape + (1, 1)) np_ans = np.tile(np_ans, batch_shape + (1, 1)) diff --git a/tensorflow/python/ops/ragged/ragged_batch_gather_with_default_op.py b/tensorflow/python/ops/ragged/ragged_batch_gather_with_default_op.py index 377fd84f96e..06690f86a50 100644 --- a/tensorflow/python/ops/ragged/ragged_batch_gather_with_default_op.py +++ b/tensorflow/python/ops/ragged/ragged_batch_gather_with_default_op.py @@ -81,8 +81,7 @@ def batch_gather_with_default(params, return_dtype=True)) # TODO(hterry): lift this restriction and support default_values of # of rank > 1 - if (default_value.shape.ndims is not 0 - and default_value.shape.ndims is not 1): + if default_value.shape.ndims not in (0, 1): raise ValueError('"default_value" must be a scalar or vector') upper_bounds = None if indices.shape.ndims is None: diff --git a/tensorflow/python/ops/random_ops.py b/tensorflow/python/ops/random_ops.py index 1af91ed0dd3..9932a76b678 100644 --- a/tensorflow/python/ops/random_ops.py +++ b/tensorflow/python/ops/random_ops.py @@ -288,8 +288,8 @@ def random_uniform(shape, shape = tensor_util.shape_tensor(shape) # In case of [0,1) floating results, minval and maxval is unused. We do an # `is` comparison here since this is cheaper than isinstance or __eq__. - minval_is_zero = minval is 0 # pylint: disable=literal-comparison - maxval_is_one = maxval is 1 # pylint: disable=literal-comparison + minval_is_zero = minval == 0 + maxval_is_one = maxval == 1 if not minval_is_zero or not maxval_is_one or dtype.is_integer: minval = ops.convert_to_tensor(minval, dtype=dtype, name="min") maxval = ops.convert_to_tensor(maxval, dtype=dtype, name="max") diff --git a/tensorflow/tools/docs/doc_generator_visitor.py b/tensorflow/tools/docs/doc_generator_visitor.py index ec2102a5935..ac5b09346ec 100644 --- a/tensorflow/tools/docs/doc_generator_visitor.py +++ b/tensorflow/tools/docs/doc_generator_visitor.py @@ -240,10 +240,9 @@ class DocGeneratorVisitor(object): # We cannot use the duplicate mechanism for some constants, since e.g., # id(c1) == id(c2) with c1=1, c2=1. This is unproblematic since constants # have no usable docstring and won't be documented automatically. - if (py_object is not None and + if (py_object not in (None, ()) not isinstance(py_object, six.integer_types + six.string_types + - (six.binary_type, six.text_type, float, complex, bool)) - and py_object is not ()): # pylint: disable=literal-comparison + (six.binary_type, six.text_type, float, complex, bool))): object_id = id(py_object) if object_id in reverse_index: master_name = reverse_index[object_id] From e474fc8ebaff4100c427a75926518fce8eb9807a Mon Sep 17 00:00:00 2001 From: Sean Settle Date: Sun, 7 Jun 2020 19:14:41 -0700 Subject: [PATCH 0094/1390] Re-generated the full Dockerfiles --- .../devel-cpu-arm64v8-jupyter.Dockerfile | 134 ++++++++++++++++++ .../arm64v8/devel-cpu-arm64v8.Dockerfile | 108 ++++++++++++++ .../devel-cpu-ppc64le-jupyter.Dockerfile | 2 +- .../ppc64le/devel-cpu-ppc64le.Dockerfile | 2 +- .../devel-gpu-ppc64le-jupyter.Dockerfile | 2 +- .../ppc64le/devel-gpu-ppc64le.Dockerfile | 2 +- 6 files changed, 246 insertions(+), 4 deletions(-) create mode 100644 tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile create mode 100644 tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile diff --git a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile new file mode 100644 index 00000000000..704b1b344aa --- /dev/null +++ b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile @@ -0,0 +1,134 @@ +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================ +# +# THIS IS A GENERATED DOCKERFILE. +# +# This file was assembled from multiple pieces, whose use is documented +# throughout. Please refer to the TensorFlow dockerfiles documentation +# for more information. + +ARG UBUNTU_VERSION=18.04 + +FROM ubuntu:${UBUNTU_VERSION} AS base + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + curl \ + git \ + libcurl3-dev \ + libfreetype6-dev \ + libhdf5-serial-dev \ + libzmq3-dev \ + pkg-config \ + rsync \ + software-properties-common \ + sudo \ + unzip \ + zip \ + zlib1g-dev \ + openjdk-8-jdk \ + openjdk-8-jre-headless \ + && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +ENV CI_BUILD_PYTHON python + +# CACHE_STOP is used to rerun future commands, otherwise cloning tensorflow will be cached and will not pull the most recent version +ARG CACHE_STOP=1 +# Check out TensorFlow source code if --build-arg CHECKOUT_TF_SRC=1 +ARG CHECKOUT_TF_SRC=0 +# In case of Python 2.7+ we need to add passwd entries for user and group id +RUN chmod a+w /etc/passwd /etc/group +RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src || true + +# See http://bugs.python.org/issue19846 +ENV LANG C.UTF-8 + +RUN apt-get update && apt-get install -y \ + python3 \ + python3-pip + +RUN python3 -m pip --no-cache-dir install --upgrade \ + pip \ + setuptools + +# Some TF tools expect a "python" binary +RUN ln -s $(which python3) /usr/local/bin/python + +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + openjdk-8-jdk \ + python3-dev \ + virtualenv \ + swig + +RUN apt-get update && apt-get install -y \ + python3-pil \ + python3-h5py \ + python3-keras-preprocessing \ + python3-matplotlib \ + python3-mock \ + python3-numpy \ + python3-scipy \ + python3-sklearn \ + python3-pandas \ + python3-portpicker + +RUN python3 -m pip --no-cache-dir install \ + enum34 + +# Build and install bazel +ENV BAZEL_VERSION 3.0.0 +WORKDIR / +RUN mkdir /bazel && \ + cd /bazel && \ + curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip && \ + unzip bazel-$BAZEL_VERSION-dist.zip && \ + bash ./compile.sh && \ + cp output/bazel /usr/local/bin/ && \ + rm -rf /bazel && \ + cd - + +COPY bashrc /etc/bash.bashrc +RUN chmod a+rwx /etc/bash.bashrc + +RUN python3 -m pip install --no-cache-dir jupyter matplotlib +# Pin ipykernel and nbformat; see https://github.com/ipython/ipykernel/issues/422 +RUN python3 -m pip install --no-cache-dir jupyter_http_over_ws ipykernel==5.1.1 nbformat==4.4.0 +RUN jupyter serverextension enable --py jupyter_http_over_ws + +RUN mkdir -p /tf/tensorflow-tutorials && chmod -R a+rwx /tf/ +RUN mkdir /.local && chmod a+rwx /.local +RUN apt-get install -y --no-install-recommends wget +# some examples require git to fetch dependencies +RUN apt-get install -y --no-install-recommends git +WORKDIR /tf/tensorflow-tutorials +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/classification.ipynb +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/overfit_and_underfit.ipynb +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/regression.ipynb +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/save_and_load.ipynb +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/text_classification.ipynb +RUN wget https://raw.githubusercontent.com/tensorflow/docs/master/site/en/tutorials/keras/text_classification_with_hub.ipynb +COPY readme-for-jupyter.md README.md +RUN apt-get autoremove -y && apt-get remove -y wget +WORKDIR /tf +EXPOSE 8888 + +RUN python3 -m ipykernel.kernelspec + +CMD ["bash", "-c", "source /etc/bash.bashrc && jupyter notebook --notebook-dir=/tf --ip 0.0.0.0 --no-browser --allow-root"] diff --git a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile new file mode 100644 index 00000000000..10e4512cf29 --- /dev/null +++ b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile @@ -0,0 +1,108 @@ +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================ +# +# THIS IS A GENERATED DOCKERFILE. +# +# This file was assembled from multiple pieces, whose use is documented +# throughout. Please refer to the TensorFlow dockerfiles documentation +# for more information. + +ARG UBUNTU_VERSION=18.04 + +FROM ubuntu:${UBUNTU_VERSION} AS base + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + curl \ + git \ + libcurl3-dev \ + libfreetype6-dev \ + libhdf5-serial-dev \ + libzmq3-dev \ + pkg-config \ + rsync \ + software-properties-common \ + sudo \ + unzip \ + zip \ + zlib1g-dev \ + openjdk-8-jdk \ + openjdk-8-jre-headless \ + && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +ENV CI_BUILD_PYTHON python + +# CACHE_STOP is used to rerun future commands, otherwise cloning tensorflow will be cached and will not pull the most recent version +ARG CACHE_STOP=1 +# Check out TensorFlow source code if --build-arg CHECKOUT_TF_SRC=1 +ARG CHECKOUT_TF_SRC=0 +# In case of Python 2.7+ we need to add passwd entries for user and group id +RUN chmod a+w /etc/passwd /etc/group +RUN test "${CHECKOUT_TF_SRC}" -eq 1 && git clone https://github.com/tensorflow/tensorflow.git /tensorflow_src || true + +# See http://bugs.python.org/issue19846 +ENV LANG C.UTF-8 + +RUN apt-get update && apt-get install -y \ + python3 \ + python3-pip + +RUN python3 -m pip --no-cache-dir install --upgrade \ + pip \ + setuptools + +# Some TF tools expect a "python" binary +RUN ln -s $(which python3) /usr/local/bin/python + +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + openjdk-8-jdk \ + python3-dev \ + virtualenv \ + swig + +RUN apt-get update && apt-get install -y \ + python3-pil \ + python3-h5py \ + python3-keras-preprocessing \ + python3-matplotlib \ + python3-mock \ + python3-numpy \ + python3-scipy \ + python3-sklearn \ + python3-pandas \ + python3-portpicker + +RUN python3 -m pip --no-cache-dir install \ + enum34 + +# Build and install bazel +ENV BAZEL_VERSION 3.0.0 +WORKDIR / +RUN mkdir /bazel && \ + cd /bazel && \ + curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip && \ + unzip bazel-$BAZEL_VERSION-dist.zip && \ + bash ./compile.sh && \ + cp output/bazel /usr/local/bin/ && \ + rm -rf /bazel && \ + cd - + +COPY bashrc /etc/bash.bashrc +RUN chmod a+rwx /etc/bash.bashrc diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le-jupyter.Dockerfile index 53ccffd1403..905faca7893 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le-jupyter.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le-jupyter.Dockerfile @@ -90,7 +90,7 @@ RUN python3 -m pip --no-cache-dir install \ portpicker \ enum34 - # Build and install bazel +# Build and install bazel ENV BAZEL_VERSION 3.0.0 WORKDIR / RUN mkdir /bazel && \ diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le.Dockerfile index 1bbe7129479..378c5f8279b 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-cpu-ppc64le.Dockerfile @@ -90,7 +90,7 @@ RUN python3 -m pip --no-cache-dir install \ portpicker \ enum34 - # Build and install bazel +# Build and install bazel ENV BAZEL_VERSION 3.0.0 WORKDIR / RUN mkdir /bazel && \ diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile index 0700a354d3c..083ce05d2b2 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le-jupyter.Dockerfile @@ -132,7 +132,7 @@ RUN python3 -m pip --no-cache-dir install \ portpicker \ enum34 - # Build and install bazel +# Build and install bazel ENV BAZEL_VERSION 3.0.0 WORKDIR / RUN mkdir /bazel && \ diff --git a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile index b6d8ff8b90e..2c13e47c257 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/ppc64le/devel-gpu-ppc64le.Dockerfile @@ -132,7 +132,7 @@ RUN python3 -m pip --no-cache-dir install \ portpicker \ enum34 - # Build and install bazel +# Build and install bazel ENV BAZEL_VERSION 3.0.0 WORKDIR / RUN mkdir /bazel && \ From 7fea4ccf1c91ff52ef1fbd36db6a1743d147e253 Mon Sep 17 00:00:00 2001 From: Sean Settle Date: Sun, 7 Jun 2020 20:34:46 -0700 Subject: [PATCH 0095/1390] Re-generated the full Dockerfiles --- .../dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile | 2 +- .../dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile | 2 +- .../partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile index 704b1b344aa..168e57d363a 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8-jupyter.Dockerfile @@ -80,7 +80,6 @@ RUN apt-get update && apt-get install -y \ RUN apt-get update && apt-get install -y \ python3-pil \ python3-h5py \ - python3-keras-preprocessing \ python3-matplotlib \ python3-mock \ python3-numpy \ @@ -90,6 +89,7 @@ RUN apt-get update && apt-get install -y \ python3-portpicker RUN python3 -m pip --no-cache-dir install \ + keras_preprocessing \ enum34 # Build and install bazel diff --git a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile index 10e4512cf29..70d6df8df14 100644 --- a/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile +++ b/tensorflow/tools/dockerfiles/dockerfiles/arm64v8/devel-cpu-arm64v8.Dockerfile @@ -80,7 +80,6 @@ RUN apt-get update && apt-get install -y \ RUN apt-get update && apt-get install -y \ python3-pil \ python3-h5py \ - python3-keras-preprocessing \ python3-matplotlib \ python3-mock \ python3-numpy \ @@ -90,6 +89,7 @@ RUN apt-get update && apt-get install -y \ python3-portpicker RUN python3 -m pip --no-cache-dir install \ + keras_preprocessing \ enum34 # Build and install bazel diff --git a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile index b833657aa69..2f923e84737 100644 --- a/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile +++ b/tensorflow/tools/dockerfiles/partials/ubuntu/bazelbuild-arm64v8.partial.Dockerfile @@ -10,7 +10,6 @@ RUN apt-get update && apt-get install -y \ RUN apt-get update && apt-get install -y \ python3-pil \ python3-h5py \ - python3-keras-preprocessing \ python3-matplotlib \ python3-mock \ python3-numpy \ @@ -20,6 +19,7 @@ RUN apt-get update && apt-get install -y \ python3-portpicker RUN python3 -m pip --no-cache-dir install \ + keras_preprocessing \ enum34 # Build and install bazel From b9db1ee4174f26a94ac332ff8f60c9e0152403a8 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Mon, 8 Jun 2020 16:53:21 +0800 Subject: [PATCH 0096/1390] sync third_party_downloads to avoid conflict --- .../micro/tools/make/third_party_downloads.inc | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/micro/tools/make/third_party_downloads.inc b/tensorflow/lite/micro/tools/make/third_party_downloads.inc index 75a51e0df10..85016dc49b6 100644 --- a/tensorflow/lite/micro/tools/make/third_party_downloads.inc +++ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc @@ -62,14 +62,14 @@ RUY_MD5="2d54f058f8f7120dfc1ecee79dbf259e" CIFAR10_DATASET_URL="https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz" CIFAR10_DATASET_MD5="c32a1d4ab5d03f1284b67883e8d87530" -IMAGE_RECOGNITION_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/models/tflite/cifar_image_recognition_model_2020_4_14.zip" -IMAGE_RECOGNITION_MODEL_MD5 := "2b886156e7ef4d6e53d0f1a4bc800e56" +IMAGE_RECOGNITION_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/models/tflite/cifar_image_recognition_model_2020_05_27.zip" +IMAGE_RECOGNITION_MODEL_MD5 := "1f4607b05ac45b8a6146fb883dbc2d7b" -PERSON_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_grayscale_2019_11_21.zip" -PERSON_MODEL_MD5 := "fe2934bd0788f1dcc7af3f0a954542ab" +PERSON_MODEL_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_grayscale_2020_05_27.zip" +PERSON_MODEL_MD5 := "55b85f76e2995153e660391d4a209ef1" -PERSON_MODEL_INT8_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_01_13.zip" -PERSON_MODEL_INT8_MD5 := "8a7d2c70325f53136faea6dde517b8cc" +PERSON_MODEL_INT8_URL := "https://storage.googleapis.com/download.tensorflow.org/data/tf_lite_micro_person_data_int8_grayscale_2020_05_27.zip" +PERSON_MODEL_INT8_MD5 := "a0ede2d058aa2a1d413893455dd55352" EMBARC_MLI_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/archive/58284867ca52d1f43b25045e8601999d7359d986.zip" EMBARC_MLI_MD5 := "2bf4982a327fdaa9d475803ce014d1ef" @@ -77,9 +77,15 @@ EMBARC_MLI_MD5 := "2bf4982a327fdaa9d475803ce014d1ef" EMBARC_MLI_PRE_COMPILED_URL := "https://github.com/foss-for-synopsys-dwc-arc-processors/embarc_mli/releases/download/Release_1.1_RC2/embARC_MLI_package.zip" EMBARC_MLI_PRE_COMPILED_MD5 := "a95ff9e0370434484f14e7e4114327f6" +ZEPHYR_URL := "https://github.com/antmicro/zephyr/archive/55e36b9.zip" +ZEPHYR_MD5 := "755622eb4812fde918a6382b65d50c3b" + XTENSA_HIFI4_URL :="https://github.com/foss-xtensa/nnlib-hifi4/raw/master/archive/xa_nnlib_04_07.zip" XTENSA_HIFI4_MD5 :="f234764928f9a42901df33a27e118c8b" +ETHOSU_URL := "https://git.mlplatform.org/ml/ethos-u/ethos-u-core-driver.git/snapshot/ethos-u-core-driver-bcb5aaa99756f1b5c1295b079ebdd60996bc75a5.tar.gz" +ETHOSU_MD5 := "d2073c8d88fc167fd5c46b5dcda58ea1" + HIMAX_WE1_SDK_URL ="https://www.himax.com.tw/we-i/himax_we1_sdk_v02.zip" HIMAX_WE1_SDK_MD5 ="9a4b2f29b16052764e437b64bdcba816" From 60f6c58f9ea92c53c7434129045bc72761a5080a Mon Sep 17 00:00:00 2001 From: Steenu Johnson Date: Mon, 8 Jun 2020 17:57:31 +0530 Subject: [PATCH 0097/1390] This commit creates a new op CSVDatasetV2 with an additional parameter exclude_cols for backwards compatibility. Other changes include: Making the new op forward compatible. Adding error info to the docstring. Adding couple of tests to verify errors are being raised. Signed-off-by: Steenu Johnson --- .../data/experimental/csv_dataset_op.cc | 28 ++++++---- .../core/ops/experimental_dataset_ops.cc | 56 ++++++++++++++++--- .../kernel_tests/csv_dataset_test.py | 31 ++++++++++ .../python/data/experimental/ops/readers.py | 49 +++++++++++----- 4 files changed, 133 insertions(+), 31 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc index 62d27294f04..d2023ecec6e 100644 --- a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc @@ -28,7 +28,9 @@ namespace { class CSVDatasetOp : public DatasetOpKernel { public: - explicit CSVDatasetOp(OpKernelConstruction* ctx) : DatasetOpKernel(ctx) { + explicit CSVDatasetOp(OpKernelConstruction* ctx) + : DatasetOpKernel(ctx), + op_version_(ctx->def().op() == "CSVDataset" ? 1 : 2) { OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_)); } @@ -62,8 +64,11 @@ class CSVDatasetOp : public DatasetOpKernel { OP_REQUIRES(ctx, select_cols_tensor->dims() == 1, errors::InvalidArgument("`select_cols` must be a vector.")); - const Tensor* exclude_cols_tensor; - OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); + const Tensor* exclude_cols_tensor = new const Tensor(); + if (op_version_ > 1) { + OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); + } + OP_REQUIRES(ctx, exclude_cols_tensor->dims() == 1, errors::InvalidArgument("`exclude_cols` must be a vector")); @@ -138,7 +143,7 @@ class CSVDatasetOp : public DatasetOpKernel { } OP_REQUIRES(ctx, select_cols.empty() || exclude_cols.empty(), errors::InvalidArgument( - "Either select_cols or exlcude_cols should be empty")); + "Either select_cols or exclude_cols should be empty")); for (int i = 1; i < exclude_cols.size(); i++) { OP_REQUIRES(ctx, exclude_cols[i - 1] < exclude_cols[i], errors::InvalidArgument( @@ -238,8 +243,8 @@ class CSVDatasetOp : public DatasetOpKernel { std::make_pair(2, buffer_size), std::make_pair(3, header), std::make_pair(4, delim), std::make_pair(5, use_quote_delim), std::make_pair(6, na_value), std::make_pair(7, select_cols), - std::make_pair(8, exclude_cols)}, // Single tensor inputs - {std::make_pair(9, record_defaults)}, // Tensor list inputs + std::make_pair(9, exclude_cols)}, // Single tensor inputs + {std::make_pair(8, record_defaults)}, // Tensor list inputs {}, output)); return Status::OK(); @@ -386,12 +391,12 @@ class CSVDatasetOp : public DatasetOpKernel { Status result; while (!end_of_record) { // Read till we reach \n, \r or EOF - bool exclude = num_excluded_parsed < excluded.size() && - excluded[num_excluded_parsed] == num_parsed; + bool explicit_exclude = num_excluded_parsed < excluded.size() && + excluded[num_excluded_parsed] == num_parsed; bool include = select_all || (num_selected_parsed < selected.size() && selected[num_selected_parsed] == num_parsed) || - (!excluded.empty() && !exclude); + (!excluded.empty() && !explicit_exclude); // Don't fail fast, so that the next call to GetNext may still return // a valid record @@ -400,7 +405,7 @@ class CSVDatasetOp : public DatasetOpKernel { num_parsed++; if (include) num_selected_parsed++; - if (exclude) num_excluded_parsed++; + if (explicit_exclude) num_excluded_parsed++; } return result; @@ -894,6 +899,8 @@ class CSVDatasetOp : public DatasetOpKernel { const io::ZlibCompressionOptions options_; }; // class Dataset + const int op_version_; + DataTypeVector output_types_; std::vector output_shapes_; }; // class CSVDatasetOp @@ -901,6 +908,7 @@ class CSVDatasetOp : public DatasetOpKernel { REGISTER_KERNEL_BUILDER(Name("CSVDataset").Device(DEVICE_CPU), CSVDatasetOp); REGISTER_KERNEL_BUILDER(Name("ExperimentalCSVDataset").Device(DEVICE_CPU), CSVDatasetOp); +REGISTER_KERNEL_BUILDER(Name("CSVDatasetV2").Device(DEVICE_CPU), CSVDatasetOp); } // namespace } // namespace experimental diff --git a/tensorflow/core/ops/experimental_dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc index 9910472a2c6..33aa416fe56 100644 --- a/tensorflow/core/ops/experimental_dataset_ops.cc +++ b/tensorflow/core/ops/experimental_dataset_ops.cc @@ -154,7 +154,6 @@ REGISTER_OP("CSVDataset") .Input("use_quote_delim: bool") .Input("na_value: string") .Input("select_cols: int64") - .Input("exclude_cols: int64") .Input("record_defaults: output_types") .Output("handle: variant") .Attr("output_types: list({float,double,int32,int64,string}) >= 1") @@ -175,10 +174,8 @@ REGISTER_OP("CSVDataset") TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); // `select_cols` must be a vector TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); - //`exclude_cols` must be a vecotr - TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused)); // `record_defaults` must be lists of scalars - for (size_t i = 9; i < c->num_inputs(); ++i) { + for (size_t i = 8; i < c->num_inputs(); ++i) { shape_inference::ShapeHandle v; TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { @@ -190,7 +187,7 @@ REGISTER_OP("CSVDataset") return shape_inference::ScalarShape(c); }); -REGISTER_OP("ExperimentalCSVDataset") +REGISTER_OP("CSVDatasetV2") .Input("filenames: string") .Input("compression_type: string") .Input("buffer_size: int64") @@ -199,8 +196,8 @@ REGISTER_OP("ExperimentalCSVDataset") .Input("use_quote_delim: bool") .Input("na_value: string") .Input("select_cols: int64") - .Input("exclude_cols: int64") .Input("record_defaults: output_types") + .Input("exclude_cols: int64") .Output("handle: variant") .Attr("output_types: list({float,double,int32,int64,string}) >= 1") .Attr("output_shapes: list(shape) >= 1") @@ -221,9 +218,52 @@ REGISTER_OP("ExperimentalCSVDataset") // `select_cols` must be a vector TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); // `exclude_cols` must be a vector - TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused)); + TF_RETURN_IF_ERROR( + c->WithRank(c->input(c->num_inputs() - 1), 1, &unused)); // `record_defaults` must be lists of scalars - for (size_t i = 9; i < c->num_inputs(); ++i) { + for (size_t i = 8; i < c->num_inputs() - 1; ++i) { + shape_inference::ShapeHandle v; + TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); + if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { + return errors::InvalidArgument( + "Shape of a default must be a length-0 or length-1 vector, or a " + "scalar."); + } + } + return shape_inference::ScalarShape(c); + }); + +REGISTER_OP("ExperimentalCSVDataset") + .Input("filenames: string") + .Input("compression_type: string") + .Input("buffer_size: int64") + .Input("header: bool") + .Input("field_delim: string") + .Input("use_quote_delim: bool") + .Input("na_value: string") + .Input("select_cols: int64") + .Input("record_defaults: output_types") + .Output("handle: variant") + .Attr("output_types: list({float,double,int32,int64,string}) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetDoNotOptimize() // TODO(b/123753214): Source dataset ops must + // disable constant folding. + .SetShapeFn([](shape_inference::InferenceContext* c) { + shape_inference::ShapeHandle unused; + // `filenames` must be a scalar or a vector. + TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused)); + // `compression_type`, `buffer_size`, `header`, `field_delim`, + // `use_quote_delim`, `na_value` must be scalars + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); + // `select_cols` must be a vector + TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused)); + // `record_defaults` must be lists of scalars + for (size_t i = 8; i < c->num_inputs(); ++i) { shape_inference::ShapeHandle v; TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(i), 1, &v)); if (c->Rank(c->input(i)) == 1 && c->Value(c->Dim(v, 0)) > 1) { diff --git a/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py index 42a96812bb4..20c983cf2a9 100644 --- a/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/csv_dataset_test.py @@ -424,6 +424,37 @@ class CsvDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): record_defaults=record_defaults, exclude_cols=[1, 2]) + @combinations.generate(test_base.default_test_combinations()) + def testCsvDataset_withSelectandExcludeCol(self): + record_defaults = [['']] + inputs = [['1,2,3', '5,6,7']] + self._test_dataset( + inputs, + expected_err_re='Either select_cols or exclude_cols should be empty', + record_defaults=record_defaults, + select_cols=[0], + exclude_cols=[1, 2]) + + @combinations.generate(test_base.default_test_combinations()) + def testCsvDataset_withExcludeColandRecordDefaultsTooLow(self): + record_defaults = [['']] + inputs = [['1,2,3', '5,6,7']] + self._test_dataset( + inputs, + expected_err_re='Expect 1 fields but have more in record', + record_defaults=record_defaults, + exclude_cols=[0]) + + @combinations.generate(test_base.default_test_combinations()) + def testCsvDataset_withExcludeColandRecordDefaultsTooHigh(self): + record_defaults = [['']]*3 + inputs = [['1,2,3', '5,6,7']] + self._test_dataset( + inputs, + expected_err_re='Expect 3 fields but have 2 in record', + record_defaults=record_defaults, + exclude_cols=[0]) + @combinations.generate(test_base.default_test_combinations()) def testCsvDataset_withMultipleNewLines(self): # In this case, we expect it to behave differently from diff --git a/tensorflow/python/data/experimental/ops/readers.py b/tensorflow/python/data/experimental/ops/readers.py index fcf92f5aaf9..14a507580ad 100644 --- a/tensorflow/python/data/experimental/ops/readers.py +++ b/tensorflow/python/data/experimental/ops/readers.py @@ -41,6 +41,7 @@ from tensorflow.python.ops import gen_experimental_dataset_ops from tensorflow.python.ops import io_ops from tensorflow.python.platform import gfile from tensorflow.python.util.tf_export import tf_export +from tensorflow.python.compat import compat _ACCEPTABLE_CSV_TYPES = (dtypes.float32, dtypes.float64, dtypes.int32, dtypes.int64, dtypes.string) @@ -661,7 +662,9 @@ class CsvDatasetV2(dataset_ops.DatasetSource): column if it is optional, or `DType` or empty `Tensor` if required. If both this and `select_columns` are specified, these must have the same lengths, and `column_defaults` is assumed to be sorted in order of - increasing column index. + increasing column index. If both this and 'exclude_cols' are specified, + the sum of lengths of record_defaults and exclude_cols should equal + the total number of columns in the CSV file. compression_type: (Optional.) A `tf.string` scalar evaluating to one of `""` (no compression), `"ZLIB"`, or `"GZIP"`. Defaults to no compression. @@ -683,6 +686,13 @@ class CsvDatasetV2(dataset_ops.DatasetSource): exclude_cols:(Optional.) A sorted list of column indices to exclude from the input data. If specified, only the complement of this set of column will be parsed. Defaults to parsing all columns. + + Raises: + InvalidArgumentError: If exclude_cols is not None and + len(exclude_cols) + len(record_defaults) does not match the total + number of columns in the file(s) + + """ self._filenames = ops.convert_to_tensor( filenames, dtype=dtypes.string, name="filenames") @@ -722,18 +732,31 @@ class CsvDatasetV2(dataset_ops.DatasetSource): ) self._element_spec = tuple( tensor_spec.TensorSpec([], d.dtype) for d in self._record_defaults) - variant_tensor = gen_experimental_dataset_ops.csv_dataset( - filenames=self._filenames, - record_defaults=self._record_defaults, - buffer_size=self._buffer_size, - header=self._header, - output_shapes=self._flat_shapes, - field_delim=self._field_delim, - use_quote_delim=self._use_quote_delim, - na_value=self._na_value, - select_cols=self._select_cols, - exclude_cols=self._exclude_cols, - compression_type=self._compression_type) + if compat.forward_compatible(2020, 6, 25): + variant_tensor = gen_experimental_dataset_ops.csv_dataset_v2( + filenames=self._filenames, + record_defaults=self._record_defaults, + buffer_size=self._buffer_size, + header=self._header, + output_shapes=self._flat_shapes, + field_delim=self._field_delim, + use_quote_delim=self._use_quote_delim, + na_value=self._na_value, + select_cols=self._select_cols, + exclude_cols=self._exclude_cols, + compression_type=self._compression_type) + else: + variant_tensor = gen_experimental_dataset_ops.csv_dataset( + filenames=self._filenames, + record_defaults=self._record_defaults, + buffer_size=self._buffer_size, + header=self._header, + output_shapes=self._flat_shapes, + field_delim=self._field_delim, + use_quote_delim=self._use_quote_delim, + na_value=self._na_value, + select_cols=self._select_cols, + compression_type=self._compression_type) super(CsvDatasetV2, self).__init__(variant_tensor) @property From bacc5e5927360ad119cdb9967d53df4c95c53fd4 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Mon, 8 Jun 2020 14:38:12 +0200 Subject: [PATCH 0098/1390] Fix build issues. --- tensorflow/lite/micro/micro_allocator.cc | 3 ++- tensorflow/lite/micro/micro_allocator_test.cc | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 8fac421750d..18c82bca57d 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -94,7 +94,7 @@ TfLiteStatus AllocateVariables( return kTfLiteOk; } - +#if !defined(__clang__) // Helper function to check flatbuffer metadata correctness. This function is // not called by default. Hence it's not linked in to the final binary code. TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, @@ -157,6 +157,7 @@ TfLiteStatus CheckOfflinePlannedOffsets(const Model* model, } return kTfLiteOk; } +#endif // A helper class to construct AllocationInfo array. This array contains the // lifetime of tensors / scratch_buffer and will be used to calculate the memory diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 67052ce12d9..ac581305340 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -238,7 +238,7 @@ TF_LITE_MICRO_TEST(TestFinishComplexTensorAllocation) { TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { int version = 1; int subgraph = 0; - int nbr_tensors = 4; + constexpr int nbr_tensors = 4; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = {version, subgraph, nbr_tensors, // header @@ -283,7 +283,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { } TF_LITE_MICRO_TEST(OfflinePlannerBasic) { - int nbr_tensors = 4; + constexpr int nbr_tensors = 4; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = {1, 0, nbr_tensors, 0, // t0 @@ -328,7 +328,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) { } TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { - int nbr_tensors = 4; + constexpr int nbr_tensors = 4; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = { 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors @@ -374,7 +374,7 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { } TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { - int nbr_tensors = 5; + constexpr int nbr_tensors = 5; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = { 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors From 0c8343f7f3a8066c507b97ff84d8b298655cc5f4 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Mon, 8 Jun 2020 11:21:49 -0500 Subject: [PATCH 0099/1390] Address review comments --- tensorflow/core/platform/tf32_utils.cc | 10 ++++++---- tensorflow/core/platform/tf32_utils.h | 2 +- tensorflow/python/framework/config.py | 7 +++++-- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/tensorflow/core/platform/tf32_utils.cc b/tensorflow/core/platform/tf32_utils.cc index 715b5996dc3..4456e768c0a 100644 --- a/tensorflow/core/platform/tf32_utils.cc +++ b/tensorflow/core/platform/tf32_utils.cc @@ -14,14 +14,16 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/platform/tf32_utils.h" +#include namespace tensorflow { -// TODO(nluehr): enable tf32 execution by default after TF32 Ampere testing. -static bool tf32_enabled = false; +// Whether TensorFloat-32 should be used where supported. +// TODO(nluehr): Maybe enable by default after TF32 Ampere testing. +static std::atomic tf32_allowed{false}; -void allow_tf32_execution(bool allow) { tf32_enabled = allow; } +void allow_tf32_execution(bool allowed) { tf32_allowed = allowed; } -bool tf32_execution_allowed() { return tf32_enabled; } +bool tf32_execution_allowed() { return tf32_allowed; } } // namespace tensorflow diff --git a/tensorflow/core/platform/tf32_utils.h b/tensorflow/core/platform/tf32_utils.h index a0ce58f9bbd..7a158d00ad3 100644 --- a/tensorflow/core/platform/tf32_utils.h +++ b/tensorflow/core/platform/tf32_utils.h @@ -18,7 +18,7 @@ limitations under the License. namespace tensorflow { -void allow_tf32_execution(bool allow); +void allow_tf32_execution(bool allowed); bool tf32_execution_allowed(); diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py index 042af4d1023..a356e6d9a16 100644 --- a/tensorflow/python/framework/config.py +++ b/tensorflow/python/framework/config.py @@ -23,6 +23,8 @@ from tensorflow.python.eager import context from tensorflow.python.util import deprecation from tensorflow.python.util.tf_export import tf_export + +# No tf_export until TF is built against CUDA11 which is required for TF32. def tensor_float32_execution_allowed(): """Get if TensorFloat-32 operations are enabled on supported hardware. @@ -31,7 +33,8 @@ def tensor_float32_execution_allowed(): """ return _pywrap_tf32_execution.is_allowed() -def allow_tensor_float_32_execution(allow): +# No tf_export until TF is built against CUDA11 which is required for TF32. +def allow_tensor_float_32_execution(allowed): """Allow use of TensorFloat-32 with float32 ops on supported hardware. TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture. @@ -47,7 +50,7 @@ def allow_tensor_float_32_execution(allow): Args: allow: whether to allow TensorFloat-32 execution """ - _pywrap_tf32_execution.allow(allow) + _pywrap_tf32_execution.allow(allowed) @tf_export('config.threading.get_intra_op_parallelism_threads') def get_intra_op_parallelism_threads(): From ed41bb08250cad9f3ddbb6c7fb83e1216ee06031 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Mon, 8 Jun 2020 17:57:27 +0100 Subject: [PATCH 0100/1390] Fix for CI failure. Change-Id: I66a5b5ab559207071ea62619e9e612fda9a73202 --- tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt | 4 ++++ .../tools/api/golden/v1/tensorflow.lite.constants.pbtxt | 4 ++++ tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt index c3199b24d98..9538fe382a0 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.lite.-ops-set.pbtxt @@ -1,6 +1,10 @@ path: "tensorflow.lite.OpsSet" tf_class { is_instance: "" + member { + name: "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" + mtype: "" + } member { name: "SELECT_TF_OPS" mtype: "" diff --git a/tensorflow/tools/api/golden/v1/tensorflow.lite.constants.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.lite.constants.pbtxt index 27c227dac64..7f62da6662a 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.lite.constants.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.lite.constants.pbtxt @@ -12,6 +12,10 @@ tf_module { name: "GRAPHVIZ_DOT" mtype: "" } + member { + name: "INT16" + mtype: "" + } member { name: "INT32" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt index c3199b24d98..9538fe382a0 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.lite.-ops-set.pbtxt @@ -1,6 +1,10 @@ path: "tensorflow.lite.OpsSet" tf_class { is_instance: "" + member { + name: "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" + mtype: "" + } member { name: "SELECT_TF_OPS" mtype: "" From ad4323e93479caf12a11b6e089174be2a7ec7462 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Thu, 7 May 2020 17:34:50 +0000 Subject: [PATCH 0101/1390] Update tf.map_fn to specify that at least one tensor must be present Signed-off-by: Yong Tang Typo fix Signed-off-by: Yong Tang --- tensorflow/python/ops/map_fn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/ops/map_fn.py b/tensorflow/python/ops/map_fn.py index 6f59bcf5599..287070af07c 100644 --- a/tensorflow/python/ops/map_fn.py +++ b/tensorflow/python/ops/map_fn.py @@ -267,7 +267,7 @@ def map_fn(fn, elems: A tensor or (possibly nested) sequence of tensors, each of which will be unstacked along their first dimension. `fn` will be applied to the nested sequence of the resulting slices. `elems` may include ragged and - sparse tensors. + sparse tensors. `elems` must consist of at least one tensor. dtype: Deprecated: Equivalent to `fn_output_signature`. parallel_iterations: (optional) The number of iterations allowed to run in parallel. When graph building, the default value is 10. While executing @@ -296,7 +296,7 @@ def map_fn(fn, TypeError: if `fn` is not callable or the structure of the output of `fn` and `fn_output_signature` do not match. ValueError: if the lengths of the output of `fn` and `fn_output_signature` - do not match. + do not match, or if the `elems` does not contain any tensor. Examples: From dbc7faeecdbdc223c67b03284ba5dc7d25668d3c Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Mon, 8 Jun 2020 20:20:22 +0100 Subject: [PATCH 0102/1390] Addressed reviewer's comment. Change-Id: I5bda332514d8070731b807b750ee7a423d6b4d78 --- tensorflow/lite/python/convert.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index b1095a469f6..52edb700195 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -94,10 +94,15 @@ class OpsSet(enum.Enum): # quantized implementations. TFLITE_BUILTINS_INT8 = "TFLITE_BUILTINS_INT8" - # Convert model using only TensorFlow Lite operations with quantized int8 weights - # and int16 activations. + # Convert model using only TensorFlow Lite operations with quantized int8 weights, + # int16 activations and int64 bias. # Specifying this will throw an error for operations that do not yet have # quantized implementations. + # This quantization mode should be used in models for super-resolution, + # audio signal processing or image de-noising. It improves accuracy + # significantly, but only slightly increases the model size. + # WARNING: These ops are currently experimental and have not yet been finalized. + # They are only compatible with CPU execution, and have not been optimized for production. EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" def __str__(self): From 13d3b343498d499f87230f2e596b738be5cf1109 Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Tue, 9 Jun 2020 09:28:23 +0800 Subject: [PATCH 0103/1390] modify example main API usage --- .../himax_we1_evb/main_functions.cc | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc index 552b52c9c51..f0c7a405974 100644 --- a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc @@ -72,21 +72,20 @@ void setup() { // incur some penalty in code space for op implementations that are not // needed by this graph. // - // tflite::ops::micro::AllOpsResolver resolver; + // tflite::AllOpsResolver resolver; // NOLINTNEXTLINE(runtime-global-variables) - static tflite::MicroOpResolver<12> micro_op_resolver; - micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_DEPTHWISE_CONV_2D, - tflite::ops::micro::Register_DEPTHWISE_CONV_2D(), - 1, 3); + static tflite::MicroMutableOpResolver<5> micro_op_resolver; + micro_op_resolver.AddBuiltin( + tflite::BuiltinOperator_DEPTHWISE_CONV_2D, + tflite::ops::micro::Register_DEPTHWISE_CONV_2D()); micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D, - tflite::ops::micro::Register_CONV_2D(), 1, 3); + tflite::ops::micro::Register_CONV_2D()); micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D, - tflite::ops::micro::Register_AVERAGE_POOL_2D(), - 1, 2); + tflite::ops::micro::Register_AVERAGE_POOL_2D()); micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE, tflite::ops::micro::Register_RESHAPE()); micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX, - tflite::ops::micro::Register_SOFTMAX(), 1, 3); + tflite::ops::micro::Register_SOFTMAX()); // Build an interpreter to run the model with. // NOLINTNEXTLINE(runtime-global-variables) From 4644b47db199ccda856c81ca88d6a3c58c41890f Mon Sep 17 00:00:00 2001 From: Eugene Kuznetsov Date: Tue, 9 Jun 2020 05:13:35 +0000 Subject: [PATCH 0104/1390] Remove duplicate macros --- .../core/kernels/non_max_suppression_op.cu.cc | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tensorflow/core/kernels/non_max_suppression_op.cu.cc b/tensorflow/core/kernels/non_max_suppression_op.cu.cc index 53559b20419..c2cae2ab212 100644 --- a/tensorflow/core/kernels/non_max_suppression_op.cu.cc +++ b/tensorflow/core/kernels/non_max_suppression_op.cu.cc @@ -28,25 +28,6 @@ limitations under the License. #include "tensorflow/core/util/gpu_launch_config.h" #include "tensorflow/stream_executor/stream_executor.h" -#define TF_RETURN_IF_CUDA_ERROR(result) \ - do { \ - cudaError_t error(result); \ - if (!SE_PREDICT_TRUE(error == cudaSuccess)) { \ - return errors::Internal("Cuda call failed with ", \ - cudaGetErrorString(error)); \ - } \ - } while (0) - -#define TF_OP_REQUIRES_CUDA_SUCCESS(context, result) \ - do { \ - cudaError_t error(result); \ - if (!SE_PREDICT_TRUE(error == cudaSuccess)) { \ - context->SetStatus(errors::Internal("Cuda call failed with", \ - cudaGetErrorString(error))); \ - return; \ - } \ - } while (0) - struct __align__(16) Box { float x1, y1, x2, y2; }; From 8681b1bf543a33c62a9bec29625594d1fd9b921a Mon Sep 17 00:00:00 2001 From: Eugene Kuznetsov Date: Tue, 9 Jun 2020 15:11:45 +0000 Subject: [PATCH 0105/1390] Adding 3d Pooling using latest MIOpen API --- tensorflow/cc/gradients/nn_grad_test.cc | 8 - tensorflow/core/kernels/cudnn_pooling_gpu.cc | 31 ++ tensorflow/python/eager/backprop_test.py | 3 - tensorflow/python/kernel_tests/BUILD | 1 - tensorflow/python/kernel_tests/pool_test.py | 4 - .../ops/parallel_for/control_flow_ops_test.py | 2 - tensorflow/stream_executor/rocm/rocm_dnn.cc | 443 +++++++++++------- tensorflow/stream_executor/rocm/rocm_dnn.h | 68 +++ 8 files changed, 360 insertions(+), 200 deletions(-) diff --git a/tensorflow/cc/gradients/nn_grad_test.cc b/tensorflow/cc/gradients/nn_grad_test.cc index 942ec08f451..f5a09e09dcd 100644 --- a/tensorflow/cc/gradients/nn_grad_test.cc +++ b/tensorflow/cc/gradients/nn_grad_test.cc @@ -259,9 +259,6 @@ TEST_F(NNGradTest, MaxPoolGradV2Helper) { RunTest(x, x_init_value, y, y_shape); } -// TODO(rocm): -// Re-enable this test once 3D pooling is supported on ROCm platform -#ifndef TENSORFLOW_USE_ROCM TEST_F(NNGradTest, MaxPool3DGradHelper) { TensorShape x_shape({1, 3, 3, 3, 1}); TensorShape y_shape({1, 1, 1, 1, 1}); @@ -274,7 +271,6 @@ TEST_F(NNGradTest, MaxPool3DGradHelper) { SetRandomValuesForMaxPooling(&x_init_value); RunTest(x, x_init_value, y, y_shape); } -#endif TEST_F(NNGradTest, AvgPoolGradHelper) { TensorShape x_shape({1, 2, 2, 1}); @@ -287,9 +283,6 @@ TEST_F(NNGradTest, AvgPoolGradHelper) { RunTest(x, x_shape, y, y_shape); } -// TODO(rocm): -// Re-enable this test once 3D pooling is supported on ROCm platform -#ifndef TENSORFLOW_USE_ROCM TEST_F(NNGradTest, AvgPool3DGradHelper) { TensorShape x_shape({1, 3, 3, 3, 1}); TensorShape y_shape({1, 1, 1, 1, 1}); @@ -300,7 +293,6 @@ TEST_F(NNGradTest, AvgPool3DGradHelper) { auto y = AvgPool3D(scope_, x, ksize, strides, "SAME"); RunTest(x, x_shape, y, y_shape); } -#endif TEST_F(NNGradTest, LRN) { TensorShape x_shape({1, 1, 2, 1}); diff --git a/tensorflow/core/kernels/cudnn_pooling_gpu.cc b/tensorflow/core/kernels/cudnn_pooling_gpu.cc index eb7d16e3074..60088133d5d 100644 --- a/tensorflow/core/kernels/cudnn_pooling_gpu.cc +++ b/tensorflow/core/kernels/cudnn_pooling_gpu.cc @@ -98,10 +98,25 @@ void DnnPooling3dOp::Compute(OpKernelContext* context, auto* stream = context->op_device_context()->stream(); OP_REQUIRES(context, stream, errors::Internal("No GPU stream available.")); +#if TENSORFLOW_USE_ROCM + static int64 PoolingScratchSize = GetDnnWorkspaceLimit( + // default value is in bytes despite the name of the environment variable + "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB + ); + + DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); + bool status = + stream + ->ThenPoolForward(pooling_desc, input_desc, input_data, output_desc, + &output_data, &scratch_allocator) + .ok(); +#else bool status = stream ->ThenPoolForward(pooling_desc, input_desc, input_data, output_desc, &output_data) .ok(); +#endif + OP_REQUIRES(context, status, errors::Internal("dnn PoolForward launch failed")); @@ -225,12 +240,28 @@ void DnnPooling3dGradOp::Compute( auto* stream = context->op_device_context()->stream(); OP_REQUIRES(context, stream, errors::Internal("No GPU stream available.")); +#if TENSORFLOW_USE_ROCM + static int64 PoolingScratchSize = GetDnnWorkspaceLimit( + // default value is in bytes despite the name of the environment variable + "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB + ); + + DnnScratchAllocator scratch_allocator(PoolingScratchSize, context); + bool status = stream + ->ThenPoolBackward(pooling_desc, orig_input_desc, + orig_input_data, orig_output_desc, + orig_output_data, output_backprop_data, + &input_backprop_data, &scratch_allocator) + .ok(); +#else bool status = stream ->ThenPoolBackward(pooling_desc, orig_input_desc, orig_input_data, orig_output_desc, orig_output_data, output_backprop_data, &input_backprop_data) .ok(); +#endif + OP_REQUIRES(context, status, errors::Internal("dnn PoolBackward launch failed")); diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py index b28aaa3a626..bb909433fa7 100644 --- a/tensorflow/python/eager/backprop_test.py +++ b/tensorflow/python/eager/backprop_test.py @@ -1455,9 +1455,6 @@ class BackpropTest(test.TestCase, parameterized.TestCase): @test_util.run_in_graph_and_eager_modes def testMaxPooling3DGradient(self): - if test.is_built_with_rocm(): - self.skipTest('Pooling with 3D tensors is not supported in ROCm') - def forward(a): r = max_pooling3d(a, pool_size=pool_size, strides=strides, padding='SAME') return r diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD index a04c874c9d6..ae4e0244357 100644 --- a/tensorflow/python/kernel_tests/BUILD +++ b/tensorflow/python/kernel_tests/BUILD @@ -2995,7 +2995,6 @@ cuda_py_test( name = "pooling_ops_3d_test", size = "medium", srcs = ["pooling_ops_3d_test.py"], - tags = ["no_rocm"], deps = [ "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", diff --git a/tensorflow/python/kernel_tests/pool_test.py b/tensorflow/python/kernel_tests/pool_test.py index 0f0eaa25402..01680f5c1f0 100644 --- a/tensorflow/python/kernel_tests/pool_test.py +++ b/tensorflow/python/kernel_tests/pool_test.py @@ -219,8 +219,6 @@ class PoolingTest(test.TestCase): strides=strides) def testPool3D(self): - if test.is_built_with_rocm(): - self.skipTest("Pooling with 3D tensors is not supported in ROCm") with self.session(use_gpu=test.is_gpu_available()): for padding in ["SAME", "VALID"]: for pooling_type in ["MAX", "AVG"]: @@ -363,8 +361,6 @@ class PoolingTest(test.TestCase): @test_util.run_deprecated_v1 def testGradient3D(self): - if test.is_built_with_rocm(): - self.skipTest("Pooling with 3D tensors is not supported in ROCm") with self.session(use_gpu=test.is_gpu_available()): for padding in ["SAME", "VALID"]: for pooling_type in ["AVG", "MAX"]: diff --git a/tensorflow/python/ops/parallel_for/control_flow_ops_test.py b/tensorflow/python/ops/parallel_for/control_flow_ops_test.py index 243471553d9..a6546fee742 100644 --- a/tensorflow/python/ops/parallel_for/control_flow_ops_test.py +++ b/tensorflow/python/ops/parallel_for/control_flow_ops_test.py @@ -488,8 +488,6 @@ class NNTest(PForTestCase): self._test_loop_fn(loop_fn, 3) def test_max_pool3d(self): - if test.is_built_with_rocm(): - self.skipTest("Pooling with 3D tensors is not supported in ROCm") with backprop.GradientTape(persistent=True) as g: x = random_ops.random_uniform([3, 3, 2, 12, 12, 3]) g.watch(x) diff --git a/tensorflow/stream_executor/rocm/rocm_dnn.cc b/tensorflow/stream_executor/rocm/rocm_dnn.cc index e0ead6d57e8..9c09784b3f1 100644 --- a/tensorflow/stream_executor/rocm/rocm_dnn.cc +++ b/tensorflow/stream_executor/rocm/rocm_dnn.cc @@ -263,7 +263,8 @@ namespace wrap { __macro(miopenFindConvolutionForwardAlgorithm) \ __macro(miopenCreateTensorDescriptor) \ __macro(miopenDestroyTensorDescriptor) \ - __macro(miopenSet2dPoolingDescriptor) \ + __macro(miopenSetNdPoolingDescriptor) \ + __macro(miopenSetPoolingIndexType) \ __macro(miopenSetLRNDescriptor) \ __macro(miopenLRNGetWorkSpaceSize) \ __macro(miopenCreateConvolutionDescriptor) \ @@ -290,7 +291,7 @@ namespace wrap { __macro(miopenSetTensorDescriptor) \ __macro(miopenGetTensorDescriptorSize) \ __macro(miopenPoolingForward) \ - __macro(miopenPoolingGetWorkSpaceSize) \ + __macro(miopenPoolingGetWorkSpaceSizeV2 \ __macro(miopenPoolingBackward) \ __macro(miopenLRNForward) \ __macro(miopenLRNBackward) \ @@ -605,6 +606,11 @@ MIOpenSupport::MIOpenSupport(GpuExecutor* parent) : parent_(parent) { // swich to Find Mode if env var TF_ROCM_USE_IMMEDIATE_MODE is set tensorflow::ReadBoolFromEnvVar("TF_ROCM_USE_IMMEDIATE_MODE", false, &use_immediate_mode_); + + bool enable_pooling_cache = false; + tensorflow::ReadBoolFromEnvVar("TF_ROCM_BW_POOL_CACHE", false, + &enable_pooling_cache); + if (enable_pooling_cache) m_pooling_cache_allowed = true; } port::Status MIOpenSupport::Init() { @@ -844,17 +850,19 @@ class ScopedPoolingDescriptor { std::transform(shape64.cbegin(), shape64.cend(), shape.begin(), &CheckedNarrowing); - if (nd != 2) { - LOG(FATAL) << "miopen requires pooling dimensions be 2" - << ToString(status); - } - - status = wrap::miopenSet2dPoolingDescriptor( + status = wrap::miopenSetNdPoolingDescriptor( handle_, (pooling_descriptor.mode() == dnn::PoolingMode::kMaximum ? miopenPoolingMax : miopenPoolingAverage), - shape[0], shape[1], padding[0], padding[1], strides[0], strides[1]); + nd, shape.data(), padding.data(), strides.data()); + + // Note: The index type has to be uint32 type for now because MIOpen + // API assumes all input indexes to be the same type. Since a tensor + // descriptor can only use int32 type, the index type here need to be + // aligned with the tensor index type of the (input) tensor descritptor + status = wrap::miopenSetPoolingIndexType(handle_, miopenIndexUint32); + if (status != miopenStatusSuccess) { LOG(FATAL) << "could not set miopen pooling descriptor: " << ToString(status); @@ -4009,10 +4017,94 @@ bool MIOpenSupport::DoPoolForward( const DeviceMemory& input_data, const dnn::BatchDescriptor& output_dimensions, DeviceMemory* output_data, ScratchAllocator* workspace_allocator) { - LOG(ERROR) << "miopen does not support pooling for dobule type yet"; + LOG(ERROR) << "miopen does not support pooling for double type yet"; return false; } +bool PoolingWorkspaceDescriptor::IsSame( + const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type) { + return dtype == _type && + input_dims == + input_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX) && + output_dims == + output_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX) && + op.mode() == pooling_dimensions.mode() && + op.window() == pooling_dimensions.window() && + op.padding() == pooling_dimensions.padding() && + op.strides() == pooling_dimensions.strides(); +} + +bool PoolingWorkspaceCache::find( + const void* p, const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type, + PoolingWorkspaceDescriptor*& pdesc) { + pdesc = 0; + auto it = cache.find(p); + if (it == cache.end()) { + return false; + } + if (!it->second.IsSame(input_dimensions, output_dimensions, + pooling_dimensions, _type)) { + return false; + } + pdesc = &it->second; + return true; +} + +void PoolingWorkspaceCache::insert( + const void* p, const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type, + std::unique_ptr>& workspace, size_t wsp_size, + hipStream_t hip_stream) { + PoolingWorkspaceDescriptor* desc = 0; + auto it = cache.find(p); + if (it != cache.end()) { + // replacing an entry with the same pointer but different attributes + // (if everything matches, the caller is expected to reuse the entry) + desc = &it->second; + hipStreamSynchronize(hip_stream); + memory_used -= desc->workspace_size; + } else { + cache[p] = PoolingWorkspaceDescriptor(); + desc = &cache[p]; + } + desc->input_dims = input_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX); + desc->output_dims = + output_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX); + desc->op = pooling_dimensions; + desc->dtype = _type; + desc->timestamp = timestamp; + timestamp++; + desc->workspace = std::move(workspace); + desc->workspace_size = wsp_size; + memory_used += wsp_size; + trim(hip_stream); +} + +void PoolingWorkspaceCache::trim(hipStream_t hip_stream) { + if (memory_used < memory_budget && cache.size() < trim_size) return; + bool must_sync = true; + while (true) { + int new_size = cache.size() - (cache.size() >> 2); + std::vector old_entries; + for (auto& x : cache) + if (x.second.timestamp + new_size < timestamp) + old_entries.push_back(x.first); + if (old_entries.empty()) break; + if (must_sync) hipStreamSynchronize(hip_stream); + must_sync = true; + for (auto x : old_entries) { + memory_used -= cache[x].workspace_size; + cache.erase(x); + } + if (memory_used < memory_budget || cache.size() < 10) break; + } +} + bool MIOpenSupport::DoPoolForward( Stream* stream, const dnn::PoolingDescriptor& pooling_dimensions, const dnn::BatchDescriptor& input_dimensions, @@ -4020,7 +4112,6 @@ bool MIOpenSupport::DoPoolForward( const dnn::BatchDescriptor& output_dimensions, DeviceMemory* output_data, ScratchAllocator* workspace_allocator) { auto miopen = miopen_->GetHandle(parent_, stream); - // Alpha is the scaling factor for input. float alpha = 1.0; // Beta is the scaling factor for output. @@ -4030,10 +4121,48 @@ bool MIOpenSupport::DoPoolForward( ScopedTensorDescriptor dest_desc{output_dimensions, miopenFloat}; ScopedPoolingDescriptor pooling_desc{pooling_dimensions}; + bool do_backward = false; + uint8* workspace = 0; + size_t workspace_size = 0; + std::unique_ptr> wsp_mem; + if (m_pooling_cache_enabled) { + do_backward = true; + auto status = wrap::miopenPoolingGetWorkSpaceSizeV2( + pooling_desc.handle(), dest_desc.handle(), &workspace_size); + if (status != miopenStatusSuccess) { + LOG(ERROR) + << "failed to obtain workspace size for backward pooling on stream: " + << ToString(status); + return false; + } + if (workspace_size != 0) { + PoolingWorkspaceDescriptor* pdesc = 0; + bool cache_hit = + m_pooling_cache_allowed && + m_pooling_cache.find(input_data.opaque(), input_dimensions, + output_dimensions, pooling_dimensions, + miopenFloat, pdesc); + if (cache_hit) { + // reusing the same buffer + workspace = reinterpret_cast( + pdesc->workspace->mutable_device_memory()->opaque()); + } else { + wsp_mem = stream->AllocateTemporaryArray(workspace_size) + .ConsumeValueOrDie(); + workspace = reinterpret_cast( + wsp_mem->mutable_device_memory()->opaque()); + m_pooling_cache.insert(input_data.opaque(), input_dimensions, + output_dimensions, pooling_dimensions, + miopenFloat, wsp_mem, workspace_size, + AsGpuStreamValue(stream)); + } + } + } + auto status = wrap::miopenPoolingForward( miopen.handle(), pooling_desc.handle(), &alpha, src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque(), - false, nullptr, 0); + do_backward, workspace, workspace_size); if (status != miopenStatusSuccess) { LOG(ERROR) << "failed to enqueue forward pooling on stream: " << ToString(status); @@ -4072,6 +4201,118 @@ bool MIOpenSupport::DoPoolForward( return true; } +template +bool MIOpenSupport::DoPoolBackwardImpl( + Stream* stream, const dnn::PoolingDescriptor& pooling_dimensions, + const dnn::BatchDescriptor& input_dimensions, + const DeviceMemory& input_data, + const dnn::BatchDescriptor& output_dimensions, + const DeviceMemory& output_data, const DeviceMemory& input_diff_data, + DeviceMemory* output_diff_data, ScratchAllocator* workspace_allocator) { + auto miopen = miopen_->GetHandle(parent_, stream); + if (m_pooling_cache_allowed) m_pooling_cache_enabled = true; + // Alpha is the scaling factor for input. + float alpha = 1.0; + // Beta is the scaling factor for output. + float beta = 0.0; + + auto type = + std::is_same::value + ? miopenFloat + : (std::is_same::value ? miopenHalf + : (miopenDataType_t)-1); + + ScopedTensorDescriptor src_desc{input_dimensions, type}; + ScopedTensorDescriptor dest_desc{output_dimensions, type}; + ScopedPoolingDescriptor pooling_desc{pooling_dimensions}; + + uint8* workspace_ptr = 0; + DeviceMemory workspace; + PoolingWorkspaceDescriptor* pdesc = 0; + + size_t workspace_size_in_bytes = 0; + auto status = wrap::miopenPoolingGetWorkSpaceSizeV2( + pooling_desc.handle(), dest_desc.handle(), &workspace_size_in_bytes); + if (status != miopenStatusSuccess) { + LOG(ERROR) + << "failed to obtain workspace size for backward pooling on stream: " + << ToString(status); + return false; + } + + // Allocate the workspace. + if (workspace_size_in_bytes > 0) { + bool cache_hit = m_pooling_cache_allowed && + m_pooling_cache.find(input_data.opaque(), input_dimensions, + output_dimensions, pooling_dimensions, + type, pdesc); + if (cache_hit) { + assert(pdesc != 0); + workspace_ptr = reinterpret_cast( + pdesc->workspace->mutable_device_memory()->opaque()); + VLOG(1) << "Pooling cache hit"; + } else { + VLOG(1) << "Pooling cache miss"; + assert(workspace_allocator); + auto allocated = + workspace_allocator->AllocateBytes(workspace_size_in_bytes); + if (!allocated.ok() || (workspace = allocated.ValueOrDie()) == nullptr) { + LOG(ERROR) << "Failed to allocate backward pooling workspace"; + return false; + } + DeviceMemory dest2; // duplicated dest from forward: + int64 dest2_size = 0; + + // miopen requires the strides and dims to be ordered as BDYX. + std::vector dims64 = + output_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX); + // miopen does not use strides and must have 4D tensor. + // std::vector dims(pooling_dimensions.ndims() + 2); + + dest2_size = sizeof(T); + for (auto& x : dims64) dest2_size *= x; + + if (dest2_size > 0) { + assert(workspace_allocator); + auto allocated = workspace_allocator->AllocateBytes(dest2_size); + if (!allocated.ok() || (dest2 = allocated.ValueOrDie()) == nullptr) { + LOG(ERROR) << "Failed to allocate backward pooling workspace"; + return false; + } + } else { + LOG(ERROR) << "Failed to calculate tensor size to chain forward and " + "backward pooling"; + } + + status = wrap::miopenPoolingForward( + miopen.handle(), pooling_desc.handle(), &alpha, src_desc.handle(), + input_data.opaque(), &beta, dest_desc.handle(), dest2.opaque(), true, + workspace.opaque(), workspace_size_in_bytes); + + if (status != miopenStatusSuccess) { + LOG(ERROR) + << "failed to enqueue forward pooling (before backward) on stream: " + << ToString(status); + return false; + } + workspace_ptr = reinterpret_cast(workspace.opaque()); + } + } + status = wrap::miopenPoolingBackward( + miopen.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(), + output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(), + src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), + output_diff_data->opaque(), workspace_ptr); + + if (status != miopenStatusSuccess) { + LOG(ERROR) << "failed to enqueue backward pooling on stream: " + << ToString(status); + return false; + } + + return true; +} + bool MIOpenSupport::DoPoolBackward( Stream* stream, const dnn::PoolingDescriptor& pooling_dimensions, const dnn::BatchDescriptor& input_dimensions, @@ -4094,91 +4335,10 @@ bool MIOpenSupport::DoPoolBackward( const DeviceMemory& input_diff_data, DeviceMemory* output_diff_data, ScratchAllocator* workspace_allocator) { - auto miopen = miopen_->GetHandle(parent_, stream); - - // Alpha is the scaling factor for input. - float alpha = 1.0; - // Beta is the scaling factor for output. - float beta = 0.0; - - ScopedTensorDescriptor src_desc{input_dimensions, miopenFloat}; - ScopedTensorDescriptor dest_desc{output_dimensions, miopenFloat}; - ScopedPoolingDescriptor pooling_desc{pooling_dimensions}; - - DeviceMemory workspace; - size_t workspace_size_in_bytes = 0; - auto status = wrap::miopenPoolingGetWorkSpaceSize(dest_desc.handle(), - &workspace_size_in_bytes); - - if (status != miopenStatusSuccess) { - LOG(ERROR) - << "failed to obtain workspace size for backward pooling on stream: " - << ToString(status); - return false; - } - - // Allocate the workspace. - if (workspace_size_in_bytes > 0) { - assert(workspace_allocator); - auto allocated = - workspace_allocator->AllocateBytes(workspace_size_in_bytes); - if (!allocated.ok() || (workspace = allocated.ValueOrDie()) == nullptr) { - LOG(ERROR) << "Failed to allocate backward pooling workspace"; - return false; - } - } - - DeviceMemory dest2; // duplicated dest from forward: - int dest2_size = 0; - - // miopen requires the strides and dims to be ordered as BDYX. - std::vector dims64 = - output_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX); - - // miopen does not use strides and must have 4D tensor. - std::vector dims(4); - - std::transform(dims64.cbegin(), dims64.cend(), dims.begin(), - &CheckedNarrowing); - - dest2_size = dims[0] * dims[1] * dims[2] * dims[3] * sizeof(float); - - if (dest2_size > 0) { - assert(workspace_allocator); - auto allocated = workspace_allocator->AllocateBytes(dest2_size); - if (!allocated.ok() || (dest2 = allocated.ValueOrDie()) == nullptr) { - LOG(ERROR) << "Failed to allocate backward pooling workspace"; - return false; - } - } else { - LOG(ERROR) << "Failed to calculate tensor size to chain forward and " - "backward pooling"; - } - - status = wrap::miopenPoolingForward( - miopen.handle(), pooling_desc.handle(), &alpha, src_desc.handle(), - input_data.opaque(), &beta, dest_desc.handle(), dest2.opaque(), true, - workspace.opaque(), workspace_size_in_bytes); - - if (status != miopenStatusSuccess) { - LOG(ERROR) - << "failed to enqueue forward pooling (before backward) on stream: " - << ToString(status); - return false; - } - - status = wrap::miopenPoolingBackward( - miopen.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(), - dest2.opaque(), dest_desc.handle(), input_diff_data.opaque(), - src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), - output_diff_data->opaque(), workspace.opaque()); - - if (status != miopenStatusSuccess) { - LOG(ERROR) << "failed to enqueue backward pooling on stream: " - << ToString(status); - return false; - } - return true; + return DoPoolBackwardImpl(stream, pooling_dimensions, input_dimensions, + input_data, output_dimensions, output_data, + input_diff_data, output_diff_data, + workspace_allocator); } bool MIOpenSupport::DoPoolBackward( @@ -4190,91 +4350,10 @@ bool MIOpenSupport::DoPoolBackward( const DeviceMemory& input_diff_data, DeviceMemory* output_diff_data, ScratchAllocator* workspace_allocator) { - auto miopen = miopen_->GetHandle(parent_, stream); - - // Alpha is the scaling factor for input. - float alpha = 1.0; - // Beta is the scaling factor for output. - float beta = 0.0; - - ScopedTensorDescriptor src_desc{input_dimensions, miopenHalf}; - ScopedTensorDescriptor dest_desc{output_dimensions, miopenHalf}; - ScopedPoolingDescriptor pooling_desc{pooling_dimensions}; - - DeviceMemory workspace; - size_t workspace_size_in_bytes = 0; - auto status = wrap::miopenPoolingGetWorkSpaceSize(dest_desc.handle(), - &workspace_size_in_bytes); - - if (status != miopenStatusSuccess) { - LOG(ERROR) - << "failed to obtain workspace size for backward pooling on stream: " - << ToString(status); - return false; - } - - // Allocate the workspace. - if (workspace_size_in_bytes > 0) { - assert(workspace_allocator); - auto allocated = - workspace_allocator->AllocateBytes(workspace_size_in_bytes); - if (!allocated.ok() || (workspace = allocated.ValueOrDie()) == nullptr) { - LOG(ERROR) << "Failed to allocate backward pooling workspace"; - return false; - } - } - - DeviceMemory dest2; // duplicated dest from forward: - int dest2_size = 0; - - // miopen requires the strides and dims to be ordered as BDYX. - std::vector dims64 = - output_dimensions.full_dims(dnn::DataLayout::kBatchDepthYX); - - // miopen does not use strides and must have 4D tensor. - std::vector dims(4); - - std::transform(dims64.cbegin(), dims64.cend(), dims.begin(), - &CheckedNarrowing); - - dest2_size = dims[0] * dims[1] * dims[2] * dims[3] * sizeof(float); - - if (dest2_size > 0) { - assert(workspace_allocator); - auto allocated = workspace_allocator->AllocateBytes(dest2_size); - if (!allocated.ok() || (dest2 = allocated.ValueOrDie()) == nullptr) { - LOG(ERROR) << "Failed to allocate backward pooling workspace"; - return false; - } - } else { - LOG(ERROR) << "Failed to calculate tensor size to chain forward and " - "backward pooling"; - } - - status = wrap::miopenPoolingForward( - miopen.handle(), pooling_desc.handle(), &alpha, src_desc.handle(), - input_data.opaque(), &beta, dest_desc.handle(), dest2.opaque(), true, - workspace.opaque(), workspace_size_in_bytes); - - if (status != miopenStatusSuccess) { - LOG(ERROR) - << "failed to enqueue forward pooling (before backward) on stream: " - << ToString(status); - return false; - } - - status = wrap::miopenPoolingBackward( - miopen.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(), - dest2.opaque(), dest_desc.handle(), input_diff_data.opaque(), - src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(), - output_diff_data->opaque(), workspace.opaque()); - - if (status != miopenStatusSuccess) { - LOG(ERROR) << "failed to enqueue backward pooling on stream: " - << ToString(status); - return false; - } - return true; + return DoPoolBackwardImpl(stream, pooling_dimensions, input_dimensions, + input_data, output_dimensions, output_data, + input_diff_data, output_diff_data, + workspace_allocator); } bool MIOpenSupport::DoNormalizeWithDimensions( diff --git a/tensorflow/stream_executor/rocm/rocm_dnn.h b/tensorflow/stream_executor/rocm/rocm_dnn.h index 40e156b5f74..4f568702d96 100644 --- a/tensorflow/stream_executor/rocm/rocm_dnn.h +++ b/tensorflow/stream_executor/rocm/rocm_dnn.h @@ -20,6 +20,7 @@ limitations under the License. #define TENSORFLOW_STREAM_EXECUTOR_ROCM_ROCM_DNN_H_ #include "absl/synchronization/mutex.h" +#include "rocm/include/miopen/miopen.h" #include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/stream_executor/dnn.h" #include "tensorflow/stream_executor/lib/status.h" @@ -38,6 +39,39 @@ class MIOpenCTCLossDescriptor; // Opaque and unique identifier for the MIOpen plugin. extern const PluginId kMIOpenPlugin; +struct PoolingWorkspaceDescriptor { + std::vector input_dims; + std::vector output_dims; + dnn::PoolingDescriptor op; + int dtype; + uint64_t timestamp; + std::unique_ptr> workspace; + size_t workspace_size; + bool IsSame(const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type); +}; + +struct PoolingWorkspaceCache { + std::map cache; + const int trim_size = 1000; + const uint64_t memory_budget = 2e7; + uint64_t timestamp = 0; + uint64_t memory_used = 0; + bool find(const void* p, const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type, + PoolingWorkspaceDescriptor*& pdesc); + void insert(const void* p, const dnn::BatchDescriptor& input_dimensions, + const dnn::BatchDescriptor& output_dimensions, + const dnn::PoolingDescriptor& pooling_dimensions, int _type, + std::unique_ptr>& workspace, + size_t wsp_size, hipStream_t hip_stream); + + private: + void trim(hipStream_t hip_stream); +}; + // miopen-library based DNN support. For details on overridden interface // functions, see dnn.h. class MIOpenSupport : public dnn::DnnSupport { @@ -664,6 +698,10 @@ class MIOpenSupport : public dnn::DnnSupport { // Provide access to the MIOpen handle. std::unique_ptr miopen_; + PoolingWorkspaceCache m_pooling_cache; + bool m_pooling_cache_allowed = false; + bool m_pooling_cache_enabled = false; + template bool DoBatchNormalizationForwardImpl( Stream* stream, dnn::DataType input_data_type, @@ -847,6 +885,36 @@ class MIOpenSupport : public dnn::DnnSupport { ScratchAllocator* scratch_allocator, std::vector* out_algorithms); + port::Status DoCtcLossImpl( + Stream* stream, const MIOpenRnnStateTensorDescriptor& probs_desc, + const DeviceMemoryBase probs_data, absl::Span labels_data, + absl::Span labels_lengths_data, + absl::Span input_lengths_data, DeviceMemoryBase costs_data, + const MIOpenRnnStateTensorDescriptor& grads_desc, + DeviceMemoryBase grads_data, const MIOpenCTCLossDescriptor& ctc_loss_desc, + DeviceMemory scratch_memory); + + port::Status DoPrepareForCtcLoss( + Stream* stream, dnn::DataType element_type, + const dnn::RnnStateTensorDescriptor& probs_desc, + const dnn::RnnStateTensorDescriptor& grads_desc, + absl::Span labels_data, + absl::Span labels_lengths_data, + absl::Span input_lengths_data, + ScratchAllocator* scratch_allocator, + DeviceMemory* scratch_memory) override; + + template + bool DoPoolBackwardImpl(Stream* stream, + const dnn::PoolingDescriptor& pooling_dimensions, + const dnn::BatchDescriptor& input_dimensions, + const DeviceMemory& input_data, + const dnn::BatchDescriptor& output_dimensions, + const DeviceMemory& output_data, + const DeviceMemory& input_diff_data, + DeviceMemory* output_diff_data, + ScratchAllocator* workspace_allocator = nullptr); + SE_DISALLOW_COPY_AND_ASSIGN(MIOpenSupport); }; From 507c7549317221bcf5b418a66fd0212cd4a7443b Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 9 Jun 2020 17:47:16 +0100 Subject: [PATCH 0106/1390] Fix for pylint errors. Change-Id: Idd96d7a41fd459c86ab0f6fbb63e5d543509145d --- tensorflow/lite/python/convert.py | 3 ++- tensorflow/lite/python/lite.py | 27 ++++++++++--------- tensorflow/lite/python/lite_test.py | 10 ++++--- tensorflow/lite/python/optimize/calibrator.py | 3 ++- .../lite/python/optimize/calibrator_test.py | 3 ++- 5 files changed, 28 insertions(+), 18 deletions(-) diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 52edb700195..68e23634b2e 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -103,7 +103,8 @@ class OpsSet(enum.Enum): # significantly, but only slightly increases the model size. # WARNING: These ops are currently experimental and have not yet been finalized. # They are only compatible with CPU execution, and have not been optimized for production. - EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" + EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = \ + "EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8" def __str__(self): return self.value diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index 26c6f0855af..bed48860b00 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -251,7 +251,8 @@ class QuantizationMode(object): self.post_training_fp16()) def activations_type(self): - return constants.INT16 if self._is_int16x8_target_required() else constants.INT8 + return constants.INT16 if self._is_int16x8_target_required() \ + else constants.INT8 def converter_flags(self, inference_ty=None, inference_input_ty=None): """Flags to the converter.""" @@ -262,7 +263,8 @@ class QuantizationMode(object): if self.training_time_int8_allow_float(): return { - "inference_type": inference_ty if inference_ty else self.activations_type(), + "inference_type": inference_ty if inference_ty else \ + self.activations_type(), "inference_input_type": inference_input_ty if inference_input_ty else constants.FLOAT, "post_training_quantize": False, # disable dynamic range quantization @@ -359,15 +361,15 @@ class QuantizationMode(object): def _is_int16x8_target_required(self): return bool( - set(self._target_spec.supported_ops).intersection([ - OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 - ])) + set(self._target_spec.supported_ops).intersection([ + OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 + ])) def _is_allow_float(self): return bool( - set(self._target_spec.supported_ops).intersection([ - OpsSet.TFLITE_BUILTINS - ])) + set(self._target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS + ])) def _any_optimization_enabled(self): return bool( @@ -441,7 +443,8 @@ class TFLiteConverterBase(object): return _get_grappler_config(optimizers) def _calibrate_quantize_model(self, result, inference_input_type, - inference_output_type, activations_type, allow_float): + inference_output_type, activations_type, + allow_float): """Calibrate and quantize the model.""" if not isinstance(self.representative_dataset, RepresentativeDataset): self.representative_dataset = RepresentativeDataset( @@ -458,8 +461,8 @@ class TFLiteConverterBase(object): return _mlir_quantize(calibrated) else: return calibrate_quantize.calibrate_and_quantize( - self.representative_dataset.input_gen, inference_input_type, - inference_output_type, allow_float, activations_type) + self.representative_dataset.input_gen, inference_input_type, + inference_output_type, allow_float, activations_type) def _is_unknown_shapes_allowed(self): # Unknown dimensions are only allowed with the new converter. @@ -1992,7 +1995,7 @@ class TocoConverter(object): @classmethod @_deprecation.deprecated( - None, "Use `lite.TFLiteConverter.from_keras_model_file` instead.") + None, "Use `lite.TFLiteConverter.from_keras_model_file` instead.") def from_keras_model_file(cls, model_file, input_arrays=None, diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index 044b1211e17..cae49cb147f 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -882,11 +882,15 @@ class FromSessionTest(TestModels, parameterized.TestCase): @parameterized.named_parameters( # Quantize model to Int8: with enable mlir - ('UseTfliteBuiltinsIntEnableMLIR', [lite.OpsSet.TFLITE_BUILTINS_INT8], True), + ('UseTfliteBuiltinsIntEnableMLIR', + [lite.OpsSet.TFLITE_BUILTINS_INT8], True), # Quantize model to Int8: with disable mlir - ('UseTfliteBuiltinsIntDisableMLIR', [lite.OpsSet.TFLITE_BUILTINS_INT8], False), + ('UseTfliteBuiltinsIntDisableMLIR', + [lite.OpsSet.TFLITE_BUILTINS_INT8], False), # Quantize model to Int16: with disable mlir - ('UseTfliteBuiltinsInt16DisableMLIR', [lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], False)) + ('UseTfliteBuiltinsInt16DisableMLIR', + [lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], + False)) def testCalibrateAndQuantizeBuiltinInt(self, supported_ops, enable_mlir): with ops.Graph().as_default(): inp, output, calibration_gen = self._getCalibrationQuantizeModel() diff --git a/tensorflow/lite/python/optimize/calibrator.py b/tensorflow/lite/python/optimize/calibrator.py index 90c43fcddfa..2b08ec690ff 100644 --- a/tensorflow/lite/python/optimize/calibrator.py +++ b/tensorflow/lite/python/optimize/calibrator.py @@ -78,7 +78,8 @@ class Calibrator(object): computation, useful when targeting an integer-only backend. If False, an error will be thrown if an operation cannot be quantized, otherwise the model will fallback to float ops. - activations_type: A tf.dtype representing the desired type for activations. + activations_type: A tf.dtype representing the desired type for + activations. resize_input: A boolean. True if the shape of the sample data is different from the input. """ diff --git a/tensorflow/lite/python/optimize/calibrator_test.py b/tensorflow/lite/python/optimize/calibrator_test.py index f778c8a555d..d79d76b09ed 100644 --- a/tensorflow/lite/python/optimize/calibrator_test.py +++ b/tensorflow/lite/python/optimize/calibrator_test.py @@ -96,7 +96,8 @@ class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase): ('UseActivationTypeInt8 - EnableMlirQuantizer', constants.INT8), # Activation type Int16 ('UseActivationTypeInt16 - DisableEnableMlirQuantizer', constants.INT16)) - def test_calibration_with_quantization_multiple_inputs(self, activations_type): + def test_calibration_with_quantization_multiple_inputs(self, + activations_type): # Load multi add model from test data. # This model has 4 inputs of size (1, 8, 8, 3). model_path = resource_loader.get_path_to_datafile( From 3fc256a0dc31eae6711d2f7680493925f0fa4091 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 9 Jun 2020 17:51:37 +0100 Subject: [PATCH 0107/1390] Fix for pylint. Change-Id: If2674380c25eb8973e73a407b75660088098e6da --- tensorflow/lite/python/lite_test.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index cae49cb147f..e6661c82894 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -883,14 +883,15 @@ class FromSessionTest(TestModels, parameterized.TestCase): @parameterized.named_parameters( # Quantize model to Int8: with enable mlir ('UseTfliteBuiltinsIntEnableMLIR', - [lite.OpsSet.TFLITE_BUILTINS_INT8], True), + [lite.OpsSet.TFLITE_BUILTINS_INT8], True), # Quantize model to Int8: with disable mlir ('UseTfliteBuiltinsIntDisableMLIR', - [lite.OpsSet.TFLITE_BUILTINS_INT8], False), + [lite.OpsSet.TFLITE_BUILTINS_INT8], False), # Quantize model to Int16: with disable mlir ('UseTfliteBuiltinsInt16DisableMLIR', - [lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], - False)) + [lite.OpsSet.\ + EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], + False)) def testCalibrateAndQuantizeBuiltinInt(self, supported_ops, enable_mlir): with ops.Graph().as_default(): inp, output, calibration_gen = self._getCalibrationQuantizeModel() From 15b97dd44a42ae1ab4cc4192a5f7c820ac90b9ac Mon Sep 17 00:00:00 2001 From: nammbash Date: Tue, 9 Jun 2020 11:25:45 -0700 Subject: [PATCH 0108/1390] CPU Feature Guard MKL brand needs to change to oneDNNchange to one dnn --- tensorflow/core/platform/cpu_feature_guard.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/platform/cpu_feature_guard.cc b/tensorflow/core/platform/cpu_feature_guard.cc index c5a5c287283..a020d3fd70e 100644 --- a/tensorflow/core/platform/cpu_feature_guard.cc +++ b/tensorflow/core/platform/cpu_feature_guard.cc @@ -138,8 +138,16 @@ void InfoAboutUnusedCPUFeatures() { CheckIfFeatureUnused(CPUFeature::FMA, "FMA", missing_instructions); #endif // __FMA__ #endif // else of if defined(_MSC_VER) && !defined(__clang__) - if (!missing_instructions.empty()) { - LOG(INFO) << "This TensorFlow binary is optimized with Intel(R) MKL-DNN " + + string intel_library_official_name( + "Intel(R) oneAPI Deep Neural Network Library (oneDNN) "); +#ifndef INTEL_MKL + intel_library_official_name = "oneAPI Deep Neural Network Library (oneDNN) "; +#endif + + if (!missing_instructions.empty()) { + LOG(INFO) << "This TensorFlow binary is optimized with " + << intel_library_official_name << "to use the following CPU instructions in performance-" << "critical operations: " << missing_instructions << std::endl << "To enable them in other operations, rebuild TensorFlow " From 47674cac85b61f7a438c6970a6a7ca49946a2622 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 9 Jun 2020 22:27:57 +0100 Subject: [PATCH 0109/1390] Fix for pylint Change-Id: If03f60a3eebc7aed61c10870c545fe6035bcb2a3 --- tensorflow/lite/python/lite.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py index bed48860b00..06796ba820b 100644 --- a/tensorflow/lite/python/lite.py +++ b/tensorflow/lite/python/lite.py @@ -361,15 +361,15 @@ class QuantizationMode(object): def _is_int16x8_target_required(self): return bool( - set(self._target_spec.supported_ops).intersection([ - OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 - ])) + set(self._target_spec.supported_ops).intersection([ + OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 + ])) def _is_allow_float(self): return bool( - set(self._target_spec.supported_ops).intersection([ - OpsSet.TFLITE_BUILTINS - ])) + set(self._target_spec.supported_ops).intersection([ + OpsSet.TFLITE_BUILTINS + ])) def _any_optimization_enabled(self): return bool( From 84afc268a77f543fe64ecb45832701278a9eb129 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Fri, 5 Jun 2020 11:47:09 +0100 Subject: [PATCH 0110/1390] Documentation on the new experimental option for 16x8. --- .../performance/post_training_quantization.md | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tensorflow/lite/g3doc/performance/post_training_quantization.md b/tensorflow/lite/g3doc/performance/post_training_quantization.md index af7d9dbf02d..c48a2820d2f 100644 --- a/tensorflow/lite/g3doc/performance/post_training_quantization.md +++ b/tensorflow/lite/g3doc/performance/post_training_quantization.md @@ -151,6 +151,49 @@ The disadvantages of float16 quantization are as follows: to float32 when run on the CPU. (Note that the GPU delegate will not perform this dequantization, since it can operate on float16 data.) +### Integer only: 16-bit activations with 8-bit weights (experimental) + +This is an experimental quantization scheme. It is similar to the "integer only" +scheme, but activations are quantized based on their range to 16-bits, weights are +quantized in 8-bit integer and bias is quantized into 64-bit integer. +This is referred to as 16x8 quantization further. + +The main advantage of this quantization is that it can improve accuracy +significantly, but only slightly increase model size. + +
+import tensorflow as tf
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter.optimizations = [tf.lite.Optimize.DEFAULT]
+converter.target_spec.supported_types = [tf.lite.constants.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8]
+tflite_quant_model = converter.convert()
+
+ +If 16x8 quantization is not supported for some operators in the model, +then the model still can be quantized, but unsupported operators kept in float. +The following option should be added to the target_spec to allow this. +
+import tensorflow as tf
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter.optimizations = [tf.lite.Optimize.DEFAULT]
+converter.target_spec.supported_types = [tf.lite.constants.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
+tf.lite.OpsSet.TFLITE_BUILTINS]
+tflite_quant_model = converter.convert()
+
+ +Examples of the use cases where accuracy improvements provided by this quantization scheme include: +* super-resolution, +* audio signal processing such as noise cancelling and beamforming, +* image de-noising, +* HDR reconstruction from a single image. + +The disadvantage of this quantization is: + +* Currently inference is noticeably slower than 8-bit full integer due to the lack of optimized kernel implementation. +* Currently it is incompatible with the existing hardware accelerated TFLite delegates. + +Note: This is an experimental feature. + ### Model accuracy Since weights are quantized post training, there could be an accuracy loss, From dcfc2175c79ee6c610770b597c8d637daa1649bc Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Tue, 9 Jun 2020 22:35:57 +0100 Subject: [PATCH 0111/1390] Small change of comment per reviewer's note. Change-Id: I1233b95282befebfa0e6c06173f5e928aef60b22 --- tensorflow/lite/python/convert.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/python/convert.py b/tensorflow/lite/python/convert.py index 68e23634b2e..ec70f793f21 100644 --- a/tensorflow/lite/python/convert.py +++ b/tensorflow/lite/python/convert.py @@ -98,9 +98,9 @@ class OpsSet(enum.Enum): # int16 activations and int64 bias. # Specifying this will throw an error for operations that do not yet have # quantized implementations. - # This quantization mode should be used in models for super-resolution, + # This quantization mode may be used in models for super-resolution, # audio signal processing or image de-noising. It improves accuracy - # significantly, but only slightly increases the model size. + # significantly, but only slightly increases the model size. # WARNING: These ops are currently experimental and have not yet been finalized. # They are only compatible with CPU execution, and have not been optimized for production. EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8 = \ From aaf693d29bb5a52d6dfd8c106e45b2ff513e6d84 Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Tue, 2 Jun 2020 08:07:53 -0700 Subject: [PATCH 0112/1390] Added code to build multiinstance/multinode container --- .../ci_build/linux/mkl/Dockerfile.devel-mkl | 16 ++++-- .../ci_build/linux/mkl/build-dev-container.sh | 26 ++++++++++ .../linux/mkl/install_openmpi_horovod.sh | 49 +++++++++++++++++++ 3 files changed, 88 insertions(+), 3 deletions(-) create mode 100755 tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl index 45ccf67d707..3893f61d940 100755 --- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl +++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl @@ -15,6 +15,11 @@ ARG CONFIG_BFLOAT16_BUILD="" ARG ENABLE_SECURE_BUILD ARG BAZEL_VERSION="" ARG ENABLE_DNNL1="" +ARG ENABLE_HOROVOD="" +ARG OPENMPI_VERSION="" +ARG OPENMPI_DOWNLOAD_URL="" + +ENV DEBIAN_FRONTEND=noninteractive # Upgrade Bazel version if argument is passed RUN if [ "${BAZEL_VERSION}" != "" ]; then \ @@ -45,9 +50,6 @@ RUN ${PYTHON} set-build-env.py -p ${TARGET_PLATFORM} -f /root/.mkl.bazelrc \ # Pull the compiler flags we just wrote into root user's .bazelrc file RUN echo "import /root/.mkl.bazelrc" >>/root/.bazelrc -# Install futures>=0.17.1 for Python2.7 compatibility mode -RUN ${PIP} install future>=0.17.1 - RUN bazel --bazelrc=/root/.bazelrc build -c opt \ tensorflow/tools/pip_package:build_pip_package && \ bazel-bin/tensorflow/tools/pip_package/build_pip_package "${WHL_DIR}" && \ @@ -55,6 +57,14 @@ RUN bazel --bazelrc=/root/.bazelrc build -c opt \ rm -rf /root/.cache # Clean up Bazel cache when done. +#Install OpenMPI/Horovod +COPY install_openmpi_horovod.sh . +RUN if [ "${ENABLE_HOROVOD}" = "yes" ]; then \ + chmod +x install_openmpi_horovod.sh && \ + ${OPENMPI_VERSION} ${OPENMPI_DOWNLOAD_URL} install_openmpi_horovod.sh && \ + rm -rf install_openmpi_horovod.sh; \ + fi + # TensorBoard EXPOSE 6006 # IPython diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh index eceef65aa38..da647153cdb 100755 --- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh +++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh @@ -64,6 +64,9 @@ ENABLE_SECURE_BUILD=${ENABLE_SECURE_BUILD:-no} BAZEL_VERSION=${BAZEL_VERSION} BUILD_PY2_CONTAINERS=${BUILD_PY2_CONTAINERS:-no} ENABLE_DNNL1=${ENABLE_DNNL1:-no} +ENABLE_HOROVOD=${ENABLE_HOROVOD:-no} +OPENMPI_VERSION=${OPENMPI_VERSION} +OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} debug "ROOT_CONTAINER=${ROOT_CONTAINER}" debug "TF_ROOT_CONTAINER_TAG=${TF_ROOT_CONTAINER_TAG}" @@ -82,6 +85,9 @@ debug "TMP_DIR=${TMP_DIR}" debug "BAZEL_VERSION=${BAZEL_VERSION}" debug "BUILD_PY2_CONTAINERS=${BUILD_PY2_CONTAINERS}" debug "ENABLE_DNNL1=${ENABLE_DNNL1}" +debug "ENABLE_HOROVOD=${ENABLE_HOROVOD}" +debug "OPENMPI_VERSION=${OPENMPI_VERSION}" +debug "OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL}" function build_container() { @@ -131,6 +137,13 @@ function build_container() TF_DOCKER_BUILD_ARGS+=("--build-arg BAZEL_VERSION=${BAZEL_VERSION}") fi + # Add build arg for installing OpenMPI/Horovod + if [[ ${ENABLE_HOROVOD} == "yes" ]]; then + TF_DOCKER_BUILD_ARGS+=("--build-arg ENABLE_HOROVOD=${ENABLE_HOROVOD}") + TF_DOCKER_BUILD_ARGS+=("--build-arg OPENMPI_VERSION=${OPENMPI_VERSION}") + TF_DOCKER_BUILD_ARGS+=("--build-arg OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL}") + fi + # Perform docker build debug "Building docker image with image name and tag: ${TEMP_IMAGE_NAME}" CMD="${DOCKER_BINARY} build ${TF_DOCKER_BUILD_ARGS[@]} --no-cache --pull -t ${TEMP_IMAGE_NAME} -f Dockerfile.devel-mkl ." @@ -188,6 +201,19 @@ function test_container() die "FAIL: MKL enabled test in ${TEMP_IMAGE_NAME}" fi + # Test to check if horovod is installed successfully + debug "Test horovod in the container..." + if [[ ${ENABLE_HOROVOD} == "yes" ]]; then + HOROVOD_TEST_CMD=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'") + ${HOROVOD_TEST_CMD} + if [[ $? == "0" ]]; then + echo "PASS: HOROVOD installation test in ${TEMP_IMAGE_NAME}" + else + die "FAIL: HOROVOD installation test in ${TEMP_IMAGE_NAME}" + fi + fi + + # Stop the running docker container sleep 1 "${DOCKER_BINARY}" stop --time=0 ${CONTAINER_ID} diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh new file mode 100755 index 00000000000..d1b297726ed --- /dev/null +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# install OpenMPI, OpenSSH and Horovod + +set -e + +apt-get clean && apt-get update -y + +# Install Open MPI +OPENMPI_VERSION=${OPENMPI_VERSION:-openmpi-2.1.1} +OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL:-https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz} +echo "Installing OpenMPI version ${OPENMPI_VERSION}..." +echo "OpenMPI Download url ${OPENMPI_DOWNLOAD_URL}..." + +mkdir /tmp/openmpi +cd /tmp/openmpi +curl -fSsL -O ${OPENMPI_DOWNLOAD_URL} +tar zxf ${OPENMPI_VERSION}.tar.gz +cd ${OPENMPI_VERSION} +./configure --enable-mpirun-prefix-by-default +make -j $(nproc) all +make install +ldconfig +cd / +rm -rf /tmp/openmpi + +# Create a wrapper for OpenMPI to allow running as root by default +mv /usr/local/bin/mpirun /usr/local/bin/mpirun.real +echo '#!/bin/bash' > /usr/local/bin/mpirun +echo 'mpirun.real --allow-run-as-root "$@"' >> /usr/local/bin/mpirun +chmod a+x /usr/local/bin/mpirun + +# Configure OpenMPI to run good defaults: +echo "btl_tcp_if_exclude = lo,docker0" >> /usr/local/etc/openmpi-mca-params.conf + +#Check mpi version +echo 'OpenMPI version:' +mpirun --version + +# Install OpenSSH for MPI to communicate between containers +apt-get install -y --no-install-recommends --fix-missing openssh-client openssh-server libnuma-dev +mkdir -p /var/run/sshd +# Allow OpenSSH to talk to containers without asking for confirmation +cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new +echo " StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new +mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config + +#Install Horovod +HOROVOD_WITH_TENSORFLOW=1 +python3 -m pip install --no-cache-dir horovod==0.19.1 From ff359d4a48aeb1905f767d32e7da1a2d01d4ce6a Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Tue, 2 Jun 2020 13:46:16 -0700 Subject: [PATCH 0113/1390] Setting default values for OpenMPI versions --- .../tools/ci_build/linux/mkl/Dockerfile.devel-mkl | 2 +- .../ci_build/linux/mkl/install_openmpi_horovod.sh | 11 +++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl index 3893f61d940..1fd54ff703f 100755 --- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl +++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl @@ -61,7 +61,7 @@ RUN bazel --bazelrc=/root/.bazelrc build -c opt \ COPY install_openmpi_horovod.sh . RUN if [ "${ENABLE_HOROVOD}" = "yes" ]; then \ chmod +x install_openmpi_horovod.sh && \ - ${OPENMPI_VERSION} ${OPENMPI_DOWNLOAD_URL} install_openmpi_horovod.sh && \ + ./install_openmpi_horovod.sh ${OPENMPI_VERSION} ${OPENMPI_DOWNLOAD_URL} && \ rm -rf install_openmpi_horovod.sh; \ fi diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index d1b297726ed..4c8b04f6024 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -5,9 +5,16 @@ set -e apt-get clean && apt-get update -y +# Set default +if [[ $# -gt 1 ]]; then + OPENMPI_VERSION="${1}" + OPENMPI_DOWNLOAD_URL="${2}" +else + OPENMPI_VERSION=openmpi-2.1.1 + OPENMPI_DOWNLOAD_URL=https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz +fi + # Install Open MPI -OPENMPI_VERSION=${OPENMPI_VERSION:-openmpi-2.1.1} -OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL:-https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz} echo "Installing OpenMPI version ${OPENMPI_VERSION}..." echo "OpenMPI Download url ${OPENMPI_DOWNLOAD_URL}..." From aa5bfd35fa5292d820493483da8540f8a6386c5f Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Wed, 3 Jun 2020 12:15:24 -0700 Subject: [PATCH 0114/1390] Added license to the shell script --- .../linux/mkl/install_openmpi_horovod.sh | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index 4c8b04f6024..0f5a670f0f2 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -1,5 +1,20 @@ #!/usr/bin/env bash -# install OpenMPI, OpenSSH and Horovod +# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +# Install OpenMPI, OpenSSH and Horovod in Intel(R) MKL support +# Usage: install_openmpi_horovod.sh [openmpi version] [openmpi download url] set -e From 615d3ce1af92614a3285807caaf42f50acd66fae Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Wed, 3 Jun 2020 13:14:39 -0700 Subject: [PATCH 0115/1390] Added install futures for backward compatibility --- tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl index 1fd54ff703f..f4ab7ba21c4 100755 --- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl +++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl @@ -50,6 +50,9 @@ RUN ${PYTHON} set-build-env.py -p ${TARGET_PLATFORM} -f /root/.mkl.bazelrc \ # Pull the compiler flags we just wrote into root user's .bazelrc file RUN echo "import /root/.mkl.bazelrc" >>/root/.bazelrc +# Install futures>=0.17.1 for Python2.7 compatibility mode +RUN ${PIP} install future>=0.17.1 + RUN bazel --bazelrc=/root/.bazelrc build -c opt \ tensorflow/tools/pip_package:build_pip_package && \ bazel-bin/tensorflow/tools/pip_package/build_pip_package "${WHL_DIR}" && \ From 7a048082c1c7aa9057c8448b77ee4cde069ec3a7 Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Wed, 3 Jun 2020 13:19:53 -0700 Subject: [PATCH 0116/1390] Removed extra line --- tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh index da647153cdb..7278724ff64 100755 --- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh +++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh @@ -202,8 +202,8 @@ function test_container() fi # Test to check if horovod is installed successfully - debug "Test horovod in the container..." if [[ ${ENABLE_HOROVOD} == "yes" ]]; then + debug "Test horovod in the container..." HOROVOD_TEST_CMD=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'") ${HOROVOD_TEST_CMD} if [[ $? == "0" ]]; then @@ -213,7 +213,6 @@ function test_container() fi fi - # Stop the running docker container sleep 1 "${DOCKER_BINARY}" stop --time=0 ${CONTAINER_ID} From 50f0ba885bc112b24b437c8a974c6a8deaace96b Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Wed, 3 Jun 2020 20:05:47 -0700 Subject: [PATCH 0117/1390] Added parameter to pass horovod version and fixed comments --- .../ci_build/linux/mkl/Dockerfile.devel-mkl | 5 ++-- .../ci_build/linux/mkl/build-dev-container.sh | 3 +++ .../linux/mkl/install_openmpi_horovod.sh | 26 ++++++++----------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl index f4ab7ba21c4..8a5a0a42050 100755 --- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl +++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl @@ -18,6 +18,7 @@ ARG ENABLE_DNNL1="" ARG ENABLE_HOROVOD="" ARG OPENMPI_VERSION="" ARG OPENMPI_DOWNLOAD_URL="" +ARG HOROVOD_VERSION="" ENV DEBIAN_FRONTEND=noninteractive @@ -60,11 +61,11 @@ RUN bazel --bazelrc=/root/.bazelrc build -c opt \ rm -rf /root/.cache # Clean up Bazel cache when done. -#Install OpenMPI/Horovod +# Install OpenMPI/Horovod COPY install_openmpi_horovod.sh . RUN if [ "${ENABLE_HOROVOD}" = "yes" ]; then \ chmod +x install_openmpi_horovod.sh && \ - ./install_openmpi_horovod.sh ${OPENMPI_VERSION} ${OPENMPI_DOWNLOAD_URL} && \ + ./install_openmpi_horovod.sh OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} HOROVOD_VERSION=${HOROVOD_VERSION} && \ rm -rf install_openmpi_horovod.sh; \ fi diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh index 7278724ff64..e9d7f1ff388 100755 --- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh +++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh @@ -67,6 +67,7 @@ ENABLE_DNNL1=${ENABLE_DNNL1:-no} ENABLE_HOROVOD=${ENABLE_HOROVOD:-no} OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} +HOROVOD_VERSION=${HOROVOD_VERSION} debug "ROOT_CONTAINER=${ROOT_CONTAINER}" debug "TF_ROOT_CONTAINER_TAG=${TF_ROOT_CONTAINER_TAG}" @@ -88,6 +89,7 @@ debug "ENABLE_DNNL1=${ENABLE_DNNL1}" debug "ENABLE_HOROVOD=${ENABLE_HOROVOD}" debug "OPENMPI_VERSION=${OPENMPI_VERSION}" debug "OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL}" +debug "HOROVOD_VERSION=${HOROVOD_VERSION}" function build_container() { @@ -142,6 +144,7 @@ function build_container() TF_DOCKER_BUILD_ARGS+=("--build-arg ENABLE_HOROVOD=${ENABLE_HOROVOD}") TF_DOCKER_BUILD_ARGS+=("--build-arg OPENMPI_VERSION=${OPENMPI_VERSION}") TF_DOCKER_BUILD_ARGS+=("--build-arg OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL}") + TF_DOCKER_BUILD_ARGS+=("--build-arg HOROVOD_VERSION=${HOROVOD_VERSION}") fi # Perform docker build diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index 0f5a670f0f2..b8d9739ceb6 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# Copyright 2016 The TensorFlow Authors. All Rights Reserved. +# Copyright 2020 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,25 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -# Install OpenMPI, OpenSSH and Horovod in Intel(R) MKL support -# Usage: install_openmpi_horovod.sh [openmpi version] [openmpi download url] +# Install OpenMPI, OpenSSH and Horovod during Intel(R) MKL container build +# Usage: install_openmpi_horovod.sh [OPENMPI_VERSION=] [OPENMPI_DOWNLOAD_URL=] [HOROVOD_VERSION=] set -e apt-get clean && apt-get update -y # Set default -if [[ $# -gt 1 ]]; then - OPENMPI_VERSION="${1}" - OPENMPI_DOWNLOAD_URL="${2}" -else - OPENMPI_VERSION=openmpi-2.1.1 - OPENMPI_DOWNLOAD_URL=https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz -fi +OPENMPI_VERSION=${OPENMPI_VERSION:-openmpi-2.1.1} +OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL:-https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz} +HOROVOD_VERSION=${HOROVOD_VERSION:-0.19.1} # Install Open MPI -echo "Installing OpenMPI version ${OPENMPI_VERSION}..." -echo "OpenMPI Download url ${OPENMPI_DOWNLOAD_URL}..." +echo "Installing OpenMPI version ${OPENMPI_VERSION} ..." +echo "OpenMPI Download url ${OPENMPI_DOWNLOAD_URL} ..." mkdir /tmp/openmpi cd /tmp/openmpi @@ -54,7 +50,7 @@ chmod a+x /usr/local/bin/mpirun # Configure OpenMPI to run good defaults: echo "btl_tcp_if_exclude = lo,docker0" >> /usr/local/etc/openmpi-mca-params.conf -#Check mpi version +# Check mpi version echo 'OpenMPI version:' mpirun --version @@ -66,6 +62,6 @@ cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.ne echo " StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config -#Install Horovod +# Install Horovod HOROVOD_WITH_TENSORFLOW=1 -python3 -m pip install --no-cache-dir horovod==0.19.1 +python3 -m pip install --no-cache-dir horovod==${HOROVOD_VERSION} From 1ed2ab4638c56b3cce6d0f85efeaea3600b75214 Mon Sep 17 00:00:00 2001 From: justkw Date: Thu, 4 Jun 2020 09:22:07 -0700 Subject: [PATCH 0118/1390] Adding parameter to use --nightly_flag to install specific packages if building the nightly build --- tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl | 5 +++-- tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh | 7 +++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl index 8a5a0a42050..a78d13c7755 100755 --- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl +++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl @@ -19,6 +19,7 @@ ARG ENABLE_HOROVOD="" ARG OPENMPI_VERSION="" ARG OPENMPI_DOWNLOAD_URL="" ARG HOROVOD_VERSION="" +ARG TF_NIGHTLY_FLAG="" ENV DEBIAN_FRONTEND=noninteractive @@ -56,8 +57,8 @@ RUN ${PIP} install future>=0.17.1 RUN bazel --bazelrc=/root/.bazelrc build -c opt \ tensorflow/tools/pip_package:build_pip_package && \ - bazel-bin/tensorflow/tools/pip_package/build_pip_package "${WHL_DIR}" && \ - ${PIP} --no-cache-dir install --upgrade "${WHL_DIR}"/tensorflow-*.whl && \ + bazel-bin/tensorflow/tools/pip_package/build_pip_package "${TF_NIGHTLY_FLAG}" "${WHL_DIR}" && \ + ${PIP} --no-cache-dir install --upgrade "${WHL_DIR}"/*.whl && \ rm -rf /root/.cache # Clean up Bazel cache when done. diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh index e9d7f1ff388..83b3ebaf9c9 100755 --- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh +++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh @@ -68,6 +68,7 @@ ENABLE_HOROVOD=${ENABLE_HOROVOD:-no} OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} HOROVOD_VERSION=${HOROVOD_VERSION} +IS_NIGHTLY=${IS_NIGHTLY:-no} debug "ROOT_CONTAINER=${ROOT_CONTAINER}" debug "TF_ROOT_CONTAINER_TAG=${TF_ROOT_CONTAINER_TAG}" @@ -90,6 +91,7 @@ debug "ENABLE_HOROVOD=${ENABLE_HOROVOD}" debug "OPENMPI_VERSION=${OPENMPI_VERSION}" debug "OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL}" debug "HOROVOD_VERSION=${HOROVOD_VERSION}" +debug "IS_NIGHTLY=${IS_NIGHTLY}" function build_container() { @@ -147,6 +149,11 @@ function build_container() TF_DOCKER_BUILD_ARGS+=("--build-arg HOROVOD_VERSION=${HOROVOD_VERSION}") fi + # Add build arg --nightly_flag for the nightly build + if [[ ${IS_NIGHTLY} == "yes" ]]; then + TF_DOCKER_BUILD_ARGS+=("--build-arg TF_NIGHTLY_FLAG=--nightly_flag") + fi + # Perform docker build debug "Building docker image with image name and tag: ${TEMP_IMAGE_NAME}" CMD="${DOCKER_BINARY} build ${TF_DOCKER_BUILD_ARGS[@]} --no-cache --pull -t ${TEMP_IMAGE_NAME} -f Dockerfile.devel-mkl ." From 83decf0d382b76771e2b1ad4fa43d208d5e40eb0 Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Mon, 8 Jun 2020 07:50:14 -0700 Subject: [PATCH 0119/1390] Support multiple OS --- .../ci_build/linux/mkl/install_openmpi_horovod.sh | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index b8d9739ceb6..6044927d2ce 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -55,7 +55,18 @@ echo 'OpenMPI version:' mpirun --version # Install OpenSSH for MPI to communicate between containers -apt-get install -y --no-install-recommends --fix-missing openssh-client openssh-server libnuma-dev +( apt-get update && apt-get install -y --no-install-recommends --fix-missing \ + libnuma-dev \ + openssh-server \ + openssh-client \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* ) || \ + ( yum -y update && yum -y install \ + numactl-devel \ + openssh-server \ + openssh-clients \ + yum clean all ) || \ + ( echo "Unsupported Linux distribution. Aborting!" && exit 1 ) mkdir -p /var/run/sshd # Allow OpenSSH to talk to containers without asking for confirmation cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new From 371c2e6f4f3f233041eda2d292a13824d98d769f Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Mon, 8 Jun 2020 08:52:49 -0700 Subject: [PATCH 0120/1390] Bug fix --- .../tools/ci_build/linux/mkl/install_openmpi_horovod.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index 6044927d2ce..276d9945ab6 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -58,13 +58,13 @@ mpirun --version ( apt-get update && apt-get install -y --no-install-recommends --fix-missing \ libnuma-dev \ openssh-server \ - openssh-client \ + openssh-client && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* ) || \ ( yum -y update && yum -y install \ numactl-devel \ openssh-server \ - openssh-clients \ + openssh-clients && \ yum clean all ) || \ ( echo "Unsupported Linux distribution. Aborting!" && exit 1 ) mkdir -p /var/run/sshd From da18384ad585b2d88a08119268c9a7134ee36bf5 Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Mon, 8 Jun 2020 09:51:25 -0700 Subject: [PATCH 0121/1390] Bug fix --- tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index 276d9945ab6..b765dbd70a6 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -58,7 +58,7 @@ mpirun --version ( apt-get update && apt-get install -y --no-install-recommends --fix-missing \ libnuma-dev \ openssh-server \ - openssh-client && \ + openssh-clients && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* ) || \ ( yum -y update && yum -y install \ From ab86bb82faabec7b1d29c61df1cae0b45d0b0e8e Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Tue, 9 Jun 2020 07:53:05 -0700 Subject: [PATCH 0122/1390] Bug fix --- .../linux/mkl/install_openmpi_horovod.sh | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh index b765dbd70a6..aec40543a17 100755 --- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh +++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh @@ -14,15 +14,14 @@ # limitations under the License. # ============================================================================== # Install OpenMPI, OpenSSH and Horovod during Intel(R) MKL container build -# Usage: install_openmpi_horovod.sh [OPENMPI_VERSION=] [OPENMPI_DOWNLOAD_URL=] [HOROVOD_VERSION=] +# Usage: install_openmpi_horovod.sh [OPENMPI_VERSION=] [OPENMPI_DOWNLOAD_URL=] +# [HOROVOD_VERSION=] set -e -apt-get clean && apt-get update -y - # Set default OPENMPI_VERSION=${OPENMPI_VERSION:-openmpi-2.1.1} -OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL:-https://www.open-mpi.org/software/ompi/v2.1/downloads/${OPENMPI_VERSION}.tar.gz} +OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL:-https://www.open-mpi.org/software/ompi/v2.1/downloads/openmpi-2.1.1.tar.gz} HOROVOD_VERSION=${HOROVOD_VERSION:-0.19.1} # Install Open MPI @@ -55,18 +54,20 @@ echo 'OpenMPI version:' mpirun --version # Install OpenSSH for MPI to communicate between containers -( apt-get update && apt-get install -y --no-install-recommends --fix-missing \ - libnuma-dev \ - openssh-server \ - openssh-clients && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* ) || \ - ( yum -y update && yum -y install \ - numactl-devel \ - openssh-server \ - openssh-clients && \ - yum clean all ) || \ - ( echo "Unsupported Linux distribution. Aborting!" && exit 1 ) +apt-get clean && apt-get update && apt-get install -y --no-install-recommends --fix-missing \ + openssh-client openssh-server libnuma-dev && \ + rm -rf /var/lib/apt/lists/* +if [[ $? == "0" ]]; then + echo "PASS: OpenSSH installation" +else + yum -y update && yum -y install numactl-devel openssh-server openssh-clients && \ + yum clean all + if [[ $? == "0" ]]; then + echo "PASS: OpenSSH installation" + else + echo "Unsupported Linux distribution. Aborting!" && exit 1 + fi +fi mkdir -p /var/run/sshd # Allow OpenSSH to talk to containers without asking for confirmation cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new From ae408bb512e614469e24ccf0db6c031f6aeac030 Mon Sep 17 00:00:00 2001 From: shwetaoj Date: Tue, 9 Jun 2020 18:00:47 -0700 Subject: [PATCH 0123/1390] remvoe trailing white space --- tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh index 83b3ebaf9c9..6e789a54e87 100755 --- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh +++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh @@ -214,7 +214,7 @@ function test_container() # Test to check if horovod is installed successfully if [[ ${ENABLE_HOROVOD} == "yes" ]]; then debug "Test horovod in the container..." - HOROVOD_TEST_CMD=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'") + HOROVOD_TEST_CMD=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'") ${HOROVOD_TEST_CMD} if [[ $? == "0" ]]; then echo "PASS: HOROVOD installation test in ${TEMP_IMAGE_NAME}" From 50403ea24fa2a4a9eda2a74cfc1879ab396eb8eb Mon Sep 17 00:00:00 2001 From: Prashant Kumar Date: Tue, 2 Jun 2020 09:40:47 +0000 Subject: [PATCH 0124/1390] [MLIR]Add conversions dot op from LHLO to Affine Add conversions from MLIR Lhlo dot op to affine loops.These conversions are run as part of -lhlo-legalize-to-affine pass. Signed-off-by: Prashant Kumar --- .../xla/tests/lhlo-legalize-to-affine.mlir | 36 ++++++++++ .../xla/transforms/lhlo_legalize_to_affine.cc | 70 ++++++++++++++++--- .../xla/transforms/map_xla_to_scalar_op.h | 24 +++++++ 3 files changed, 122 insertions(+), 8 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-affine.mlir b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-affine.mlir index aaf65b5a38a..483204cf0d5 100644 --- a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-affine.mlir +++ b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-affine.mlir @@ -143,3 +143,39 @@ func @int_sub_op(%lhs: memref<7xi32>, %rhs: memref<7xi32>, : (memref<7xi32>, memref<7xi32>, memref<7xi32>) -> () return } + +// Dot tests. +// CHECK-LABEL: func @float_dot_op +func @float_dot_op(%lhs: memref<7x3xf32>, %rhs: + memref<3x4xf32>, %result: memref<7x4xf32> ) -> () { + // CHECK-NEXT: affine.for %[[I:.*]] = 0 to 7 { + // CHECK-NEXT: affine.for %[[J:.*]] = 0 to 4 { + // CHECK-NEXT: affine.for %[[K:.*]] = 0 to 3 { + // CHECK-NEXT: %[[LHS:.*]] = affine.load %{{.*}}[%[[I]], %[[K]]] : memref<7x3xf32> + // CHECK-NEXT: %[[RHS:.*]] = affine.load %{{.*}}[%[[K]], %[[J]]] : memref<3x4xf32> + // CHECK-NEXT: %[[RESULT:.*]] = affine.load %{{.*}}[%[[I]], %[[J]]] : memref<7x4xf32> + // CHECK-NEXT: %[[MULT:.*]] = mulf %[[LHS]], %[[RHS]] : f32 + // CHECK-NEXT: %[[ADD:.*]] = addf %[[MULT]], %[[RESULT]] : f32 + // CHECK-NEXT: affine.store %[[ADD]], %{{.*}}[%[[I]], %[[J]]] : memref<7x4xf32> + // CHECK: return + "xla_lhlo.dot"(%lhs, %rhs, %result) : + (memref<7x3xf32>, memref<3x4xf32>, memref<7x4xf32>) -> () + return +} +// CHECK-LABEL: func @int_dot_op +func @int_dot_op(%lhs: memref<7x3xi32>, %rhs: + memref<3x4xi32>, %result: memref<7x4xi32> ) -> () { + // CHECK-NEXT: affine.for %[[I:.*]] = 0 to 7 { + // CHECK-NEXT: affine.for %[[J:.*]] = 0 to 4 { + // CHECK-NEXT: affine.for %[[K:.*]] = 0 to 3 { + // CHECK-NEXT: %[[LHS:.*]] = affine.load %{{.*}}[%[[I]], %[[K]]] : memref<7x3xi32> + // CHECK-NEXT: %[[RHS:.*]] = affine.load %{{.*}}[%[[K]], %[[J]]] : memref<3x4xi32> + // CHECK-NEXT: %[[RESULT:.*]] = affine.load %{{.*}}[%[[I]], %[[J]]] : memref<7x4xi32> + // CHECK-NEXT: %[[MULT:.*]] = muli %[[LHS]], %[[RHS]] : i32 + // CHECK-NEXT: %[[ADD:.*]] = addi %[[MULT]], %[[RESULT]] : i32 + // CHECK-NEXT: affine.store %[[ADD]], %{{.*}}[%[[I]], %[[J]]] : memref<7x4xi32> + // CHECK: return + "xla_lhlo.dot"(%lhs, %rhs, %result) : + (memref<7x3xi32>, memref<3x4xi32>, memref<7x4xi32>) -> () + return +} diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_affine.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_affine.cc index f7f5537f882..4cc77292494 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_affine.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_affine.cc @@ -17,13 +17,13 @@ limitations under the License. #include "absl/memory/memory.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project -#include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project -#include "mlir/IR/Attributes.h" // from @llvm-project -#include "mlir/IR/Location.h" // from @llvm-project -#include "mlir/IR/MLIRContext.h" // from @llvm-project -#include "mlir/IR/PatternMatch.h" // from @llvm-project -#include "mlir/IR/StandardTypes.h" // from @llvm-project -#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Location.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/StandardTypes.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/mlir/xla/ir/lhlo_ops.h" #include "tensorflow/compiler/mlir/xla/transforms/map_xla_to_scalar_op.h" @@ -31,6 +31,59 @@ namespace mlir { namespace xla_lhlo { namespace { +struct DotOpConverter : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + // Supports only rank-2 tensors for LHS and RHS. + LogicalResult matchAndRewrite(DotOp op, + PatternRewriter& rewriter) const override { + Value lhs = op.lhs(); + Value rhs = op.rhs(); + MemRefType lhs_type = lhs.getType().cast(); + MemRefType rhs_type = rhs.getType().cast(); + Type element_type = lhs_type.getElementType(); + ArrayRef shape_lhs = lhs_type.getShape(); + ArrayRef shape_rhs = rhs_type.getShape(); + + if ((lhs_type.getRank() != 2) || (rhs_type.getRank() != 2)) { + return failure(); + } + SmallVector lhs_indices, rhs_indices, result_indices; + const auto& loc = op.getLoc(); + + // Create the canonical ijk form of matmul. + auto forOp = rewriter.create(loc, 0, shape_lhs[0]); + lhs_indices.push_back(forOp.getInductionVar()); + result_indices.push_back(forOp.getInductionVar()); + + rewriter.setInsertionPointToStart(forOp.getBody()); + forOp = rewriter.create(loc, 0, shape_rhs.back()); + result_indices.push_back(forOp.getInductionVar()); + rhs_indices.resize(2); + rhs_indices[1] = forOp.getInductionVar(); + + rewriter.setInsertionPointToStart(forOp.getBody()); + forOp = rewriter.create(loc, 0, shape_rhs.front()); + lhs_indices.push_back(forOp.getInductionVar()); + rhs_indices[0] = forOp.getInductionVar(); + + // Construct the innermost loop body. + rewriter.setInsertionPointToStart(forOp.getBody()); + auto l = rewriter.create(loc, lhs, lhs_indices); + auto r = rewriter.create(loc, rhs, rhs_indices); + auto result = + rewriter.create(loc, op.output(), result_indices); + Value op_result = xla_lhlo::XlaOpToStdScalarOp::map( + op, element_type, {l, r, result}, &rewriter); + if (op_result == nullptr) { + return failure(); + } + rewriter.create(loc, op_result, op.output(), result_indices); + rewriter.eraseOp(op); + return success(); + } +}; + template struct BinaryOpConverter : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -77,7 +130,8 @@ void populateLHLOToAffineConversionPattern(MLIRContext* context, BinaryOpConverter, BinaryOpConverter, BinaryOpConverter, - BinaryOpConverter>(context); + BinaryOpConverter, + DotOpConverter>(context); // clang-format on } diff --git a/tensorflow/compiler/mlir/xla/transforms/map_xla_to_scalar_op.h b/tensorflow/compiler/mlir/xla/transforms/map_xla_to_scalar_op.h index c317dc36b3c..07d4d3dd138 100644 --- a/tensorflow/compiler/mlir/xla/transforms/map_xla_to_scalar_op.h +++ b/tensorflow/compiler/mlir/xla/transforms/map_xla_to_scalar_op.h @@ -287,6 +287,30 @@ inline Value MapLhloOpToStdScalarOp( return nullptr; } +template <> +inline Value MapLhloOpToStdScalarOp( + Location loc, ArrayRef result_types, ArrayRef args, + OpBuilder* b) { + // Dot Op converter from lhlo to affine only accepts float and integer types. + const auto& lhs = args[0]; + const auto& rhs = args[1]; + const auto& result = args[2]; + Type element_type = lhs.getType(); + if (element_type.isa()) { + Value float_mul = MapLhloOpToStdScalarOpImpl{}( + loc, result_types, {lhs, rhs}, b); + return MapLhloOpToStdScalarOpImpl{}( + loc, result_types, {float_mul, result}, b); + } + if (element_type.isa()) { + Value int_mul = MapLhloOpToStdScalarOpImpl{}( + loc, result_types, {lhs, rhs}, b); + return MapLhloOpToStdScalarOpImpl{}( + loc, result_types, {int_mul, result}, b); + } + return nullptr; +} + template <> inline Value MapLhloOpToStdScalarOp( Location loc, ArrayRef result_types, ArrayRef args, From 77e5fb550de9aeed3a3454de06f1a7571d5e5ba3 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Wed, 10 Jun 2020 10:47:54 +0200 Subject: [PATCH 0125/1390] Fix compile errors by using the new MicroAllocator::Create() --- tensorflow/lite/micro/micro_allocator_test.cc | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index a71115dc67b..f52742ca723 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -288,9 +288,9 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator allocator(&context, model, arena, arena_size, - micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( + &context, model, arena, arena_size, micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); // Since all of the tensors are online planned and the model structure is // identical to that in TestAllocationForModelsWithBranches, @@ -337,9 +337,9 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator allocator(&context, model, arena, arena_size, - micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( + &context, model, arena, arena_size, micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); @@ -382,9 +382,9 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator allocator(&context, model, arena, arena_size, - micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( + &context, model, arena, arena_size, micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); @@ -430,9 +430,9 @@ TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator allocator(&context, model, arena, arena_size, - micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator.FinishTensorAllocation()); + tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( + &context, model, arena, arena_size, micro_test::reporter); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); From 764e2641e05e58ccd7ffdfcff0307f21fbb3fed0 Mon Sep 17 00:00:00 2001 From: Niranjan Hasabnis Date: Wed, 10 Jun 2020 09:15:38 -0700 Subject: [PATCH 0126/1390] Addressing review comments --- tensorflow/core/kernels/BUILD | 1 - tensorflow/core/kernels/mkl_conv_ops_test.cc | 23 +-- tensorflow/core/kernels/mkl_relu_op_test.cc | 200 ++++++++----------- tensorflow/core/util/mkl_util.h | 15 ++ 4 files changed, 103 insertions(+), 136 deletions(-) diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index a68eeb4479b..5abedf441a1 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -8124,7 +8124,6 @@ tf_cc_test_mkl( "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", "//tensorflow/core:test", - "//tensorflow/core:test_main", "//tensorflow/core:testlib", ], ) diff --git a/tensorflow/core/kernels/mkl_conv_ops_test.cc b/tensorflow/core/kernels/mkl_conv_ops_test.cc index 9d11b0fb006..7e7b78e004d 100644 --- a/tensorflow/core/kernels/mkl_conv_ops_test.cc +++ b/tensorflow/core/kernels/mkl_conv_ops_test.cc @@ -103,20 +103,6 @@ static Tensor GetFilterSizesTensor(const Conv2DDimensions& dims) { dims.input_depth, dims.filter_count}); } -#if defined(INTEL_MKL_DNN_ONLY) -static Tensor NonMklTensor() { - MklDnnShape non_mkl_shape; - non_mkl_shape.SetMklTensor(false); - - auto size = static_cast(non_mkl_shape.GetSerializeBufferSize()); - Tensor tensor(DT_UINT8, {size}); - - non_mkl_shape.SerializeMklDnnShape(tensor.flat().data(), - size * sizeof(uint8)); - return tensor; -} -#endif - static Graph* DefaultConv2D(const Conv2DDimensions& dims) { auto* graph = new Graph(OpRegistry::Global()); @@ -148,7 +134,8 @@ static Graph* MklConv2D(const Conv2DDimensions& dims) { Node* input = test::graph::Constant(graph, input_t, "input"); Node* filter = test::graph::Constant(graph, filter_t, "filter"); - Node* not_mkl_shape = test::graph::Constant(graph, NonMklTensor(), "not_mkl"); + Node* not_mkl_shape = + test::graph::Constant(graph, GetMklMetaTensor(), "not_mkl"); Node* conv2d; TF_CHECK_OK(NodeBuilder(graph->NewName("mkl_conv_2d"), "_MklConv2D") @@ -207,7 +194,8 @@ static Graph* MklConv2DBwdInput(const Conv2DDimensions& dims) { Node* out_backprop = test::graph::Constant(graph, out_backprop_t, "out_backprop"); - Node* not_mkl_shape = test::graph::Constant(graph, NonMklTensor(), "not_mkl"); + Node* not_mkl_shape = + test::graph::Constant(graph, GetMklMetaTensor(), "not_mkl"); Node* conv2d_bwd_input; TF_CHECK_OK(NodeBuilder(graph->NewName("conv_2d_bwd_input"), @@ -271,7 +259,8 @@ static Graph* MklConv2DBwdFilter(const Conv2DDimensions& dims) { Node* out_backprop = test::graph::Constant(graph, out_backprop_t, "out_backprop"); - Node* not_mkl_shape = test::graph::Constant(graph, NonMklTensor(), "not_mkl"); + Node* not_mkl_shape = + test::graph::Constant(graph, GetMklMetaTensor(), "not_mkl"); Node* conv2d_bwd_filter; TF_CHECK_OK(NodeBuilder(graph->NewName("conv_2d_bwd_filter"), diff --git a/tensorflow/core/kernels/mkl_relu_op_test.cc b/tensorflow/core/kernels/mkl_relu_op_test.cc index 7a3dffef0de..30f75cd23df 100644 --- a/tensorflow/core/kernels/mkl_relu_op_test.cc +++ b/tensorflow/core/kernels/mkl_relu_op_test.cc @@ -13,10 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#undef INTEL_MKL - #ifdef INTEL_MKL +#include "absl/strings/match.h" #include "tensorflow/cc/ops/const_op.h" #include "tensorflow/cc/ops/nn_ops.h" #include "tensorflow/cc/ops/standard_ops.h" @@ -27,6 +26,10 @@ limitations under the License. #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/kernels/ops_testutil.h" #include "tensorflow/core/kernels/ops_util.h" +#include "tensorflow/core/platform/cpu_info.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/stacktrace_handler.h" +#include "tensorflow/core/platform/str_util.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/test_benchmark.h" #include "tensorflow/core/public/session.h" @@ -37,133 +40,72 @@ limitations under the License. // Compare performance of default Tensorflow convolution kernels (Eigen) with // MKL kernels on CPU. -// Before running these benchmarks configure OpenMP environment variables: -// export KMP_BLOCKTIME=0 -// export OMP_NUM_THREADS=${num_threads} - namespace tensorflow { -static Tensor NonMklTensor() { - MklDnnShape non_mkl_shape; - non_mkl_shape.SetMklTensor(false); - auto size = static_cast(non_mkl_shape.GetSerializeBufferSize()); - Tensor tensor(DT_UINT8, {size}); +static Graph* Activation(const string& op_name, const string& kind, + const TensorShape& shape) { + auto* graph = new Graph(OpRegistry::Global()); + const string node_name = kind + "_" + op_name; + const bool isForwardOp = !tensorflow::str_util::EndsWith(op_name, "Grad"); + const bool isDefault = (kind == "Default"); - non_mkl_shape.SerializeMklDnnShape(tensor.flat().data(), - size * sizeof(uint8)); - return tensor; + Tensor input_t(DT_FLOAT, shape); + input_t.flat().setRandom(); + Node* input = test::graph::Constant(graph, input_t, "input"); + Node* not_mkl_shape = + test::graph::Constant(graph, GetMklMetaTensor(), "not_mkl"); + + if (isForwardOp) { + // Default forward op. + if (isDefault) { + TF_CHECK_OK(NodeBuilder(graph->NewName(node_name), op_name) + .Input(input) + .Attr("T", DT_FLOAT) + .Finalize(graph, nullptr)); + return graph; + } + // MKL forward op. + TF_CHECK_OK(NodeBuilder(graph->NewName(node_name), "_Mkl" + op_name) + .Input(input) + .Input(not_mkl_shape) + .Attr("T", DT_FLOAT) + .Attr("_kernel", "MklLayoutDependentOp") + .Finalize(graph, nullptr)); + return graph; + } + + // Default backward op. + Tensor grad_t(DT_FLOAT, shape); + grad_t.flat().setRandom(); + Node* grad = test::graph::Constant(graph, grad_t, "grad"); + if (isDefault) { + TF_CHECK_OK(NodeBuilder(graph->NewName(node_name), op_name) + .Input(grad) + .Input(input) + .Attr("T", DT_FLOAT) + .Finalize(graph, nullptr)); + return graph; + } + + // MKL backward op. + TF_CHECK_OK(NodeBuilder(graph->NewName(node_name), "_Mkl" + op_name) + .Input(grad) + .Input(input) + .Input(not_mkl_shape) + .Input(not_mkl_shape) + .Attr("T", DT_FLOAT) + .Attr("_kernel", "MklLayoutDependentOp") + .Finalize(graph, nullptr)); + return graph; } -static Tensor GetRandomTensor(const TensorShape& shape) { - Tensor tensor(DT_FLOAT, TensorShape(shape)); - tensor.flat() = tensor.flat().setRandom(); - return tensor; -} - -#define CREATE_DEFAULT_FWD_OP(NODE_NAME, OP_NAME) \ - static Graph* NODE_NAME(const TensorShape& shape) { \ - auto* graph = new Graph(OpRegistry::Global()); \ - Tensor input_t = GetRandomTensor(shape); \ - Node* input = test::graph::Constant(graph, input_t, "input"); \ - Node* op; \ - TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ - .Input(input) \ - .Attr("T", DT_FLOAT) \ - .Finalize(graph, &op)); \ - return graph; \ - } -CREATE_DEFAULT_FWD_OP(Default_Tanh, Tanh) -CREATE_DEFAULT_FWD_OP(Default_Elu, Elu) -CREATE_DEFAULT_FWD_OP(Default_Relu, Relu) -CREATE_DEFAULT_FWD_OP(Default_Relu6, Relu6) -CREATE_DEFAULT_FWD_OP(Default_LeakyRelu, LeakyRelu) - -#define CREATE_DEFAULT_BWD_OP(NODE_NAME, OP_NAME) \ - static Graph* NODE_NAME(const TensorShape& shape) { \ - auto* graph = new Graph(OpRegistry::Global()); \ - Tensor input_t = GetRandomTensor(shape); \ - Node* input = test::graph::Constant(graph, input_t, "input"); \ - Tensor grad_t = GetRandomTensor(shape); \ - Node* grad = test::graph::Constant(graph, grad_t, "grad"); \ - Node* op; \ - TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ - .Input(grad) \ - .Input(input) \ - .Attr("T", DT_FLOAT) \ - .Finalize(graph, &op)); \ - return graph; \ - } -CREATE_DEFAULT_BWD_OP(Default_TanhGrad, TanhGrad) -CREATE_DEFAULT_BWD_OP(Default_EluGrad, EluGrad) -CREATE_DEFAULT_BWD_OP(Default_ReluGrad, ReluGrad) -CREATE_DEFAULT_BWD_OP(Default_Relu6Grad, Relu6Grad) -CREATE_DEFAULT_BWD_OP(Default_LeakyReluGrad, LeakyReluGrad) - -#define CREATE_MKL_FWD_OP(NODE_NAME, OP_NAME) \ - static Graph* NODE_NAME(const TensorShape& shape) { \ - auto* graph = new Graph(OpRegistry::Global()); \ - \ - Tensor input_t = GetRandomTensor(shape); \ - Node* input = test::graph::Constant(graph, input_t, "input"); \ - \ - Node* not_mkl_shape = \ - test::graph::Constant(graph, NonMklTensor(), "not_mkl"); \ - \ - Node* op; \ - TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ - .Input(input) \ - .Input(not_mkl_shape) \ - .Attr("T", DT_FLOAT) \ - .Attr("_kernel", "MklLayoutDependentOp") \ - .Finalize(graph, &op)); \ - \ - return graph; \ - } - -CREATE_MKL_FWD_OP(Mkl_Tanh, _MklTanh) -CREATE_MKL_FWD_OP(Mkl_Elu, _MklElu) -CREATE_MKL_FWD_OP(Mkl_Relu, _MklRelu) -CREATE_MKL_FWD_OP(Mkl_Relu6, _MklRelu6) -CREATE_MKL_FWD_OP(Mkl_LeakyRelu, _MklLeakyRelu) - -#define CREATE_MKL_BWD_OP(NODE_NAME, OP_NAME) \ - static Graph* NODE_NAME(const TensorShape& shape) { \ - auto* graph = new Graph(OpRegistry::Global()); \ - \ - Tensor input_t = GetRandomTensor(shape); \ - Node* input = test::graph::Constant(graph, input_t, "input"); \ - Tensor grad_t = GetRandomTensor(shape); \ - Node* grad = test::graph::Constant(graph, grad_t, "grad"); \ - \ - Node* not_mkl_shape = \ - test::graph::Constant(graph, NonMklTensor(), "not_mkl"); \ - \ - Node* op; \ - TF_CHECK_OK(NodeBuilder(graph->NewName(#NODE_NAME), #OP_NAME) \ - .Input(grad) \ - .Input(input) \ - .Input(not_mkl_shape) \ - .Input(not_mkl_shape) \ - .Attr("T", DT_FLOAT) \ - .Attr("_kernel", "MklLayoutDependentOp") \ - .Finalize(graph, &op)); \ - \ - return graph; \ - } - -CREATE_MKL_BWD_OP(Mkl_TanhGrad, _MklTanhGrad) -CREATE_MKL_BWD_OP(Mkl_EluGrad, _MklEluGrad) -CREATE_MKL_BWD_OP(Mkl_ReluGrad, _MklReluGrad) -CREATE_MKL_BWD_OP(Mkl_Relu6Grad, _MklRelu6Grad) -CREATE_MKL_BWD_OP(Mkl_LeakyReluGrad, _MklLeakyReluGrad) - #define BM_Activation(op, kind, A, B, C, D, type) \ static void BM_##op##_##kind##_##type##_##A##_##B##_##C##_##D(int iters) { \ int64 num_computed_elements = (A) * (B) * (C) * (D); \ int64 flops_per_iter = num_computed_elements; \ testing::ItemsProcessed(static_cast(iters) * flops_per_iter); \ \ - test::Benchmark(#type, kind##_##op({A, B, C, D})).Run(iters); \ + test::Benchmark(#type, Activation(#op, #kind, {A, B, C, D})).Run(iters); \ } \ BENCHMARK(BM_##op##_##kind##_##type##_##A##_##B##_##C##_##D) @@ -190,4 +132,26 @@ TEST_ALL_SIZES(LeakyReluGrad) } // namespace tensorflow +// -------------------------------------------------------------------------- + +GTEST_API_ int main(int argc, char** argv) { + // Sets OpenMP environment variables. + // TODO(intel-tf): Remove this when OpenMP is removed. + tensorflow::setenv("KMP_BLOCKTIME", "0", true /*overwrite*/); + tensorflow::setenv("OMP_NUM_THREADS", + std::to_string(tensorflow::port::MaxParallelism()).c_str(), + true /*overwrite*/); + + tensorflow::testing::InstallStacktraceHandler(); + ::testing::InitGoogleTest(&argc, argv); + for (int i = 1; i < argc; i++) { + if (absl::StartsWith(argv[i], "--benchmarks=")) { + const char* pattern = argv[i] + strlen("--benchmarks="); + tensorflow::testing::Benchmark::Run(pattern); + return 0; + } + } + return RUN_ALL_TESTS(); +} + #endif // INTEL_MKL diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h index 7f6272b09c1..884f23b23c7 100644 --- a/tensorflow/core/util/mkl_util.h +++ b/tensorflow/core/util/mkl_util.h @@ -1028,6 +1028,21 @@ inline void ForwardMklMetaDataInToOut(OpKernelContext* context, } } +// ------------------------------------------------------------------- +// Common utility functions used by MKL unit tests + +inline Tensor GetMklMetaTensor() { + MklDnnShape non_mkl_shape; + non_mkl_shape.SetMklTensor(false); + + auto size = static_cast(non_mkl_shape.GetSerializeBufferSize()); + Tensor tensor(DT_UINT8, {size}); + + non_mkl_shape.SerializeMklDnnShape(tensor.flat().data(), + size * sizeof(uint8)); + return tensor; +} + // ------------------------------------------------------------------- /// Return MKL-DNN data type (memory::data_type) for input type T From 565490c3387dceab0c1eb0c9480a28215a30f779 Mon Sep 17 00:00:00 2001 From: Niranjan Hasabnis Date: Wed, 10 Jun 2020 12:45:28 -0700 Subject: [PATCH 0127/1390] Fixing unit test failure in eigen build --- tensorflow/core/kernels/BUILD | 1 + tensorflow/core/kernels/mkl_relu_op_test.cc | 22 --------------------- 2 files changed, 1 insertion(+), 22 deletions(-) diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 5abedf441a1..a68eeb4479b 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -8124,6 +8124,7 @@ tf_cc_test_mkl( "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", "//tensorflow/core:test", + "//tensorflow/core:test_main", "//tensorflow/core:testlib", ], ) diff --git a/tensorflow/core/kernels/mkl_relu_op_test.cc b/tensorflow/core/kernels/mkl_relu_op_test.cc index 30f75cd23df..0949b30cab7 100644 --- a/tensorflow/core/kernels/mkl_relu_op_test.cc +++ b/tensorflow/core/kernels/mkl_relu_op_test.cc @@ -132,26 +132,4 @@ TEST_ALL_SIZES(LeakyReluGrad) } // namespace tensorflow -// -------------------------------------------------------------------------- - -GTEST_API_ int main(int argc, char** argv) { - // Sets OpenMP environment variables. - // TODO(intel-tf): Remove this when OpenMP is removed. - tensorflow::setenv("KMP_BLOCKTIME", "0", true /*overwrite*/); - tensorflow::setenv("OMP_NUM_THREADS", - std::to_string(tensorflow::port::MaxParallelism()).c_str(), - true /*overwrite*/); - - tensorflow::testing::InstallStacktraceHandler(); - ::testing::InitGoogleTest(&argc, argv); - for (int i = 1; i < argc; i++) { - if (absl::StartsWith(argv[i], "--benchmarks=")) { - const char* pattern = argv[i] + strlen("--benchmarks="); - tensorflow::testing::Benchmark::Run(pattern); - return 0; - } - } - return RUN_ALL_TESTS(); -} - #endif // INTEL_MKL From eef7be73987ab556e954260cbef3e82ccbcee8e8 Mon Sep 17 00:00:00 2001 From: Niranjan Hasabnis Date: Wed, 10 Jun 2020 12:52:23 -0700 Subject: [PATCH 0128/1390] Adding comment back --- tensorflow/core/kernels/mkl_relu_op_test.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/core/kernels/mkl_relu_op_test.cc b/tensorflow/core/kernels/mkl_relu_op_test.cc index 0949b30cab7..6c0ad6facab 100644 --- a/tensorflow/core/kernels/mkl_relu_op_test.cc +++ b/tensorflow/core/kernels/mkl_relu_op_test.cc @@ -39,6 +39,9 @@ limitations under the License. // Compare performance of default Tensorflow convolution kernels (Eigen) with // MKL kernels on CPU. +// Before running these benchmarks configure OpenMP environment variables: +// export KMP_BLOCKTIME=0 +// export OMP_NUM_THREADS=${num_threads} namespace tensorflow { From 02c6b9edaf14fb7ca35c2b4f0bc80b74cc8a551a Mon Sep 17 00:00:00 2001 From: jerryyin Date: Wed, 10 Jun 2020 21:11:38 +0000 Subject: [PATCH 0129/1390] [ROCm][mlir] Disable mlir saved model test --- .../mlir/tensorflow/tests/tf_saved_model/build_defs.bzl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/build_defs.bzl b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/build_defs.bzl index 594afa10453..95ad05aa1e6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/build_defs.bzl +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/build_defs.bzl @@ -4,8 +4,6 @@ load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "lit_test") def tf_saved_model_test(name, data, tags = None): """Create a SavedModel test.""" - if tags == None: - tags = ["no_rocm"] native.py_binary( name = name, testonly = 1, @@ -26,5 +24,5 @@ def tf_saved_model_test(name, data, tags = None): name = name + ".py", data = [name] + data, driver = "@llvm-project//mlir:run_lit.sh", - tags = tags, + tags = tags + ["no_rocm"], ) From 3b95c2c54df8a7bc3641871197262841b803f8cd Mon Sep 17 00:00:00 2001 From: "902449@58880@bigcat_chen@ASIC" Date: Thu, 11 Jun 2020 13:54:16 +0800 Subject: [PATCH 0130/1390] Correcting for PR comments --- .../himax_we1_evb/detection_responder.cc | 1 - .../himax_we1_evb/image_provider.cc | 10 +- .../himax_we1_evb/main_functions.cc | 126 ------------------ .../lite/micro/himax_we1_evb/debug_log.cc | 5 +- 4 files changed, 6 insertions(+), 136 deletions(-) delete mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc index a353dc8a9b8..ae5de962fd3 100644 --- a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc @@ -22,7 +22,6 @@ limitations under the License. // should implement their own versions of this function. void RespondToDetection(tflite::ErrorReporter* error_reporter, int8_t person_score, int8_t no_person_score) { - if (person_score > no_person_score) { hx_drv_led_on(HX_DRV_LED_GREEN); } else { diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc index 727d93c61d1..4a3ab5775be 100644 --- a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc +++ b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc @@ -21,14 +21,12 @@ limitations under the License. hx_drv_sensor_image_config_t g_pimg_config; - TfLiteStatus GetImage(tflite::ErrorReporter* error_reporter, int image_width, int image_height, int channels, int8_t* image_data) { static bool is_initialized = false; if (!is_initialized) { - if(hx_drv_sensor_initial(&g_pimg_config)!= HX_DRV_LIB_PASS) - { + if (hx_drv_sensor_initial(&g_pimg_config) != HX_DRV_LIB_PASS) { return kTfLiteError; } is_initialized = true; @@ -36,9 +34,9 @@ TfLiteStatus GetImage(tflite::ErrorReporter* error_reporter, int image_width, hx_drv_sensor_capture(&g_pimg_config); - hx_drv_image_rescale((uint8_t*)g_pimg_config.raw_address, g_pimg_config.img_width, g_pimg_config.img_height, - image_data, image_width, image_height); - + hx_drv_image_rescale((uint8_t*)g_pimg_config.raw_address, + g_pimg_config.img_width, g_pimg_config.img_height, + image_data, image_width, image_height); return kTfLiteOk; } diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc b/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc deleted file mode 100644 index f0c7a405974..00000000000 --- a/tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/main_functions.cc +++ /dev/null @@ -1,126 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -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 - - http://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. -==============================================================================*/ - -#include "tensorflow/lite/micro/examples/person_detection_experimental/main_functions.h" - -#include "tensorflow/lite/micro/examples/person_detection_experimental/detection_responder.h" -#include "tensorflow/lite/micro/examples/person_detection_experimental/image_provider.h" -#include "tensorflow/lite/micro/examples/person_detection_experimental/model_settings.h" -#include "tensorflow/lite/micro/examples/person_detection_experimental/person_detect_model_data.h" -#include "tensorflow/lite/micro/kernels/micro_ops.h" -#include "tensorflow/lite/micro/micro_error_reporter.h" -#include "tensorflow/lite/micro/micro_interpreter.h" -#include "tensorflow/lite/micro/micro_mutable_op_resolver.h" -#include "tensorflow/lite/schema/schema_generated.h" -#include "tensorflow/lite/version.h" - -// Globals, used for compatibility with Arduino-style sketches. -namespace { -tflite::ErrorReporter* error_reporter = nullptr; -const tflite::Model* model = nullptr; -tflite::MicroInterpreter* interpreter = nullptr; -TfLiteTensor* input = nullptr; - -// In order to use optimized tensorflow lite kernels, a signed int8 quantized -// model is preferred over the legacy unsigned model format. This means that -// throughout this project, input images must be converted from unisgned to -// signed format. The easiest and quickest way to convert from unsigned to -// signed 8-bit integers is to subtract 128 from the unsigned value to get a -// signed value. - -// An area of memory to use for input, output, and intermediate arrays. -constexpr int kTensorArenaSize = 125 * 1024; -#pragma Bss(".tensor_arena") -static uint8_t tensor_arena[kTensorArenaSize]; -#pragma Bss() -} // namespace - -// The name of this function is important for Arduino compatibility. -void setup() { - // Set up logging. Google style is to avoid globals or statics because of - // lifetime uncertainty, but since this has a trivial destructor it's okay. - // NOLINTNEXTLINE(runtime-global-variables) - static tflite::MicroErrorReporter micro_error_reporter; - error_reporter = µ_error_reporter; - - // Map the model into a usable data structure. This doesn't involve any - // copying or parsing, it's a very lightweight operation. - model = tflite::GetModel(g_person_detect_model_data); - if (model->version() != TFLITE_SCHEMA_VERSION) { - TF_LITE_REPORT_ERROR(error_reporter, - "Model provided is schema version %d not equal " - "to supported version %d.", - model->version(), TFLITE_SCHEMA_VERSION); - return; - } - - // Pull in only the operation implementations we need. - // This relies on a complete list of all the ops needed by this graph. - // An easier approach is to just use the AllOpsResolver, but this will - // incur some penalty in code space for op implementations that are not - // needed by this graph. - // - // tflite::AllOpsResolver resolver; - // NOLINTNEXTLINE(runtime-global-variables) - static tflite::MicroMutableOpResolver<5> micro_op_resolver; - micro_op_resolver.AddBuiltin( - tflite::BuiltinOperator_DEPTHWISE_CONV_2D, - tflite::ops::micro::Register_DEPTHWISE_CONV_2D()); - micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D, - tflite::ops::micro::Register_CONV_2D()); - micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D, - tflite::ops::micro::Register_AVERAGE_POOL_2D()); - micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE, - tflite::ops::micro::Register_RESHAPE()); - micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX, - tflite::ops::micro::Register_SOFTMAX()); - - // Build an interpreter to run the model with. - // NOLINTNEXTLINE(runtime-global-variables) - static tflite::MicroInterpreter static_interpreter( - model, micro_op_resolver, tensor_arena, kTensorArenaSize, error_reporter); - interpreter = &static_interpreter; - - // Allocate memory from the tensor_arena for the model's tensors. - TfLiteStatus allocate_status = interpreter->AllocateTensors(); - if (allocate_status != kTfLiteOk) { - TF_LITE_REPORT_ERROR(error_reporter, "AllocateTensors() failed"); - return; - } - - // Get information about the memory area to use for the model's input. - input = interpreter->input(0); -} - -// The name of this function is important for Arduino compatibility. -void loop() { - // Get image from provider. - if (kTfLiteOk != GetImage(error_reporter, kNumCols, kNumRows, kNumChannels, - input->data.int8)) { - TF_LITE_REPORT_ERROR(error_reporter, "Image capture failed."); - } - - // Run the model on this input and make sure it succeeds. - if (kTfLiteOk != interpreter->Invoke()) { - TF_LITE_REPORT_ERROR(error_reporter, "Invoke failed."); - } - - TfLiteTensor* output = interpreter->output(0); - - // Process the inference results. - int8_t person_score = output->data.uint8[kPersonIndex]; - int8_t no_person_score = output->data.uint8[kNotAPersonIndex]; - RespondToDetection(error_reporter, person_score, no_person_score); -} diff --git a/tensorflow/lite/micro/himax_we1_evb/debug_log.cc b/tensorflow/lite/micro/himax_we1_evb/debug_log.cc index 32af2625630..36ac3f3fa03 100644 --- a/tensorflow/lite/micro/himax_we1_evb/debug_log.cc +++ b/tensorflow/lite/micro/himax_we1_evb/debug_log.cc @@ -20,12 +20,11 @@ limitations under the License. #include "tensorflow/lite/micro/debug_log.h" #include "hx_drv_tflm.h" - extern "C" void DebugLog(const char* s) { static bool is_initialized = false; if (!is_initialized) { - hx_drv_uart_initial(); - is_initialized = true; + hx_drv_uart_initial(); + is_initialized = true; } hx_drv_uart_print("%s", s); From b34536e5784f6d87a8905e73021f75a5cf98077a Mon Sep 17 00:00:00 2001 From: stjohnso98 <44154075+stjohnso98@users.noreply.github.com> Date: Thu, 11 Jun 2020 14:44:04 +0530 Subject: [PATCH 0131/1390] Update tensorflow/core/kernels/data/experimental/csv_dataset_op.cc Co-authored-by: Rachel Lim --- .../kernels/data/experimental/csv_dataset_op.cc | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc index d2023ecec6e..906e57066a2 100644 --- a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc @@ -64,10 +64,15 @@ class CSVDatasetOp : public DatasetOpKernel { OP_REQUIRES(ctx, select_cols_tensor->dims() == 1, errors::InvalidArgument("`select_cols` must be a vector.")); - const Tensor* exclude_cols_tensor = new const Tensor(); - if (op_version_ > 1) { - OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); - } +std::vector exclude_cols; +if (op_version_ > 1) { + const Tensor* exclude_cols_tensor; + OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); + exclude_cols.reserve(exclude_cols_tensor->NumElements()); + for (int i = 0; i < exclude_cols_tensor->NumElements(); ++i) { + exclude_cols.push_back(exclude_cols_tensor->flat()(i)); + } +} OP_REQUIRES(ctx, exclude_cols_tensor->dims() == 1, errors::InvalidArgument("`exclude_cols` must be a vector")); From 5cbd0bcf412c56c6610c24ae12c83840dc9724a6 Mon Sep 17 00:00:00 2001 From: Tare Gaskin Date: Thu, 11 Jun 2020 01:03:25 +0000 Subject: [PATCH 0132/1390] [-Wsign-compare] batch resolution 1 --- .../quantization/import_quant_stats_pass.cc | 4 +-- .../lite/quantization/quantization_config.cc | 4 +-- .../lite/quantization/quantization_driver.cc | 4 +-- .../lite/quantization/quantization_utils.cc | 10 +++---- .../mlir/tensorflow/utils/dump_mlir_util.cc | 2 +- tensorflow/compiler/mlir/xla/ir/chlo_ops.cc | 2 +- tensorflow/compiler/mlir/xla/ir/hlo_ops.cc | 6 ++-- tensorflow/compiler/xla/window_util.cc | 2 +- tensorflow/core/kernels/batch_kernels.cc | 6 ++-- .../core/kernels/data/prefetch_autotuner.cc | 4 +-- tensorflow/core/kernels/quantization_utils.h | 2 +- tensorflow/core/platform/s3/s3_file_system.cc | 2 +- .../core/profiler/utils/derived_timeline.cc | 2 +- .../core/profiler/utils/derived_timeline.h | 2 +- .../core/profiler/utils/xplane_utils.cc | 2 +- tensorflow/core/util/bcast.h | 4 +-- .../convert_trivial_tile_to_concat.cc | 2 +- .../convert_trivial_transpose_to_reshape.cc | 2 +- .../toco/graph_transformations/dequantize.cc | 2 +- .../graph_transformations/drop_fake_quant.cc | 2 +- ...int8_weights_safe_for_fast_int8_kernels.cc | 2 +- .../fuse_broadcast_into_following_binary.cc | 2 +- .../group_bidirectional_sequence_ops.cc | 4 +-- .../graph_transformations/hardcode_min_max.cc | 2 +- .../identify_nearest_upsample.cc | 2 +- .../merge_reshape_into_preceding_transpose.cc | 4 +-- .../propagate_array_data_types.cc | 2 +- .../propagate_fake_quant_num_bits.cc | 2 +- .../propagate_fixed_sizes.cc | 28 +++++++++---------- .../remove_successive_transpose.cc | 10 +++---- .../remove_trivial_passthrough.cc | 2 +- .../reorder_elementwise_unary.cc | 4 +-- .../reorder_reshape_transpose.cc | 12 ++++---- .../resolve_batch_normalization.cc | 10 +++---- .../resolve_constant_concatenation.cc | 2 +- .../resolve_constant_pack.cc | 2 +- .../resolve_constant_slice.cc | 2 +- .../resolve_constant_transpose.cc | 2 +- .../resolve_constant_unary.cc | 4 +-- .../unpartition_embedding_lookup.cc | 4 +-- tensorflow/lite/toco/model_cmdline_flags.cc | 8 +++--- tensorflow/lite/toco/toco_cmdline_flags.cc | 2 +- 42 files changed, 89 insertions(+), 89 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc index d924a3e82ac..5419a0d5e1b 100644 --- a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc +++ b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc @@ -76,7 +76,7 @@ class ImportQuantStatsPass // If the index is out of range, this method returns false. Otherwise it // returns true if the value is a float tensor. bool IsQuantizableResult(Operation *op, int index) { - if (index < 0 || index >= op->getNumResults()) return false; + if (index < 0 || index >= static_cast(op->getNumResults())) return false; Value res = op->getResult(index); return res.getType().isa() && res.getType().cast().getElementType().isa(); @@ -158,7 +158,7 @@ void ImportQuantStatsPass::ImportAsStatsOps(OpBuilder b, Operation *op, InsertStatsOpAtResult(b, op->getResult(index), layer_stats, axis_stats, axis); } else { - for (int i = 0; i < op->getNumResults(); ++i) { + for (int i = 0; i < static_cast(op->getNumResults()); ++i) { if (IsQuantizableResult(op, i)) { InsertStatsOpAtResult(b, op->getResult(i), layer_stats, axis_stats, axis); diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc index 6b897bd5608..c4cf6e71cf3 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc @@ -48,7 +48,7 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_mins; if (!min_values.empty()) { std::vector node_mins_str = absl::StrSplit(min_values, ','); - for (int i = 0; i < node_mins_str.size(); i++) { + for (size_t i = 0; i < node_mins_str.size(); i++) { double value; if (!absl::SimpleAtod(node_mins_str[i], &value)) { return true; @@ -60,7 +60,7 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_maxs; if (!max_values.empty()) { std::vector node_maxs_str = absl::StrSplit(max_values, ','); - for (int i = 0; i < node_maxs_str.size(); i++) { + for (size_t i = 0; i < node_maxs_str.size(); i++) { double value; if (!absl::SimpleAtod(node_maxs_str[i], &value)) { llvm::errs() << "Unexpected mins: " << node_maxs_str[i] << "\n"; diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc index 2964a3e79f8..fc11604ef8a 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc @@ -294,7 +294,7 @@ class QuantizationDriver { return; if (current_op == op) llvm::errs() << "===>>>"; llvm::errs() << op->getName() << " : ("; - for (auto i = 0; i < op->getNumOperands(); ++i) { + for (size_t i = 0; i < op->getNumOperands(); ++i) { if (auto params = GetOperandQuantState(op, i).params) params.print(llvm::errs()); else @@ -303,7 +303,7 @@ class QuantizationDriver { llvm::errs() << ","; } llvm::errs() << ") -> ("; - for (auto i = 0; i < op->getNumResults(); ++i) { + for (size_t i = 0; i < op->getNumResults(); ++i) { if (auto params = GetResultQuantState(op, i).params) params.print(llvm::errs()); else diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc index 3d50f280d0f..b9ca5329519 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc @@ -54,7 +54,7 @@ static Type GetQuantizedType(Builder builder, Type input_type, } else if (min.size() == max.size()) { auto shape = input_type.dyn_cast(); if (!shape || shape.getRank() <= quant_dim || - min.size() != shape.getDimSize(quant_dim)) { + static_cast(min.size()) != shape.getDimSize(quant_dim)) { return {}; } // TODO(b/141508873): the quantization dim is set to the last dimension. @@ -75,7 +75,7 @@ TypeAttr RescaleQuantizedType(Type input, Attribute factor) { if (auto qtype = ele_type.dyn_cast()) { ArrayRef scales = qtype.getScales(); // Broadcasting hasn't been implemented yet. - if (scales.size() != factor_values.getNumElements()) return {}; + if (static_cast(scales.size()) != factor_values.getNumElements()) return {}; SmallVector new_scales; new_scales.reserve(scales.size()); auto scales_iter = scales.begin(); @@ -269,7 +269,7 @@ Type GetUniformQuantizedPerAxisTypeForWeight(ElementsAttr attr, int quant_dim, bool narrow_range) { Builder builder(attr.getContext()); auto shape = attr.getType().cast().getShape(); - if (shape.size() <= quant_dim) return {}; + if (static_cast(shape.size()) <= quant_dim) return {}; // `symmetric` can only be used when it is `signed` and `narrow_range`. if (symmetric && (!is_signed || !narrow_range)) return {}; @@ -334,7 +334,7 @@ quant::QuantizedType GetUniformQuantizedTypeForBias( const std::vector& op_types) { if (op_types.empty()) return {}; - int axis_size = 1; + size_t axis_size = 1; int32_t quant_dim = -1; Type expressed_type; // Requires all the op types are valid UniformQuantizedTypes or @@ -368,7 +368,7 @@ quant::QuantizedType GetUniformQuantizedTypeForBias( scales[index_scale.index()] *= index_scale.value(); } } else if (auto type = op_type.dyn_cast()) { - for (int index = 0; index != axis_size; ++index) { + for (size_t index = 0; index != axis_size; ++index) { scales[index] *= type.getScale(); } } diff --git a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc index 797687ea658..b5a6c922707 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc @@ -41,7 +41,7 @@ std::string MakeUniqueFilename(string name) { static NameCounts& instance = *new NameCounts; // Remove illegal characters from `name`. - for (int i = 0; i < name.size(); ++i) { + for (size_t i = 0; i < name.size(); ++i) { char ch = name[i]; if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' || ch == '\\') { diff --git a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc index 26db4549a2a..f5b895f0c76 100644 --- a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc +++ b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc @@ -49,7 +49,7 @@ static Type GetBroadcastType(Type x, Type y, Type element_type, if (shape_x.size() == shape_y.size()) { llvm::SmallVector out_shape(shape_x.size()); - for (int i = 0; i < shape_x.size(); i++) { + for (size_t i = 0; i < shape_x.size(); i++) { auto x_val = shape_x[i]; auto y_val = shape_y[i]; if (x_val == -1 || y_val == -1) { diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc index d20f1713eba..569e45912a2 100644 --- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc +++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc @@ -143,7 +143,7 @@ DenseIntElementsAttr BuildConvPaddingAttrs( int rank = padding_low.size(); SmallVector padding; - for (unsigned i = 0; i < rank; ++i) { + for (unsigned i = 0; i < static_cast(rank); ++i) { padding.push_back(GetPaddingValue(padding_attr, {i, 0}) + padding_low[i]); padding.push_back(GetPaddingValue(padding_attr, {i, 1}) + padding_high[i]); } @@ -853,7 +853,7 @@ static Attribute foldConcatenateHelper(ConcatenateOp* op, auto shape = type.getShape(); size_t top_size = 1; - for (int i = 0; i < axis; i++) { + for (size_t i = 0; i < axis; i++) { top_size = top_size * shape[i]; } @@ -1118,7 +1118,7 @@ static LogicalResult Verify(MapOp op) { // increasing. auto values = op.dimensions().getValues(); auto dimensions = std::vector{values.begin(), values.end()}; - for (int i = 0; i < dimensions.size(); ++i) { + for (int i = 0; static_cast(i) < dimensions.size(); ++i) { if (dimensions[i] != i) return op.emitOpError() << "requires monotonically increasing dimension " "numbers, but got: " diff --git a/tensorflow/compiler/xla/window_util.cc b/tensorflow/compiler/xla/window_util.cc index a58179c3ee0..e33d0b6d1dc 100644 --- a/tensorflow/compiler/xla/window_util.cc +++ b/tensorflow/compiler/xla/window_util.cc @@ -42,7 +42,7 @@ Window MakeWindow(absl::Span sizes, absl::Span strides) { Window window; CHECK_EQ(sizes.size(), strides.size()); - for (auto nb = 0; nb < sizes.size(); ++nb) { + for (auto nb = 0; static_cast(nb) < sizes.size(); ++nb) { auto* dimension = window.add_dimensions(); dimension->set_size(sizes[nb]); dimension->set_stride(strides[nb]); diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc index 151f2367c95..ee271f1a123 100644 --- a/tensorflow/core/kernels/batch_kernels.cc +++ b/tensorflow/core/kernels/batch_kernels.cc @@ -486,18 +486,18 @@ class BatchResource : public ResourceBase { std::map> split_tensors; DCHECK_EQ(batch->task(0).context->num_outputs(), combined_outputs.size()); - if (combined_outputs.size() != batch->task(0).context->num_outputs()) { + if (static_cast(combined_outputs.size()) != batch->task(0).context->num_outputs()) { return errors::Internal("Wrong number of batched output tensors"); } // Generate 'split_tensors' and populate the context outputs. - for (int i = 0; i < combined_outputs.size(); ++i) { + for (size_t i = 0; i < combined_outputs.size(); ++i) { const Tensor& output_tensor = combined_outputs[i]; if (output_tensor.shape().dims() == 0) { return errors::FailedPrecondition( "Batched output tensor has 0 dimensions"); } - if (output_tensor.shape().dim_size(0) != batch->size() + padding_size) { + if (output_tensor.shape().dim_size(0) != static_cast(batch->size() + padding_size)) { return errors::FailedPrecondition( "Batched output tensor's 0th dimension does not equal the sum of " "the 0th dimension sizes of the input tensors"); diff --git a/tensorflow/core/kernels/data/prefetch_autotuner.cc b/tensorflow/core/kernels/data/prefetch_autotuner.cc index a3bb1acc352..a3fd9919d6b 100644 --- a/tensorflow/core/kernels/data/prefetch_autotuner.cc +++ b/tensorflow/core/kernels/data/prefetch_autotuner.cc @@ -40,13 +40,13 @@ void PrefetchAutotuner::RecordConsumption(size_t current_buffer_size) { case Mode::kDisabled: return; case Mode::kUpswing: - if (current_buffer_size == buffer_limit_) { + if (static_cast(current_buffer_size) == buffer_limit_) { mode_ = Mode::kDownswing; } return; case Mode::kDownswing: if (current_buffer_size == 0) { - if (buffer_limit_ >= kBufferLimitThreshold) { + if (buffer_limit_ >= static_cast(kBufferLimitThreshold)) { buffer_limit_ += kBufferLimitThreshold; } else { buffer_limit_ *= 2; diff --git a/tensorflow/core/kernels/quantization_utils.h b/tensorflow/core/kernels/quantization_utils.h index 315616f3fb3..06c901967b0 100644 --- a/tensorflow/core/kernels/quantization_utils.h +++ b/tensorflow/core/kernels/quantization_utils.h @@ -268,7 +268,7 @@ inline void RequantizeManyInNewRangeReference(const qint32* input, int64 count, // that could be easily adapted for a SIMD implementation. It should also be // possible to perform all the calculations in 32-bit rather than 64, but // that's not been implemented yet. - for (size_t index = 0; index < count; ++index) { + for (size_t index = 0; static_cast(index) < count; ++index) { const int64 input_value = static_cast(input[index]); const int64 fp_value = ((input_value * range_scale_fp) >> 32) + input_offset_fp; diff --git a/tensorflow/core/platform/s3/s3_file_system.cc b/tensorflow/core/platform/s3/s3_file_system.cc index 1726c9fbc6c..45d648abcc0 100644 --- a/tensorflow/core/platform/s3/s3_file_system.cc +++ b/tensorflow/core/platform/s3/s3_file_system.cc @@ -906,7 +906,7 @@ Status S3FileSystem::MultiPartCopy(const Aws::String& source, // wait on the mutex until notify is called // then check the finished parts as there could be false notifications multi_part_copy_cv.wait(lock, [&finishedPartStates, num_parts] { - return finishedPartStates.size() == num_parts; + return static_cast(finishedPartStates.size()) == num_parts; }); } // check if there was any error for any part diff --git a/tensorflow/core/profiler/utils/derived_timeline.cc b/tensorflow/core/profiler/utils/derived_timeline.cc index 112c0977763..3d03fc22c16 100644 --- a/tensorflow/core/profiler/utils/derived_timeline.cc +++ b/tensorflow/core/profiler/utils/derived_timeline.cc @@ -130,7 +130,7 @@ void DerivedXLineBuilder::ExpandOrAddLevelEvent(const XEvent& event, } void DerivedXLineBuilder::ResetLastEvents(int level) { - for (int i = level; i < last_event_by_level_.size(); ++i) { + for (int i = level; i < static_cast(last_event_by_level_.size()); ++i) { last_event_by_level_[i] = absl::nullopt; } if (level == 0) ResetDependentLines(); diff --git a/tensorflow/core/profiler/utils/derived_timeline.h b/tensorflow/core/profiler/utils/derived_timeline.h index cd4da7996c5..92489399b8f 100644 --- a/tensorflow/core/profiler/utils/derived_timeline.h +++ b/tensorflow/core/profiler/utils/derived_timeline.h @@ -37,7 +37,7 @@ class DerivedXLineBuilder { std::vector dependent_lines); void ExpandOrAddEvents(const std::vector& event_per_level) { - for (int level = 0; level < event_per_level.size(); ++level) { + for (size_t level = 0; level < event_per_level.size(); ++level) { ExpandOrAddLevelEvent(event_per_level[level], level); } } diff --git a/tensorflow/core/profiler/utils/xplane_utils.cc b/tensorflow/core/profiler/utils/xplane_utils.cc index 7f5221c5391..1fe476ce79c 100644 --- a/tensorflow/core/profiler/utils/xplane_utils.cc +++ b/tensorflow/core/profiler/utils/xplane_utils.cc @@ -266,7 +266,7 @@ void SortXSpace(XSpace* space) { // smaller than these value. void NormalizeTimestamps(XPlane* plane, uint64 start_time_ns) { for (XLine& line : *plane->mutable_lines()) { - if (line.timestamp_ns() >= start_time_ns) { + if (line.timestamp_ns() >= static_cast(start_time_ns)) { line.set_timestamp_ns(line.timestamp_ns() - start_time_ns); } } diff --git a/tensorflow/core/util/bcast.h b/tensorflow/core/util/bcast.h index 7bb8ea18ad3..075de84964e 100644 --- a/tensorflow/core/util/bcast.h +++ b/tensorflow/core/util/bcast.h @@ -139,7 +139,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N], if (x[i] != x[0]) { all_equal = false; } - if (x[i].size() > largest_rank) { + if (static_cast(x[i].size()) > largest_rank) { largest_rank = x[i].size(); } } @@ -176,7 +176,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N], // 1-extend and align all vectors. for (int i = 0; i < N; ++i) { - if (copy[i].size() < largest_rank) { + if (static_cast(copy[i].size()) < largest_rank) { copy[i].resize(largest_rank, 1); } } diff --git a/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc b/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc index 46288d2a1ed..c19ccf676c9 100644 --- a/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc +++ b/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc @@ -52,7 +52,7 @@ namespace toco { // It then just becomes a concat along that dimension. int non_one_dims = 0; int concat_axis = 0; - for (int i = 0; i < multiples.size(); ++i) { + for (size_t i = 0; i < multiples.size(); ++i) { if (multiples[i] != 1) { ++non_one_dims; concat_axis = i; diff --git a/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc b/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc index 2b5aaea2b23..fa8a69a1e7a 100644 --- a/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc +++ b/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc @@ -31,7 +31,7 @@ bool TransposeAffectsMemoryOrder(std::vector perm, // just the shape) then the flat buffer representation shouldn't change. std::vector old_major_index_ordering; std::vector new_major_index_ordering; - for (int i = 0; i < in_shape.size(); i++) { + for (int i = 0; static_cast(i) < in_shape.size(); i++) { if (in_shape[i] != 1) { old_major_index_ordering.push_back(i); } diff --git a/tensorflow/lite/toco/graph_transformations/dequantize.cc b/tensorflow/lite/toco/graph_transformations/dequantize.cc index cc5dddbb40e..c87c305a70d 100644 --- a/tensorflow/lite/toco/graph_transformations/dequantize.cc +++ b/tensorflow/lite/toco/graph_transformations/dequantize.cc @@ -35,7 +35,7 @@ void DequantizeBuffer(Array* array) { auto& new_data = array->GetMutableBuffer().data; new_data.resize(old_data.size()); const auto& qparams = array->GetQuantizationParams(); - for (int i = 0; i < old_data.size(); i++) { + for (size_t i = 0; i < old_data.size(); i++) { new_data[i] = qparams.scale * (old_data[i] - qparams.zero_point); } } diff --git a/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc b/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc index bb8679bced8..3a0b4d0103f 100644 --- a/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc +++ b/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc @@ -45,7 +45,7 @@ namespace toco { } // Drop min/max inputs - for (int i = 1; i < fakequant_op->inputs.size(); i++) { + for (size_t i = 1; i < fakequant_op->inputs.size(); i++) { if (CountOpsWithInput(*model, fakequant_op->inputs[i]) == 1) { model->EraseArray(fakequant_op->inputs[i]); } diff --git a/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc b/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc index 918bb489995..ce4574cdfbf 100644 --- a/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc +++ b/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc @@ -166,7 +166,7 @@ namespace toco { int index_of_previous_bad_value = 0; bool changed = false; - for (int i = 0; i < buffer_data.size(); i++) { + for (size_t i = 0; i < buffer_data.size(); i++) { if (buffer_data[i] == 0) { count_bad++; if (count_bad > 1) { diff --git a/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc b/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc index ba3e277f676..2c5c2cbb5f1 100644 --- a/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc +++ b/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc @@ -34,7 +34,7 @@ bool IsBroadcastingOp(const Model& model, Operator* op) { // Concatenation of identical inputs is usually a broadcast. if (op->type == OperatorType::kConcatenation) { // Verify that all inputs are the same. - for (int i = 1; i < op->inputs.size(); ++i) { + for (size_t i = 1; i < op->inputs.size(); ++i) { if (op->inputs[i] != op->inputs[0]) { return false; } diff --git a/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc b/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc index fa252b1a61b..a6d95ec43b1 100644 --- a/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc +++ b/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc @@ -125,7 +125,7 @@ bool CheckTwoUnidirectionalSequenceOpsAreValid( return false; // Make sure the inputs datatype matches. - for (int i = 0; i < fw_sequence_op->inputs.size(); ++i) { + for (size_t i = 0; i < fw_sequence_op->inputs.size(); ++i) { const auto& fw_input_array_name = fw_sequence_op->inputs[i]; const auto& bw_input_array_name = bw_sequence_op->inputs[i]; if (model.HasArray(fw_input_array_name) && @@ -137,7 +137,7 @@ bool CheckTwoUnidirectionalSequenceOpsAreValid( } // Make sure the outputs datatype matches. - for (int i = 0; i < fw_sequence_op->outputs.size(); ++i) { + for (size_t i = 0; i < fw_sequence_op->outputs.size(); ++i) { const auto& fw_output_array_name = fw_sequence_op->outputs[i]; const auto& bw_output_array_name = bw_sequence_op->outputs[i]; if (model.HasArray(fw_output_array_name) && diff --git a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc index 171d522daa7..4250668bcf5 100644 --- a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc +++ b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc @@ -405,7 +405,7 @@ bool HardcodeMinMaxForPack(Model* model, Operator* op) { } const auto& first_input_minmax = first_input_array.GetMinMax(); - for (int i = 1; i < op->inputs.size(); i++) { + for (size_t i = 1; i < op->inputs.size(); i++) { const auto& input_array = model->GetArray(op->inputs[i]); if (!input_array.minmax) { return false; diff --git a/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc b/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc index 2ab6692a3a8..08894c93a5b 100644 --- a/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc +++ b/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc @@ -199,7 +199,7 @@ std::vector>::iterator FindOperator( shape_array.data_type = ArrayDataType::kInt32; auto& shape_buffer = shape_array.GetMutableBuffer(); // This is what imagined as the original shape. - for (int i = 0; i < imagined_original_shape.size(); ++i) { + for (size_t i = 0; i < imagined_original_shape.size(); ++i) { shape_buffer.data.push_back(imagined_original_shape.at(i)); } diff --git a/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc b/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc index 80170fe8bcb..a76ae1a0635 100644 --- a/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc @@ -70,7 +70,7 @@ std::vector ReshapeToTranspose(const Model& model, std::vector not_one_indices; // Separate into one indices and not one indices. - for (int i = 0; i < in_shape.size(); i++) { + for (size_t i = 0; i < in_shape.size(); i++) { if (in_shape[i] == 1) { one_indices.push_back(i); } else { @@ -167,7 +167,7 @@ std::vector ReshapeToTranspose(const Model& model, // Combine the permutations. const auto& transpose_perm = transpose_op->perm; - for (int i = 0; i < merged_perm.size(); i++) { + for (size_t i = 0; i < merged_perm.size(); i++) { merged_perm[i] = transpose_perm[merged_perm[i]]; } diff --git a/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc index 49d59de860b..2f316934311 100644 --- a/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc +++ b/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc @@ -170,7 +170,7 @@ void SetDataTypeForAllOutputs(Model* model, Operator* op, if (unsupported_op->output_data_types.size() < op->outputs.size()) { return ::tensorflow::Status::OK(); } - for (int i = 0; i < op->outputs.size(); ++i) { + for (size_t i = 0; i < op->outputs.size(); ++i) { const string& output = op->outputs[i]; const ArrayDataType data_type = unsupported_op->output_data_types[i]; model->GetArray(output).data_type = data_type; diff --git a/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc b/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc index 1ed618879c1..94779f54af2 100644 --- a/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc +++ b/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc @@ -149,7 +149,7 @@ bool RecursivelyBackwardPropagateDataType(GraphTransformation* transformation, ArrayDataType new_data_type, const MinMax& new_minmax) { bool did_change = false; - for (int input_index = 0; input_index < op->inputs.size(); ++input_index) { + for (size_t input_index = 0; input_index < op->inputs.size(); ++input_index) { const auto& input = op->inputs[input_index]; auto& input_array = model->GetArray(input); diff --git a/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc index 006e624eb7a..520cd8b495a 100644 --- a/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc +++ b/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc @@ -431,7 +431,7 @@ void ProcessTensorFlowReshapeOperator(Model* model, bool has_wildcard = false; int wildcard_index = 0; int product_non_wildcard_dims = 1; - for (int i = 0; i < shape_data.size(); i++) { + for (size_t i = 0; i < shape_data.size(); i++) { if (shape_data[i] == -1) { CHECK(!has_wildcard); has_wildcard = true; @@ -574,7 +574,7 @@ void ProcessTensorFlowReductionOperator(Model* model, Operator* op) { std::set true_indices; const auto& reduction_indices = reduction_indices_array.GetBuffer().data; - for (int i = 0; i < reduction_indices.size(); ++i) { + for (size_t i = 0; i < reduction_indices.size(); ++i) { const int32 reduction_index = reduction_indices[i]; if (reduction_index < -input_rank || reduction_index >= input_rank) { CHECK(false) << "Invalid reduction dimension " << reduction_index @@ -627,7 +627,7 @@ void ProcessSliceOperator(Model* model, SliceOperator* op) { CHECK_EQ(op->begin.size(), op->size.size()); std::vector output_dims; - for (int i = 0; i < op->begin.size(); ++i) { + for (size_t i = 0; i < op->begin.size(); ++i) { int size = op->size[i]; if (size == -1) { size = input_array.shape().dims(i) - op->begin[i]; @@ -883,7 +883,7 @@ void ProcessTensorFlowSplitVOperator(Model* model, CHECK_EQ(op->outputs.size(), op->num_split); - for (int i = 0; i < op->outputs.size(); ++i) { + for (size_t i = 0; i < op->outputs.size(); ++i) { const auto& output = op->outputs[i]; Shape output_shape = input_shape; (*output_shape.mutable_dims())[axis] = size_splits_vector.at(i); @@ -1514,7 +1514,7 @@ void ProcessPadOperator(Model* model, PadOperator* op) { std::vector& dims = *output_shape.mutable_dims(); CHECK_EQ(op->left_padding.size(), dims.size()); - for (int i = 0; i < op->left_padding.size(); ++i) { + for (size_t i = 0; i < op->left_padding.size(); ++i) { dims[i] += op->left_padding[i] + op->right_padding[i]; } @@ -1540,7 +1540,7 @@ void ProcessPadV2Operator(Model* model, PadV2Operator* op) { std::vector& dims = *output_shape.mutable_dims(); CHECK_EQ(op->left_padding.size(), dims.size()); - for (int i = 0; i < op->left_padding.size(); ++i) { + for (size_t i = 0; i < op->left_padding.size(); ++i) { dims[i] += op->left_padding[i] + op->right_padding[i]; } @@ -1683,7 +1683,7 @@ void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) { CHECK_LE(op->strides.size(), num_input_axes) << "StridedSlice op with output \"" << op->outputs[0] << "\", requires no more than " << num_input_axes << " strides"; - for (int i = 0; i < op->strides.size(); i++) { + for (size_t i = 0; i < op->strides.size(); i++) { CHECK_NE(op->strides[i], 0) << "Strides must be non-zero. Axis " << i << " has stride=" << op->strides[i] << "."; } @@ -1814,7 +1814,7 @@ void ProcessTransposeOperator(Model* model, TransposeOperator* op) { << "Transpose permutation input " << op->inputs[1] << " must be same length as input dimensions"; std::vector* output_dims = output_array.mutable_shape()->mutable_dims(); - for (int i = 0; i < perm.size(); i++) { + for (size_t i = 0; i < perm.size(); i++) { int axis = perm[i]; CHECK_GE(axis, 0); CHECK_LT(axis, input_shape.dimensions_count()); @@ -1856,8 +1856,8 @@ void ProcessArgMinMaxOperator(Model* model, Op* op) { std::vector output_dims; output_dims.reserve(input_dims.size() - 1); - for (int i = 0; i < input_dims.size(); ++i) { - if (i != axis) { + for (size_t i = 0; i < input_dims.size(); ++i) { + if ( static_cast(i) != axis) { output_dims.push_back(input_dims[i]); } } @@ -1938,7 +1938,7 @@ void ProcessTileOperator(Model* model, TensorFlowTileOperator* op) { auto* mutable_dims = output_array.mutable_shape()->mutable_dims(); mutable_dims->resize(multiples.size()); - for (int i = 0; i < mutable_dims->size(); ++i) { + for (size_t i = 0; i < mutable_dims->size(); ++i) { (*mutable_dims)[i] = input_shape.dims(i) * multiples[i]; } } @@ -2010,8 +2010,8 @@ void ProcessUnpackOperator(Model* model, UnpackOperator* op) { std::vector output_dims; output_dims.reserve(input_dims.size() - 1); - for (int i = 0; i < input_dims.size(); ++i) { - if (i != op->axis) { + for (size_t i = 0; i < input_dims.size(); ++i) { + if ( static_cast(i) != op->axis) { output_dims.push_back(input_dims[i]); } } @@ -2399,7 +2399,7 @@ void ProcessScatterNdOperator(Model* model, ScatterNdOperator* op) { if (unsupported_op->output_shapes.size() < op->outputs.size()) { return ::tensorflow::Status::OK(); } - for (int i = 0; i < op->outputs.size(); ++i) { + for (size_t i = 0; i < op->outputs.size(); ++i) { const string& output = op->outputs[i]; model->GetArray(output).copy_shape(unsupported_op->output_shapes.at(i)); } diff --git a/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc b/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc index 6eccda04c18..1cb3a300127 100644 --- a/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc @@ -31,12 +31,12 @@ bool TransformsToIdentity(std::vector const& perm1, // perm1 is the order of the indices after first transpose. When perm1 is // reordered according to perm2, if the result is simple increasing sequence // i.e., range(0, perm1.size()), then the two transposes cancel each other. - for (int i = 0; i < perm1.size(); ++i) { - if (perm1[i] < 0 || perm1[i] >= perm1.size() || perm2[i] < 0 || - perm2[i] >= perm1.size()) { + for (size_t i = 0; i < perm1.size(); ++i) { + if (perm1[i] < 0 || perm1[i] >= static_cast(perm1.size()) || perm2[i] < 0 || + perm2[i] >= static_cast(perm1.size())) { return false; } - if (perm1[perm2[i]] != i) { + if (perm1[perm2[i]] != static_cast(i)) { return false; } } @@ -46,7 +46,7 @@ bool TransformsToIdentity(std::vector const& perm1, void ReplaceOpInputsWith(Model* model, const string& lookfor, const string& replacewith) { for (const auto& op : model->operators) { - for (int i = 0; i < op->inputs.size(); ++i) { + for (size_t i = 0; i < op->inputs.size(); ++i) { if (op->inputs[i] == lookfor) { op->inputs[i] = replacewith; } diff --git a/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc b/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc index bd529bd9ecd..eeb8751bf86 100644 --- a/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc +++ b/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc @@ -82,7 +82,7 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation, // We call 'main input' the unique nonconstant input array if there is one, // or else the 0-th input. int count_nonconstant_input_arrays = 0; - for (int i = 0; i < passthru_op->inputs.size(); i++) { + for (size_t i = 0; i < passthru_op->inputs.size(); i++) { if (!model->GetArray(passthru_op->inputs[i]).buffer) { count_nonconstant_input_arrays++; if (count_nonconstant_input_arrays == 1) { diff --git a/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc b/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc index 17a5e9a1d6a..38edff76d55 100644 --- a/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc +++ b/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc @@ -127,9 +127,9 @@ bool IsMoveOperator(OperatorType optype) { move_op->outputs[0] = output_name; } else { // The intermediate array is now the output array. - for (int i = 0; i < model->operators.size(); i++) { + for (size_t i = 0; i < model->operators.size(); i++) { Operator* consumer = model->operators[i].get(); - for (int j = 0; j < consumer->inputs.size(); j++) { + for (size_t j = 0; j < consumer->inputs.size(); j++) { if (consumer->inputs[j] == output_name) { consumer->inputs[j] = intermediate_name; } diff --git a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc index 0fbcf9f73b1..b2d184cdc31 100644 --- a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc @@ -60,7 +60,7 @@ std::vector ComputeNewPerm(std::vector input_dims, std::vector perm) { // These are the major axis of the input. std::vector input_indices; - for (int i = 0; i < input_dims.size(); i++) { + for (size_t i = 0; i < input_dims.size(); i++) { if (input_dims[i] != 1) { input_indices.push_back(i); } @@ -69,7 +69,7 @@ std::vector ComputeNewPerm(std::vector input_dims, // This maps which indices of the input produced the intermediate indices for // non-unary dimensions. std::unordered_map intermediate_to_input_indices_map; - for (int i = 0; i < intermediate_dims.size(); i++) { + for (size_t i = 0; i < intermediate_dims.size(); i++) { if (intermediate_dims[i] != 1) { intermediate_to_input_indices_map[i] = input_indices[intermediate_to_input_indices_map.size()]; @@ -80,14 +80,14 @@ std::vector ComputeNewPerm(std::vector input_dims, // major indices. std::vector new_perm; new_perm.reserve(input_dims.size()); - for (int i = 0; i < perm.size(); i++) { + for (size_t i = 0; i < perm.size(); i++) { if (intermediate_dims[perm[i]] == 1) continue; new_perm.push_back(intermediate_to_input_indices_map[perm[i]]); } // Fill the rest of the transpose in with the ones. - for (int index = 0; index < input_dims.size(); index++) { + for (size_t index = 0; index < input_dims.size(); index++) { if (input_dims[index] == 1) { new_perm.push_back(index); } @@ -193,9 +193,9 @@ std::vector ComputeNewPerm(std::vector input_dims, DeleteArrayIfUnused(intermediate_name, model); } else { // The intermediate array is now the output array. - for (int i = 0; i < model->operators.size(); i++) { + for (size_t i = 0; i < model->operators.size(); i++) { Operator* consumer = model->operators[i].get(); - for (int j = 0; j < consumer->inputs.size(); j++) { + for (size_t j = 0; j < consumer->inputs.size(); j++) { if (consumer->inputs[j] == output_name) { consumer->inputs[j] = intermediate_name; } diff --git a/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc b/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc index 6e5815ee94d..545c53fb31a 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc @@ -124,11 +124,11 @@ namespace toco { const auto& offset_float_data = offset_array.GetBuffer().data; - CHECK(mul_float_data.size() == buffer_size); - CHECK(add_float_data.size() == buffer_size); - CHECK(mean_float_data.size() == buffer_size); - CHECK(multiplier_float_data.size() == buffer_size); - CHECK(offset_float_data.size() == buffer_size); + CHECK(static_cast(mul_float_data.size()) == buffer_size); + CHECK(static_cast(add_float_data.size()) == buffer_size); + CHECK(static_cast(mean_float_data.size()) == buffer_size); + CHECK(static_cast(multiplier_float_data.size()) == buffer_size); + CHECK(static_cast(offset_float_data.size()) == buffer_size); for (int i = 0; i < buffer_size; i++) { mul_float_data[i] = multiplier_float_data[i]; diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc index 7c9aa025f64..20e805a29e0 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc @@ -64,7 +64,7 @@ void CopyTensorSegments(const std::vector& input_arrays, // Copy the data from input_arrays to concatenated_array_buffer. T* dest_ptr = concatenated_array_buffer.data(); for (int s = 0; s < total_copy_steps; s++) { - for (int i = 0; i < input_arrays.size(); i++) { + for (size_t i = 0; i < input_arrays.size(); i++) { std::copy(src_ptr[i], src_ptr[i] + array_copy_size[i], dest_ptr); src_ptr[i] += array_copy_size[i]; dest_ptr += array_copy_size[i]; diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc index 0df35509d3d..c6dc093ba00 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc @@ -36,7 +36,7 @@ void Pack(Model* model, PackOperator const& op) { // Pack inputs into buffer CHECK_EQ(op.axis, 0) << "Packing only supported along first axis"; int dst_offset = 0; - for (int i = 0; i < op.inputs.size(); i++) { + for (size_t i = 0; i < op.inputs.size(); i++) { // Append array data to output for each input array const auto& input_array = model->GetArray(op.inputs[i]); int input_size = RequiredBufferSizeForShape(input_array.shape()); diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc index fd71fb1873a..34a1a1ce899 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc @@ -50,7 +50,7 @@ bool Slice(SliceOperator const& op, Array const& input_array, CHECK_LE(size.size(), 4); std::vector begin = op.begin; std::vector end; - for (int i = 0; i < begin.size(); ++i) { + for (size_t i = 0; i < begin.size(); ++i) { int dim_size = size[i]; if (dim_size == -1) { // -1 means the rest of the dimension. diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc index 7ceffe6307e..a822f7b79e3 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc @@ -40,7 +40,7 @@ void Transpose(Model* model, const Array& input_array, CHECK(input_shape.dimensions_count() == output_shape.dimensions_count()); const int dim = input_shape.dimensions_count(); CHECK_LE(dim, 4); - CHECK(perm.size() >= dim); + CHECK(static_cast(perm.size()) >= dim); for (int i = 0; i < dim; i++) { CHECK(perm[i] >= 0 && perm[i] < dim); CHECK(input_shape.dims(perm[i]) == output_shape.dims(i)); diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc index 197e17eee16..4d6cd188729 100644 --- a/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc +++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc @@ -62,7 +62,7 @@ void ReduceGeneric(bool keep_dims, const std::vector& axes, } std::vector output_indices(input_shape.dimensions_count()); - for (int input_offset = 0; input_offset < input.size(); ++input_offset) { + for (size_t input_offset = 0; input_offset < input.size(); ++input_offset) { std::vector input_indices = ReverseOffset(input_shape, input_offset); // Calculate the output location by squashing input indices to 0 // in reduced axes. @@ -319,7 +319,7 @@ bool CopyMinMaxFromFirstInput(const Operator& op, Model* model) { } else if (unary_op->type == OperatorType::kRelu6 || unary_op->type == OperatorType::kRelu1 || unary_op->type == OperatorType::kRelu) { - for (size_t i = 0; i < output_buffer_size; ++i) { + for (int i = 0; i < output_buffer_size; ++i) { const float value = (*input_float_data)[i]; float new_value = 0.0f; switch (unary_op->type) { diff --git a/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc b/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc index 1f7035c21e2..84d5922aae8 100644 --- a/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc +++ b/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc @@ -57,10 +57,10 @@ namespace toco { // Split up the DynamicStitch inputs into the indices and data. std::vector stitch_indices_inputs; std::vector stitch_data_inputs; - for (size_t i = 0; i < stitch_op->num_partitions; ++i) { + for (int i = 0; i < stitch_op->num_partitions; ++i) { stitch_indices_inputs.push_back(stitch_op->inputs[i]); } - for (size_t i = stitch_op->num_partitions; i < stitch_op->num_partitions * 2; + for (int i = stitch_op->num_partitions; i < stitch_op->num_partitions * 2; ++i) { stitch_data_inputs.push_back(stitch_op->inputs[i]); } diff --git a/tensorflow/lite/toco/model_cmdline_flags.cc b/tensorflow/lite/toco/model_cmdline_flags.cc index 2434481272f..351884fbf1e 100644 --- a/tensorflow/lite/toco/model_cmdline_flags.cc +++ b/tensorflow/lite/toco/model_cmdline_flags.cc @@ -263,7 +263,7 @@ void ReadModelFlagsFromCommandLineFlags( QCHECK(uses_multi_input_flags); std::vector mean_values = absl::StrSplit(parsed_model_flags.mean_values.value(), ','); - QCHECK(mean_values.size() == model_flags->input_arrays_size()); + QCHECK(static_cast(mean_values.size()) == model_flags->input_arrays_size()); for (size_t i = 0; i < mean_values.size(); ++i) { char* last = nullptr; model_flags->mutable_input_arrays(i)->set_mean_value( @@ -280,7 +280,7 @@ void ReadModelFlagsFromCommandLineFlags( QCHECK(uses_multi_input_flags); std::vector std_values = absl::StrSplit(parsed_model_flags.std_values.value(), ','); - QCHECK(std_values.size() == model_flags->input_arrays_size()); + QCHECK( static_cast(std_values.size()) == model_flags->input_arrays_size()); for (size_t i = 0; i < std_values.size(); ++i) { char* last = nullptr; model_flags->mutable_input_arrays(i)->set_std_value( @@ -298,7 +298,7 @@ void ReadModelFlagsFromCommandLineFlags( QCHECK(uses_multi_input_flags); std::vector input_data_types = absl::StrSplit(parsed_model_flags.input_data_types.value(), ','); - QCHECK(input_data_types.size() == model_flags->input_arrays_size()); + QCHECK(static_cast(input_data_types.size()) == model_flags->input_arrays_size()); for (size_t i = 0; i < input_data_types.size(); ++i) { IODataType type; QCHECK(IODataType_Parse(input_data_types[i], &type)); @@ -321,7 +321,7 @@ void ReadModelFlagsFromCommandLineFlags( QCHECK(uses_multi_input_flags); std::vector input_shapes = absl::StrSplit(parsed_model_flags.input_shapes.value(), ':'); - QCHECK(input_shapes.size() == model_flags->input_arrays_size()); + QCHECK(static_cast(input_shapes.size()) == model_flags->input_arrays_size()); for (size_t i = 0; i < input_shapes.size(); ++i) { auto* shape = model_flags->mutable_input_arrays(i)->mutable_shape(); shape->clear_dims(); diff --git a/tensorflow/lite/toco/toco_cmdline_flags.cc b/tensorflow/lite/toco/toco_cmdline_flags.cc index c133db8f2a4..9697a1ecbbd 100644 --- a/tensorflow/lite/toco/toco_cmdline_flags.cc +++ b/tensorflow/lite/toco/toco_cmdline_flags.cc @@ -320,7 +320,7 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags, std::vector input_types = absl::StrSplit(parsed_toco_flags.input_types.value(), ','); QCHECK(!input_types.empty()); - for (int i = 1; i < input_types.size(); i++) { + for (size_t i = 1; i < input_types.size(); i++) { QCHECK_EQ(input_types[i], input_types[0]); } toco::IODataType input_type; From ceed71957998d0d69c1a47e21647461282b60b48 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Thu, 11 Jun 2020 11:13:15 -0700 Subject: [PATCH 0133/1390] Change varible names --- tensorflow/core/kernels/relu_op_gpu.cu.cc | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/relu_op_gpu.cu.cc b/tensorflow/core/kernels/relu_op_gpu.cu.cc index ca1a4235f3a..b17303238f4 100644 --- a/tensorflow/core/kernels/relu_op_gpu.cu.cc +++ b/tensorflow/core/kernels/relu_op_gpu.cu.cc @@ -35,7 +35,7 @@ namespace tensorflow { typedef Eigen::GpuDevice GPUDevice; -static constexpr int VectorSize = 8; +static constexpr int VectorSizeElements = 8; namespace functor { // This kernel computes ReluGrad by processing one half2, two fp16, at a time. @@ -98,10 +98,11 @@ __global__ void ReluGradHalfKernelVector( const Eigen::half* __restrict__ gradient, const Eigen::half* __restrict__ feature, Eigen::half* __restrict__ backprop, int32 count) { - int32 half8_count = count / VectorSize; + int32 half8_count = count / VectorSizeElements; int32 index = blockIdx.x * blockDim.x + threadIdx.x; if (index < half8_count) { + // Cast to xx_h8 for vector load and store. float4 gradient_h8 = reinterpret_cast(gradient)[index]; float4 feature_h8 = reinterpret_cast(feature)[index]; float4* p_backprop_h8 = reinterpret_cast(backprop) + index; @@ -115,7 +116,7 @@ __global__ void ReluGradHalfKernelVector( #if __CUDA_ARCH__ >= 530 const half2 kZeroH2 = __float2half2_rn(0.f); #endif - for (int i = 0; i < VectorSize / 2; i++) { + for (int i = 0; i < VectorSizeElements / 2; i++) { #if __CUDA_ARCH__ >= 530 // mask = (feature > 0) half2 mask_h2 = __hgt2(feature_h2[i], kZeroH2); @@ -136,19 +137,19 @@ __global__ void ReluGradHalfKernelVector( *p_backprop_h8 = backprop_h8; } - int remaining_count = (count % VectorSize); + int remaining_count = (count % VectorSizeElements); if (index < remaining_count) { // Use first threads to process the remaining elements. - Eigen::half grad_h = gradient[half8_count * VectorSize + index]; - Eigen::half feature_h = feature[half8_count * VectorSize + index]; + Eigen::half grad_h = gradient[half8_count * VectorSizeElements + index]; + Eigen::half feature_h = feature[half8_count * VectorSizeElements + index]; float grad_f = static_cast(grad_h); float feature_f = static_cast(feature_h); float backprop_f = (feature_f > 0) ? grad_f : 0; Eigen::half backprop_h(backprop_f); - backprop[half8_count * VectorSize + index] = backprop_h; + backprop[half8_count * VectorSizeElements + index] = backprop_h; } } @@ -176,7 +177,7 @@ struct ReluGrad { constexpr int32 kThreadInBlock = 512; if (count == 0) return; if (aligned) { - int32 half8_count = Eigen::divup(count, VectorSize); + int32 half8_count = Eigen::divup(count, VectorSizeElements); int32 kBlock = Eigen::divup(half8_count, kThreadInBlock); TF_CHECK_OK(GpuLaunchKernel( ReluGradHalfKernelVector, kBlock, kThreadInBlock, From 013ddd96d0ce111ca5ec1422b2899b66ec41a036 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Tue, 9 Jun 2020 15:35:47 -0500 Subject: [PATCH 0134/1390] Use CUDNN_TENSOR_OP_MATH to enable tensor cores. --- tensorflow/stream_executor/cuda/cuda_dnn.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc index fa06d410323..28ec6a842bb 100644 --- a/tensorflow/stream_executor/cuda/cuda_dnn.cc +++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc @@ -723,7 +723,7 @@ class CudnnConvolutionDescriptor { #if CUDNN_VERSION >= 7000 cudnnMathType_t math_type = #if CUDNN_VERSION >= 8000 - (use_tensor_op_math ? CUDNN_DEFAULT_MATH : CUDNN_FMA_MATH); + (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH); #else (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH); #endif @@ -1179,8 +1179,7 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor { cudnnMathType_t math_type; if (use_tensor_ops) { - math_type = - CUDNN_VERSION >= 8000 ? CUDNN_DEFAULT_MATH : CUDNN_TENSOR_OP_MATH; + math_type = CUDNN_TENSOR_OP_MATH; } else { math_type = CUDNN_VERSION >= 8000 ? CUDNN_FMA_MATH : CUDNN_DEFAULT_MATH; } From da0d85808e77aa62287ba22e822a9e83866a43a4 Mon Sep 17 00:00:00 2001 From: Nathan Luehr Date: Thu, 11 Jun 2020 15:35:04 -0500 Subject: [PATCH 0135/1390] Make python names consistent --- tensorflow/python/framework/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py index a356e6d9a16..1ff2fa613da 100644 --- a/tensorflow/python/framework/config.py +++ b/tensorflow/python/framework/config.py @@ -34,7 +34,7 @@ def tensor_float32_execution_allowed(): return _pywrap_tf32_execution.is_allowed() # No tf_export until TF is built against CUDA11 which is required for TF32. -def allow_tensor_float_32_execution(allowed): +def allow_tensor_float32_execution(allowed): """Allow use of TensorFloat-32 with float32 ops on supported hardware. TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture. From d03b86fe4462ebad1f73d460c2aceab47372b239 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Thu, 11 Jun 2020 23:05:40 +0100 Subject: [PATCH 0136/1390] Extended test to the case when new converter is enabled. Change-Id: I83f20d025027ad1266f99f9d79932cab4f1a9ed5 --- tensorflow/lite/python/lite_test.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index e6661c82894..478840c5549 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -891,7 +891,11 @@ class FromSessionTest(TestModels, parameterized.TestCase): ('UseTfliteBuiltinsInt16DisableMLIR', [lite.OpsSet.\ EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], - False)) + False), + ('UseTfliteBuiltinsInt16EnableMLIR', + [lite.OpsSet.\ + EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8], + True)) def testCalibrateAndQuantizeBuiltinInt(self, supported_ops, enable_mlir): with ops.Graph().as_default(): inp, output, calibration_gen = self._getCalibrationQuantizeModel() From 4bd42cdd0e68cebe8b280b323cde4d01a9d2bf3a Mon Sep 17 00:00:00 2001 From: Reed Date: Thu, 11 Jun 2020 17:16:28 -0700 Subject: [PATCH 0137/1390] Use float_32 instead of float32 in function names --- tensorflow/python/framework/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py index 1ff2fa613da..bbaa2ca8248 100644 --- a/tensorflow/python/framework/config.py +++ b/tensorflow/python/framework/config.py @@ -25,7 +25,7 @@ from tensorflow.python.util.tf_export import tf_export # No tf_export until TF is built against CUDA11 which is required for TF32. -def tensor_float32_execution_allowed(): +def tensor_float_32_execution_allowed(): """Get if TensorFloat-32 operations are enabled on supported hardware. Returns: @@ -34,7 +34,7 @@ def tensor_float32_execution_allowed(): return _pywrap_tf32_execution.is_allowed() # No tf_export until TF is built against CUDA11 which is required for TF32. -def allow_tensor_float32_execution(allowed): +def allow_tensor_float_32_execution(allowed): """Allow use of TensorFloat-32 with float32 ops on supported hardware. TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture. From 77e46ebf9745d798863b5c7ed26d6bf077700008 Mon Sep 17 00:00:00 2001 From: Kaixi Hou Date: Mon, 30 Mar 2020 17:25:38 -0700 Subject: [PATCH 0138/1390] conv3d ndhwc plumbing --- tensorflow/core/kernels/conv_grad_ops_3d.cc | 148 ++++++++++++++++---- tensorflow/core/kernels/conv_ops_3d.cc | 88 +++++++++--- 2 files changed, 193 insertions(+), 43 deletions(-) diff --git a/tensorflow/core/kernels/conv_grad_ops_3d.cc b/tensorflow/core/kernels/conv_grad_ops_3d.cc index 2183d0e0885..fc2d58ec94f 100644 --- a/tensorflow/core/kernels/conv_grad_ops_3d.cc +++ b/tensorflow/core/kernels/conv_grad_ops_3d.cc @@ -50,6 +50,7 @@ using stream_executor::dnn::DimIndex; #include "tensorflow/stream_executor/gpu/gpu_asm_opts.h" #include "tensorflow/stream_executor/gpu/redzone_allocator.h" #include "tensorflow/stream_executor/tf_allocator_adapter.h" +#include "third_party/gpus/cudnn/cudnn.h" #endif // GOOGLE_CUDA namespace { @@ -1264,26 +1265,56 @@ class Conv3DBackpropInputOp : public OpKernel { CHECK(padding_rows >= 0 && padding_cols >= 0 && padding_planes >= 0) << "Negative paddings: (" << padding_rows << ", " << padding_cols << ", " << padding_planes << ")"; + +#if GOOGLE_CUDA + const bool compute_in_nhwc = CUDNN_VERSION >= 8000 && + DataTypeToEnum::value == DT_HALF; +#else + // fast NDHWC implementation is a CUDA only feature + const bool compute_in_nhwc = false; +#endif + const TensorFormat compute_data_format = + (compute_in_nhwc && data_format_ == FORMAT_NHWC) ? FORMAT_NHWC + : FORMAT_NCHW; + + VLOG(3) << "Compute Conv3DBackpropInput with cuDNN:" + << " data_format=" << ToString(data_format_) + << " compute_data_format=" << ToString(compute_data_format); + + constexpr auto kComputeInNHWC = + std::make_tuple(se::dnn::DataLayout::kBatchYXDepth, + se::dnn::FilterLayout::kOutputYXInput); + constexpr auto kComputeInNCHW = + std::make_tuple(se::dnn::DataLayout::kBatchDepthYX, + se::dnn::FilterLayout::kOutputInputYX); + + se::dnn::DataLayout compute_data_layout; + se::dnn::FilterLayout filter_layout; + + std::tie(compute_data_layout, filter_layout) = + compute_data_format == FORMAT_NHWC ? kComputeInNHWC : kComputeInNCHW; + se::dnn::BatchDescriptor input_desc(3); input_desc.set_count(dims.batch_size) .set_spatial_dim(DimIndex::X, compatible_input_shape.dim_size(4)) .set_spatial_dim(DimIndex::Y, compatible_input_shape.dim_size(3)) .set_spatial_dim(DimIndex::Z, compatible_input_shape.dim_size(2)) .set_feature_map_count(dims.in_depth) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::BatchDescriptor output_desc(3); output_desc.set_count(dims.batch_size) .set_spatial_dim(DimIndex::X, dims.output_size(2)) .set_spatial_dim(DimIndex::Y, dims.output_size(1)) .set_spatial_dim(DimIndex::Z, dims.output_size(0)) .set_feature_map_count(dims.out_depth) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::FilterDescriptor filter_desc(3); filter_desc.set_spatial_dim(DimIndex::X, dims.filter_size(2)) .set_spatial_dim(DimIndex::Y, dims.filter_size(1)) .set_spatial_dim(DimIndex::Z, dims.filter_size(0)) .set_input_feature_map_count(filter_shape.dim_size(3)) - .set_output_feature_map_count(filter_shape.dim_size(4)); + .set_output_feature_map_count(filter_shape.dim_size(4)) + .set_layout(filter_layout); se::dnn::ConvolutionDescriptor conv_desc(3); conv_desc.set_dilation_rate(DimIndex::X, dims.dilation(2)) .set_dilation_rate(DimIndex::Y, dims.dilation(1)) @@ -1298,21 +1329,33 @@ class Conv3DBackpropInputOp : public OpKernel { // Shape: out, in, z, y, x. Tensor transformed_filter; + auto dst_format = + compute_data_format == FORMAT_NCHW ? FORMAT_OIHW: FORMAT_OHWI; + TensorShape dst_shape = + dst_format == FORMAT_OIHW + ? TensorShape({filter_shape.dim_size(4), filter_shape.dim_size(3), + dims.filter_size(0), + dims.filter_size(1), + dims.filter_size(2)}) + : TensorShape({filter_shape.dim_size(4), + dims.filter_size(0), + dims.filter_size(1), + dims.filter_size(2), + filter_shape.dim_size(3)}); OP_REQUIRES_OK( context, context->allocate_temp( DataTypeToEnum::value, - TensorShape({filter_shape.dim_size(4), - filter_shape.dim_size(3), dims.filter_size(0), - dims.filter_size(1), dims.filter_size(2)}), + dst_shape, &transformed_filter)); + functor::TransformFilter()( - context->eigen_device(), FORMAT_OIHW, + context->eigen_device(), dst_format, To32Bit(filter.tensor()), To32Bit(transformed_filter.tensor())); // Shape: batch, filters, z, y, x. Tensor transformed_out_backprop; - if (data_format_ == FORMAT_NHWC) { + if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { TensorShape nchw_shape = {dims.batch_size, dims.out_depth, dims.output_size(0), dims.output_size(1), dims.output_size(2)}; @@ -1333,8 +1376,15 @@ class Conv3DBackpropInputOp : public OpKernel { Tensor pre_transformed_in_backprop; OP_REQUIRES_OK( context, - context->allocate_temp(DataTypeToEnum::value, compatible_input_shape, - &pre_transformed_in_backprop)); + context->allocate_temp( + DataTypeToEnum::value, + ShapeFromFormat( + compute_data_format, compatible_input_shape.dim_size(0), + {{compatible_input_shape.dim_size(2), + compatible_input_shape.dim_size(3), + compatible_input_shape.dim_size(4)}}, + compatible_input_shape.dim_size(1)), + &pre_transformed_in_backprop)); auto out_backprop_ptr = AsDeviceMemory(transformed_out_backprop.template flat().data(), @@ -1355,7 +1405,7 @@ class Conv3DBackpropInputOp : public OpKernel { dims.batch_size, dims.in_depth, {{dims.input_size(0), dims.input_size(1), dims.input_size(2)}}, - FORMAT_NCHW, + compute_data_format, dims.out_depth, {{dims.filter_size(0), dims.filter_size(1), dims.filter_size(2)}}, {{dims.dilation(0), dims.dilation(1), dims.dilation(2)}}, @@ -1500,8 +1550,11 @@ class Conv3DBackpropInputOp : public OpKernel { OP_REQUIRES_OK(context, context->allocate_temp( DataTypeToEnum::value, - {dims.batch_size, dims.in_depth, dims.input_size(0), - dims.input_size(1), dims.input_size(2)}, + ShapeFromFormat( + compute_data_format, dims.batch_size, + {{dims.input_size(0), dims.input_size(1), + dims.input_size(2)}}, + dims.in_depth), &in_backprop_remove_padding)); // Remove the padding for odd spatial dimensions. @@ -1510,12 +1563,13 @@ class Conv3DBackpropInputOp : public OpKernel { To32Bit(const_cast(pre_transformed_in_backprop) .tensor()), {{0, 0, 0}}, {{-planes_odd, -rows_odd, -cols_odd}}, - To32Bit(in_backprop_remove_padding.tensor()), FORMAT_NCHW); + To32Bit(in_backprop_remove_padding.tensor()), + compute_data_format); pre_transformed_in_backprop = in_backprop_remove_padding; } - if (data_format_ == FORMAT_NHWC) { + if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { auto toConstTensor = [](const Tensor& x) -> const Tensor { return x; }; functor::NCHWToNHWC()( context->eigen_device(), @@ -1723,6 +1777,35 @@ class Conv3DBackpropFilterOp : public OpKernel { CHECK(padding_rows >= 0 && padding_cols >= 0 && padding_planes >= 0) << "Negative paddings: (" << padding_rows << ", " << padding_cols << ", " << padding_planes << ")"; + +#if GOOGLE_CUDA + const bool compute_in_nhwc = CUDNN_VERSION >= 8000 && + DataTypeToEnum::value == DT_HALF; +#else + // fast NDHWC implementation is a CUDA only feature + const bool compute_in_nhwc = false; +#endif + const TensorFormat compute_data_format = + (compute_in_nhwc && data_format_ == FORMAT_NHWC) ? FORMAT_NHWC + : FORMAT_NCHW; + + VLOG(3) << "Compute Conv3DBackpropFilter with cuDNN:" + << " data_format=" << ToString(data_format_) + << " compute_data_format=" << ToString(compute_data_format); + + constexpr auto kComputeInNHWC = + std::make_tuple(se::dnn::DataLayout::kBatchYXDepth, + se::dnn::FilterLayout::kOutputYXInput); + constexpr auto kComputeInNCHW = + std::make_tuple(se::dnn::DataLayout::kBatchDepthYX, + se::dnn::FilterLayout::kOutputInputYX); + + se::dnn::DataLayout compute_data_layout; + se::dnn::FilterLayout filter_layout; + + std::tie(compute_data_layout, filter_layout) = + compute_data_format == FORMAT_NHWC ? kComputeInNHWC : kComputeInNCHW; + se::dnn::BatchDescriptor input_desc(3); input_desc.set_count(dims.batch_size) .set_spatial_dim(DimIndex::X, @@ -1732,20 +1815,21 @@ class Conv3DBackpropFilterOp : public OpKernel { .set_spatial_dim(DimIndex::Z, GetTensorDim(compatible_input, data_format_, '0')) .set_feature_map_count(dims.in_depth) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::BatchDescriptor output_desc(3); output_desc.set_count(dims.batch_size) .set_spatial_dim(DimIndex::X, dims.output_size(2)) .set_spatial_dim(DimIndex::Y, dims.output_size(1)) .set_spatial_dim(DimIndex::Z, dims.output_size(0)) .set_feature_map_count(dims.out_depth) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::FilterDescriptor filter_desc(3); filter_desc.set_spatial_dim(DimIndex::X, dims.filter_size(2)) .set_spatial_dim(DimIndex::Y, dims.filter_size(1)) .set_spatial_dim(DimIndex::Z, dims.filter_size(0)) .set_input_feature_map_count(filter_shape.dim_size(3)) - .set_output_feature_map_count(filter_shape.dim_size(4)); + .set_output_feature_map_count(filter_shape.dim_size(4)) + .set_layout(filter_layout); se::dnn::ConvolutionDescriptor conv_desc(3); conv_desc.set_dilation_rate(DimIndex::X, dims.dilation(2)) .set_dilation_rate(DimIndex::Y, dims.dilation(1)) @@ -1757,17 +1841,30 @@ class Conv3DBackpropFilterOp : public OpKernel { .set_zero_padding(DimIndex::Y, padding_rows / 2) .set_zero_padding(DimIndex::Z, padding_planes / 2) .set_group_count(dims.in_depth / filter_shape.dim_size(3)); + Tensor pre_transformed_filter_backprop; + auto dst_format = + compute_data_format == FORMAT_NCHW ? FORMAT_OIHW: FORMAT_OHWI; + TensorShape dst_shape = + dst_format == FORMAT_OIHW + ? TensorShape({filter_shape.dim_size(4), filter_shape.dim_size(3), + dims.filter_size(0), + dims.filter_size(1), + dims.filter_size(2)}) + : TensorShape({filter_shape.dim_size(4), + dims.filter_size(0), + dims.filter_size(1), + dims.filter_size(2), + filter_shape.dim_size(3)}); OP_REQUIRES_OK( context, context->allocate_temp( DataTypeToEnum::value, - TensorShape({filter_shape.dim_size(4), - filter_shape.dim_size(3), dims.filter_size(0), - dims.filter_size(1), dims.filter_size(2)}), + dst_shape, &pre_transformed_filter_backprop)); Tensor transformed_out_backprop; - if (data_format_ == FORMAT_NHWC) { + if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { + VLOG(4) << "Convert the `out_backprop` tensor from NDHWC to NCDHW."; TensorShape nchw_shape = {dims.batch_size, dims.out_depth, dims.output_size(0), dims.output_size(1), dims.output_size(2)}; @@ -1785,7 +1882,8 @@ class Conv3DBackpropFilterOp : public OpKernel { transformed_out_backprop = out_backprop; } Tensor transformed_input; - if (data_format_ == FORMAT_NHWC) { + if (data_format_ == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { + VLOG(4) << "Convert the `input` tensor from NDHWC to NCDHW."; TensorShape nchw_shape = { dims.batch_size, dims.in_depth, compatible_input.dim_size(1), compatible_input.dim_size(2), compatible_input.dim_size(3)}; @@ -1823,7 +1921,7 @@ class Conv3DBackpropFilterOp : public OpKernel { dims.batch_size, dims.in_depth, {{dims.input_size(0), dims.input_size(1), dims.input_size(2)}}, - FORMAT_NCHW, + compute_data_format, dims.out_depth, {{dims.filter_size(0), dims.filter_size(1), dims.filter_size(2)}}, {{dims.dilation(0), dims.dilation(1), dims.dilation(2)}}, @@ -1947,7 +2045,7 @@ class Conv3DBackpropFilterOp : public OpKernel { auto toConstTensor = [](const Tensor& x) -> const Tensor { return x; }; functor::ReverseTransformFilter()( - context->eigen_device(), /*src_filter_format=*/FORMAT_OIHW, + context->eigen_device(), /*src_filter_format=*/dst_format, toConstTensor(pre_transformed_filter_backprop).template tensor(), filter_backprop->tensor()); } diff --git a/tensorflow/core/kernels/conv_ops_3d.cc b/tensorflow/core/kernels/conv_ops_3d.cc index 69e6fba4192..c1fe6c690cd 100644 --- a/tensorflow/core/kernels/conv_ops_3d.cc +++ b/tensorflow/core/kernels/conv_ops_3d.cc @@ -43,6 +43,7 @@ using stream_executor::dnn::DimIndex; #include "tensorflow/stream_executor/gpu/asm_compiler.h" #include "tensorflow/stream_executor/gpu/redzone_allocator.h" #include "tensorflow/stream_executor/tf_allocator_adapter.h" +#include "third_party/gpus/cudnn/cudnn.h" #endif // GOOGLE_CUDA namespace tensorflow { @@ -201,7 +202,23 @@ struct LaunchConvOp { } } - if (data_format == FORMAT_NHWC) { +#if GOOGLE_CUDA + const bool compute_in_nhwc = CUDNN_VERSION >= 8000 && + DataTypeToEnum::value == DT_HALF; +#else + // fast NHWC implementation is a CUDA only feature + const bool compute_in_nhwc = false; +#endif + const TensorFormat compute_data_format = + (compute_in_nhwc && data_format == FORMAT_NHWC) ? FORMAT_NHWC + : FORMAT_NCHW; + + VLOG(3) << "Compute Conv3D with cuDNN:" + << " data_format=" << ToString(data_format) + << " compute_data_format=" << ToString(compute_data_format); + + if (data_format == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { + VLOG(4) << "Convert the input tensor from NDHWC to NCDHW."; const TensorShape nchw_shape = ShapeFromFormat( FORMAT_NCHW, in_batch, {{in_planes, in_rows, in_cols}}, in_depth); if (in_depth > 1) { @@ -219,8 +236,26 @@ struct LaunchConvOp { } else { CHECK(input.CopyFrom(input, nchw_shape)); } + } else { + CHECK(data_format == compute_data_format) // Crash OK + << "Illegal data and compute format pair:" + << " data_format=" << ToString(data_format) + << " compute_data_format=" << ToString(compute_data_format); } + constexpr auto kComputeInNHWC = + std::make_tuple(se::dnn::DataLayout::kBatchYXDepth, + se::dnn::FilterLayout::kOutputYXInput); + constexpr auto kComputeInNCHW = + std::make_tuple(se::dnn::DataLayout::kBatchDepthYX, + se::dnn::FilterLayout::kOutputInputYX); + + se::dnn::DataLayout compute_data_layout; + se::dnn::FilterLayout filter_layout; + + std::tie(compute_data_layout, filter_layout) = + compute_data_format == FORMAT_NHWC ? kComputeInNHWC : kComputeInNCHW; + CHECK(pad_rows >= 0 && pad_cols >= 0 && pad_planes >= 0) << "Negative paddings: (" << pad_rows << ", " << pad_cols << ", " << pad_planes << ")"; @@ -230,20 +265,21 @@ struct LaunchConvOp { .set_spatial_dim(DimIndex::X, in_cols) .set_spatial_dim(DimIndex::Y, in_rows) .set_spatial_dim(DimIndex::Z, in_planes) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::BatchDescriptor output_desc(3); output_desc.set_count(in_batch) .set_spatial_dim(DimIndex::X, out_cols) .set_spatial_dim(DimIndex::Y, out_rows) .set_spatial_dim(DimIndex::Z, out_planes) .set_feature_map_count(out_depth) - .set_layout(se::dnn::DataLayout::kBatchDepthYX); + .set_layout(compute_data_layout); se::dnn::FilterDescriptor filter_desc(3); filter_desc.set_spatial_dim(DimIndex::X, filter_cols) .set_spatial_dim(DimIndex::Y, filter_rows) .set_spatial_dim(DimIndex::Z, filter_planes) .set_input_feature_map_count(filter_depth) - .set_output_feature_map_count(out_depth); + .set_output_feature_map_count(out_depth) + .set_layout(filter_layout); se::dnn::ConvolutionDescriptor conv_desc(3); conv_desc.set_dilation_rate(DimIndex::X, dilations[2]) .set_dilation_rate(DimIndex::Y, dilations[1]) @@ -257,25 +293,42 @@ struct LaunchConvOp { .set_group_count(in_depth / filter_depth); Tensor transformed_filter; + auto dst_format = + compute_data_format == FORMAT_NCHW ? FORMAT_OIHW: FORMAT_OHWI; + VLOG(4) << "Transform filter tensor from " << ToString(FORMAT_HWIO) + << " to " << ToString(dst_format); + TensorShape dst_shape = + dst_format == FORMAT_OIHW + ? TensorShape({filter.dim_size(4), filter.dim_size(3), + filter.dim_size(0), filter.dim_size(1), + filter.dim_size(2)}) + : TensorShape({filter.dim_size(4), filter.dim_size(0), + filter.dim_size(1), filter.dim_size(2), + filter.dim_size(3)}); OP_REQUIRES_OK( ctx, ctx->allocate_temp(DataTypeToEnum::value, - TensorShape({out_depth, in_depth, filter_planes, - filter_rows, filter_cols}), + dst_shape, &transformed_filter)); // filter: [x, y, z, in, out] - // t_filter: [out, in, x, y, z] + // t_filter: [out, in, x, y, z] (NCDHW) or + // t_filter: [out, x, y, z, in] (NDHWC) functor::TransformFilter()( - ctx->eigen_device(), FORMAT_OIHW, + ctx->eigen_device(), dst_format, To32Bit(filter.tensor()), To32Bit(transformed_filter.tensor())); Tensor transformed_output; - OP_REQUIRES_OK( - ctx, ctx->allocate_temp( - DataTypeToEnum::value, - ShapeFromFormat(FORMAT_NCHW, in_batch, - {{out_planes, out_rows, out_cols}}, out_depth), - &transformed_output)); + if (data_format != compute_data_format) { + VLOG(4) << "Allocate temporary memory for output in compute data format"; + OP_REQUIRES_OK( + ctx, ctx->allocate_temp( + DataTypeToEnum::value, + ShapeFromFormat(FORMAT_NCHW, in_batch, + {{out_planes, out_rows, out_cols}}, out_depth), + &transformed_output)); + } else { + transformed_output = *output; + } auto input_ptr = AsDeviceMemory(input.template flat().data(), input.template flat().size()); @@ -295,7 +348,7 @@ struct LaunchConvOp { in_batch, in_depth, {{in_planes, in_rows, in_cols}}, - FORMAT_NCHW, + compute_data_format, out_depth, {{filter_planes, filter_rows, filter_cols}}, {{dilations[0], dilations[1], dilations[2]}}, @@ -455,15 +508,14 @@ struct LaunchConvOp { ") filter shape(", filter.shape().DebugString(), ")")); } - if (data_format == FORMAT_NHWC) { + if (data_format == FORMAT_NHWC && compute_data_format == FORMAT_NCHW) { + VLOG(4) << "Convert the output tensor back from NCDHW to NDHWC."; // t_output: [b, out, x, y, z] // output: [b, x, y, z, out] functor::NCHWToNHWC()( ctx->eigen_device(), const_cast(transformed_output).tensor(), output->tensor()); - } else { - *output = transformed_output; } } }; From 744010d48dbe7c3877fec417797fb128882b6ec5 Mon Sep 17 00:00:00 2001 From: Steenu Johnson Date: Fri, 12 Jun 2020 11:39:54 +0530 Subject: [PATCH 0139/1390] Forward compatiblity and backwards compatibility changes. Changing forward compatibility date to be 3 weeks from the current commit. Using V2 version when user specifies exclude_cols. Making op_version in kernel more specific to CSVDatasetV2. Serialization done only for version 2 op. Removing clashes from the previous commit and fixing indentation errors. Signed-off-by: Steenu Johnson --- .../data/experimental/csv_dataset_op.cc | 81 +++++++++++-------- .../python/data/experimental/ops/readers.py | 2 +- 2 files changed, 47 insertions(+), 36 deletions(-) diff --git a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc index 906e57066a2..b1c5937747d 100644 --- a/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc +++ b/tensorflow/core/kernels/data/experimental/csv_dataset_op.cc @@ -30,7 +30,7 @@ class CSVDatasetOp : public DatasetOpKernel { public: explicit CSVDatasetOp(OpKernelConstruction* ctx) : DatasetOpKernel(ctx), - op_version_(ctx->def().op() == "CSVDataset" ? 1 : 2) { + op_version_(ctx->def().op() == "CSVDatasetV2" ? 2 : 1) { OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_)); } @@ -64,18 +64,17 @@ class CSVDatasetOp : public DatasetOpKernel { OP_REQUIRES(ctx, select_cols_tensor->dims() == 1, errors::InvalidArgument("`select_cols` must be a vector.")); -std::vector exclude_cols; -if (op_version_ > 1) { - const Tensor* exclude_cols_tensor; - OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); - exclude_cols.reserve(exclude_cols_tensor->NumElements()); - for (int i = 0; i < exclude_cols_tensor->NumElements(); ++i) { - exclude_cols.push_back(exclude_cols_tensor->flat()(i)); - } -} - - OP_REQUIRES(ctx, exclude_cols_tensor->dims() == 1, - errors::InvalidArgument("`exclude_cols` must be a vector")); + std::vector exclude_cols; + if (op_version_ > 1) { + const Tensor* exclude_cols_tensor; + OP_REQUIRES_OK(ctx, ctx->input("exclude_cols", &exclude_cols_tensor)); + OP_REQUIRES(ctx, exclude_cols_tensor->dims() == 1, + errors::InvalidArgument("`exclude_cols` must be a vector")); + exclude_cols.reserve(exclude_cols_tensor->NumElements()); + for (int i = 0; i < exclude_cols_tensor->NumElements(); ++i) { + exclude_cols.push_back(exclude_cols_tensor->flat()(i)); + } + } int64 buffer_size = 0; OP_REQUIRES_OK( @@ -141,11 +140,6 @@ if (op_version_ > 1) { ctx, select_cols.empty() || select_cols.front() >= 0, errors::InvalidArgument("select_cols should be non-negative indices")); - std::vector exclude_cols; - exclude_cols.reserve(exclude_cols_tensor->NumElements()); - for (int i = 0; i < exclude_cols_tensor->NumElements(); ++i) { - exclude_cols.push_back(exclude_cols_tensor->flat()(i)); - } OP_REQUIRES(ctx, select_cols.empty() || exclude_cols.empty(), errors::InvalidArgument( "Either select_cols or exclude_cols should be empty")); @@ -158,12 +152,12 @@ if (op_version_ > 1) { ctx, exclude_cols.empty() || exclude_cols.front() >= 0, errors::InvalidArgument("exclude_cols should be non-negative indices")); - *output = - new Dataset(ctx, std::move(filenames), header, - std::move(compression_type), zlib_compression_options, - output_types_, output_shapes_, std::move(record_defaults), - std::move(select_cols), std::move(exclude_cols), - use_quote_delim, delim[0], std::move(na_value)); + *output = new Dataset(ctx, std::move(filenames), header, + std::move(compression_type), zlib_compression_options, + output_types_, output_shapes_, + std::move(record_defaults), std::move(select_cols), + std::move(exclude_cols), use_quote_delim, delim[0], + std::move(na_value), op_version_); } private: @@ -175,7 +169,7 @@ if (op_version_ > 1) { const std::vector& output_shapes, std::vector record_defaults, std::vector select_cols, std::vector exclude_cols, bool use_quote_delim, char delim, - string na_value) + string na_value, int op_version) : DatasetBase(DatasetContext(ctx)), filenames_(std::move(filenames)), header_(header), @@ -187,6 +181,7 @@ if (op_version_ > 1) { use_quote_delim_(use_quote_delim), delim_(delim), na_value_(std::move(na_value)), + op_version_(op_version), use_compression_(!compression_type.empty()), compression_type_(std::move(compression_type)), options_(options) {} @@ -242,16 +237,31 @@ if (op_version_ > 1) { TF_RETURN_IF_ERROR(b->AddVector(select_cols_, &select_cols)); TF_RETURN_IF_ERROR(b->AddVector(exclude_cols_, &exclude_cols)); - TF_RETURN_IF_ERROR(b->AddDataset( - this, - {std::make_pair(0, filenames), std::make_pair(1, compression_type), - std::make_pair(2, buffer_size), std::make_pair(3, header), - std::make_pair(4, delim), std::make_pair(5, use_quote_delim), - std::make_pair(6, na_value), std::make_pair(7, select_cols), - std::make_pair(9, exclude_cols)}, // Single tensor inputs - {std::make_pair(8, record_defaults)}, // Tensor list inputs - {}, - output)); + if (op_version_ > 1) { + TF_RETURN_IF_ERROR(b->AddDataset( + this, + {std::make_pair(0, filenames), std::make_pair(1, compression_type), + std::make_pair(2, buffer_size), std::make_pair(3, header), + std::make_pair(4, delim), std::make_pair(5, use_quote_delim), + std::make_pair(6, na_value), std::make_pair(7, select_cols), + std::make_pair(9, exclude_cols)}, // Single tensor inputs + {std::make_pair(8, record_defaults)}, // Tensor list inputs + {}, + output)); + } else { + TF_RETURN_IF_ERROR(b->AddDataset( + this, + { + std::make_pair(0, filenames), + std::make_pair(1, compression_type), + std::make_pair(2, buffer_size), std::make_pair(3, header), + std::make_pair(4, delim), std::make_pair(5, use_quote_delim), + std::make_pair(6, na_value), std::make_pair(7, select_cols), + }, // Single tensor inputs + {std::make_pair(8, record_defaults)}, // Tensor list inputs + {}, + output)); + } return Status::OK(); } @@ -899,6 +909,7 @@ if (op_version_ > 1) { const bool use_quote_delim_; const char delim_; const tstring na_value_; + const int op_version_; const bool use_compression_; const tstring compression_type_; const io::ZlibCompressionOptions options_; diff --git a/tensorflow/python/data/experimental/ops/readers.py b/tensorflow/python/data/experimental/ops/readers.py index 14a507580ad..eed955e51b0 100644 --- a/tensorflow/python/data/experimental/ops/readers.py +++ b/tensorflow/python/data/experimental/ops/readers.py @@ -732,7 +732,7 @@ class CsvDatasetV2(dataset_ops.DatasetSource): ) self._element_spec = tuple( tensor_spec.TensorSpec([], d.dtype) for d in self._record_defaults) - if compat.forward_compatible(2020, 6, 25): + if compat.forward_compatible(2020, 7, 3) or exclude_cols is not None: variant_tensor = gen_experimental_dataset_ops.csv_dataset_v2( filenames=self._filenames, record_defaults=self._record_defaults, From bb5e2e63d9bfc197f632c0e28446d47ffdba7661 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Mon, 18 May 2020 16:15:38 +0200 Subject: [PATCH 0140/1390] TFLu: Add support for unknown output dimensions Some models (e.g. Convnet) with broadcasting have unspecified output dimensions for some operators, where the dimensions can be deduced from the input dimensions. Change-Id: Ica7d3715c4f4cfa0ce05580afe5f25bb2cfce981 --- tensorflow/lite/micro/kernels/add.cc | 15 ++++++++- tensorflow/lite/micro/kernels/cmsis-nn/add.cc | 15 ++++++++- tensorflow/lite/micro/kernels/cmsis-nn/mul.cc | 15 ++++++++- tensorflow/lite/micro/kernels/mul.cc | 15 ++++++++- tensorflow/lite/micro/memory_helpers.cc | 31 +++++++++++++++++++ tensorflow/lite/micro/memory_helpers.h | 5 +++ 6 files changed, 92 insertions(+), 4 deletions(-) diff --git a/tensorflow/lite/micro/kernels/add.cc b/tensorflow/lite/micro/kernels/add.cc index 42609301232..feb5f2c0b29 100644 --- a/tensorflow/lite/micro/kernels/add.cc +++ b/tensorflow/lite/micro/kernels/add.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" #include "tensorflow/lite/kernels/op_macros.h" +#include "tensorflow/lite/micro/memory_helpers.h" namespace tflite { namespace ops { @@ -94,6 +95,18 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteAddParams* params, return kTfLiteOk; } +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + const TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->dims->size == 0) { + return AllocateOutputDimensionsFromInput(context, input1, input2, output); + } + + return kTfLiteOk; +} + void EvalAdd(TfLiteContext* context, TfLiteNode* node, TfLiteAddParams* params, const OpData* data, const TfLiteTensor* input1, const TfLiteTensor* input2, TfLiteTensor* output) { @@ -190,7 +203,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TfLiteRegistration* Register_ADD() { static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - /*prepare=*/nullptr, + /*prepare=*/add::Prepare, /*invoke=*/add::Eval, /*profiling_string=*/nullptr, /*builtin_code=*/0, diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/add.cc b/tensorflow/lite/micro/kernels/cmsis-nn/add.cc index 6dbe4a618ab..d1f3eae233a 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/add.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/add.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" #include "tensorflow/lite/kernels/op_macros.h" +#include "tensorflow/lite/micro/memory_helpers.h" namespace tflite { namespace ops { @@ -94,6 +95,18 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteAddParams* params, return kTfLiteOk; } +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + const TfLiteTensor* input1 = GetInput(context, node, kInputTensor1); + const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->dims->size == 0) { + return AllocateOutputDimensionsFromInput(context, input1, input2, output); + } + + return kTfLiteOk; +} + void EvalAdd(TfLiteContext* context, TfLiteNode* node, TfLiteAddParams* params, const OpData* data, const TfLiteTensor* input1, const TfLiteTensor* input2, TfLiteTensor* output) { @@ -199,7 +212,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TfLiteRegistration* Register_ADD() { static TfLiteRegistration r = {nullptr /* Init */, nullptr /* Free */, - nullptr /* Prepare */, add::Eval}; + add::Prepare, add::Eval}; return &r; } diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc b/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc index d746166ebd9..6d0592d6c09 100644 --- a/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc +++ b/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h" #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/memory_helpers.h" namespace tflite { namespace ops { @@ -64,6 +65,18 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node, return kTfLiteOk; } +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + const TfLiteTensor* input1 = GetInput(context, node, kInput1Tensor); + const TfLiteTensor* input2 = GetInput(context, node, kInput2Tensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->dims->size == 0) { + return AllocateOutputDimensionsFromInput(context, input1, input2, output); + } + + return kTfLiteOk; +} + void EvalQuantized(TfLiteContext* context, TfLiteNode* node, TfLiteMulParams* params, OpData* data, const TfLiteTensor* input1, const TfLiteTensor* input2, @@ -167,7 +180,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TfLiteRegistration* Register_MUL() { static TfLiteRegistration r = {nullptr /* Init */, nullptr /* Free */, - nullptr /* Prepare */, mul::Eval}; + mul::Prepare, mul::Eval}; return &r; } diff --git a/tensorflow/lite/micro/kernels/mul.cc b/tensorflow/lite/micro/kernels/mul.cc index fb47728a1a4..b9c989e1f87 100644 --- a/tensorflow/lite/micro/kernels/mul.cc +++ b/tensorflow/lite/micro/kernels/mul.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h" #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/micro/memory_helpers.h" namespace tflite { namespace ops { @@ -65,6 +66,18 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node, return kTfLiteOk; } +TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { + const TfLiteTensor* input1 = GetInput(context, node, kInput1Tensor); + const TfLiteTensor* input2 = GetInput(context, node, kInput2Tensor); + TfLiteTensor* output = GetOutput(context, node, kOutputTensor); + + if (output->dims->size == 0) { + return AllocateOutputDimensionsFromInput(context, input1, input2, output); + } + + return kTfLiteOk; +} + void EvalQuantized(TfLiteContext* context, TfLiteNode* node, TfLiteMulParams* params, OpData* data, const TfLiteTensor* input1, const TfLiteTensor* input2, @@ -161,7 +174,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TfLiteRegistration* Register_MUL() { static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - /*prepare=*/nullptr, + /*prepare=*/mul::Prepare, /*invoke=*/mul::Eval, /*profiling_string=*/nullptr, /*builtin_code=*/0, diff --git a/tensorflow/lite/micro/memory_helpers.cc b/tensorflow/lite/micro/memory_helpers.cc index 37c78162b62..a0c0d85429a 100644 --- a/tensorflow/lite/micro/memory_helpers.cc +++ b/tensorflow/lite/micro/memory_helpers.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/lite/core/api/error_reporter.h" #include "tensorflow/lite/core/api/flatbuffer_conversions.h" #include "tensorflow/lite/schema/schema_generated.h" +#include "tensorflow/lite/kernels/internal/tensor_ctypes.h" namespace tflite { @@ -101,4 +102,34 @@ TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, return kTfLiteOk; } +TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, const TfLiteTensor* input1, + const TfLiteTensor* input2, TfLiteTensor* output) { + int size = 1, i = 0; + const TfLiteTensor* input = nullptr; + + TF_LITE_ENSURE(context, input1->dims != nullptr); + TF_LITE_ENSURE(context, input2->dims != nullptr); + TF_LITE_ENSURE(context, output->dims->size == 0); + + input = input1->dims->size > input2->dims->size ? input1 : input2; + TF_LITE_ENSURE(context, output->type == input->type); + + const int dimensions_count = tflite::GetTensorShape(input).DimensionsCount(); + for (i = 0; i < dimensions_count; i++) { + size *= input->dims->data[i]; + } + output->bytes += size; + + TF_LITE_ENSURE_STATUS(context->AllocatePersistentBuffer( + context, TfLiteIntArrayGetSizeInBytes(size), + reinterpret_cast(&output->dims))); + + output->dims->size = input->dims->size; + for (i = 0; i < dimensions_count; i++) { + output->dims->data[i] = input->dims->data[i]; + } + + return kTfLiteOk; +} + } // namespace tflite diff --git a/tensorflow/lite/micro/memory_helpers.h b/tensorflow/lite/micro/memory_helpers.h index f52da062271..1b6fa48dd99 100644 --- a/tensorflow/lite/micro/memory_helpers.h +++ b/tensorflow/lite/micro/memory_helpers.h @@ -42,6 +42,11 @@ TfLiteStatus BytesRequiredForTensor(const tflite::Tensor& flatbuffer_tensor, size_t* bytes, size_t* type_size, ErrorReporter* error_reporter); +// Deduce output dimensions from input and allocate given size. +// Useful for operators with two inputs where the largest input should equal the output dimension. +TfLiteStatus AllocateOutputDimensionsFromInput(TfLiteContext* context, const TfLiteTensor* input1, + const TfLiteTensor* input2, TfLiteTensor* output); + } // namespace tflite #endif // TENSORFLOW_LITE_MICRO_MEMORY_HELPERS_H_ From 2f9642602d9d3da721d603e583b2569c765704a1 Mon Sep 17 00:00:00 2001 From: Jens Elofsson Date: Fri, 12 Jun 2020 10:49:57 +0200 Subject: [PATCH 0141/1390] Adapt to changes in micro_allocator. --- tensorflow/lite/micro/micro_allocator.cc | 20 ++++--- tensorflow/lite/micro/micro_allocator.h | 3 +- tensorflow/lite/micro/micro_allocator_test.cc | 52 ++++++++++++++----- 3 files changed, 55 insertions(+), 20 deletions(-) diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc index 7cd40e54435..16b3b986a52 100644 --- a/tensorflow/lite/micro/micro_allocator.cc +++ b/tensorflow/lite/micro/micro_allocator.cc @@ -375,10 +375,9 @@ TfLiteStatus CreatePlan(ErrorReporter* error_reporter, planner->AddBuffer(error_reporter, aligned_bytes_required, current->first_created, current->last_used)); } else { - TF_LITE_ENSURE_STATUS( - planner->AddBuffer(error_reporter, aligned_bytes_required, - current->first_created, current->last_used, - current->offline_offset)); + TF_LITE_ENSURE_STATUS(planner->AddBuffer( + error_reporter, aligned_bytes_required, current->first_created, + current->last_used, current->offline_offset)); } } } @@ -647,7 +646,7 @@ TfLiteStatus MicroAllocator::FinishModelAllocation(const Model* model, const SubGraph* subgraph = GetSubGraphFromModel(model); TFLITE_DCHECK(subgraph != nullptr); - TF_LITE_ENSURE_STATUS(CommitStaticMemoryPlan(subgraph, context)); + TF_LITE_ENSURE_STATUS(CommitStaticMemoryPlan(model, subgraph, context)); TF_LITE_ENSURE_STATUS(AllocateVariables(subgraph->tensors(), context->tensors, memory_allocator_)); @@ -874,7 +873,8 @@ const SubGraph* MicroAllocator::GetSubGraphFromModel(const Model* model) { return (*subgraphs)[0]; } -TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(const SubGraph* subgraph, +TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(const Model* model, + const SubGraph* subgraph, TfLiteContext* context) { // Create static memory plan // 1. Calculate AllocationInfo to know the lifetime of each tensor/buffer. @@ -891,7 +891,13 @@ TfLiteStatus MicroAllocator::CommitStaticMemoryPlan(const SubGraph* subgraph, AllocationInfoBuilder builder(error_reporter_, &tmp_allocator); TF_LITE_ENSURE_STATUS( builder.Init(subgraph->tensors()->size(), scratch_buffer_count_)); - TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph, context->tensors)); + + int32_t* offline_planner_offsets = nullptr; + TF_LITE_ENSURE_STATUS( + builder.GetOfflinePlannedOffsets(model, &offline_planner_offsets)); + TF_LITE_ENSURE_STATUS(builder.AddTensors(subgraph, offline_planner_offsets, + context->tensors)); + TF_LITE_ENSURE_STATUS(builder.AddScratchBuffers(scratch_buffer_handles_)); const AllocationInfo* allocation_info = builder.Finish(); diff --git a/tensorflow/lite/micro/micro_allocator.h b/tensorflow/lite/micro/micro_allocator.h index 7fc091196a5..9cfc1793fc7 100644 --- a/tensorflow/lite/micro/micro_allocator.h +++ b/tensorflow/lite/micro/micro_allocator.h @@ -185,7 +185,8 @@ class MicroAllocator { // Commits a memory plan for all non-persistent buffer allocations in the // 'head' section of the memory arena. - virtual TfLiteStatus CommitStaticMemoryPlan(const SubGraph* subgraph, + virtual TfLiteStatus CommitStaticMemoryPlan(const Model* model, + const SubGraph* subgraph, TfLiteContext* context); // A simple memory allocator that always allocate from the arena tail or head. diff --git a/tensorflow/lite/micro/micro_allocator_test.cc b/tensorflow/lite/micro/micro_allocator_test.cc index 04f4732b9d3..f3f3f32611e 100644 --- a/tensorflow/lite/micro/micro_allocator_test.cc +++ b/tensorflow/lite/micro/micro_allocator_test.cc @@ -312,6 +312,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { int version = 1; int subgraph = 0; constexpr int nbr_tensors = 4; + tflite::testing::MockOpResolver mock_resolver; + tflite::NodeAndRegistration* node_and_registration; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = {version, subgraph, nbr_tensors, // header @@ -340,9 +342,14 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( - &context, model, arena, arena_size, micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); + tflite::MicroAllocator* allocator = + tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator->StartModelAllocation(model, &context, mock_resolver, + &node_and_registration)); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, + allocator->FinishModelAllocation(model, &context)); // Since all of the tensors are online planned and the model structure is // identical to that in TestAllocationForModelsWithBranches, @@ -357,6 +364,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerBranchesAllOnline) { TF_LITE_MICRO_TEST(OfflinePlannerBasic) { constexpr int nbr_tensors = 4; + tflite::testing::MockOpResolver mock_resolver; + tflite::NodeAndRegistration* node_and_registration; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = {1, 0, nbr_tensors, 0, // t0 @@ -389,9 +398,14 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( - &context, model, arena, arena_size, micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); + tflite::MicroAllocator* allocator = + tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator->StartModelAllocation(model, &context, mock_resolver, + &node_and_registration)); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, + allocator->FinishModelAllocation(model, &context)); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); @@ -402,6 +416,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerBasic) { TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { constexpr int nbr_tensors = 4; + tflite::testing::MockOpResolver mock_resolver; + tflite::NodeAndRegistration* node_and_registration; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = { 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors @@ -434,9 +450,14 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( - &context, model, arena, arena_size, micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); + tflite::MicroAllocator* allocator = + tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator->StartModelAllocation(model, &context, mock_resolver, + &node_and_registration)); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, + allocator->FinishModelAllocation(model, &context)); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); @@ -448,6 +469,8 @@ TF_LITE_MICRO_TEST(OfflinePlannerOverlappingAllocation) { TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { constexpr int nbr_tensors = 5; + tflite::testing::MockOpResolver mock_resolver; + tflite::NodeAndRegistration* node_and_registration; const int32_t metadata_buffer[tflite::testing::kOfflinePlannerHeaderSize + nbr_tensors] = { 1, 0, nbr_tensors, // header: version, subgraph, nbr tensors @@ -482,9 +505,14 @@ TF_LITE_MICRO_TEST(OfflinePlannerOfflineOnline) { TfLiteContext context; constexpr size_t arena_size = 4096; uint8_t arena[arena_size]; - tflite::MicroAllocator* allocator = tflite::MicroAllocator::Create( - &context, model, arena, arena_size, micro_test::reporter); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, allocator->FinishTensorAllocation()); + tflite::MicroAllocator* allocator = + tflite::MicroAllocator::Create(arena, arena_size, micro_test::reporter); + + TF_LITE_MICRO_EXPECT_EQ( + kTfLiteOk, allocator->StartModelAllocation(model, &context, mock_resolver, + &node_and_registration)); + TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, + allocator->FinishModelAllocation(model, &context)); uint8_t* start = context.tensors[0].data.uint8; TF_LITE_MICRO_EXPECT_EQ(0, context.tensors[0].data.uint8 - start); From 3626d6cf24924957b001ac92db934967139cc3c5 Mon Sep 17 00:00:00 2001 From: Michael137 Date: Sun, 14 Jun 2020 11:22:21 -0400 Subject: [PATCH 0142/1390] Document usage of the profiling_output_csv_file option --- tensorflow/lite/tools/benchmark/README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/lite/tools/benchmark/README.md b/tensorflow/lite/tools/benchmark/README.md index a43383cff9d..d76e8d4d031 100644 --- a/tensorflow/lite/tools/benchmark/README.md +++ b/tensorflow/lite/tools/benchmark/README.md @@ -48,6 +48,12 @@ and the following optional parameters: 'enable_op_profiling'. Note, the platform-wide tracing might not work if the tool runs as a commandline native binary. For example, on Android, the ATrace-based tracing only works when the tool is launched as an APK. +* `profiling_output_csv_file`: `str` (default="") \ + File path to export profile data to as CSV. + The results are printed to `stdout` if option is not set. + Requires `enable_op_profiling` to be `true` and the path + to include the name of the output CSV; otherwise + results are printed to `stdout`. ### TFLite delegate parameters The tool supports all runtime/delegate parameters introduced by From ec0217f4ac194741cf6c566f01747a5eb771edc6 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 11:29:41 -0400 Subject: [PATCH 0143/1390] Update quantization_config.cc --- .../compiler/mlir/lite/quantization/quantization_config.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc index c4cf6e71cf3..634d212409e 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc @@ -60,7 +60,7 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_maxs; if (!max_values.empty()) { std::vector node_maxs_str = absl::StrSplit(max_values, ','); - for (size_t i = 0; i < node_maxs_str.size(); i++) { + for (int i : llvm::seq(node_maxs_str.size())) { double value; if (!absl::SimpleAtod(node_maxs_str[i], &value)) { llvm::errs() << "Unexpected mins: " << node_maxs_str[i] << "\n"; From fc205f91b67f38b02e026245fa39f3f61981e605 Mon Sep 17 00:00:00 2001 From: Chen Lei Date: Mon, 15 Jun 2020 00:07:37 +0800 Subject: [PATCH 0144/1390] Fix wrong comments Reference [Probot: Stale](https://github.com/probot/stale) --- .github/stale.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/stale.yml b/.github/stale.yml index e1184ce37b4..5f8dd12f477 100644 --- a/.github/stale.yml +++ b/.github/stale.yml @@ -23,7 +23,7 @@ daysUntilStale: 7 # Number of days of inactivity before a stale Issue or Pull Request is closed daysUntilClose: 7 -# Issues or Pull Requests with these labels will never be considered stale. Set to `[]` to disable +# Only issues or pull requests with all of these labels are check if stale. Defaults to `[]` (disabled) onlyLabels: - stat:awaiting response # Comment to post when marking as stale. Set to `false` to disable From 02090cac6a7f1958920155522d912de0b2769301 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 12:20:43 -0400 Subject: [PATCH 0145/1390] Update import_quant_stats_pass.cc --- .../compiler/mlir/lite/quantization/import_quant_stats_pass.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc index 5419a0d5e1b..e00a088c38c 100644 --- a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc +++ b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc @@ -158,7 +158,7 @@ void ImportQuantStatsPass::ImportAsStatsOps(OpBuilder b, Operation *op, InsertStatsOpAtResult(b, op->getResult(index), layer_stats, axis_stats, axis); } else { - for (int i = 0; i < static_cast(op->getNumResults()); ++i) { + for (int i = 0, e = op->getNumResults(); i < e; ++i) { if (IsQuantizableResult(op, i)) { InsertStatsOpAtResult(b, op->getResult(i), layer_stats, axis_stats, axis); From 0052918f2b9332e7eabe4b2ababbbdb464889cec Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 12:23:40 -0400 Subject: [PATCH 0146/1390] Update quantization_config.cc --- .../compiler/mlir/lite/quantization/quantization_config.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc index 634d212409e..b299fa8f4c2 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc @@ -48,7 +48,7 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_mins; if (!min_values.empty()) { std::vector node_mins_str = absl::StrSplit(min_values, ','); - for (size_t i = 0; i < node_mins_str.size(); i++) { + for (int i : llvm::seq(node_mins_str.size())) { double value; if (!absl::SimpleAtod(node_mins_str[i], &value)) { return true; From 09f45f4f5edba341b21c89266431c7dd9e950af8 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 12:23:47 -0400 Subject: [PATCH 0147/1390] Update quantization_config.cc From 79eead7a46147cced45480d7da02e798e6e4ba56 Mon Sep 17 00:00:00 2001 From: Vo Van Nghia Date: Sun, 14 Jun 2020 23:41:06 +0700 Subject: [PATCH 0148/1390] Add GCS Path Parser --- .../experimental/filesystem/plugins/gcs/BUILD | 1 + .../filesystem/plugins/gcs/gcs_filesystem.cc | 40 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD index c7a2d68ac3e..c9fee433589 100644 --- a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD +++ b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD @@ -27,5 +27,6 @@ cc_library( "//tensorflow/c:tf_status", "//tensorflow/c/experimental/filesystem:filesystem_interface", "@com_github_googlecloudplatform_google_cloud_cpp//:storage_client", + "@com_google_absl//absl/strings", ], ) diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc index 7918e7a7310..11641fac53b 100644 --- a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc +++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include "absl/strings/string_view.h" #include "google/cloud/storage/client.h" #include "tensorflow/c/experimental/filesystem/filesystem_interface.h" #include "tensorflow/c/tf_status.h" @@ -35,6 +36,45 @@ static inline void TF_SetStatusFromGCSStatus( static void* plugin_memory_allocate(size_t size) { return calloc(1, size); } static void plugin_memory_free(void* ptr) { free(ptr); } +static void ParseGCSPath(absl::string_view fname, bool object_empty_ok, char** bucket, + char** object, TF_Status* status) { + size_t scheme_end = fname.find("://") + 2; + if (fname.substr(0, scheme_end + 1) != "gs://") { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "GCS path doesn't start with 'gs://'."); + return; + } + + size_t bucket_end = fname.find("/", scheme_end + 1); + if (bucket_end == absl::string_view::npos) { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "GCS path doesn't contain a bucket name."); + return; + } + absl::string_view bucket_view = + fname.substr(scheme_end + 1, bucket_end - scheme_end - 1); + *bucket = + static_cast(plugin_memory_allocate(bucket_view.length() + 1)); + memcpy(*bucket, bucket_view.data(), bucket_view.length()); + (*bucket)[bucket_view.length()] = '\0'; + + absl::string_view object_view = fname.substr(bucket_end + 1); + if (object_view == "") { + if (object_empty_ok) { + *object = nullptr; + return; + } else { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "GCS path doesn't contain an object name."); + return; + } + } + *object = + static_cast(plugin_memory_allocate(object_view.length() + 1)); + // object_view.data() is a null-terminated string_view because fname is. + strcpy(*object, object_view.data()); +} + // SECTION 1. Implementation for `TF_RandomAccessFile` // ---------------------------------------------------------------------------- namespace tf_random_access_file { From 2f515039033a27b6253ae788e79c7ead32265007 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 14:36:49 -0400 Subject: [PATCH 0149/1390] Update quantization_driver.cc --- .../compiler/mlir/lite/quantization/quantization_driver.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc index fc11604ef8a..a9f4eb78431 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc @@ -294,7 +294,7 @@ class QuantizationDriver { return; if (current_op == op) llvm::errs() << "===>>>"; llvm::errs() << op->getName() << " : ("; - for (size_t i = 0; i < op->getNumOperands(); ++i) { + for (int i = 0, e = op->getNumOperands(); i < e; ++i) { if (auto params = GetOperandQuantState(op, i).params) params.print(llvm::errs()); else @@ -303,7 +303,7 @@ class QuantizationDriver { llvm::errs() << ","; } llvm::errs() << ") -> ("; - for (size_t i = 0; i < op->getNumResults(); ++i) { + for (int i = 0, e = op->getNumResults(); i < e; ++i) { if (auto params = GetResultQuantState(op, i).params) params.print(llvm::errs()); else From 5ac4a4a3ea522b2d0b6b7fa0058e7bc1bb5ba6d3 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 14:47:27 -0400 Subject: [PATCH 0150/1390] Update quantization_utils.cc --- .../compiler/mlir/lite/quantization/quantization_utils.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc index b9ca5329519..57b24eb8772 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc @@ -368,7 +368,7 @@ quant::QuantizedType GetUniformQuantizedTypeForBias( scales[index_scale.index()] *= index_scale.value(); } } else if (auto type = op_type.dyn_cast()) { - for (size_t index = 0; index != axis_size; ++index) { + for (int index = 0, e = axis_size; index != e; ++index) { scales[index] *= type.getScale(); } } From eabb1453f2efc62a648096c21b2ed993079d8c51 Mon Sep 17 00:00:00 2001 From: nammbash Date: Sun, 14 Jun 2020 11:47:44 -0700 Subject: [PATCH 0151/1390] Fix and Refactor NonAVX512 CPU platform --- tensorflow/core/graph/mkl_graph_util.h | 49 ++++++++++++++++++-------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/tensorflow/core/graph/mkl_graph_util.h b/tensorflow/core/graph/mkl_graph_util.h index cd09ac522d7..3c4c186b791 100644 --- a/tensorflow/core/graph/mkl_graph_util.h +++ b/tensorflow/core/graph/mkl_graph_util.h @@ -126,22 +126,19 @@ inline string GetMklEagerOpName(const string& name) { } #ifdef ENABLE_INTEL_MKL_BFLOAT16 -static inline bool CheckBfloat16Support(DataType T) { - static absl::once_flag cpu_bfloat16_warn_once_flag; - // Restrict bfloat16 ops to platforms with at least AVX512 support, fall back - // to Eigen implementation otherwise. - if (!(port::TestCPUFeature(port::CPUFeature::AVX512F)) && T == DT_BFLOAT16) { - absl::call_once(cpu_bfloat16_warn_once_flag, [] { - LOG(ERROR) - << "oneDNN BFloat16 support are only on platforms with AVX512. " - "Falling back to default implementation if present."; - }); - return false; - } - return true; +static inline bool IsBF16SupportedByOneDNNOnThisCPU() { + return port::TestCPUFeature(port::CPUFeature::AVX512F); } #endif +static inline void BF16UnsupportedWarning() { + static absl::once_flag cpu_bfloat16_warn_once_flag; + absl::call_once(cpu_bfloat16_warn_once_flag, [] { + LOG(ERROR) << "oneDNN BFloat16 support are only on platforms with AVX512. " + "Falling back to default implementation if present."; + }); +} + // Check whether opname with type T is registered as MKL operator // that can accept input tensors in MKL layout. // @@ -159,7 +156,18 @@ static inline bool IsMklLayoutDependentOp(const string& op_name, DataType T) { #ifdef ENABLE_INTEL_MKL_BFLOAT16 // Restrict regular ops to FLOAT and BFLOAT16 if (kernel.find(kMklLayoutDependentOpLabelPattern) != string::npos) { - return (T == DT_FLOAT || CheckBfloat16Support(T)); + if (T == DT_FLOAT) return true; + if (T == DT_BFLOAT16) { + if (IsBF16SupportedByOneDNNOnThisCPU()) { + return true; + } else { + // Restrict bfloat16 ops to platforms with at least AVX512 support, fall + // back to Eigen implementation otherwise. + BF16UnsupportedWarning(); + return false; + } + } + return false; } #else // Restrict regular ops to FLOAT @@ -216,7 +224,18 @@ static inline bool IsMklNameChangeOp(const string& op_name, DataType T) { isTypeAllowed = (T == DT_COMPLEX128 || T == DT_COMPLEX64 || T == DT_DOUBLE || T == DT_FLOAT); #ifdef ENABLE_INTEL_MKL_BFLOAT16 - isTypeAllowed = (isTypeAllowed || CheckBfloat16Support(T)); + if (!isTypeAllowed) { + if (T == DT_BFLOAT16) { + if (IsBF16SupportedByOneDNNOnThisCPU()) { + isTypeAllowed = true; + } else { + // Restrict bfloat16 ops to platforms with at least AVX512 support, + // fall back to Eigen implementation otherwise. + BF16UnsupportedWarning(); + isTypeAllowed = false; + } + } + } #endif return isTypeAllowed; } From 5394d892605fd158ad6a9c366e88e95b675a0227 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 14:48:54 -0400 Subject: [PATCH 0152/1390] Update dump_mlir_util.cc --- tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc index b5a6c922707..febf2bc096d 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc @@ -41,7 +41,7 @@ std::string MakeUniqueFilename(string name) { static NameCounts& instance = *new NameCounts; // Remove illegal characters from `name`. - for (size_t i = 0; i < name.size(); ++i) { + for (int i = 0, e = name.size(); i < e; ++i) { char ch = name[i]; if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' || ch == '\\') { From 1bccc9e1959fca38ce94a3cb1cfe7be1b6d4050c Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 14:52:29 -0400 Subject: [PATCH 0153/1390] Update chlo_ops.cc --- tensorflow/compiler/mlir/xla/ir/chlo_ops.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc index f5b895f0c76..3408f3ed0cc 100644 --- a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc +++ b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc @@ -49,7 +49,7 @@ static Type GetBroadcastType(Type x, Type y, Type element_type, if (shape_x.size() == shape_y.size()) { llvm::SmallVector out_shape(shape_x.size()); - for (size_t i = 0; i < shape_x.size(); i++) { + for (int i = 0, e = shape_x.size(); i < e; i++) { auto x_val = shape_x[i]; auto y_val = shape_y[i]; if (x_val == -1 || y_val == -1) { From 340053608bd1ac4168dc1be35e019cc1ac9d595a Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 14:58:56 -0400 Subject: [PATCH 0154/1390] Update hlo_ops.cc --- tensorflow/compiler/mlir/xla/ir/hlo_ops.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc index 569e45912a2..7f313b56925 100644 --- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc +++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc @@ -143,7 +143,7 @@ DenseIntElementsAttr BuildConvPaddingAttrs( int rank = padding_low.size(); SmallVector padding; - for (unsigned i = 0; i < static_cast(rank); ++i) { + for (unsigned i = 0, e = rank; i < e; ++i) { padding.push_back(GetPaddingValue(padding_attr, {i, 0}) + padding_low[i]); padding.push_back(GetPaddingValue(padding_attr, {i, 1}) + padding_high[i]); } @@ -853,7 +853,7 @@ static Attribute foldConcatenateHelper(ConcatenateOp* op, auto shape = type.getShape(); size_t top_size = 1; - for (size_t i = 0; i < axis; i++) { + for (int i = 0, e = axis; i < e; i++) { top_size = top_size * shape[i]; } @@ -1118,7 +1118,7 @@ static LogicalResult Verify(MapOp op) { // increasing. auto values = op.dimensions().getValues(); auto dimensions = std::vector{values.begin(), values.end()}; - for (int i = 0; static_cast(i) < dimensions.size(); ++i) { + for (int i = 0, e = dimensions.size(); i < e; ++i) { if (dimensions[i] != i) return op.emitOpError() << "requires monotonically increasing dimension " "numbers, but got: " From 9f535c3290cc3a8bdc503a4c48d3b0640b9f4798 Mon Sep 17 00:00:00 2001 From: tg-at-google Date: Sun, 14 Jun 2020 15:22:51 -0400 Subject: [PATCH 0155/1390] Update quantization_config.cc --- .../mlir/lite/quantization/quantization_config.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc index b299fa8f4c2..cdff93502f2 100644 --- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc +++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc @@ -48,9 +48,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_mins; if (!min_values.empty()) { std::vector node_mins_str = absl::StrSplit(min_values, ','); - for (int i : llvm::seq(node_mins_str.size())) { + for (const std::string&node_min : node_mins_str.size()) { double value; - if (!absl::SimpleAtod(node_mins_str[i], &value)) { + if (!absl::SimpleAtod(node_min, &value)) { return true; } node_mins.push_back(value); @@ -60,9 +60,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names, std::vector node_maxs; if (!max_values.empty()) { std::vector node_maxs_str = absl::StrSplit(max_values, ','); - for (int i : llvm::seq(node_maxs_str.size())) { + for (const std::string&node_max : node_maxs_str.size()) { double value; - if (!absl::SimpleAtod(node_maxs_str[i], &value)) { + if (!absl::SimpleAtod(node_max, &value)) { llvm::errs() << "Unexpected mins: " << node_maxs_str[i] << "\n"; return true; } From 6bb481c58b27b3fc99c0d8fc71b9d58f13e8b0ba Mon Sep 17 00:00:00 2001 From: Amedeo Cavallo Date: Mon, 15 Jun 2020 11:56:50 +0200 Subject: [PATCH 0156/1390] C linkage for stm32l4HAL target C linkage for output retargeting on stm32l4HAL target --- tensorflow/lite/micro/stm32f4HAL/debug_log.cc | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorflow/lite/micro/stm32f4HAL/debug_log.cc b/tensorflow/lite/micro/stm32f4HAL/debug_log.cc index 4be3b40e782..90dd7cfd787 100644 --- a/tensorflow/lite/micro/stm32f4HAL/debug_log.cc +++ b/tensorflow/lite/micro/stm32f4HAL/debug_log.cc @@ -22,6 +22,10 @@ limitations under the License. extern UART_HandleTypeDef DEBUG_UART_HANDLE; +#ifdef __cplusplus + extern "C" { +#endif + #ifdef __GNUC__ int __io_putchar(int ch) { HAL_UART_Transmit(&DEBUG_UART_HANDLE, (uint8_t *)&ch, 1, HAL_MAX_DELAY); @@ -36,4 +40,8 @@ int fputc(int ch, FILE *f) { } #endif /* __GNUC__ */ -extern "C" void DebugLog(const char *s) { fprintf(stderr, "%s", s); } +void DebugLog(const char *s) { fprintf(stderr, "%s", s); } + +#ifdef __cplusplus +} +#endif From f926d8c10efb07176ae559d0e098cdfdb4d03219 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 03:38:44 -0700 Subject: [PATCH 0157/1390] Introduces a new experimental package that: - Defines a schema for configuring delegates - Defines a C++ plugin mechanism using the schema, so that code can support configuring arbitrary delegates without a build-time dependency PiperOrigin-RevId: 316433209 Change-Id: Id538d0d3885bba9dd4094892915e9b6b736efd7d --- .../acceleration/configuration/BUILD | 160 ++++++++++++++ .../configuration/configuration.proto | 208 ++++++++++++++++++ .../configuration/delegate_registry.cc | 60 +++++ .../configuration/delegate_registry.h | 95 ++++++++ .../acceleration/configuration/gpu_plugin.cc | 62 ++++++ .../configuration/hexagon_plugin.cc | 73 ++++++ .../configuration/nnapi_plugin.cc | 93 ++++++++ .../configuration/nnapi_plugin_test.cc | 175 +++++++++++++++ .../configuration/proto_to_flatbuffer.cc | 58 +++++ .../configuration/proto_to_flatbuffer.h | 32 +++ 10 files changed, 1016 insertions(+) create mode 100644 tensorflow/lite/experimental/acceleration/configuration/BUILD create mode 100644 tensorflow/lite/experimental/acceleration/configuration/configuration.proto create mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h create mode 100644 tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc create mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h diff --git a/tensorflow/lite/experimental/acceleration/configuration/BUILD b/tensorflow/lite/experimental/acceleration/configuration/BUILD new file mode 100644 index 00000000000..38d28d5cc2e --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/BUILD @@ -0,0 +1,160 @@ +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== + +load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_java_library", "flatc_path") + +package( + default_visibility = [ + "//visibility:public", + ], + licenses = ["notice"], # Apache 2.0 +) + +genrule( + name = "configuration_schema", + srcs = ["configuration.proto"], + outs = ["configuration.fbs"], + # We rename the namespace since otherwise the proto classes and flatbuffer + # classes would have the same names. + cmd = """ + $(location {}) --proto -o $(@D) $(location :configuration.proto) + perl -p -i -e 's/tflite.proto/tflite/' $(@D)/configuration.fbs + """.format(flatc_path), + tools = [ + flatc_path, + ], +) + +genrule( + name = "configuration_fbs_contents_cc", + srcs = ["configuration.fbs"], + outs = ["configuration_fbs_contents-inl.h"], + cmd = """ + echo 'constexpr char configuration_fbs_contents[] = R"Delimiter(' > $(@) + cat < $(<) >> $(@) + echo ')Delimiter";' >> $(@) + """, +) + +proto_library( + name = "configuration_proto", + srcs = [ + "configuration.proto", + ], +) + +cc_proto_library( + name = "configuration_cc_proto", + deps = [":configuration_proto"], +) + +java_lite_proto_library( + name = "configuration_java_proto_lite", + deps = [":configuration_proto"], +) + +flatbuffer_cc_library( + name = "configuration_fbs", + srcs = [":configuration.fbs"], +) + +flatbuffer_java_library( + name = "configuration_fbs_java", + srcs = [":configuration.fbs"], +) + +cc_library( + name = "proto_to_flatbuffer", + srcs = [ + "configuration_fbs_contents-inl.h", + "proto_to_flatbuffer.cc", + ], + hdrs = ["proto_to_flatbuffer.h"], + deps = [ + ":configuration_cc_proto", + ":configuration_fbs", + "//tensorflow/core/platform:protobuf", + "//tensorflow/lite:minimal_logging", + "@flatbuffers", + ], +) + +cc_library( + name = "delegate_registry", + srcs = ["delegate_registry.cc"], + hdrs = ["delegate_registry.h"], + deps = [ + ":configuration_fbs", + "//tensorflow/lite/c:common", + "@com_google_absl//absl/synchronization", + ], +) + +cc_library( + name = "nnapi_plugin", + srcs = ["nnapi_plugin.cc"], + deps = [ + ":configuration_fbs", + ":delegate_registry", + "//tensorflow/lite/delegates/nnapi:nnapi_delegate", + "@com_google_absl//absl/memory", + ], + alwayslink = 1, # For registration to always run. +) + +cc_test( + name = "nnapi_plugin_test", + srcs = ["nnapi_plugin_test.cc"], + deps = [ + ":configuration_fbs", + ":delegate_registry", + ":nnapi_plugin", + "//tensorflow/lite:framework", + "//tensorflow/lite/c:common", + "//tensorflow/lite/delegates/nnapi:nnapi_delegate", + "//tensorflow/lite/delegates/nnapi:nnapi_delegate_mock_test", + "//tensorflow/lite/kernels:test_util", + "@com_google_googletest//:gtest_main", + "@flatbuffers", + ], +) + +cc_library( + name = "hexagon_plugin", + srcs = ["hexagon_plugin.cc"], + deps = [ + ":configuration_fbs", + ":delegate_registry", + "@com_google_absl//absl/memory", + ] + select({ + "//tensorflow:android": [ + "//tensorflow/lite/delegates/hexagon:hexagon_delegate", + ], + "//conditions:default": [], + }), + alwayslink = 1, # For registration to always run. +) + +cc_library( + name = "gpu_plugin", + srcs = ["gpu_plugin.cc"], + deps = [ + ":configuration_fbs", + ":delegate_registry", + "//tensorflow/lite/delegates/gpu:delegate", + "@com_google_absl//absl/memory", + ], + alwayslink = 1, # For registration to always run. +) diff --git a/tensorflow/lite/experimental/acceleration/configuration/configuration.proto b/tensorflow/lite/experimental/acceleration/configuration/configuration.proto new file mode 100644 index 00000000000..e1c49f02856 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/configuration.proto @@ -0,0 +1,208 @@ +// Copyright 2020 The TensorFlow Authors. All Rights Reserved. +// +// 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 +// +// http://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. + +// This schema defines how to configure TFLite for delegation. These +// definitions can be used in multiple ways: as output of a compatibility list, +// in benchmarking tools and to decouple delegate instantiation from code. +// +// The schema is work-in-progress, covering the most broadly used delegates and +// options. + +syntax = "proto2"; + +package tflite.proto; + +// ExecutionPreference is used to match accelerators against the preferences of +// the current application or usecase. Some of the values here can appear both +// in the compatibility list and as input, some only as input. +// +// These are separate from NNAPIExecutionPreference - the compatibility list +// design doesn't assume a one-to-one mapping between which usecases +// compatibility list entries have been developed for and what settings are used +// for NNAPI. +enum ExecutionPreference { + // Match any selected preference. Whitelist (semantically - value is same as + // on input). + ANY = 0; + // Match low latency preference. Both compatibility list and input. + LOW_LATENCY = 1; + // Math low power preference. Both compatibility list and input. + LOW_POWER = 2; + // Never accelerate. Can be used for input to compatibility list or for + // standalone Acceleration configuration. + FORCE_CPU = 3; +} + +// TFLite delegate to use. +enum Delegate { + NONE = 0; + NNAPI = 1; + GPU = 2; + HEXAGON = 3; + XNNPACK = 4; + // TODO(b/157893534): Support exposing edgetpu tflite delegate creation + // options. + EDGETPU = 5; +} + +enum NNAPIExecutionPreference { + // Undefined. + UNDEFINED = 0; + // Prefer executing in a way that minimizes battery drain. + NNAPI_LOW_POWER = 1; + // Prefer returning a single answer as fast as possible, even if this causes + // more power consumption. + NNAPI_FAST_SINGLE_ANSWER = 2; + // Prefer maximizing the throughput of successive frames, for example when + // processing successive frames coming from the camera. + NNAPI_SUSTAINED_SPEED = 3; +} + +// One possible acceleration configuration. +message ComputeSettings { + // Which preference to use this accelerator for. + optional ExecutionPreference preference = 1; + // How to configure TFLite + optional TFLiteSettings tflite_settings = 2; + // Identifiers to use for instrumentation and telemetry. + optional string model_namespace_for_statistics = 3; + optional string model_identifier_for_statistics = 4; +} + +// NNAPI delegate settings. +message NNAPISettings { + // Which instance (NNAPI accelerator) to use. One driver may provide several + // accelerators (though a driver may also hide several back-ends behind one + // name, at the choice of the driver vendor). + // Note that driver introspection is only available in Android Q and later. + optional string accelerator_name = 1; + + // NNAPI model compilation caching settings to be passed to + // tflite::StatefulNnApiDelegate + optional string cache_directory = 2; + optional string model_token = 3; + + // NNAPI execution preference to pass. See + // https://developer.android.com/ndk/reference/group/neural-networks.html + optional NNAPIExecutionPreference execution_preference = 4; + + // Number of instances to cache for the same model (for input size + // changes). This is mandatory for getting reasonable performance in that + // case. + optional int32 no_of_nnapi_instances_to_cache = 5; + + // Whether to automatically fall back to TFLite CPU path. + optional FallbackSettings fallback_settings = 6; + + // Whether to allow use of NNAPI CPU (nnapi-reference accelerator) on Android + // 10+ when an accelerator name is not specified. The NNAPI CPU typically + // performs less well than the TfLite built-in kernels; but allowing allows a + // model to be partially accelerated which may be a win. + optional bool allow_nnapi_cpu_on_android_10_plus = 7; +} + +// Which GPU backend to select. Default behaviour on Android is to try OpenCL +// and if it's not available fall back to OpenGL. +enum GPUBackend { + UNSET = 0; + OPENCL = 1; + OPENGL = 2; + // Not yet supported. + // VULKAN = 3; + // METAL = 4; +} + +// GPU Delegate settings. +// +// See +// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/gpu/delegate.h +message GPUSettings { + optional bool is_precision_loss_allowed = 1; + optional bool enable_quantized_inference = 2 [default = true]; + optional GPUBackend force_backend = 3; + // TODO(b/152019007): add remaining options. +} + +// Hexagon Delegate settings. +// +// See +// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/hexagon/hexagon_delegate.h +message HexagonSettings { + optional int32 debug_level = 1; + optional int32 powersave_level = 2; + optional bool print_graph_profile = 3; + optional bool print_graph_debug = 4; +} + +// XNNPack Delegate settings. +// +// See +// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h +message XNNPackSettings { + optional int32 num_threads = 1; +} + +message CPUSettings { + optional int32 num_threads = 1; +} + +// How to configure TFLite. +message TFLiteSettings { + // Which delegate to use. + optional Delegate delegate = 1; + + // How to configure the chosen delegate. + // (In principle we would like to use 'oneof', but flatc turns that into an + // nested anonymous table rather than a union. See + // https://github.com/google/flatbuffers/issues/4628). + optional NNAPISettings nnapi_settings = 2; + optional GPUSettings gpu_settings = 3; + optional HexagonSettings hexagon_settings = 4; + optional XNNPackSettings xnnpack_settings = 5; + + // How to configure CPU execution. + optional CPUSettings cpu_settings = 6; + + // Shared delegation settings. + optional int32 max_delegated_partitions = 7; +} + +// Whether to automatically fallback to TFLite CPU path on delegation errors. +// +// Typically fallback is enabled in production use but disabled in tests and +// benchmarks to ensure they test the intended path. +message FallbackSettings { + // Whether to allow automatically falling back to TfLite CPU path on + // compilation failure. Default is not allowing automatic fallback. + // + // This is useful in naive production usecases where the caller would prefer + // for the model to run even if it's not accelerated. More advanced users will + // implement fallback themselves; e.g., by using a different model on CPU. + // + // Note that compilation errors may occur either at initial + // ModifyGraphWithDelegate() time, or when calling AllocateTensors() after + // resizing. + optional bool allow_automatic_fallback_on_compilation_error = 7; + // Whether to allow automatically falling back to TfLite CPU path on + // execution error. Default is not allowing automatic fallback. + // + // Experimental, use with care (only when you have complete control over the + // client code). + // + // The caveat above for compilation error holds. Additionally, execution-time + // errors are harder to handle automatically as they require invalidating the + // TfLite interpreter which most client code has not been designed to deal + // with. + optional bool allow_automatic_fallback_on_execution_error = 8; +} diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc new file mode 100644 index 00000000000..b8d80342d5f --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc @@ -0,0 +1,60 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h" + +#include "absl/synchronization/mutex.h" + +namespace tflite { +namespace delegates { + +void DelegatePluginRegistry::RegisterImpl( + const std::string& name, + std::function< + std::unique_ptr(const TFLiteSettings&)> + creator_function) { + absl::MutexLock lock(&mutex_); + factories_[name] = creator_function; +} + +std::unique_ptr DelegatePluginRegistry::CreateImpl( + const std::string& name, const TFLiteSettings& settings) { + absl::MutexLock lock(&mutex_); + auto it = factories_.find(name); + if (it != factories_.end()) { + return it->second(settings); + } else { + return nullptr; + } +} + +DelegatePluginRegistry* DelegatePluginRegistry::GetSingleton() { + static auto* instance = new DelegatePluginRegistry(); + return instance; +} + +std::unique_ptr DelegatePluginRegistry::CreateByName( + const std::string& name, const TFLiteSettings& settings) { + auto* const instance = DelegatePluginRegistry::GetSingleton(); + return instance->CreateImpl(name, settings); +} + +DelegatePluginRegistry::Register::Register(const std::string& name, + CreatorFunction creator_function) { + auto* const instance = DelegatePluginRegistry::GetSingleton(); + instance->RegisterImpl(name, creator_function); +} + +} // namespace delegates +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h new file mode 100644 index 00000000000..c86759dcc3f --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h @@ -0,0 +1,95 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_ +#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_ + +#include +#include + +#include "absl/synchronization/mutex.h" +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" + +// Defines an interface for TFLite delegate plugins. +// +// The acceleration library aims to support all TFLite delegates based on +// configuration expressed as data (flatbuffers). However, consumers tend to +// care about size and also use a subset of delegates. Hence we don't want to +// statically build against all delegates. +// +// This interface allows plugins to handle specific delegates. +// +// Goal of this interface is not to abstract away all the differences between +// delegates. The goal is only to avoid static linking. +// +// Note to implementers: this interface may change if new delegates don't fit +// into the same design. +namespace tflite { +namespace delegates { + +// Same w/ Interpreter::TfLiteDelegatePtr to avoid pulling +// tensorflow/lite/interpreter.h dependency +using TfLiteDelegatePtr = + std::unique_ptr; + +class DelegatePluginInterface { + public: + virtual TfLiteDelegatePtr Create() = 0; + virtual int GetDelegateErrno(TfLiteDelegate* from_delegate) = 0; + virtual ~DelegatePluginInterface() = default; +}; + +// A stripped-down registry that allows delegate plugins to be created by name. +// +// Limitations: +// - Doesn't allow deregistration. +// - Doesn't check for duplication registration. +// +class DelegatePluginRegistry { + public: + typedef std::function( + const TFLiteSettings&)> + CreatorFunction; + // Returns a DelegatePluginInterface registered with `name` or nullptr if no + // matching plugin found. + // TFLiteSettings is per-plugin, so that the corresponding delegate options + // data lifetime is maintained. + static std::unique_ptr CreateByName( + const std::string& name, const TFLiteSettings& settings); + + // Struct to be statically allocated for registration. + struct Register { + Register(const std::string& name, CreatorFunction creator_function); + }; + + private: + void RegisterImpl(const std::string& name, CreatorFunction creator_function); + std::unique_ptr CreateImpl( + const std::string& name, const TFLiteSettings& settings); + static DelegatePluginRegistry* GetSingleton(); + std::unordered_map factories_; + absl::Mutex mutex_; +}; + +} // namespace delegates +} // namespace tflite + +#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f) \ + static auto* g_delegate_plugin_##name##_ = \ + new DelegatePluginRegistry::Register(#name, f); +#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(name, f) \ + TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f); + +#endif // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_ diff --git a/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc new file mode 100644 index 00000000000..25b8171c5ea --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc @@ -0,0 +1,62 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include + +#include "absl/memory/memory.h" +#include "tensorflow/lite/delegates/gpu/delegate.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" +#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h" + +namespace tflite { +namespace delegates { +class GpuPlugin : public DelegatePluginInterface { + public: + TfLiteDelegatePtr Create() override { + return TfLiteDelegatePtr(TfLiteGpuDelegateV2Create(&options_), + TfLiteGpuDelegateV2Delete); + } + int GetDelegateErrno(TfLiteDelegate* from_delegate) override { return 0; } + static std::unique_ptr New( + const TFLiteSettings& acceleration) { + return absl::make_unique(acceleration); + } + explicit GpuPlugin(const TFLiteSettings& tflite_settings) + : options_(TfLiteGpuDelegateOptionsV2Default()) { + const auto* gpu_settings = tflite_settings.gpu_settings(); + if (gpu_settings) { + options_.inference_priority1 = + gpu_settings->is_precision_loss_allowed() + ? TFLITE_GPU_INFERENCE_PRIORITY_MIN_LATENCY + : TFLITE_GPU_INFERENCE_PRIORITY_MAX_PRECISION; + if (gpu_settings->enable_quantized_inference()) { + options_.experimental_flags |= + TFLITE_GPU_EXPERIMENTAL_FLAGS_ENABLE_QUANT; + } + if (gpu_settings->force_backend() == GPUBackend_OPENCL) { + options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_CL_ONLY; + } else if (gpu_settings->force_backend() == GPUBackend_OPENGL) { + options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_GL_ONLY; + } + } + } + + private: + TfLiteGpuDelegateOptionsV2 options_; +}; + +TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(GpuPlugin, GpuPlugin::New); + +} // namespace delegates +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc new file mode 100644 index 00000000000..7f2674604b0 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc @@ -0,0 +1,73 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include + +#include "absl/memory/memory.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" +#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h" + +#if defined(__ARM_ARCH) +#include "tensorflow/lite/delegates/hexagon/hexagon_delegate.h" +#endif + +namespace tflite { +namespace delegates { +class HexagonPlugin : public DelegatePluginInterface { + public: + TfLiteDelegatePtr Create() override { +#if defined(__ARM_ARCH) + TfLiteHexagonInit(); + auto* delegate_ptr = TfLiteHexagonDelegateCreate(&options_); + TfLiteDelegatePtr delegate(delegate_ptr, [](TfLiteDelegate* delegate) { + TfLiteHexagonDelegateDelete(delegate); + TfLiteHexagonTearDown(); + }); + return delegate; +#else // !defined(__ARM_ARCH) + return TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {}); +#endif // defined(__ARM_ARCH) + } + int GetDelegateErrno(TfLiteDelegate* /* from_delegate */) override { + return 0; + } + static std::unique_ptr New( + const TFLiteSettings& tflite_settings) { + return absl::make_unique(tflite_settings); + } + explicit HexagonPlugin(const TFLiteSettings& tflite_settings) { + const HexagonSettings* settings = tflite_settings.hexagon_settings(); +#if defined(__ARM_ARCH) + options_ = TfLiteHexagonDelegateOptions({0}); + if (settings) { + options_.debug_level = settings->debug_level(); + options_.powersave_level = settings->powersave_level(); + options_.print_graph_profile = settings->print_graph_profile(); + options_.print_graph_debug = settings->print_graph_debug(); + } +#else + (void)settings; +#endif + } + + private: +#if defined(__ARM_ARCH) + TfLiteHexagonDelegateOptions options_; +#endif +}; + +TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(HexagonPlugin, HexagonPlugin::New); + +} // namespace delegates +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc new file mode 100644 index 00000000000..7301983a815 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc @@ -0,0 +1,93 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include + +#include "absl/memory/memory.h" +#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" +#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h" + +namespace tflite { +namespace delegates { + +inline tflite::StatefulNnApiDelegate::Options::ExecutionPreference +ConvertExecutionPrefence( + NNAPIExecutionPreference from_compatibility_preference) { + using TflitePreference = + tflite::StatefulNnApiDelegate::Options::ExecutionPreference; + switch (from_compatibility_preference) { + case NNAPIExecutionPreference_NNAPI_LOW_POWER: + return TflitePreference::kLowPower; + case NNAPIExecutionPreference_NNAPI_FAST_SINGLE_ANSWER: + return TflitePreference::kFastSingleAnswer; + case NNAPIExecutionPreference_NNAPI_SUSTAINED_SPEED: + return TflitePreference::kSustainedSpeed; + default: + return TflitePreference::kUndefined; + } +} + +class NnapiPlugin : public DelegatePluginInterface { + public: + TfLiteDelegatePtr Create() override { + auto nnapi_delegate = + absl::make_unique(options_); + return TfLiteDelegatePtr( + nnapi_delegate.release(), [](TfLiteDelegate* delegate) { + delete reinterpret_cast(delegate); + }); + } + int GetDelegateErrno(TfLiteDelegate* from_delegate) override { + auto nnapi_delegate = + reinterpret_cast(from_delegate); + return nnapi_delegate->GetNnApiErrno(); + } + static std::unique_ptr New( + const TFLiteSettings& tflite_settings) { + return absl::make_unique(tflite_settings); + } + explicit NnapiPlugin(const TFLiteSettings& tflite_settings) { + const NNAPISettings* nnapi_settings = tflite_settings.nnapi_settings(); + if (!nnapi_settings) return; + if (nnapi_settings->accelerator_name() && + nnapi_settings->accelerator_name()->Length() != 0) { + accelerator_ = nnapi_settings->accelerator_name()->str(); + options_.accelerator_name = accelerator_.c_str(); + } + if (nnapi_settings->cache_directory() && + nnapi_settings->cache_directory()->Length() != 0) { + cache_dir_ = nnapi_settings->cache_directory()->str(); + options_.cache_dir = cache_dir_.c_str(); + } + if (nnapi_settings->model_token() && + nnapi_settings->model_token()->Length() != 0) { + model_token_ = nnapi_settings->model_token()->str(); + options_.model_token = model_token_.c_str(); + } + options_.execution_preference = + ConvertExecutionPrefence(nnapi_settings->execution_preference()); + options_.disallow_nnapi_cpu = + !nnapi_settings->allow_nnapi_cpu_on_android_10_plus(); + } + + private: + std::string accelerator_, cache_dir_, model_token_; + tflite::StatefulNnApiDelegate::Options options_; +}; + +TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(NnapiPlugin, NnapiPlugin::New); + +} // namespace delegates +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc new file mode 100644 index 00000000000..4f9f5dd08c1 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc @@ -0,0 +1,175 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include + +#include +#include "flatbuffers/flatbuffers.h" // from @flatbuffers +#include "tensorflow/lite/c/common.h" +#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h" +#include "tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" +#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h" +#include "tensorflow/lite/interpreter.h" +#include "tensorflow/lite/kernels/test_util.h" + +// Tests for checking that the NNAPI Delegate plugin correctly handles all the +// options from the flatbuffer. +// +// Checking done at NNAPI call level, as that is where we have a mockable +// layer. +namespace tflite { +namespace { + +using delegate::nnapi::NnApiMock; + +class SingleAddOpModel : tflite::SingleOpModel { + public: + void Build() { + int input = AddInput({tflite::TensorType_FLOAT32, {1, 2, 2}}); + int constant = AddConstInput({tflite::TensorType_FLOAT32, {1, 2, 2}}, + {1.0f, 1.0f, 1.0f, 1.0f}); + AddOutput({tflite::TensorType_FLOAT32, {}}); + + SetBuiltinOp(tflite::BuiltinOperator_ADD, tflite::BuiltinOptions_AddOptions, + tflite::CreateAddOptions(builder_).Union()); + BuildInterpreter({GetShape(input), GetShape(constant)}); + } + + tflite::Interpreter* Interpreter() const { return interpreter_.get(); } +}; + +class NNAPIPluginTest : public ::testing::Test { + protected: + NNAPIPluginTest() : delegate_(nullptr, [](TfLiteDelegate*) {}) {} + void SetUp() override { + nnapi_ = const_cast(NnApiImplementation()); + nnapi_mock_ = absl::make_unique(nnapi_); + nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices = + [](const ANeuralNetworksModel* model, + const ANeuralNetworksDevice* const* devices, uint32_t numDevices, + bool* supportedOps) -> int { + supportedOps[0] = true; + return 0; + }; + model_.Build(); + } + template + void CheckExecutionPreference() { + // Note - this uses a template since the NNAPI functions are C function + // pointers rather than lambdas so can't capture variables. + nnapi_->ANeuralNetworksCompilation_setPreference = + [](ANeuralNetworksCompilation* compilation, int32_t preference) { + return preference - output; + }; + CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0, input)); + // Since delegation succeeds, the model becomes immutable and hence can't + // reuse it. + SingleAddOpModel model; + model.Build(); + EXPECT_EQ(model.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteOk) + << " given input: " << input << " expected output: " << output; + } + + void CreateDelegate(flatbuffers::Offset settings) { + settings_ = flatbuffers::GetTemporaryPointer( + fbb_, CreateTFLiteSettings(fbb_, tflite::Delegate_NNAPI, settings)); + + plugin_ = delegates::DelegatePluginRegistry::CreateByName("NnapiPlugin", + *settings_); + delegate_ = plugin_->Create(); + } + + NnApi* nnapi_; + std::unique_ptr nnapi_mock_; + SingleAddOpModel model_; + flatbuffers::FlatBufferBuilder fbb_; + const TFLiteSettings* settings_ = nullptr; + delegates::TfLiteDelegatePtr delegate_; + std::unique_ptr plugin_; +}; + +TEST_F(NNAPIPluginTest, PassesAcceleratorName) { + // Fails with non-existent "foo". + CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("foo"))); + EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteDelegateError); + + // Succeeds with "test-device" supported by the mock. + CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("test-device"))); + EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteOk); +} + +TEST_F(NNAPIPluginTest, PassesExecutionPreference) { + CheckExecutionPreference(); + CheckExecutionPreference(); + CheckExecutionPreference(); + CheckExecutionPreference(); +} + +TEST_F(NNAPIPluginTest, PassesCachingParameters) { + nnapi_->ANeuralNetworksCompilation_setCaching = + [](ANeuralNetworksCompilation* compilation, const char* cacheDir, + const uint8_t* token) -> int { + if (std::string(cacheDir) != "d") return 1; + // Token is hashed with other bits, just check that it's not empty. + if (std::string(reinterpret_cast(token)).empty()) return 2; + return 0; + }; + CreateDelegate(CreateNNAPISettings(fbb_, 0, fbb_.CreateString("d"), + fbb_.CreateString("t"))); + EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteOk); +} + +TEST_F(NNAPIPluginTest, PassesFalseNNAPICpuFlag) { + CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0, + NNAPIExecutionPreference_UNDEFINED, 0, 0, + /* allow CPU */ false)); + nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices = + [](const ANeuralNetworksModel* model, + const ANeuralNetworksDevice* const* devices, uint32_t numDevices, + bool* supportedOps) -> int { + supportedOps[0] = true; + // Since no CPU, should only pass one device. + return numDevices - 1; + }; + EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteOk); +} + +TEST_F(NNAPIPluginTest, PassesTrueNNAPICpuFlag) { + CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0, + NNAPIExecutionPreference_UNDEFINED, 0, 0, + /* allow CPU */ true)); + nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices = + [](const ANeuralNetworksModel* model, + const ANeuralNetworksDevice* const* devices, uint32_t numDevices, + bool* supportedOps) -> int { + supportedOps[0] = true; + // With CPU allowed, should pass two devices. + return numDevices - 2; + }; + EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()), + kTfLiteOk); +} + +} // namespace +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc new file mode 100644 index 00000000000..709bb70ca70 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc @@ -0,0 +1,58 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#include "tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h" + +#include + +#include "flatbuffers/idl.h" // from @flatbuffers +#include "flatbuffers/util.h" // from @flatbuffers +#include "tensorflow/core/platform/protobuf.h" +#include "tensorflow/lite/minimal_logging.h" + +namespace tflite { + +namespace { +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_fbs_contents-inl.h" +} + +const ComputeSettings* ConvertFromProto( + flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings) { + std::string json; + tensorflow::protobuf::util::JsonPrintOptions options; + options.preserve_proto_field_names = true; + options.always_print_primitive_fields = true; // For catching problems. + auto status = tensorflow::protobuf::util::MessageToJsonString(proto_settings, + &json, options); + if (!status.ok()) { + TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to convert to Json: %s", + status.ToString().c_str()); + return nullptr; + } + if (!parser->Parse(configuration_fbs_contents)) { + TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse schema: %s", + parser->error_.c_str()); + return nullptr; + } + parser->SetRootType("tflite.ComputeSettings"); + if (!parser->Parse(json.c_str())) { + TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse json: %s", + parser->error_.c_str()); + return nullptr; + } + return flatbuffers::GetRoot( + parser->builder_.GetBufferPointer()); +} + +} // namespace tflite diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h new file mode 100644 index 00000000000..3b69e8465a5 --- /dev/null +++ b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h @@ -0,0 +1,32 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ +#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_ +#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_ + +#include "flatbuffers/idl.h" // from @flatbuffers +#include "tensorflow/lite/experimental/acceleration/configuration/configuration.pb.h" +#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h" + +namespace tflite { + +// Converts the protobuf version ComputeSettings to the flatbuffer version, via +// json. The parser is used for state - the returned pointer is valid only as +// long as the parser is kept alive and unmutated. +const ComputeSettings* ConvertFromProto( + flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings); + +} // namespace tflite + +#endif // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_ From af52bd27dce7a568ff111b87f922582f6dbeff7e Mon Sep 17 00:00:00 2001 From: Jacques Pienaar Date: Mon, 15 Jun 2020 06:25:37 -0700 Subject: [PATCH 0158/1390] Update ::OperandAdaptor to ::Adaptor The Adaptor can refer to attributes too, so update naming. Follow up from https://reviews.llvm.org/D81741 PiperOrigin-RevId: 316451631 Change-Id: If2882e8ef2e75f70ae2c9193b4e8286ab3b0326f --- .../compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc | 4 ++-- .../compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc index 8846d7918c7..99d2c08aa98 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc @@ -35,7 +35,7 @@ struct StaticMemRefCastOpConverter auto loc = op->getLoc(); auto cast_op = cast(op); - StaticMemRefCastOp::OperandAdaptor operands_adaptor(operands); + StaticMemRefCastOp::Adaptor operands_adaptor(operands); MemRefDescriptor sourceMemRef(operands_adaptor.operand()); MemRefType targetMemRefType = @@ -86,7 +86,7 @@ struct DynamicMemRefCastOpConverter auto loc = op->getLoc(); auto cast_op = cast(op); - DynamicMemRefCastOp::OperandAdaptor operands_adaptor(operands); + DynamicMemRefCastOp::Adaptor operands_adaptor(operands); MemRefDescriptor sourceMemRef(operands_adaptor.operand()); MemRefType targetMemRefType = diff --git a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc index 0e4842537ef..ad78a01100b 100644 --- a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc +++ b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc @@ -415,7 +415,7 @@ class LhloBroadcastInDimConverter LogicalResult matchAndRewrite( xla_lhlo::BroadcastInDimOp op, ArrayRef args, ConversionPatternRewriter& rewriter) const final { - xla_lhlo::BroadcastInDimOp::OperandAdaptor operand_adaptor(args); + xla_lhlo::BroadcastInDimOp::Adaptor operand_adaptor(args); auto result_type = operand_adaptor.output().getType().cast(); auto result_shape = result_type.getShape(); @@ -476,7 +476,7 @@ class LhloBroadcastInDimConverter std::pair> InsertReshapeIfNecessary( xla_lhlo::BroadcastInDimOp op, ArrayRef args, ConversionPatternRewriter& rewriter) const { - xla_lhlo::BroadcastInDimOp::OperandAdaptor operand_adaptor(args); + xla_lhlo::BroadcastInDimOp::Adaptor operand_adaptor(args); Value operand = operand_adaptor.operand(); auto operand_type = operand_adaptor.operand().getType().cast(); auto operand_shape = operand_type.getShape(); From fb173c22bd9e35df96828ee2a6273506902d456d Mon Sep 17 00:00:00 2001 From: Guillaume Klein Date: Mon, 15 Jun 2020 15:51:55 +0200 Subject: [PATCH 0159/1390] Fix docstring format of tf.executing_eagerly --- tensorflow/python/eager/context.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py index b01f0795c72..4560c3b634e 100644 --- a/tensorflow/python/eager/context.py +++ b/tensorflow/python/eager/context.py @@ -1877,9 +1877,8 @@ def executing_eagerly(): True False - Inside `tf.function` after + Inside `tf.function` after `tf.config.run_functions_eagerly(True)` is called: - `tf.config.run_functions_eagerly(True)` is called: >>> tf.config.run_functions_eagerly(True) >>> @tf.function ... def fn(): From 08311077edd04cb9fb90939bfb0e5caf3059ccb2 Mon Sep 17 00:00:00 2001 From: Guillaume Klein Date: Mon, 15 Jun 2020 15:57:20 +0200 Subject: [PATCH 0160/1390] Copy reverse_sequence docstring to _v2 to remove deprecation notices --- tensorflow/python/ops/array_ops.py | 41 +++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index c77977bf7d2..6ca0eb975ac 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -4535,6 +4535,46 @@ def reverse_sequence_v2(input, seq_axis=None, batch_axis=None, name=None): + """Reverses variable length slices. + + This op first slices `input` along the dimension `batch_axis`, and for + each slice `i`, reverses the first `seq_lengths[i]` elements along the + dimension `seq_axis`. + + The elements of `seq_lengths` must obey `seq_lengths[i] <= + input.dims[seq_axis]`, and `seq_lengths` must be a vector of length + `input.dims[batch_axis]`. + + The output slice `i` along dimension `batch_axis` is then given by + input slice `i`, with the first `seq_lengths[i]` slices along + dimension `seq_axis` reversed. + + Example usage: + + >>> seq_lengths = [7, 2, 3, 5] + >>> input = [[1, 2, 3, 4, 5, 0, 0, 0], [1, 2, 0, 0, 0, 0, 0, 0], + ... [1, 2, 3, 4, 0, 0, 0, 0], [1, 2, 3, 4, 5, 6, 7, 8]] + >>> output = tf.reverse_sequence(input, seq_lengths, seq_axis=1, batch_axis=0) + >>> output + + + Args: + input: A `Tensor`. The input to reverse. + seq_lengths: A `Tensor`. Must be one of the following types: `int32`, + `int64`. 1-D with length `input.dims(batch_axis)` and `max(seq_lengths) <= + input.dims(seq_axis)` + seq_axis: An `int`. The dimension which is partially reversed. + batch_axis: An optional `int`. Defaults to `0`. The dimension along which + reversal is performed. + name: A name for the operation (optional). + + Returns: + A Tensor. Has the same type as input. + """ return gen_array_ops.reverse_sequence( input=input, seq_lengths=seq_lengths, @@ -4542,7 +4582,6 @@ def reverse_sequence_v2(input, batch_dim=batch_axis, name=name) -reverse_sequence_v2.__doc__ = reverse_sequence.__doc__ # pylint: enable=redefined-builtin From eabae7b8e94da529e9da72bbacd534904c5e4b79 Mon Sep 17 00:00:00 2001 From: Kibeom Kim Date: Mon, 15 Jun 2020 07:00:00 -0700 Subject: [PATCH 0161/1390] Set async as default for TFRT runtime PiperOrigin-RevId: 316456716 Change-Id: Ifdffb129c4609108285892a7ead5758e67802641 --- tensorflow/python/eager/context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py index b01f0795c72..68234985d15 100644 --- a/tensorflow/python/eager/context.py +++ b/tensorflow/python/eager/context.py @@ -427,7 +427,7 @@ class Context(object): raise ValueError( "execution_mode should be None/SYNC/ASYNC. Got %s" % execution_mode) if execution_mode is None: - execution_mode = SYNC + execution_mode = ASYNC if is_tfrt_enabled() else SYNC self._default_is_async = execution_mode == ASYNC self._lazy_remote_inputs_copy = None self._use_tfrt = is_tfrt_enabled() From 83b4360a3fd7955bf807e12c05734c388b67c2f2 Mon Sep 17 00:00:00 2001 From: Tres Popp Date: Mon, 15 Jun 2020 07:19:59 -0700 Subject: [PATCH 0162/1390] Add shape constraints to CHLO->HLO lowering. PiperOrigin-RevId: 316459663 Change-Id: Ifff45b67a039c5a8e7cf8fa1bedf187c33900091 --- .../chlo_legalize_to_hlo_broadcasts.mlir | 52 ++++++---- .../tests/legalize-tf-binary-elementwise.mlir | 99 +++++++++---------- .../xla/transforms/chlo_legalize_to_hlo.cc | 19 +++- 3 files changed, 95 insertions(+), 75 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir b/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir index 107a668c0a7..65285021fd4 100644 --- a/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir +++ b/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir @@ -1,4 +1,4 @@ -// RUN: xla-opt -test-xla-chlo-legalize-to-hlo -split-input-file -verify-diagnostics %s -o - | FileCheck %s +// RUN: xla-opt -test-xla-chlo-legalize-to-hlo -cse -split-input-file -verify-diagnostics %s -o - | FileCheck %s // Check the non-broadcast case for each registered op, then just check a // representative op for detailed broadcast semantics. @@ -14,14 +14,18 @@ func @addWithoutBroadcast(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor< // CHECK-SAME: %[[ARG0:.+]]: tensor // CHECK-SAME: %[[ARG1:.+]]: tensor func @dynamicBroadcast(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] - // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) - // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] - // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT:.+]] = xla_hlo.add %[[ARG0_B]], %[[ARG1_B]] - // CHECK: return %[[RESULT]] : tensor + // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] + // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] + // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] + // CHECK-NEXT: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] + // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] + // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK-NEXT: %[[RESULT:.+]] = xla_hlo.add %[[ARG0_B]], %[[ARG1_B]] + // CHECK-NEXT: shape.assuming_yield %[[RESULT]] + // CHECK-NEXT: } + // CHECK-NEXT: return %[[FINAL_RESULT]] : tensor %0 = xla_chlo.broadcast_add %arg0, %arg1 : (tensor, tensor) -> tensor return %0 : tensor } @@ -31,14 +35,18 @@ func @dynamicBroadcast(%arg0: tensor, %arg1: tensor) -> tensor // CHECK-SAME: %[[ARG1:.+]]: tensor func @dynamicBroadcastComplex(%arg0: tensor, %arg1: tensor) -> tensor> { - // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] - // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) - // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] - // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK-DAG: %[[RESULT:.+]] = "xla_hlo.complex"(%[[ARG0_B]], %[[ARG1_B]]) : (tensor, tensor) -> tensor> - // CHECK: return %[[RESULT]] : tensor> + // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] + // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] + // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] + // CHECK-NEXT: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] + // CHECK-NEXT: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] + // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor + // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor + // CHECK-NEXT: %[[RESULT:.+]] = "xla_hlo.complex"(%[[ARG0_B]], %[[ARG1_B]]) : (tensor, tensor) -> tensor> + // CHECK-NEXT: shape.assuming_yield %[[RESULT]] + // CHECK-NEXT: } + // CHECK-NEXT: return %[[FINAL_RESULT]] : tensor> %0 = xla_chlo.broadcast_complex %arg0, %arg1 : (tensor, tensor) -> tensor> return %0 : tensor> } @@ -50,12 +58,16 @@ func @dynamicBroadcastComplex(%arg0: tensor, %arg1: tensor) -> t func @dynamicBroadcastCompare(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] + // CHECK: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] + // CHECK: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK-DAG: %[[RESULT:.+]] = "xla_hlo.compare"(%[[ARG0_B]], %[[ARG1_B]]) {comparison_direction = "EQ"} : (tensor, tensor) -> tensor - // CHECK: return %[[RESULT]] : tensor + // CHECK: %[[RESULT:.+]] = "xla_hlo.compare"(%[[ARG0_B]], %[[ARG1_B]]) {comparison_direction = "EQ"} : (tensor, tensor) -> tensor + // CHECK: shape.assuming_yield %[[RESULT]] + // CHECK-NEXT: } + // CHECK: return %[[FINAL_RESULT]] : tensor %0 = xla_chlo.broadcast_compare %arg0, %arg1 {comparison_direction = "EQ"} : (tensor, tensor) -> tensor return %0 : tensor } diff --git a/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir b/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir index 2153258993a..3d270a52f48 100644 --- a/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir +++ b/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir @@ -1,7 +1,7 @@ // Note that binary elementwise tests are run with chlo legalization enabled // (unlike the rest), since this is the primary use case for such ops and // verification of shapes and broadcasts is desired. -// RUN: tf-opt "-xla-legalize-tf=allow-partial-conversion legalize-chlo=true" %s | FileCheck %s +// RUN: tf-opt "-xla-legalize-tf=allow-partial-conversion legalize-chlo=true" -canonicalize %s | FileCheck %s //===----------------------------------------------------------------------===// // Binary op legalizations. @@ -24,13 +24,8 @@ func @add(%arg0: tensor<2xi32>) -> tensor<2xi32> { // patterns unambiguous and more interesting (once broadcastable trait is // fixed upstream). func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi32> { - // CHECK: %[[UNUSED_LHS_SHAPE:.+]] = shape.const_shape [1] - // CHECK: %[[UNUSED_RHS_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK: xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] + // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-NEXT: xla_hlo.add %[[LHS_BCAST]], %arg1 %0 = "tf.Add"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32> return %0: tensor<1x2xi32> } @@ -39,26 +34,26 @@ func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2x // TODO(laurenzo): Change this to a (4x1x1 + 1x4x4x4) shaped add once upstream // broadcastable bug is fixed (helps make the CHECK matching unambiguous) func @broadcast_multi_dim_add(%arg0: tensor<4x1x1xi32>, %arg1: tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> { - // CHECK: %[[UNUSED_LHS_SHAPE:.+]] = shape.const_shape [4, 1, 1] - // CHECK: %[[UNUSED_RHS_SHAPE:.+]] = shape.const_shape [4, 4, 4, 4] - // CHECK: %[[RESULT_SHAPE:.+]] = shape.const_shape [4, 4, 4, 4] - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>} - // CHECK: xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] + // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>} + // CHECK-NEXT: xla_hlo.add %[[LHS_BCAST]], %arg1 %0 = "tf.Add"(%arg0, %arg1) : (tensor<4x1x1xi32>, tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> return %0: tensor<4x4x4x4xi32> } // CHECK-LABEL: func @add_dynamic func @add_dynamic(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 - // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK: xla_hlo.add %4, %5 : tensor + // CHECK-DAG: %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1 + // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]] + // CHECK-NEXT: shape.assuming %[[WITNESS:.+]] + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 + // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) + // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-NEXT: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK-NEXT: %[[RESULT:.+]] = xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] : tensor + // CHECK-NEXT: shape.assuming_yield %[[RESULT]] %0 = "tf.Add"(%arg0, %arg1) : (tensor, tensor) -> tensor return %0: tensor } @@ -80,21 +75,21 @@ func @shift_left(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> { // CHECK-LABEL: func @div_unranked func @div_unranked(%arg0: tensor<*xi32>, %arg1: tensor) -> tensor { - // CHECK: tf.Div + // CHECK-NEXT: tf.Div %0 = "tf.Div"(%arg0, %arg1) : (tensor<*xi32>, tensor) -> tensor return %0: tensor } // CHECK-LABEL: func @maximum func @maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> { - // CHECK: xla_hlo.maximum %arg0, %arg1 : tensor<4xf32> + // CHECK-NEXT: xla_hlo.maximum %arg0, %arg1 : tensor<4xf32> %0 = "tf.Maximum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> return %0 : tensor<4xf32> } // CHECK-LABEL: func @minimum func @minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> { - // CHECK: xla_hlo.minimum %arg0, %arg1 : tensor<4xf32> + // CHECK-NEXT: xla_hlo.minimum %arg0, %arg1 : tensor<4xf32> %0 = "tf.Minimum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> return %0 : tensor<4xf32> } @@ -200,26 +195,25 @@ func @equal(%arg0: tensor<2xi32>) -> tensor<2xi1> { // CHECK-LABEL: func @equal_dynamic func @equal_dynamic(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1] - // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1] + // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]] + // CHECK-NEXT: shape.assuming %[[WITNESS]] -> (tensor) { + // CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0 + // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE1]], %[[RHS_SHAPE]]) + // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-NEXT: %[[RESULT:.+]] = "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} + // CHECK-NEXT: shape.assuming_yield %[[RESULT]] %0 = "tf.Equal"(%arg0, %arg1) : (tensor, tensor<1xi32>) -> tensor return %0: tensor } // CHECK-LABEL: func @equal_broadcast func @equal_broadcast(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.const_shape [1] - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK-DAG: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %arg1) {comparison_direction = "EQ"} %0 = "tf.Equal"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> return %0: tensor<1x2xi1> } @@ -281,26 +275,25 @@ func @greater(%arg0: tensor<2xi32>) -> tensor<2xi1> { // CHECK-LABEL: func @broadcast_greater func @broadcast_greater(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.const_shape [1] - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK-DAG: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} + // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %arg1) {comparison_direction = "GT"} %0 = "tf.Greater"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> return %0: tensor<1x2xi1> } // CHECK-LABEL: func @greater_dynamic func @greater_dynamic(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 - // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) - // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 + // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]] + // CHECK-NEXT: shape.assuming %[[WITNESS]] + // CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE1:.+]] = shape.shape_of %arg1 + // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE1]], %[[RHS_SHAPE1]]) + // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} %0 = "tf.Greater"(%arg0, %arg1) : (tensor, tensor) -> tensor return %0: tensor } diff --git a/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc b/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc index e5a79616d5b..97afa9617c4 100644 --- a/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc +++ b/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc @@ -112,6 +112,19 @@ struct ConvertRankedDynamicBroadcastBinaryOp // Compute result shape. auto loc = op.getLoc(); + + // Insert a constraint on the shapes being broadcastable and insert all + // future code into an assuming block reliant on the constraint. + Value lhs_shape = rewriter.create(loc, lhs); + Value rhs_shape = rewriter.create(loc, rhs); + auto broadcastable_cstr = + rewriter.create(loc, lhs_shape, rhs_shape); + auto assuming_op = rewriter.create( + loc, ArrayRef{result_type}, broadcastable_cstr.result()); + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.createBlock(&assuming_op.doRegion()); + int64_t result_rank = std::max(lhs_type.getRank(), rhs_type.getRank()); Value result_extents = xla::ComputeBinaryElementwiseBroadcastingResultExtents(loc, lhs, rhs, @@ -140,8 +153,10 @@ struct ConvertRankedDynamicBroadcastBinaryOp rewriter.getI64TensorAttr(rhs_broadcast_dimensions)); // And generate the final non-broadcasted binary op. - rewriter.replaceOp(op, {Adaptor::CreateOp(op, result_type, broadcasted_lhs, - broadcasted_rhs, rewriter)}); + Value final_result = Adaptor::CreateOp(op, result_type, broadcasted_lhs, + broadcasted_rhs, rewriter); + rewriter.create(loc, final_result); + rewriter.replaceOp(op, {assuming_op.getResult(0)}); return success(); } }; From 1a0909a9f44cd60419aa9ffa0c5710957c83901d Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 07:50:42 -0700 Subject: [PATCH 0163/1390] Add shape constraints to CHLO->HLO lowering. PiperOrigin-RevId: 316464141 Change-Id: If31be3a8f1644335897feb7c693eed02ce52f029 --- .../chlo_legalize_to_hlo_broadcasts.mlir | 52 ++++------ .../tests/legalize-tf-binary-elementwise.mlir | 99 ++++++++++--------- .../xla/transforms/chlo_legalize_to_hlo.cc | 19 +--- 3 files changed, 75 insertions(+), 95 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir b/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir index 65285021fd4..107a668c0a7 100644 --- a/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir +++ b/tensorflow/compiler/mlir/xla/tests/chlo_legalize_to_hlo_broadcasts.mlir @@ -1,4 +1,4 @@ -// RUN: xla-opt -test-xla-chlo-legalize-to-hlo -cse -split-input-file -verify-diagnostics %s -o - | FileCheck %s +// RUN: xla-opt -test-xla-chlo-legalize-to-hlo -split-input-file -verify-diagnostics %s -o - | FileCheck %s // Check the non-broadcast case for each registered op, then just check a // representative op for detailed broadcast semantics. @@ -14,18 +14,14 @@ func @addWithoutBroadcast(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor< // CHECK-SAME: %[[ARG0:.+]]: tensor // CHECK-SAME: %[[ARG1:.+]]: tensor func @dynamicBroadcast(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] - // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] - // CHECK-NEXT: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] - // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) - // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] - // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK-NEXT: %[[RESULT:.+]] = xla_hlo.add %[[ARG0_B]], %[[ARG1_B]] - // CHECK-NEXT: shape.assuming_yield %[[RESULT]] - // CHECK-NEXT: } - // CHECK-NEXT: return %[[FINAL_RESULT]] : tensor + // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] + // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] + // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] + // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK-DAG: %[[RESULT:.+]] = xla_hlo.add %[[ARG0_B]], %[[ARG1_B]] + // CHECK: return %[[RESULT]] : tensor %0 = xla_chlo.broadcast_add %arg0, %arg1 : (tensor, tensor) -> tensor return %0 : tensor } @@ -35,18 +31,14 @@ func @dynamicBroadcast(%arg0: tensor, %arg1: tensor) -> tensor // CHECK-SAME: %[[ARG1:.+]]: tensor func @dynamicBroadcastComplex(%arg0: tensor, %arg1: tensor) -> tensor> { - // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] - // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] - // CHECK-NEXT: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] - // CHECK-NEXT: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) - // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] - // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK-NEXT: %[[RESULT:.+]] = "xla_hlo.complex"(%[[ARG0_B]], %[[ARG1_B]]) : (tensor, tensor) -> tensor> - // CHECK-NEXT: shape.assuming_yield %[[RESULT]] - // CHECK-NEXT: } - // CHECK-NEXT: return %[[FINAL_RESULT]] : tensor> + // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] + // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] + // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] + // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor + // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor + // CHECK-DAG: %[[RESULT:.+]] = "xla_hlo.complex"(%[[ARG0_B]], %[[ARG1_B]]) : (tensor, tensor) -> tensor> + // CHECK: return %[[RESULT]] : tensor> %0 = xla_chlo.broadcast_complex %arg0, %arg1 : (tensor, tensor) -> tensor> return %0 : tensor> } @@ -58,16 +50,12 @@ func @dynamicBroadcastComplex(%arg0: tensor, %arg1: tensor) -> t func @dynamicBroadcastCompare(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK-DAG: %[[ARG0_S:.+]] = shape.shape_of %[[ARG0]] // CHECK-DAG: %[[ARG1_S:.+]] = shape.shape_of %[[ARG1]] - // CHECK: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[ARG0_S]], %[[ARG1_S]] - // CHECK: %[[FINAL_RESULT:.+]] = shape.assuming %[[WITNESS]] - // CHECK: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) + // CHECK-DAG: %[[RESULT_S:.+]] = "shape.broadcast"(%[[ARG0_S]], %[[ARG1_S]]) // CHECK: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_S]] // CHECK-DAG: %[[ARG0_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor, tensor<2xindex>) -> tensor // CHECK-DAG: %[[ARG1_B:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor, tensor<2xindex>) -> tensor - // CHECK: %[[RESULT:.+]] = "xla_hlo.compare"(%[[ARG0_B]], %[[ARG1_B]]) {comparison_direction = "EQ"} : (tensor, tensor) -> tensor - // CHECK: shape.assuming_yield %[[RESULT]] - // CHECK-NEXT: } - // CHECK: return %[[FINAL_RESULT]] : tensor + // CHECK-DAG: %[[RESULT:.+]] = "xla_hlo.compare"(%[[ARG0_B]], %[[ARG1_B]]) {comparison_direction = "EQ"} : (tensor, tensor) -> tensor + // CHECK: return %[[RESULT]] : tensor %0 = xla_chlo.broadcast_compare %arg0, %arg1 {comparison_direction = "EQ"} : (tensor, tensor) -> tensor return %0 : tensor } diff --git a/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir b/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir index 3d270a52f48..2153258993a 100644 --- a/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir +++ b/tensorflow/compiler/mlir/xla/tests/legalize-tf-binary-elementwise.mlir @@ -1,7 +1,7 @@ // Note that binary elementwise tests are run with chlo legalization enabled // (unlike the rest), since this is the primary use case for such ops and // verification of shapes and broadcasts is desired. -// RUN: tf-opt "-xla-legalize-tf=allow-partial-conversion legalize-chlo=true" -canonicalize %s | FileCheck %s +// RUN: tf-opt "-xla-legalize-tf=allow-partial-conversion legalize-chlo=true" %s | FileCheck %s //===----------------------------------------------------------------------===// // Binary op legalizations. @@ -24,8 +24,13 @@ func @add(%arg0: tensor<2xi32>) -> tensor<2xi32> { // patterns unambiguous and more interesting (once broadcastable trait is // fixed upstream). func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi32> { - // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-NEXT: xla_hlo.add %[[LHS_BCAST]], %arg1 + // CHECK: %[[UNUSED_LHS_SHAPE:.+]] = shape.const_shape [1] + // CHECK: %[[UNUSED_RHS_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK: xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] %0 = "tf.Add"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi32> return %0: tensor<1x2xi32> } @@ -34,26 +39,26 @@ func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2x // TODO(laurenzo): Change this to a (4x1x1 + 1x4x4x4) shaped add once upstream // broadcastable bug is fixed (helps make the CHECK matching unambiguous) func @broadcast_multi_dim_add(%arg0: tensor<4x1x1xi32>, %arg1: tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> { - // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>} - // CHECK-NEXT: xla_hlo.add %[[LHS_BCAST]], %arg1 + // CHECK: %[[UNUSED_LHS_SHAPE:.+]] = shape.const_shape [4, 1, 1] + // CHECK: %[[UNUSED_RHS_SHAPE:.+]] = shape.const_shape [4, 4, 4, 4] + // CHECK: %[[RESULT_SHAPE:.+]] = shape.const_shape [4, 4, 4, 4] + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>} + // CHECK: xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] %0 = "tf.Add"(%arg0, %arg1) : (tensor<4x1x1xi32>, tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> return %0: tensor<4x4x4x4xi32> } // CHECK-LABEL: func @add_dynamic func @add_dynamic(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1 - // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]] - // CHECK-NEXT: shape.assuming %[[WITNESS:.+]] - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 - // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) - // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-NEXT: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} - // CHECK-NEXT: %[[RESULT:.+]] = xla_hlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] : tensor - // CHECK-NEXT: shape.assuming_yield %[[RESULT]] + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 + // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK: xla_hlo.add %4, %5 : tensor %0 = "tf.Add"(%arg0, %arg1) : (tensor, tensor) -> tensor return %0: tensor } @@ -75,21 +80,21 @@ func @shift_left(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> { // CHECK-LABEL: func @div_unranked func @div_unranked(%arg0: tensor<*xi32>, %arg1: tensor) -> tensor { - // CHECK-NEXT: tf.Div + // CHECK: tf.Div %0 = "tf.Div"(%arg0, %arg1) : (tensor<*xi32>, tensor) -> tensor return %0: tensor } // CHECK-LABEL: func @maximum func @maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> { - // CHECK-NEXT: xla_hlo.maximum %arg0, %arg1 : tensor<4xf32> + // CHECK: xla_hlo.maximum %arg0, %arg1 : tensor<4xf32> %0 = "tf.Maximum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> return %0 : tensor<4xf32> } // CHECK-LABEL: func @minimum func @minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> { - // CHECK-NEXT: xla_hlo.minimum %arg0, %arg1 : tensor<4xf32> + // CHECK: xla_hlo.minimum %arg0, %arg1 : tensor<4xf32> %0 = "tf.Minimum"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> return %0 : tensor<4xf32> } @@ -195,25 +200,26 @@ func @equal(%arg0: tensor<2xi32>) -> tensor<2xi1> { // CHECK-LABEL: func @equal_dynamic func @equal_dynamic(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1] - // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]] - // CHECK-NEXT: shape.assuming %[[WITNESS]] -> (tensor) { - // CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0 - // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE1]], %[[RHS_SHAPE]]) - // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-NEXT: %[[RESULT:.+]] = "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} - // CHECK-NEXT: shape.assuming_yield %[[RESULT]] + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1] + // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} %0 = "tf.Equal"(%arg0, %arg1) : (tensor, tensor<1xi32>) -> tensor return %0: tensor } // CHECK-LABEL: func @equal_broadcast func @equal_broadcast(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %arg1) {comparison_direction = "EQ"} + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.const_shape [1] + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK-DAG: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "EQ"} %0 = "tf.Equal"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> return %0: tensor<1x2xi1> } @@ -275,25 +281,26 @@ func @greater(%arg0: tensor<2xi32>) -> tensor<2xi1> { // CHECK-LABEL: func @broadcast_greater func @broadcast_greater(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { - // CHECK-NEXT: %[[LHS_BCAST:.+]] = "xla_hlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<1> : tensor<1xi64>} - // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %arg1) {comparison_direction = "GT"} + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.const_shape [1] + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK-DAG: %[[RESULT_SHAPE:.+]] = shape.const_shape [1, 2] + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<1> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} + // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} %0 = "tf.Greater"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> return %0: tensor<1x2xi1> } // CHECK-LABEL: func @greater_dynamic func @greater_dynamic(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 - // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]] - // CHECK-NEXT: shape.assuming %[[WITNESS]] - // CHECK-DAG: %[[LHS_SHAPE1:.+]] = shape.shape_of %arg0 - // CHECK-DAG: %[[RHS_SHAPE1:.+]] = shape.shape_of %arg1 - // CHECK-NEXT: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE1]], %[[RHS_SHAPE1]]) - // CHECK-NEXT: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] - // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} - // CHECK-NEXT: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} + // CHECK-DAG: %[[LHS_SHAPE:.+]] = shape.shape_of %arg0 + // CHECK-DAG: %[[RHS_SHAPE:.+]] = shape.shape_of %arg1 + // CHECK-DAG: %[[RESULT_SHAPE:.+]] = "shape.broadcast"(%[[LHS_SHAPE]], %[[RHS_SHAPE]]) + // CHECK-DAG: %[[RESULT_EXTENTS:.+]] = shape.to_extent_tensor %[[RESULT_SHAPE]] + // CHECK-DAG: %[[LHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK-DAG: %[[RHS_BCAST:.+]] = "xla_hlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) {broadcast_dimensions = dense<0> : tensor<1xi64>} + // CHECK: "xla_hlo.compare"(%[[LHS_BCAST]], %[[RHS_BCAST]]) {comparison_direction = "GT"} %0 = "tf.Greater"(%arg0, %arg1) : (tensor, tensor) -> tensor return %0: tensor } diff --git a/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc b/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc index 97afa9617c4..e5a79616d5b 100644 --- a/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc +++ b/tensorflow/compiler/mlir/xla/transforms/chlo_legalize_to_hlo.cc @@ -112,19 +112,6 @@ struct ConvertRankedDynamicBroadcastBinaryOp // Compute result shape. auto loc = op.getLoc(); - - // Insert a constraint on the shapes being broadcastable and insert all - // future code into an assuming block reliant on the constraint. - Value lhs_shape = rewriter.create(loc, lhs); - Value rhs_shape = rewriter.create(loc, rhs); - auto broadcastable_cstr = - rewriter.create(loc, lhs_shape, rhs_shape); - auto assuming_op = rewriter.create( - loc, ArrayRef{result_type}, broadcastable_cstr.result()); - - OpBuilder::InsertionGuard guard(rewriter); - rewriter.createBlock(&assuming_op.doRegion()); - int64_t result_rank = std::max(lhs_type.getRank(), rhs_type.getRank()); Value result_extents = xla::ComputeBinaryElementwiseBroadcastingResultExtents(loc, lhs, rhs, @@ -153,10 +140,8 @@ struct ConvertRankedDynamicBroadcastBinaryOp rewriter.getI64TensorAttr(rhs_broadcast_dimensions)); // And generate the final non-broadcasted binary op. - Value final_result = Adaptor::CreateOp(op, result_type, broadcasted_lhs, - broadcasted_rhs, rewriter); - rewriter.create(loc, final_result); - rewriter.replaceOp(op, {assuming_op.getResult(0)}); + rewriter.replaceOp(op, {Adaptor::CreateOp(op, result_type, broadcasted_lhs, + broadcasted_rhs, rewriter)}); return success(); } }; From 06b1b45e42aefe9dbfbedeca8abff47c20b28a21 Mon Sep 17 00:00:00 2001 From: Shanqing Cai Date: Mon, 15 Jun 2020 08:47:31 -0700 Subject: [PATCH 0164/1390] [tfdbg2] A few fixes and improvements to example debug_mnist_v2 1. Change the default `--dump_tensor_debug_mode` flag value to `FULL_HEALTH`, a mode more suitable for numerical instability debugging for the particular bug in this example than the previous default value `NO_TENSOR`. 2. Change the default `--dump_circular_buffer_size` value to -1, to accommodate the possibility of longer runs where user would want to see the debug data in its entirety. 3. Rename a few weight variables. They were previously named in a confusing way. 4. Change "logits" to "probs", as they are generated by a `softmax` operation and hence are more accurately described as probability scores. PiperOrigin-RevId: 316473164 Change-Id: I4eb13f9581a4d4e550b3b3a5cd132eeffc7dd043 --- .../debug/examples/v2/debug_mnist_v2.py | 34 +++++++++++++------ 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/tensorflow/python/debug/examples/v2/debug_mnist_v2.py b/tensorflow/python/debug/examples/v2/debug_mnist_v2.py index 539be3cd54f..f00e54500fa 100644 --- a/tensorflow/python/debug/examples/v2/debug_mnist_v2.py +++ b/tensorflow/python/debug/examples/v2/debug_mnist_v2.py @@ -98,16 +98,18 @@ def parse_args(): parser.add_argument( "--dump_tensor_debug_mode", type=str, - default="NO_TENSOR", + default="FULL_HEALTH", help="Mode for dumping tensor values. Options: NO_TENSOR, CURT_HEALTH, " - "CONCISE_HEALTH, SHAPE, FULL_TENSOR. This is relevant only when " + "CONCISE_HEALTH, SHAPE, FULL_HEALTH. This is relevant only when " "--dump_dir is set.") # TODO(cais): Add more tensor debug mode strings once they are supported. parser.add_argument( "--dump_circular_buffer_size", type=int, - default=1000, + default=-1, help="Size of the circular buffer used to dump execution events. " + "A value <= 0 disables the circular-buffer behavior and causes " + "all instrumented tensor values to be dumped. " "This is relevant only when --dump_dir is set.") parser.add_argument( "--use_random_config_path", @@ -178,9 +180,9 @@ def main(_): return activations # init model - hidden = get_dense_weights(IMAGE_SIZE**2, HIDDEN_SIZE) - logits = get_dense_weights(HIDDEN_SIZE, NUM_LABELS) - variables = hidden + logits + hidden_weights = get_dense_weights(IMAGE_SIZE**2, HIDDEN_SIZE) + output_weights = get_dense_weights(HIDDEN_SIZE, NUM_LABELS) + variables = hidden_weights + output_weights @tf.function def model(x): @@ -193,15 +195,25 @@ def main(_): Returns: A (?, 10) tensor containing the class scores for each example. """ - hidden_act = dense_layer(hidden, x) - logits_act = dense_layer(logits, hidden_act, tf.identity) + hidden_act = dense_layer(hidden_weights, x) + logits_act = dense_layer(output_weights, hidden_act, tf.identity) y = tf.nn.softmax(logits_act) return y @tf.function - def loss(logits, labels): - """Calculates cross entropy loss.""" - diff = -(labels * tf.math.log(logits)) + def loss(probs, labels): + """Calculates cross entropy loss. + + Args: + probs: Class probabilities predicted by the model. The shape is expected + to be (?, 10). + labels: Truth labels for the classes, as one-hot encoded vectors. The + shape is expected to be the same as `probs`. + + Returns: + A scalar loss tensor. + """ + diff = -labels * tf.math.log(probs) loss = tf.reduce_mean(diff) return loss From 81601f9bb422d18368f1ded540e68911bcebb7f8 Mon Sep 17 00:00:00 2001 From: Andrew Audibert Date: Mon, 15 Jun 2020 08:51:03 -0700 Subject: [PATCH 0165/1390] [tf.data service] Increase default number of uncompress threads to 4. A single thread may not be able to uncompress data as quickly as it is requested. PiperOrigin-RevId: 316473727 Change-Id: Ic041b70f2d3081d9333f3272f493d52ec3704d91 --- .../data/experimental/ops/data_service_ops.py | 3 ++- .../kernel_tests/data_service_ops_test.py | 23 +++++++++++-------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/tensorflow/python/data/experimental/ops/data_service_ops.py b/tensorflow/python/data/experimental/ops/data_service_ops.py index 39790d843ba..dd81614fa45 100644 --- a/tensorflow/python/data/experimental/ops/data_service_ops.py +++ b/tensorflow/python/data/experimental/ops/data_service_ops.py @@ -241,7 +241,8 @@ def _distribute(processing_mode, # TODO(b/157105111): Make this an autotuned parallel map when we have a way # to limit memory usage. dataset = dataset.map( - lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec)) + lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec), + num_parallel_calls=4) # Disable autosharding for shared jobs. if job_name: diff --git a/tensorflow/python/data/kernel_tests/data_service_ops_test.py b/tensorflow/python/data/kernel_tests/data_service_ops_test.py index d316009ce0c..440a4f46a20 100644 --- a/tensorflow/python/data/kernel_tests/data_service_ops_test.py +++ b/tensorflow/python/data/kernel_tests/data_service_ops_test.py @@ -201,13 +201,18 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): self._new_worker = server_lib.WorkerServer( port=port, master_address=self._master._address, protocol=PROTOCOL) - # The dataset starts over now that we read from the new worker. - for i in range(num_elements): + # There may have been some elements prefetched from the first worker + # before it was stopped. + while True: + val = next(iterator).numpy() + if val == 0: + break + + # The dataset starts over now that we read from the new worker. + # TODO(b/157086991): Iterate until end of sequence when we support + # detecting lost workers. + for i in range(1, num_elements // 2): val = next(iterator).numpy() - if val == midpoint and i != midpoint: - # There may have been one last element prefetched from the first worker - # before it was stopped. - val = next(iterator).numpy() self.assertEqual(i, val) @combinations.generate(test_base.eager_only_combinations()) @@ -291,7 +296,7 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.eager_only_combinations()) def testSharedJobNameRepeat(self): - num_elements = 10 + num_elements = 100 num_repetitions = 3 master_address = self.create_cluster(1) ds = dataset_ops.Dataset.range(num_elements) @@ -302,9 +307,9 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): results = [] iter1 = iter(ds1) iter2 = iter(ds2) - for _ in range(((num_elements * num_repetitions) // 2) - 1): + for _ in range(((num_elements * num_repetitions) // 3)): results.append(next(iter1).numpy()) - for _ in range(((num_elements * num_repetitions) // 2) - 1): + for _ in range(((num_elements * num_repetitions) // 3)): results.append(next(iter2).numpy()) for elem in iter1: results.append(elem.numpy()) From 7d36cebea4e10f108117d5c5423100cfd8f24509 Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Mon, 15 Jun 2020 08:59:00 -0700 Subject: [PATCH 0166/1390] Fork the keras related saver_test to keras. PiperOrigin-RevId: 316475172 Change-Id: I9b38d7624ed01a724c15ae571c5be1a7f0e049ad --- tensorflow/python/keras/tests/BUILD | 23 +++ tensorflow/python/keras/tests/saver_test.py | 158 ++++++++++++++++++++ tensorflow/python/training/saver_test.py | 121 --------------- 3 files changed, 181 insertions(+), 121 deletions(-) create mode 100644 tensorflow/python/keras/tests/saver_test.py diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD index d270e6f638c..4bb7d5358e5 100644 --- a/tensorflow/python/keras/tests/BUILD +++ b/tensorflow/python/keras/tests/BUILD @@ -335,6 +335,29 @@ tf_py_test( ], ) +cuda_py_test( + name = "saver_test", + size = "medium", + srcs = ["saver_test.py"], + python_version = "PY3", + deps = [ + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:nn_grad", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:saver", + "//tensorflow/python:training_lib", + "//tensorflow/python:training_util", + "//tensorflow/python/eager:context", + "//tensorflow/python/keras/engine", + "//tensorflow/python/keras/layers:core", + "//tensorflow/python/training/tracking", + "//tensorflow/python/training/tracking:util", + ], +) + tf_py_test( name = "temporal_sample_weights_correctness_test", srcs = ["temporal_sample_weights_correctness_test.py"], diff --git a/tensorflow/python/keras/tests/saver_test.py b/tensorflow/python/keras/tests/saver_test.py new file mode 100644 index 00000000000..f425414a932 --- /dev/null +++ b/tensorflow/python/keras/tests/saver_test.py @@ -0,0 +1,158 @@ +# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================= +"""Tests for tensorflow.python.training.saver.py.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os + +from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import errors +from tensorflow.python.framework import ops as ops_lib +from tensorflow.python.keras.engine import training +from tensorflow.python.keras.layers import core +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.platform import test +from tensorflow.python.training import adam +from tensorflow.python.training import saver as saver_module +from tensorflow.python.training import training_util +from tensorflow.python.training.tracking import tracking as trackable_tracking +from tensorflow.python.training.tracking import util as trackable_utils + + +class NonLayerTrackable(trackable_tracking.AutoTrackable): + + def __init__(self): + super(NonLayerTrackable, self).__init__() + self.a_variable = trackable_utils.add_variable( + self, name="a_variable", shape=[]) + + +class MyModel(training.Model): + """A concrete Model for testing.""" + + def __init__(self): + super(MyModel, self).__init__() + self._named_dense = core.Dense(1, use_bias=True) + self._second = core.Dense(1, use_bias=False) + # We can still track Trackables which aren't Layers. + self._non_layer = NonLayerTrackable() + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +class TrackableCompatibilityTests(test.TestCase): + + def _initialized_model(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + train_op = optimizer.minimize( + functools.partial(model, input_value), + global_step=optimizer_step) + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + # A regular variable, a slot variable, and a non-slot Optimizer variable + # with known values to check when loading. + self.evaluate(model._named_dense.bias.assign([1.])) + self.evaluate(optimizer.get_slot( + var=model._named_dense.bias, name="m").assign([2.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(3.)) + return root_trackable + + def _set_sentinels(self, root_trackable): + self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) + self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, name="m") + .assign([102.])) + beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(103.)) + + def _check_sentinels(self, root_trackable): + self.assertAllEqual( + [1.], self.evaluate(root_trackable.model._named_dense.bias)) + self.assertAllEqual([2.], self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, name="m"))) + beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() + self.assertAllEqual(3., self.evaluate(beta1_power)) + + def testLoadFromObjectBasedGraph(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + save_graph = ops_lib.Graph() + with save_graph.as_default(), self.session(graph=save_graph) as sess: + root = self._initialized_model() + object_saver = trackable_utils.Checkpoint(root=root) + save_path = object_saver.save(file_prefix=checkpoint_prefix) + + # An incompatible object-based checkpoint to check error messages + var = resource_variable_ops.ResourceVariable(1., name="a") + self.evaluate(var.initializer) + second_saver = trackable_utils.Checkpoint(v=var) + second_path = second_saver.save(file_prefix=os.path.join( + checkpoint_directory, "second")) + + restore_graph = ops_lib.Graph() + with restore_graph.as_default(), self.session( + graph=restore_graph) as sess: + root = self._initialized_model() + self._set_sentinels(root) + saver = saver_module.Saver() + saver.restore(sess=sess, save_path=save_path) + self._check_sentinels(root) + before_second_restore_ops = restore_graph.get_operations() + # Test that multiple restores do not pollute the graph + saver.restore(sess=sess, save_path=save_path) + self.assertEqual(before_second_restore_ops, + restore_graph.get_operations()) + with self.assertRaisesRegexp(errors.NotFoundError, + "Could not find some variables"): + saver.restore(sess=sess, save_path=second_path) + + def testLoadFromObjectBasedEager(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + save_graph = ops_lib.Graph() + with save_graph.as_default(), self.session(graph=save_graph): + root = self._initialized_model() + object_saver = trackable_utils.Checkpoint(root=root) + save_path = object_saver.save(file_prefix=checkpoint_prefix) + + with context.eager_mode(): + root = self._initialized_model() + self._set_sentinels(root) + saver = saver_module.Saver( + root.model.variables + root.optimizer.variables()) + saver.restore(sess=None, save_path=save_path) + self._check_sentinels(root) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py index 2c8bdadd5d7..5c87be37e4c 100644 --- a/tensorflow/python/training/saver_test.py +++ b/tensorflow/python/training/saver_test.py @@ -18,7 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import functools import glob import math import os @@ -48,8 +47,6 @@ from tensorflow.python.framework import graph_io from tensorflow.python.framework import meta_graph from tensorflow.python.framework import ops as ops_lib from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import training -from tensorflow.python.keras.layers import core from tensorflow.python.lib.io import file_io from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops @@ -74,10 +71,7 @@ from tensorflow.python.training import py_checkpoint_reader from tensorflow.python.training import queue_runner_impl from tensorflow.python.training import saver as saver_module from tensorflow.python.training import saver_test_utils -from tensorflow.python.training import training_util from tensorflow.python.training.tracking import base as trackable_base -from tensorflow.python.training.tracking import tracking as trackable_tracking -from tensorflow.python.training.tracking import util as trackable_utils from tensorflow.python.util import compat @@ -3024,29 +3018,6 @@ class _OwnsMirroredVariables(trackable_base.Trackable): return self.non_dep_variable.name -class NonLayerTrackable(trackable_tracking.AutoTrackable): - - def __init__(self): - super(NonLayerTrackable, self).__init__() - self.a_variable = trackable_utils.add_variable( - self, name="a_variable", shape=[]) - - -class MyModel(training.Model): - """A concrete Model for testing.""" - - def __init__(self): - super(MyModel, self).__init__() - self._named_dense = core.Dense(1, use_bias=True) - self._second = core.Dense(1, use_bias=False) - # We can still track Trackables which aren't Layers. - self._non_layer = NonLayerTrackable() - - def call(self, values): - ret = self._second(self._named_dense(values)) - return ret - - class TrackableCompatibilityTests(test.TestCase): # TODO(allenl): Track down python3 reference cycles in these tests. @@ -3112,46 +3083,6 @@ class TrackableCompatibilityTests(test.TestCase): saver.restore(sess, save_path) self.assertEqual(1, v.eval_count) - def _initialized_model(self): - input_value = constant_op.constant([[3.]]) - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - optimizer_step = training_util.get_or_create_global_step() - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, optimizer_step=optimizer_step) - train_op = optimizer.minimize( - functools.partial(model, input_value), - global_step=optimizer_step) - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - # A regular variable, a slot variable, and a non-slot Optimizer variable - # with known values to check when loading. - self.evaluate(model._named_dense.bias.assign([1.])) - self.evaluate(optimizer.get_slot( - var=model._named_dense.bias, name="m").assign([2.])) - beta1_power, _ = optimizer._get_beta_accumulators() - self.evaluate(beta1_power.assign(3.)) - return root_trackable - - def _set_sentinels(self, root_trackable): - self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) - self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, name="m") - .assign([102.])) - beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() - self.evaluate(beta1_power.assign(103.)) - - def _check_sentinels(self, root_trackable): - self.assertAllEqual( - [1.], self.evaluate(root_trackable.model._named_dense.bias)) - self.assertAllEqual([2.], self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, name="m"))) - beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() - self.assertAllEqual(3., self.evaluate(beta1_power)) - def testVariableNotFoundErrorRaised(self): # Restore does some tricky exception handling to figure out if it should # load an object-based checkpoint. Tests that the exception handling isn't @@ -3199,58 +3130,6 @@ class TrackableCompatibilityTests(test.TestCase): "a mismatch between the current graph and the graph"): a_saver.restore(sess=sess, save_path=save_path) - def testLoadFromObjectBasedGraph(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - - save_graph = ops_lib.Graph() - with save_graph.as_default(), self.session(graph=save_graph) as sess: - root = self._initialized_model() - object_saver = trackable_utils.Checkpoint(root=root) - save_path = object_saver.save(file_prefix=checkpoint_prefix) - - # An incompatible object-based checkpoint to check error messages - var = resource_variable_ops.ResourceVariable(1., name="a") - self.evaluate(var.initializer) - second_saver = trackable_utils.Checkpoint(v=var) - second_path = second_saver.save(file_prefix=os.path.join( - checkpoint_directory, "second")) - - restore_graph = ops_lib.Graph() - with restore_graph.as_default(), self.session( - graph=restore_graph) as sess: - root = self._initialized_model() - self._set_sentinels(root) - saver = saver_module.Saver() - saver.restore(sess=sess, save_path=save_path) - self._check_sentinels(root) - before_second_restore_ops = restore_graph.get_operations() - # Test that multiple restores do not pollute the graph - saver.restore(sess=sess, save_path=save_path) - self.assertEqual(before_second_restore_ops, - restore_graph.get_operations()) - with self.assertRaisesRegexp(errors.NotFoundError, - "Could not find some variables"): - saver.restore(sess=sess, save_path=second_path) - - def testLoadFromObjectBasedEager(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - - save_graph = ops_lib.Graph() - with save_graph.as_default(), self.session(graph=save_graph): - root = self._initialized_model() - object_saver = trackable_utils.Checkpoint(root=root) - save_path = object_saver.save(file_prefix=checkpoint_prefix) - - with context.eager_mode(): - root = self._initialized_model() - self._set_sentinels(root) - saver = saver_module.Saver( - root.model.variables + root.optimizer.variables()) - saver.restore(sess=None, save_path=save_path) - self._check_sentinels(root) - if __name__ == "__main__": test.main() From 01b30fa03f636396f977a76628123199b772463a Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Mon, 15 Jun 2020 17:30:29 +0100 Subject: [PATCH 0167/1390] Fix for the linter. Change-Id: Ie1185cc2cca9157655b22e1d3bb49ddc017a8f0e --- tensorflow/lite/tools/optimize/quantization_utils.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/lite/tools/optimize/quantization_utils.cc b/tensorflow/lite/tools/optimize/quantization_utils.cc index cdf2743585e..cdc794c20c4 100644 --- a/tensorflow/lite/tools/optimize/quantization_utils.cc +++ b/tensorflow/lite/tools/optimize/quantization_utils.cc @@ -92,7 +92,6 @@ void GetSymmetricQuantizationParams( min = std::min(min, 0.0f); max = std::max(max, 0.0f); const float scale = std::max(std::abs(max), std::abs(min)) / half_quant_range; - int64_t zero_point = 0; quantization_params->min = std::vector(1, min); quantization_params->max = std::vector(1, max); quantization_params->scale = std::vector(1, scale); From 16ac7c04d4d8727ae917c20678eb776fb1a4bacb Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Mon, 15 Jun 2020 09:35:04 -0700 Subject: [PATCH 0168/1390] Fork keras related tracking test to keras/tests PiperOrigin-RevId: 316482123 Change-Id: I20645bbfdd926e2c83136ee27c6ef9325cb1f438 --- tensorflow/python/keras/tests/BUILD | 65 ++ .../python/keras/tests/tracking_test.py | 610 ++++++++++++ .../python/keras/tests/tracking_util_test.py | 926 ++++++++++++++++++ tensorflow/python/training/tracking/BUILD | 7 - .../training/tracking/data_structures_test.py | 503 +--------- .../python/training/tracking/tracking_test.py | 63 -- .../python/training/tracking/util_test.py | 861 ---------------- 7 files changed, 1602 insertions(+), 1433 deletions(-) create mode 100644 tensorflow/python/keras/tests/tracking_test.py create mode 100644 tensorflow/python/keras/tests/tracking_util_test.py diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD index 4bb7d5358e5..ad52d33abc6 100644 --- a/tensorflow/python/keras/tests/BUILD +++ b/tensorflow/python/keras/tests/BUILD @@ -370,6 +370,71 @@ tf_py_test( ], ) +tf_py_test( + name = "tracking_test", + srcs = ["tracking_test.py"], + python_version = "PY3", + tags = [ + "no_windows", + "nomac", + ], + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:math_ops", + "//tensorflow/python:variables", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:test", + "//tensorflow/python/keras/engine", + "//tensorflow/python/keras/layers:core", + "//tensorflow/python/keras/layers:normalization", + "//tensorflow/python/module", + "//tensorflow/python/training/tracking", + "//tensorflow/python/training/tracking:data_structures", + "//tensorflow/python/training/tracking:util", + "@absl_py//absl/testing:parameterized", + ], +) + +tf_py_test( + name = "tracking_util_test", + srcs = ["tracking_util_test.py"], + python_version = "PY3", + tags = ["notsan"], # b/74395663 + deps = [ + "//tensorflow/compiler/tests:xla_test", + "//tensorflow/python:checkpoint_management", + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:extra_py_tests_deps", + "//tensorflow/python:framework_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", + "//tensorflow/python:platform", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:saver", + "//tensorflow/python:state_ops", + "//tensorflow/python:template", + "//tensorflow/python:training_util", + "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", + "//tensorflow/python/eager:backprop", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:def_function", + "//tensorflow/python/keras/engine", + "//tensorflow/python/keras/layers:core", + "//tensorflow/python/keras/optimizer_v2", + "//tensorflow/python/training/tracking", + "//tensorflow/python/training/tracking:graph_view", + "//tensorflow/python/training/tracking:util", + "@absl_py//absl/testing:parameterized", + ], +) + py_library( name = "get_config_samples", srcs = ["get_config_samples.py"], diff --git a/tensorflow/python/keras/tests/tracking_test.py b/tensorflow/python/keras/tests/tracking_test.py new file mode 100644 index 00000000000..b5ce6911d92 --- /dev/null +++ b/tensorflow/python/keras/tests/tracking_test.py @@ -0,0 +1,610 @@ +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from absl.testing import parameterized +import numpy +import six + +from tensorflow.python.eager import context +from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.keras.engine import sequential +from tensorflow.python.keras.engine import training +from tensorflow.python.keras.layers import core +from tensorflow.python.keras.layers import normalization +from tensorflow.python.module import module +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import variables +from tensorflow.python.training.tracking import base +from tensorflow.python.training.tracking import data_structures +from tensorflow.python.training.tracking import tracking +from tensorflow.python.training.tracking import util + + +class HasList(training.Model): + + def __init__(self): + super(HasList, self).__init__() + self.layer_list = data_structures.List([core.Dense(3)]) + self.layer_list.append(core.Dense(4)) + self.layer_list.extend( + [core.Dense(5), + core.Dense(6, kernel_regularizer=math_ops.reduce_sum)]) + self.layer_list += [ + core.Dense(7, bias_regularizer=math_ops.reduce_sum), + core.Dense(8) + ] + self.layer_list += ( + data_structures.List([core.Dense(9)]) + data_structures.List( + [core.Dense(10)])) + self.layer_list.extend( + data_structures.List( + list([core.Dense(11)]) + [core.Dense(12)])) + self.layers_with_updates = data_structures.List( + (normalization.BatchNormalization(),)) + + def call(self, x): + aggregation = 0. + for l in self.layer_list: + x = l(x) + aggregation += math_ops.reduce_sum(x) + bn, = self.layers_with_updates + return bn(x) / aggregation + + +class ListTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes + @test_util.run_v1_only("b/120545219") + def testTracking(self): + model = HasList() + output = model(array_ops.ones([32, 2])) + self.assertAllEqual([32, 12], output.shape) + self.assertEqual(11, len(model.layers)) + self.assertEqual(10, len(model.layer_list.layers)) + six.assertCountEqual( + self, + model.layers, + model.layer_list.layers + model.layers_with_updates) + for index in range(10): + self.assertEqual(3 + index, model.layer_list.layers[index].units) + self.assertEqual(2, len(model._checkpoint_dependencies)) + self.assertIs(model.layer_list, model._checkpoint_dependencies[0].ref) + self.assertIs(model.layers_with_updates, + model._checkpoint_dependencies[1].ref) + self.assertEqual( + 10, len(model._checkpoint_dependencies[0].ref._checkpoint_dependencies)) + self.evaluate([v.initializer for v in model.variables]) + self.evaluate(model.variables[0].assign([[1., 2., 3.], [4., 5., 6.]])) + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + self.evaluate(model.variables[0].assign(array_ops.zeros([2, 3]))) + model.load_weights(save_path) + self.assertAllEqual([[1., 2., 3.], [4., 5., 6.]], + self.evaluate(model.variables[0])) + v = variables.Variable(1.) + model.var_list = [v] + self.assertIn(v, model.variables) + self.assertIn(v, model.trainable_variables) + self.assertNotIn(v, model.non_trainable_variables) + self.assertIn(model.layer_list[0].trainable_weights[0], + model.trainable_weights) + + def testSubModelTracking(self): + model = training.Model() + model.v = variables.Variable(1.) + self.assertIn(model.v, model.trainable_weights) + model2 = training.Model() + model2.m = [model] + self.assertIn(model.v, model2.trainable_weights) + + def testSubSequentialTracking(self): + + class _Subclassed(training.Model): + + def __init__(self, wrapped): + super(_Subclassed, self).__init__() + self._wrapped = wrapped + + def call(self, x): + return self._wrapped(x) + + model = sequential.Sequential() + layer = core.Dense(1) + model.add(layer) + model2 = _Subclassed(model) + model2(array_ops.ones([1, 2])) + model2.m = [model] + self.assertIn(layer.kernel, model2.trainable_weights) + + def testLayerTrackedThroughSequential(self): + class AttrDict(dict): + + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + def ffnet(layer_sizes, name): + ff = sequential.Sequential(name=name) + for i, width in enumerate(layer_sizes): + ff.add(core.Dense( + width, + activation=("relu" if i < len(layer_sizes)-1 else None))) + return ff + + class MyModel2(training.Model): + + def __init__(self, config, name="my_model_2"): + super(MyModel2, self).__init__(name=name) + self._num_tokens = config.num_tokens + + # list of sub-models + self._ffnet = [ffnet(config.module_layers + (self._num_tokens,), "ff")] + + def null_input(self): + return array_ops.zeros([1, self._num_tokens], dtype=dtypes.float32) + + def call(self, input_, module_index=None): + return self._ffnet[0](input_) + + m2 = MyModel2(AttrDict( + num_tokens=5, + module_layers=(50, 30))) + + # Construct + m2(m2.null_input()) + self.assertLen(m2.trainable_variables, 6) + + @test_util.run_v1_only("b/120545219") + def testUpdatesForwarded(self): + with context.graph_mode(): + model = HasList() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertGreater(len(model.layers_with_updates[0].updates), 0) + self.assertEqual(set(model.layers_with_updates[0].updates), + set(model.updates)) + + with context.eager_mode(): + model = HasList() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertEqual(0, len(model.updates)) + + @test_util.run_in_graph_and_eager_modes + @test_util.run_v1_only("b/120545219") + def testLossesForwarded(self): + model = HasList() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertEqual(2, len(model.losses)) + + def testModelContainersCompareEqual(self): + class HasEqualContainers(training.Model): + + def __init__(self): + super(HasEqualContainers, self).__init__() + self.l1 = [] + self.l2 = [] + + model = HasEqualContainers() + first_layer = HasEqualContainers() + model.l1.append(first_layer) + second_layer = HasEqualContainers() + model.l2.append(second_layer) + self.assertEqual([first_layer, second_layer], model.layers) + + @test_util.run_in_graph_and_eager_modes + def testTensorConversion(self): + + class ListToTensor(training.Model): + + def __init__(self): + super(ListToTensor, self).__init__() + self.l = [1., 2., 3.] + + self.assertAllEqual( + [1., 2., 3.], + self.evaluate(constant_op.constant(ListToTensor().l))) + + self.assertAllEqual( + [1., 2., 3.], + self.evaluate(array_ops.pack(ListToTensor().l))) + + +class ListWrapperTest(test.TestCase): + + def testLayerCollectionWithExternalMutation(self): + l = [] + l_wrapper = data_structures.ListWrapper(l) + layer = core.Dense(1) + l.append(layer) + self.assertEqual([layer], l_wrapper.layers) + + +class HasMapping(training.Model): + + def __init__(self): + super(HasMapping, self).__init__() + self.layer_dict = data_structures.Mapping(output=core.Dense(7)) + self.layer_dict["norm"] = data_structures.List() + self.layer_dict["dense"] = data_structures.List() + self.layer_dict["dense"].extend( + [core.Dense(5), + core.Dense(6, kernel_regularizer=math_ops.reduce_sum)]) + self.layer_dict["norm"].append( + normalization.BatchNormalization()) + self.layer_dict["norm"].append( + normalization.BatchNormalization()) + + def call(self, x): + aggregation = 0. + for norm, dense in zip(self.layer_dict["norm"], self.layer_dict["dense"]): + x = norm(dense(x)) + aggregation += math_ops.reduce_sum(x) + return self.layer_dict["output"](x) / aggregation + + +class MappingTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes + def testTracking(self): + model = HasMapping() + output = model(array_ops.ones([32, 2])) + self.assertAllEqual([32, 7], output.shape.as_list()) + self.assertEqual(5, len(model.layers)) + six.assertCountEqual(self, model.layers, model.layer_dict.layers) + self.assertEqual(1, len(model._checkpoint_dependencies)) + self.assertIs(model.layer_dict, model._checkpoint_dependencies[0].ref) + self.evaluate([v.initializer for v in model.variables]) + test_var = model.layer_dict["output"].kernel + self.evaluate(test_var.assign(array_ops.ones([6, 7]))) + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + self.evaluate(test_var.assign(array_ops.zeros([6, 7]))) + model.load_weights(save_path) + self.assertAllEqual(numpy.ones([6, 7]), + self.evaluate(test_var)) + + def testLayerCollectionWithExternalMutation(self): + d = {} + root = tracking.AutoTrackable() + root.wrapper = d + self.assertEqual([], root.wrapper.layers) + self.assertEqual([], root.wrapper.trainable_weights) + layer1 = core.Dense(1) + layer2 = core.Dense(1) + d["a"] = layer1 + d["b"] = layer2 + self.assertEqual([layer1, layer2], root.wrapper.layers) + # The layers have still not created variables + self.assertEqual([], root.wrapper.trainable_weights) + + def testDictWrapperBadKeys(self): + a = tracking.AutoTrackable() + a.d = {} + a.d[1] = data_structures.List() + model = training.Model() + model.sub = a + save_path = os.path.join(self.get_temp_dir(), "ckpt") + with self.assertRaisesRegexp(ValueError, "non-string key"): + model.save_weights(save_path) + + def testDictWrapperNoDependency(self): + a = tracking.AutoTrackable() + a.d = data_structures.NoDependency({}) + a.d[1] = [3] + self.assertEqual([a], util.list_objects(a)) + model = training.Model() + model.sub = a + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + model.load_weights(save_path) + + def testNonStringKeyNotTrackableValue(self): + a = tracking.AutoTrackable() + a.d = {} + a.d["a"] = [3] + a.d[1] = data_structures.NoDependency([3]) + self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) + model = training.Model() + model.sub = a + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + model.load_weights(save_path) + + def testNonAppendNotTrackable(self): + # Non-append mutations (deleting or overwriting values) are OK when the + # values aren't tracked. + a = tracking.AutoTrackable() + a.d = {} + a.d["a"] = [3] + a.d[1] = 3 + a.d[1] = 2 + self.assertEqual(2, a.d[1]) + del a.d[1] + a.d[2] = data_structures.NoDependency(tracking.AutoTrackable()) + second = tracking.AutoTrackable() + a.d[2] = data_structures.NoDependency(second) + self.assertIs(second, a.d[2]) + self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) + model = training.Model() + model.sub = a + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + model.load_weights(save_path) + + def testPopNoSave(self): + model = training.Model() + model.d = {} + model.d["a"] = [] + model.d.pop("a") + save_path = os.path.join(self.get_temp_dir(), "ckpt") + with self.assertRaisesRegexp(ValueError, "Unable to save"): + model.save_weights(save_path) + + def testExternalModificationNoSave(self): + model = training.Model() + external_reference = {} + model.d = external_reference + external_reference["a"] = [] + save_path = os.path.join(self.get_temp_dir(), "ckpt") + with self.assertRaisesRegexp(ValueError, "modified outside the wrapper"): + model.save_weights(save_path) + + def testOverwriteCanStillSave(self): + model = training.Model() + model.d = {} + model.d["a"] = {} + model.d["a"] = {} + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + + def testIter(self): + model = training.Model() + model.d = {1: 3} + model.d[1] = 3 + self.assertEqual([1], list(model.d)) + new_dict = {} + # This update() is super tricky. If the dict wrapper subclasses dict, + # CPython will access its storage directly instead of calling any + # methods/properties on the object. So the options are either not to + # subclass dict (in which case update will call normal iter methods, but the + # object won't pass isinstance checks) or to subclass dict and keep that + # storage updated (no shadowing all its methods like ListWrapper). + new_dict.update(model.d) + self.assertEqual({1: 3}, new_dict) + + +class HasTuple(training.Model): + + def __init__(self): + super(HasTuple, self).__init__() + self.layer_list = ( + core.Dense(3), core.Dense(4), + core.Dense(5, kernel_regularizer=math_ops.reduce_sum)) + self.layers_with_updates = (normalization.BatchNormalization(),) + + def call(self, x): + aggregation = 0. + for l in self.layer_list: + x = l(x) + aggregation += math_ops.reduce_sum(x) + bn, = self.layers_with_updates + return bn(x) / aggregation + + +class TupleTests(test.TestCase, parameterized.TestCase): + + @test_util.run_in_graph_and_eager_modes + def testTracking(self): + model = HasTuple() + output = model(array_ops.ones([32, 2])) + self.assertAllEqual([32, 5], output.shape.as_list()) + self.assertLen(model.layers, 4) + self.assertLen(model.layer_list.layers, 3) + six.assertCountEqual( + self, + model.layers, + tuple(model.layer_list.layers) + model.layers_with_updates) + self.assertEqual(3, model.layer_list.layers[0].units) + self.assertEqual(4, model.layer_list.layers[1].units) + self.assertEqual(5, model.layer_list.layers[2].units) + self.assertLen(model._checkpoint_dependencies, 2) + self.assertIs(model.layer_list, model._checkpoint_dependencies[0].ref) + self.assertIs(model.layers_with_updates, + model._checkpoint_dependencies[1].ref) + self.assertLen( + model._checkpoint_dependencies[0].ref._checkpoint_dependencies, 3) + self.evaluate([v.initializer for v in model.variables]) + self.evaluate(model.variables[0].assign([[1., 2., 3.], [4., 5., 6.]])) + save_path = os.path.join(self.get_temp_dir(), "ckpt") + model.save_weights(save_path) + self.evaluate(model.variables[0].assign(array_ops.zeros([2, 3]))) + model.load_weights(save_path) + self.assertAllEqual([[1., 2., 3.], [4., 5., 6.]], + self.evaluate(model.variables[0])) + v = variables.Variable(1.) + model.var_list = (v,) + self.assertIn(id(v), [id(obj) for obj in model.variables]) + self.assertIn(id(v), [id(obj) for obj in model.trainable_variables]) + self.assertNotIn(id(v), [id(obj) for obj in model.non_trainable_variables]) + self.assertIn(id(model.layer_list[0].trainable_weights[0]), + [id(obj) for obj in model.trainable_weights]) + + @parameterized.named_parameters( + ("Module", module.Module), + ("Model", training.Model), + ) + def testSubModelTracking(self, module_subclass): + model = module_subclass() + model.v = variables.Variable(1.) + self.assertIn(model.v, model.trainable_variables) + model2 = module_subclass() + model2.m = (model,) + self.assertIn(model.v, model2.trainable_variables) + + def testSubSequentialTracking(self): + + class _Subclassed(training.Model): + + def __init__(self, wrapped): + super(_Subclassed, self).__init__() + self._wrapped = wrapped + + def call(self, x): + return self._wrapped(x) + + model = sequential.Sequential() + layer = core.Dense(1) + model.add(layer) + model2 = _Subclassed(model) + model2(array_ops.ones([1, 2])) + model2.m = (model,) + self.assertIn(layer.kernel, model2.trainable_weights) + + def testUpdatesForwarded(self): + with ops.Graph().as_default(): + model = HasTuple() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertNotEmpty(model.layers_with_updates[0].updates) + self.assertEqual(set(model.layers_with_updates[0].updates), + set(model.updates)) + + model = HasTuple() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertEmpty(model.updates) + + @test_util.run_in_graph_and_eager_modes + def testLossesForwarded(self): + model = HasTuple() + model_input = array_ops.ones([32, 2]) + model(model_input) + self.assertLen(model.losses, 1) + + def testModelContainersCompareEqual(self): + class HasEqualContainers(training.Model): + + def __init__(self): + super(HasEqualContainers, self).__init__() + self.l1 = () + self.l2 = () + + model = HasEqualContainers() + first_layer = HasEqualContainers() + model.l1 = (first_layer,) + second_layer = HasEqualContainers() + model.l2 = (second_layer,) + self.assertEqual((first_layer,), model.l1) + d = {model.l1: 1, model.l2: 2} + self.assertEqual(1, d[model.l1]) + self.assertEqual(1, d[(first_layer,)]) + self.assertEqual(2, d[model.l2]) + self.assertEqual(2, d[(second_layer,)]) + self.assertEqual([first_layer, second_layer], model.layers) + + @test_util.run_in_graph_and_eager_modes + def testTensorConversion(self): + + class TupleToTensor(training.Model): + + def __init__(self): + super(TupleToTensor, self).__init__() + self.l = (1., 2., 3.) + + self.assertAllEqual( + (1., 2., 3.), + self.evaluate(constant_op.constant(TupleToTensor().l))) + + self.assertAllEqual( + (1., 2., 3.), + self.evaluate(array_ops.pack(TupleToTensor().l))) + + +class InterfaceTests(test.TestCase): + + def testNoDependency(self): + root = tracking.AutoTrackable() + hasdep = tracking.AutoTrackable() + root.hasdep = hasdep + nodep = tracking.AutoTrackable() + root.nodep = data_structures.NoDependency(nodep) + self.assertEqual(1, len(root._checkpoint_dependencies)) + self.assertIs(root._checkpoint_dependencies[0].ref, root.hasdep) + self.assertIs(root.hasdep, hasdep) + self.assertIs(root.nodep, nodep) + + class NoDependencyModel(training.Model): + + @base.no_automatic_dependency_tracking + def __init__(self): + super(NoDependencyModel, self).__init__() + self.a = [] + self.b = tracking.AutoTrackable() + + nodeps = NoDependencyModel() + self.assertEqual([nodeps], util.list_objects(nodeps)) + + @test_util.run_in_graph_and_eager_modes + def testDictionariesBasic(self): + a = training.Model() + b = training.Model() + a.attribute = {"b": b} + c = training.Model() + a.attribute["c"] = [] + a.attribute["c"].append(c) + a_deps = util.list_objects(a) + self.assertIn(b, a_deps) + self.assertIn(c, a_deps) + self.assertIs(b, a.attribute["b"]) + six.assertCountEqual( + self, + ["b", "c"], + [dep.name for dep in a.attribute._checkpoint_dependencies]) + self.assertEqual([b, c], a.layers) + self.assertEqual([b, c], a.attribute.layers) + self.assertEqual([c], a.attribute["c"].layers) + checkpoint = util.Checkpoint(a=a) + save_path = checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) + with self.cached_session(): + checkpoint.restore(save_path).assert_consumed().initialize_or_restore() + + @test_util.run_in_graph_and_eager_modes + def testNoDepList(self): + a = training.Model() + a.l1 = data_structures.NoDependency([]) + a.l1.insert(1, 0) + self.assertIsInstance(a.l1, list) + checkpoint = util.Checkpoint(a=a) + checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) + a.l2 = [] + a.l2.insert(1, module.Module()) + with self.assertRaisesRegexp(ValueError, "A list element was replaced"): + checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/keras/tests/tracking_util_test.py b/tensorflow/python/keras/tests/tracking_util_test.py new file mode 100644 index 00000000000..ee5d7428fcc --- /dev/null +++ b/tensorflow/python/keras/tests/tracking_util_test.py @@ -0,0 +1,926 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os +import weakref + +from absl.testing import parameterized +import six + +from tensorflow.python.eager import backprop +from tensorflow.python.eager import context +from tensorflow.python.eager import def_function +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.keras.engine import input_layer +from tensorflow.python.keras.engine import sequential +from tensorflow.python.keras.engine import training +from tensorflow.python.keras.layers import core +from tensorflow.python.keras.optimizer_v2 import adam +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.ops import template +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables as variables_lib +from tensorflow.python.platform import test +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import checkpoint_management +from tensorflow.python.training import saver as saver_lib +from tensorflow.python.training import training_util +from tensorflow.python.training.tracking import graph_view +from tensorflow.python.training.tracking import tracking +from tensorflow.python.training.tracking import util as trackable_utils + + +# pylint: disable=not-callable +class MyModel(training.Model): + """A concrete Model for testing.""" + + def __init__(self): + super(MyModel, self).__init__() + self._named_dense = core.Dense(1, use_bias=True) + self._second = core.Dense(1, use_bias=False) + # We can still track Trackables which aren't Layers. + self._non_layer = NonLayerTrackable() + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +class NonLayerTrackable(tracking.AutoTrackable): + + def __init__(self): + super(NonLayerTrackable, self).__init__() + self.a_variable = trackable_utils.add_variable( + self, name="a_variable", shape=[]) + + +class InterfaceTests(test.TestCase): + + def testLayerDeduplication(self): + model = training.Model() + layer_one = core.Dense(1) + layer_two = core.Dense(1) + model.other_path = [layer_one, layer_two] + model.l2 = layer_two + model.l1 = layer_one + self.assertEqual([layer_one, layer_two], model.layers) + + def testSaveWithOnlyKerasSession(self): + + with ops.Graph().as_default(): + inp = input_layer.Input([1]) + dense = core.Dense(1)(inp) + model = training.Model(inp, dense) + model.compile(optimizer="sgd", loss="mse") + model.fit([1.], [2.]) + checkpoint = trackable_utils.Checkpoint(model=model) + checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) + + def testObjectMetadata(self): + with context.eager_mode(): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + dense = core.Dense(1) + checkpoint = trackable_utils.Checkpoint(dense=dense) + dense(constant_op.constant([[1.]])) + save_path = checkpoint.save(checkpoint_prefix) + + objects = trackable_utils.object_metadata(save_path) + all_variable_names = [] + for obj in objects.nodes: + for attribute in obj.attributes: + all_variable_names.append(attribute.full_name) + self.assertIn("dense/kernel", all_variable_names) + + +class CheckpointingTests(parameterized.TestCase, test.TestCase): + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testNamingWithOptimizer(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + # A nuisance Model using the same optimizer. Its slot variables should not + # go in the checkpoint, since it is never depended on. + other_model = MyModel() + optimizer = adam.Adam(0.001) + step = training_util.get_or_create_global_step() + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, step=step) + + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + train_op = control_flow_ops.group( + optimizer.apply_gradients(zip(gradients, variables)), + step.assign_add(1)) + + with backprop.GradientTape() as tape: + loss = other_model(input_value) + variables = other_model.trainable_variables + gradients = tape.gradient(loss, variables) + optimizer.apply_gradients(zip(gradients, variables)) + + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + named_variables, serialized_graph, _ = graph_view.ObjectGraphView( + root_trackable).serialize_object_graph() + expected_slot_keys = ( + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", + "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v", + ) + expected_checkpoint_names = ( + # Created in the root node, so no prefix. + "step", + "model/_second/kernel", + "model/_named_dense/kernel", + "model/_named_dense/bias", + # non-Layer dependency of the model + "model/_non_layer/a_variable", + "optimizer/learning_rate", + "optimizer/beta_1", + "optimizer/beta_2", + "optimizer/iter", + "optimizer/decay", + ) + expected_slot_keys + suffix = "/.ATTRIBUTES/VARIABLE_VALUE" + expected_checkpoint_names = [ + name + suffix for name in expected_checkpoint_names] + named_variables = {v.name: v for v in named_variables} + six.assertCountEqual(self, expected_checkpoint_names, + named_variables.keys()) + # Check that we've mapped to the right variable objects (not exhaustive) + self.assertEqual( + "global_step", + named_variables["step" + suffix].full_name) + self.assertEqual( + "my_model/dense_1/kernel", + named_variables["model/_second/kernel" + suffix].full_name) + self.assertEqual( + "my_model/dense/kernel", + named_variables["model/_named_dense/kernel" + suffix].full_name) + self.assertEqual("Adam/beta_1", + named_variables["optimizer/beta_1" + suffix].full_name) + self.assertEqual("Adam/beta_2", + named_variables["optimizer/beta_2" + suffix].full_name) + # Spot check the generated protocol buffers. + self.assertEqual("optimizer", + serialized_graph.nodes[0].children[1].local_name) + optimizer_node = serialized_graph.nodes[ + serialized_graph.nodes[0].children[1].node_id] + children = [node.local_name for node in optimizer_node.children] + six.assertCountEqual( + self, + # hyper variable dependencies + ["beta_1", "beta_2", "iter", "decay", "learning_rate"], + children) + serialized_slot_keys = [] + for slot in optimizer_node.slot_variables: + for attribute in ( + serialized_graph.nodes[slot.slot_variable_node_id].attributes): + serialized_slot_keys.append(attribute.checkpoint_key) + six.assertCountEqual( + self, + [key + suffix for key in expected_slot_keys], + serialized_slot_keys) + + @test_util.run_in_graph_and_eager_modes + def testSaveRestore(self): + model = MyModel() + optimizer = adam.Adam(0.001) + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + input_value = constant_op.constant([[3.]]) + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + train_op = optimizer.apply_gradients(zip(gradients, variables)) + self.assertFalse(root_trackable.save_counter.trainable) + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + prefix = os.path.join(self.get_temp_dir(), "ckpt") + self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.])) + m_bias_slot = optimizer.get_slot(model._named_dense.variables[1], "m") + self.evaluate(state_ops.assign(m_bias_slot, [1.5])) + save_path = root_trackable.save(file_prefix=prefix) + self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.])) + self.evaluate(state_ops.assign(root_trackable.save_counter, 3)) + optimizer_variables = self.evaluate( + sorted(optimizer.variables(), key=lambda v: v.name)) + self.evaluate(state_ops.assign(m_bias_slot, [-2.])) + # Immediate restoration + status = root_trackable.restore(save_path=save_path).assert_consumed() + status.run_restore_ops() + self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1])) + self.assertAllEqual(1, self.evaluate(root_trackable.save_counter)) + self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) + if not context.executing_eagerly(): + return # Restore-on-create is only supported when executing eagerly + on_create_model = MyModel() + on_create_optimizer = adam.Adam(0.001) + on_create_root = trackable_utils.Checkpoint( + optimizer=on_create_optimizer, model=on_create_model) + # Deferred restoration + status = on_create_root.restore(save_path=save_path) + status.assert_nontrivial_match() + status.assert_existing_objects_matched() + with self.assertRaises(AssertionError): + status.assert_consumed() + on_create_model(constant_op.constant([[3.]])) # create variables + self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) + self.assertAllEqual([42.], + self.evaluate( + on_create_model._named_dense.variables[1])) + on_create_m_bias_slot = on_create_optimizer.get_slot( + on_create_model._named_dense.variables[1], "m") + status.assert_existing_objects_matched() + if not context.executing_eagerly(): + with self.assertRaises(AssertionError): + status.assert_consumed() + # Optimizer slot variables are created when the original variable is + # restored. + self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot)) + dummy_var = resource_variable_ops.ResourceVariable([1.]) + on_create_optimizer.minimize(loss=dummy_var.read_value, + var_list=[dummy_var]) + status.assert_existing_objects_matched() + status.assert_consumed() + self.assertAllEqual( + optimizer_variables, + # Creation order is different, so .variables() needs to be re-sorted. + self.evaluate(sorted(optimizer.variables(), key=lambda v: v.name))) + + # TODO(allenl): Debug garbage created by this test in python3. + def testDeferredRestorationUsageEager(self): + """An idiomatic eager execution example.""" + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + model = MyModel() + optimizer = adam.Adam(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + root.restore(checkpoint_management.latest_checkpoint( + checkpoint_directory)) + for _ in range(num_training_steps): + # TODO(allenl): Use a Dataset and serialize/checkpoint it. + input_value = constant_op.constant([[3.]]) + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + optimizer.apply_gradients(zip(gradients, variables)) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer.iterations.numpy()) + + def testUsageGraph(self): + """Expected usage when graph building.""" + with context.graph_mode(): + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + with ops.Graph().as_default(): + model = MyModel() + optimizer = adam.Adam(0.001) + root = trackable_utils.CheckpointV1( + optimizer=optimizer, model=model) + input_value = constant_op.constant([[3.]]) + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + train_op = optimizer.apply_gradients(zip(gradients, variables)) + + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + with self.session(graph=ops.get_default_graph()) as session: + status = root.restore(save_path=checkpoint_path) + status.initialize_or_restore(session=session) + if checkpoint_path is None: + self.assertEqual(0, training_continuation) + with self.assertRaises(AssertionError): + status.assert_consumed() + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + else: + status.assert_consumed() + status.assert_existing_objects_matched() + for _ in range(num_training_steps): + session.run(train_op) + root.save(file_prefix=checkpoint_prefix, session=session) + self.assertEqual((training_continuation + 1) * num_training_steps, + session.run(root.optimizer.iterations)) + self.assertEqual(training_continuation + 1, + session.run(root.save_counter)) + + @test_util.run_in_graph_and_eager_modes + def testAgnosticUsage(self): + """Graph/eager agnostic usage.""" + # Does create garbage when executing eagerly due to ops.Graph() creation. + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + def _train_fn(model, input_value): + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + return optimizer.apply_gradients(zip(gradients, variables)) + for training_continuation in range(3): + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.Adam(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + manager = checkpoint_management.CheckpointManager( + root, checkpoint_directory, max_to_keep=1) + status = root.restore(save_path=manager.latest_checkpoint) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial(_train_fn, model, input_value) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + for _ in range(num_training_steps): + train_fn() + manager.save() + self.assertEqual((training_continuation + 1) * num_training_steps, + self.evaluate(root.optimizer.iterations)) + self.assertEqual(training_continuation + 1, + self.evaluate(root.save_counter)) + + def testPartialRestoreWarningObject(self): + with context.eager_mode(): + optimizer = adam.Adam(0.0) + original_root = trackable_utils.Checkpoint(v1=variables_lib.Variable(2.), + v2=variables_lib.Variable(3.), + optimizer=optimizer) + # Create a slot variable to save + optimizer.minimize(original_root.v1.read_value, [original_root.v1]) + prefix = os.path.join(self.get_temp_dir(), "ckpt") + save_path = original_root.save(prefix) + partial_root = trackable_utils.Checkpoint(v1=variables_lib.Variable(0.)) + weak_partial_root = weakref.ref(partial_root) + weak_v1 = weakref.ref(partial_root.v1) + partial_root.restore(save_path) + self.assertEqual(2., partial_root.v1.numpy()) + with test.mock.patch.object(logging, "warning") as mock_log: + del partial_root + self.assertIsNone(weak_partial_root()) + self.assertIsNone(weak_v1()) + messages = str(mock_log.call_args_list) + self.assertIn("(root).v2'", messages) + self.assertIn("(root).optimizer's state 'm' for (root).v1", messages) + self.assertNotIn("(root).v1'", messages) + self.assertIn("expect_partial()", messages) + + # pylint: disable=cell-var-from-loop + @test_util.run_in_graph_and_eager_modes + @test_util.run_v1_only("b/120545219") + def testWithDefun(self): + num_training_steps = 2 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + with test_util.device(use_gpu=True): + model = MyModel() + # Don't actually train so we can test variable values + optimizer = adam.Adam(0.) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + def train_fn(): + @def_function.function + def _call_model(x): + return model(x) + with backprop.GradientTape() as tape: + loss = _call_model(constant_op.constant([[3.]])) + gradients = tape.gradient(loss, model.variables) + return optimizer.apply_gradients(zip(gradients, model.variables)) + if not context.executing_eagerly(): + train_fn = functools.partial( + self.evaluate, train_fn()) + status.initialize_or_restore() + for _ in range(num_training_steps): + train_fn() + if training_continuation > 0: + status.assert_consumed() + self.assertAllClose([[42.]], self.evaluate(model.variables[0])) + else: + self.evaluate(model.variables[0].assign([[42.]])) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + self.evaluate(optimizer.iterations)) + self.assertEqual(training_continuation + 1, + self.evaluate(root.save_counter)) + # pylint: enable=cell-var-from-loop + + def testAnonymousVarsInInit(self): + + class Model(training.Model): + + def __init__(self): + super(Model, self).__init__() + self.w = resource_variable_ops.ResourceVariable(0.0) + self.b = resource_variable_ops.ResourceVariable(0.0) + self.vars = [self.w, self.b] + + def call(self, x): + return x * self.w + self.b + + with context.eager_mode(): + model = Model() + optimizer = adam.Adam(learning_rate=0.05) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + checkpoint = trackable_utils.Checkpoint( + model=model, optimizer=optimizer) + for _ in range(2): + checkpoint.save(checkpoint_prefix) + with backprop.GradientTape() as tape: + loss = (constant_op.constant(1.) + - model(constant_op.constant(1.))) ** 2 + grad = tape.gradient(loss, model.vars) + optimizer.apply_gradients( + [(g, v) for g, v in zip(grad, model.vars)]) + + @test_util.run_in_graph_and_eager_modes + def testDeferredSlotRestoration(self): + checkpoint_directory = self.get_temp_dir() + + root = trackable_utils.Checkpoint() + root.var = trackable_utils.add_variable( + root, name="var", initializer=0.) + optimizer = adam.Adam(0.1) + variables = [root.var] + gradients = [1.] + train_op = optimizer.apply_gradients(zip(gradients, variables)) + # Note that `optimizer` has not been added as a dependency of + # `root`. Create a one-off grouping so that slot variables for `root.var` + # get initialized too. + self.evaluate(trackable_utils.gather_initializers( + trackable_utils.Checkpoint(root=root, optimizer=optimizer))) + self.evaluate(train_op) + self.evaluate(state_ops.assign(root.var, 12.)) + no_slots_path = root.save(os.path.join(checkpoint_directory, "no_slots")) + root.optimizer = optimizer + self.evaluate(state_ops.assign(root.var, 13.)) + self.evaluate(state_ops.assign( + optimizer.get_slot(slot_name="m", var=root.var), + 14.)) + slots_path = root.save(os.path.join(checkpoint_directory, "with_slots")) + new_root = trackable_utils.Checkpoint() + # Load the slot-containing checkpoint (deferred), then immediately overwrite + # the non-slot variable (also deferred). + slot_status = new_root.restore(slots_path) + no_slot_status = new_root.restore(no_slots_path) + with self.assertRaises(AssertionError): + no_slot_status.assert_consumed() + new_root.var = trackable_utils.add_variable( + new_root, name="var", shape=[]) + no_slot_status.assert_consumed() + no_slot_status.run_restore_ops() + self.assertEqual(12., self.evaluate(new_root.var)) + new_root.optimizer = adam.Adam(0.1) + slot_status.assert_existing_objects_matched() + if not context.executing_eagerly(): + with self.assertRaisesRegexp(AssertionError, "Unresolved object"): + slot_status.assert_consumed() + self.assertEqual(12., self.evaluate(new_root.var)) + if context.executing_eagerly(): + # Slot variables are only created with restoring initializers when + # executing eagerly. + self.assertEqual(14., self.evaluate( + new_root.optimizer.get_slot(slot_name="m", var=new_root.var))) + else: + # Slot variables are not created eagerly when graph building. + with self.assertRaises(KeyError): + new_root.optimizer.get_slot(slot_name="m", var=new_root.var) + variables = [new_root.var] + gradients = [1.] + train_op = new_root.optimizer.apply_gradients(zip(gradients, variables)) + # The slot variable now exists; restore() didn't create it, but we should + # now have a restore op for it. + slot_status.run_restore_ops() + if not context.executing_eagerly(): + # The train op hasn't run when graph building, so the slot variable has + # its restored value. It has run in eager, so the value will be different. + self.assertEqual(14., self.evaluate( + new_root.optimizer.get_slot(slot_name="m", var=new_root.var))) + self.evaluate(train_op) + slot_status.assert_consumed() + + def testManySavesGraph(self): + """Saves after the first should not modify the graph.""" + with context.graph_mode(): + graph = ops.Graph() + with graph.as_default(), self.session(graph): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + obj = trackable_utils.Checkpoint() + obj.var = variables_lib.Variable(0., name="v") + obj.opt = adam.Adam(0.1) + variables = [obj.var] + gradients = [1.] + obj.opt.apply_gradients(zip(gradients, variables)) + self.evaluate(trackable_utils.gather_initializers(obj)) + obj.save(checkpoint_prefix) + graph.finalize() + obj.save(checkpoint_prefix) + + def testManyRestoresGraph(self): + """Restores after the first should not modify the graph.""" + with context.graph_mode(): + graph = ops.Graph() + with graph.as_default(), self.session(graph): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + obj = trackable_utils.Checkpoint() + obj.var = variables_lib.Variable(0., name="v") + obj.opt = adam.Adam(0.1) + variables = [obj.var] + gradients = [1.] + obj.opt.apply_gradients(zip(gradients, variables)) + self.evaluate(trackable_utils.gather_initializers(obj)) + save_path = obj.save(checkpoint_prefix) + obj.restore(save_path) + graph.finalize() + obj.restore(save_path) + + @test_util.run_in_graph_and_eager_modes + def test_sequential(self): + model = sequential.Sequential() + checkpoint = trackable_utils.Checkpoint(model=model) + model.add(core.Dense(4)) + second_dense = core.Dense(5) + model.add(second_dense) + model(constant_op.constant([[1.]])) + checkpoint.restore(None).initialize_or_restore() + self.evaluate(second_dense.bias.assign( + constant_op.constant([1., 2., 3., 4., 5.]))) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = checkpoint.save(checkpoint_prefix) + self.evaluate(second_dense.bias.assign( + constant_op.constant([5., 6., 7., 8., 9.]))) + checkpoint.restore(save_path).assert_consumed().run_restore_ops() + self.assertAllEqual([1., 2., 3., 4., 5.], self.evaluate(second_dense.bias)) + + deferred_sequential = sequential.Sequential() + deferred_sequential_checkpoint = trackable_utils.Checkpoint( + model=deferred_sequential) + status = deferred_sequential_checkpoint.restore(save_path) + deferred_sequential.add(core.Dense(4)) + deferred_second_dense = core.Dense(5) + deferred_sequential.add(deferred_second_dense) + deferred_sequential(constant_op.constant([[1.]])) + status.run_restore_ops() + self.assertAllEqual([1., 2., 3., 4., 5.], + self.evaluate(deferred_second_dense.bias)) + + @test_util.run_in_graph_and_eager_modes + def test_initialize_if_not_restoring(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + optimizer_only_prefix = os.path.join(checkpoint_directory, "opt") + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.Adam(0.001) + root = trackable_utils.Checkpoint( + model=model) # Do not save the optimizer with the checkpoint. + optimizer_checkpoint = trackable_utils.Checkpoint( + optimizer=optimizer) + + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + input_value = constant_op.constant([[3.]]) + def train_fn(): + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + return optimizer.apply_gradients(zip(gradients, variables)) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + # TODO(tanzheny): Add hyper variables to .variables(), and set them with + # set_weights etc. + variables_not_in_the_variables_property = [ + obj for obj in optimizer._hyper.values() + if isinstance(obj, variables_lib.Variable)] + self.evaluate([v.initializer for v + in optimizer.variables() + + variables_not_in_the_variables_property]) + train_fn() + model_save_path = root.save(file_prefix=checkpoint_prefix) + self.evaluate(optimizer.beta_1.assign(42.)) + optimizer_save_path = optimizer_checkpoint.save(optimizer_only_prefix) + del train_fn + + # Restore into a graph with the optimizer + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.Adam(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + status = root.restore(save_path=model_save_path) + input_value = constant_op.constant([[3.]]) + def train_fn1(): + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + return optimizer.apply_gradients(zip(gradients, variables)) + if not context.executing_eagerly(): + train_fn1 = functools.partial(self.evaluate, train_fn1()) + status.initialize_or_restore() + train_fn1() + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + with self.assertRaises(AssertionError): + status.assert_consumed() + del train_fn1 + + # Make sure initialization doesn't clobber later restores + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.Adam(0.001, beta_1=1.0) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + opt_root = trackable_utils.Checkpoint( + optimizer=optimizer) + status = root.restore(save_path=model_save_path) + init_only_optimizer_status = opt_root.restore(save_path=None) + optimizer_status = opt_root.restore(save_path=optimizer_save_path) + input_value = constant_op.constant([[3.]]) + def train_fn2(): + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + return optimizer.apply_gradients(zip(gradients, variables)) + if not context.executing_eagerly(): + train_fn2 = functools.partial(self.evaluate, train_fn2()) + optimizer_status.run_restore_ops() + status.initialize_or_restore() + init_only_optimizer_status.initialize_or_restore() + train_fn2() + self.assertEqual(42., self.evaluate(optimizer.beta_1)) + + +class _ManualScope(tracking.AutoTrackable): + + def __call__(self): + with variable_scope.variable_scope("ManualScope") as vs: + self.variable_scope = vs + with trackable_utils.capture_dependencies(template=self): + return self._build() + + def _build(self): + return variable_scope.get_variable(name="in_manual_scope", shape=[]) + + +class TemplateTests(parameterized.TestCase, test.TestCase): + + @test_util.run_in_graph_and_eager_modes + def test_trackable_save_restore(self): + + def _templated(): + v = variable_scope.get_variable( + "v", shape=[1], initializer=init_ops.zeros_initializer(), + use_resource=True) + v2 = variable_scope.get_variable( + "v2", shape=[1], initializer=init_ops.zeros_initializer(), + use_resource=True) + manual = _ManualScope() + return v, v + 1., v2, manual, manual() + + save_template = template.make_template("s1", _templated) + v1_save, _, v2_save, manual_scope, manual_scope_v = save_template() + six.assertCountEqual( + self, + [id(v1_save), id(v2_save), id(manual_scope), + id(manual_scope_v), id(save_template)], + map(id, trackable_utils.list_objects(save_template))) + manual_dep, = manual_scope._checkpoint_dependencies + self.assertEqual("in_manual_scope", manual_dep.name) + self.assertIs(manual_scope_v, manual_dep.ref) + optimizer = adam.Adam(0.0) + save_root = trackable_utils.Checkpoint( + my_template=save_template, optimizer=optimizer) + optimizer.minimize(v1_save.read_value, + var_list=[v1_save]) + self.evaluate([v.initializer for v in save_template.variables]) + optimizer_variables = optimizer.variables() + list( + optimizer._hyper.values()) + self.evaluate([v.initializer for v in optimizer_variables]) + self.evaluate(v1_save.assign([12.])) + self.evaluate(v2_save.assign([14.])) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + save_path = save_root.save(checkpoint_prefix) + + load_template = template.make_template("s2", _templated) + load_optimizer = adam.Adam(0.0) + load_root = trackable_utils.Checkpoint( + my_template=load_template, optimizer=load_optimizer) + status = load_root.restore(save_path) + var, var_plus_one, var2, _, _ = load_template() + load_optimizer.minimize(var.read_value, var_list=[var]) + self.assertLen(load_template._checkpoint_dependencies, 3) + self.assertEqual("v", load_template._checkpoint_dependencies[0].name) + self.assertEqual("v2", load_template._checkpoint_dependencies[1].name) + self.assertEqual("ManualScope", + load_template._checkpoint_dependencies[2].name) + status.assert_consumed().run_restore_ops() + self.assertAllEqual([12.], self.evaluate(var)) + self.assertAllEqual([13.], self.evaluate(var_plus_one)) + self.assertAllEqual([14.], self.evaluate(var2)) + + +class CheckpointCompatibilityTests(test.TestCase): + + def _initialized_model(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + optimizer = adam.Adam(0.001) + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + with backprop.GradientTape() as tape: + loss = model(input_value) + variables = model.trainable_variables + gradients = tape.gradient(loss, variables) + train_op = optimizer.apply_gradients(zip(gradients, variables)) + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + # A regular variable, a slot variable, and a non-slot Optimizer variable + # with known values to check when loading. + self.evaluate(model._named_dense.bias.assign([1.])) + self.evaluate(optimizer.get_slot( + var=model._named_dense.bias, slot_name="m").assign([2.])) + self.evaluate(optimizer.beta_1.assign(3.)) + return root_trackable + + def _set_sentinels(self, root_trackable): + self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) + self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, slot_name="m") + .assign([102.])) + self.evaluate(root_trackable.optimizer.beta_1.assign(103.)) + + def _check_sentinels(self, root_trackable): + self.assertAllEqual( + [1.], self.evaluate(root_trackable.model._named_dense.bias)) + self.assertAllEqual([2.], self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, slot_name="m"))) + self.assertAllEqual(3., + self.evaluate(root_trackable.optimizer.beta_1)) + + def _write_name_based_checkpoint(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph) as session: + root = self._initialized_model() + name_saver = saver_lib.Saver() + return name_saver.save( + sess=session, + save_path=checkpoint_prefix, + global_step=root.optimizer.iterations) + + @test_util.run_in_graph_and_eager_modes + def testLoadFromNameBasedSaver(self): + """Save a name-based checkpoint, load it using the object-based API.""" + with test_util.device(use_gpu=True): + save_path = self._write_name_based_checkpoint() + root = self._initialized_model() + self._set_sentinels(root) + with self.assertRaises(AssertionError): + self._check_sentinels(root) + object_saver = trackable_utils.TrackableSaver( + graph_view.ObjectGraphView(root)) + self._set_sentinels(root) + status = object_saver.restore(save_path) + if context.executing_eagerly(): + self._check_sentinels(root) + if context.executing_eagerly(): + status.assert_consumed() + status.assert_existing_objects_matched() + status.assert_nontrivial_match() + else: + # When graph building, we haven't read any keys, so we don't know + # whether the restore will be complete. + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_consumed() + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_existing_objects_matched() + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_nontrivial_match() + status.run_restore_ops() + self._check_sentinels(root) + self._set_sentinels(root) + status = object_saver.restore(save_path) + status.initialize_or_restore() + status.assert_nontrivial_match() + self._check_sentinels(root) + # Check that there is no error when keys are missing from the name-based + # checkpoint. + root.not_in_name_checkpoint = resource_variable_ops.ResourceVariable([1.]) + status = object_saver.restore(save_path) + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + + def testSaveGraphLoadEager(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph): + root = self._initialized_model() + save_path = root.save(file_prefix=checkpoint_prefix) + with context.eager_mode(): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed() + self._check_sentinels(root) + + def testSaveEagerLoadGraph(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.eager_mode(): + root = self._initialized_model() + save_path = root.save(file_prefix=checkpoint_prefix) + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed().run_restore_ops() + self._check_sentinels(root) + + def testIgnoreSaveCounter(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with self.cached_session() as session: + # Create and save a model using Saver() before using a Checkpoint. This + # generates a snapshot without the Checkpoint's `save_counter`. + model = sequential.Sequential() + model.add(core.Flatten(input_shape=(1,))) + model.add(core.Dense(1)) + name_saver = saver_lib.Saver(model.trainable_variables) + save_path = name_saver.save( + sess=session, save_path=checkpoint_prefix, global_step=1) + # Checkpoint.restore must successfully load that checkpoint. + ckpt = trackable_utils.Checkpoint(model=model) + status = ckpt.restore(save_path) + status.assert_existing_objects_matched() + # It should, however, refuse to load a checkpoint where an unrelated + # `save_counter` variable is missing. + model.layers[1].var = variables_lib.Variable(0., name="save_counter") + status = ckpt.restore(save_path) + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + + +if __name__ == "__main__": + ops.enable_eager_execution() + test.main() diff --git a/tensorflow/python/training/tracking/BUILD b/tensorflow/python/training/tracking/BUILD index f893e29feab..36ca3cf4b66 100644 --- a/tensorflow/python/training/tracking/BUILD +++ b/tensorflow/python/training/tracking/BUILD @@ -98,8 +98,6 @@ tf_py_test( "//tensorflow/python:math_ops", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", - "//tensorflow/python/keras:engine", - "//tensorflow/python/keras/layers", ], ) @@ -149,7 +147,6 @@ py_library( "//tensorflow/python:variables", "//tensorflow/python/eager:context", "//tensorflow/python/eager:def_function", - "//tensorflow/python/keras:backend", "//tensorflow/python/training/saving:checkpoint_options", "//tensorflow/python/training/saving:functional_saver", "//tensorflow/python/training/saving:saveable_object_util", @@ -188,10 +185,6 @@ tf_py_test( "//tensorflow/python/eager:context", "//tensorflow/python/eager:def_function", "//tensorflow/python/eager:test", - "//tensorflow/python/keras:backend", - "//tensorflow/python/keras:engine", - "//tensorflow/python/keras/layers", - "//tensorflow/python/keras/optimizer_v2", "//tensorflow/python/training/saving:checkpoint_options", "@absl_py//absl/testing:parameterized", "@six_archive//:six", diff --git a/tensorflow/python/training/tracking/data_structures_test.py b/tensorflow/python/training/tracking/data_structures_test.py index 79c88d6873a..be795601678 100644 --- a/tensorflow/python/training/tracking/data_structures_test.py +++ b/tensorflow/python/training/tracking/data_structures_test.py @@ -23,26 +23,16 @@ import os import pickle from absl.testing import parameterized -import numpy -import six from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape -from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import sequential -from tensorflow.python.keras.engine import training -from tensorflow.python.keras.layers import core -from tensorflow.python.keras.layers import normalization from tensorflow.python.layers import core as non_keras_core from tensorflow.python.module import module from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variables from tensorflow.python.training.tracking import data_structures @@ -52,184 +42,13 @@ from tensorflow.python.util import nest from tensorflow.python.util import serialization -class HasList(training.Model): - - def __init__(self): - super(HasList, self).__init__() - self.layer_list = data_structures.List([core.Dense(3)]) - self.layer_list.append(core.Dense(4)) - self.layer_list.extend( - [core.Dense(5), - core.Dense(6, kernel_regularizer=math_ops.reduce_sum)]) - self.layer_list += [ - core.Dense(7, bias_regularizer=math_ops.reduce_sum), - core.Dense(8) - ] - self.layer_list += ( - data_structures.List([core.Dense(9)]) + data_structures.List( - [core.Dense(10)])) - self.layer_list.extend( - data_structures.List( - list([core.Dense(11)]) + [core.Dense(12)])) - self.layers_with_updates = data_structures.List( - (normalization.BatchNormalization(),)) - - def call(self, x): - aggregation = 0. - for l in self.layer_list: - x = l(x) - aggregation += math_ops.reduce_sum(x) - bn, = self.layers_with_updates - return bn(x) / aggregation - - class ListTests(test.TestCase): - @test_util.run_in_graph_and_eager_modes - @test_util.run_v1_only("b/120545219") - def testTracking(self): - model = HasList() - output = model(array_ops.ones([32, 2])) - self.assertAllEqual([32, 12], output.shape) - self.assertEqual(11, len(model.layers)) - self.assertEqual(10, len(model.layer_list.layers)) - six.assertCountEqual( - self, - model.layers, - model.layer_list.layers + model.layers_with_updates) - for index in range(10): - self.assertEqual(3 + index, model.layer_list.layers[index].units) - self.assertEqual(2, len(model._checkpoint_dependencies)) - self.assertIs(model.layer_list, model._checkpoint_dependencies[0].ref) - self.assertIs(model.layers_with_updates, - model._checkpoint_dependencies[1].ref) - self.assertEqual( - 10, len(model._checkpoint_dependencies[0].ref._checkpoint_dependencies)) - self.evaluate([v.initializer for v in model.variables]) - self.evaluate(model.variables[0].assign([[1., 2., 3.], [4., 5., 6.]])) - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - self.evaluate(model.variables[0].assign(array_ops.zeros([2, 3]))) - model.load_weights(save_path) - self.assertAllEqual([[1., 2., 3.], [4., 5., 6.]], - self.evaluate(model.variables[0])) - v = variables.Variable(1.) - model.var_list = [v] - self.assertIn(v, model.variables) - self.assertIn(v, model.trainable_variables) - self.assertNotIn(v, model.non_trainable_variables) - self.assertIn(model.layer_list[0].trainable_weights[0], - model.trainable_weights) - - def testSubModelTracking(self): - model = training.Model() - model.v = variables.Variable(1.) - self.assertIn(model.v, model.trainable_weights) - model2 = training.Model() - model2.m = [model] - self.assertIn(model.v, model2.trainable_weights) - - def testSubSequentialTracking(self): - - class _Subclassed(training.Model): - - def __init__(self, wrapped): - super(_Subclassed, self).__init__() - self._wrapped = wrapped - - def call(self, x): - return self._wrapped(x) - - model = sequential.Sequential() - layer = core.Dense(1) - model.add(layer) - model2 = _Subclassed(model) - model2(array_ops.ones([1, 2])) - model2.m = [model] - self.assertIn(layer.kernel, model2.trainable_weights) - - def testLayerTrackedThroughSequential(self): - class AttrDict(dict): - - def __init__(self, *args, **kwargs): - super(AttrDict, self).__init__(*args, **kwargs) - self.__dict__ = self - - def ffnet(layer_sizes, name): - ff = sequential.Sequential(name=name) - for i, width in enumerate(layer_sizes): - ff.add(core.Dense( - width, - activation=("relu" if i < len(layer_sizes)-1 else None))) - return ff - - class MyModel2(training.Model): - - def __init__(self, config, name="my_model_2"): - super(MyModel2, self).__init__(name=name) - self._num_tokens = config.num_tokens - - # list of sub-models - self._ffnet = [ffnet(config.module_layers + (self._num_tokens,), "ff")] - - def null_input(self): - return array_ops.zeros([1, self._num_tokens], dtype=dtypes.float32) - - def call(self, input_, module_index=None): - return self._ffnet[0](input_) - - m2 = MyModel2(AttrDict( - num_tokens=5, - module_layers=(50, 30))) - - # Construct - m2(m2.null_input()) - self.assertLen(m2.trainable_variables, 6) - def testJSONSerialization(self): obj = tracking.AutoTrackable() obj.l = [1] json.dumps(obj.l, default=serialization.get_json_type) - @test_util.run_v1_only("b/120545219") - def testUpdatesForwarded(self): - with context.graph_mode(): - model = HasList() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertGreater(len(model.layers_with_updates[0].updates), 0) - self.assertEqual(set(model.layers_with_updates[0].updates), - set(model.updates)) - - with context.eager_mode(): - model = HasList() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertEqual(0, len(model.updates)) - - @test_util.run_in_graph_and_eager_modes - @test_util.run_v1_only("b/120545219") - def testLossesForwarded(self): - model = HasList() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertEqual(2, len(model.losses)) - - def testModelContainersCompareEqual(self): - class HasEqualContainers(training.Model): - - def __init__(self): - super(HasEqualContainers, self).__init__() - self.l1 = [] - self.l2 = [] - - model = HasEqualContainers() - first_layer = HasEqualContainers() - model.l1.append(first_layer) - second_layer = HasEqualContainers() - model.l2.append(second_layer) - self.assertEqual([first_layer, second_layer], model.layers) - def testNotTrackable(self): class NotTrackable(object): pass @@ -245,23 +64,6 @@ class ListTests(test.TestCase): with self.assertRaises(AttributeError): data_structures.List().pop() - @test_util.run_in_graph_and_eager_modes - def testTensorConversion(self): - - class ListToTensor(training.Model): - - def __init__(self): - super(ListToTensor, self).__init__() - self.l = [1., 2., 3.] - - self.assertAllEqual( - [1., 2., 3.], - self.evaluate(constant_op.constant(ListToTensor().l))) - - self.assertAllEqual( - [1., 2., 3.], - self.evaluate(array_ops.pack(ListToTensor().l))) - def testNesting(self): with context.graph_mode(): inner = data_structures.List() @@ -315,8 +117,7 @@ class ListTests(test.TestCase): self.assertEqual(l[:-1], [v1, v2, v3]) def testHash(self): - has_sequences = set([data_structures.List(), - data_structures.List()]) + has_sequences = {data_structures.List(), data_structures.List()} self.assertEqual(2, len(has_sequences)) self.assertNotIn(data_structures.List(), has_sequences) @@ -454,13 +255,6 @@ class ListWrapperTest(test.TestCase): l.append(1) self.assertEqual([1], l_wrapper) - def testLayerCollectionWithExternalMutation(self): - l = [] - l_wrapper = data_structures.ListWrapper(l) - layer = core.Dense(1) - l.append(layer) - self.assertEqual([layer], l_wrapper.layers) - def testNotHashable(self): with self.assertRaises(TypeError): hash(data_structures.ListWrapper()) @@ -538,50 +332,8 @@ class ListWrapperTest(test.TestCase): return l._checkpoint_dependencies # pylint: disable=protected-access -class HasMapping(training.Model): - - def __init__(self): - super(HasMapping, self).__init__() - self.layer_dict = data_structures.Mapping(output=core.Dense(7)) - self.layer_dict["norm"] = data_structures.List() - self.layer_dict["dense"] = data_structures.List() - self.layer_dict["dense"].extend( - [core.Dense(5), - core.Dense(6, kernel_regularizer=math_ops.reduce_sum)]) - self.layer_dict["norm"].append( - normalization.BatchNormalization()) - self.layer_dict["norm"].append( - normalization.BatchNormalization()) - - def call(self, x): - aggregation = 0. - for norm, dense in zip(self.layer_dict["norm"], self.layer_dict["dense"]): - x = norm(dense(x)) - aggregation += math_ops.reduce_sum(x) - return self.layer_dict["output"](x) / aggregation - - class MappingTests(test.TestCase): - @test_util.run_in_graph_and_eager_modes - def testTracking(self): - model = HasMapping() - output = model(array_ops.ones([32, 2])) - self.assertAllEqual([32, 7], output.shape.as_list()) - self.assertEqual(5, len(model.layers)) - six.assertCountEqual(self, model.layers, model.layer_dict.layers) - self.assertEqual(1, len(model._checkpoint_dependencies)) - self.assertIs(model.layer_dict, model._checkpoint_dependencies[0].ref) - self.evaluate([v.initializer for v in model.variables]) - test_var = model.layer_dict["output"].kernel - self.evaluate(test_var.assign(array_ops.ones([6, 7]))) - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - self.evaluate(test_var.assign(array_ops.zeros([6, 7]))) - model.load_weights(save_path) - self.assertAllEqual(numpy.ones([6, 7]), - self.evaluate(test_var)) - def testJSONSerialization(self): obj = tracking.AutoTrackable() obj.d = {"a": 2} @@ -605,20 +357,6 @@ class MappingTests(test.TestCase): with self.assertRaises(TypeError): mapping[1] = data_structures.List() - def testLayerCollectionWithExternalMutation(self): - d = {} - root = tracking.AutoTrackable() - root.wrapper = d - self.assertEqual([], root.wrapper.layers) - self.assertEqual([], root.wrapper.trainable_weights) - layer1 = core.Dense(1) - layer2 = core.Dense(1) - d["a"] = layer1 - d["b"] = layer2 - self.assertEqual([layer1, layer2], root.wrapper.layers) - # The layers have still not created variables - self.assertEqual([], root.wrapper.trainable_weights) - def testHashing(self): has_mappings = set([data_structures.Mapping(), data_structures.Mapping()]) @@ -633,101 +371,6 @@ class MappingTests(test.TestCase): with self.assertRaisesRegexp(TypeError, "unhashable"): set([a.d]) - def testDictWrapperBadKeys(self): - a = tracking.AutoTrackable() - a.d = {} - a.d[1] = data_structures.List() - model = training.Model() - model.sub = a - save_path = os.path.join(self.get_temp_dir(), "ckpt") - with self.assertRaisesRegexp(ValueError, "non-string key"): - model.save_weights(save_path) - - def testDictWrapperNoDependency(self): - a = tracking.AutoTrackable() - a.d = data_structures.NoDependency({}) - a.d[1] = [3] - self.assertEqual([a], util.list_objects(a)) - model = training.Model() - model.sub = a - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - model.load_weights(save_path) - - def testNonStringKeyNotTrackableValue(self): - a = tracking.AutoTrackable() - a.d = {} - a.d["a"] = [3] - a.d[1] = data_structures.NoDependency([3]) - self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) - model = training.Model() - model.sub = a - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - model.load_weights(save_path) - - def testNonAppendNotTrackable(self): - # Non-append mutations (deleting or overwriting values) are OK when the - # values aren't tracked. - a = tracking.AutoTrackable() - a.d = {} - a.d["a"] = [3] - a.d[1] = 3 - a.d[1] = 2 - self.assertEqual(2, a.d[1]) - del a.d[1] - a.d[2] = data_structures.NoDependency(tracking.AutoTrackable()) - second = tracking.AutoTrackable() - a.d[2] = data_structures.NoDependency(second) - self.assertIs(second, a.d[2]) - self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) - model = training.Model() - model.sub = a - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - model.load_weights(save_path) - - def testPopNoSave(self): - model = training.Model() - model.d = {} - model.d["a"] = [] - model.d.pop("a") - save_path = os.path.join(self.get_temp_dir(), "ckpt") - with self.assertRaisesRegexp(ValueError, "Unable to save"): - model.save_weights(save_path) - - def testExternalModificationNoSave(self): - model = training.Model() - external_reference = {} - model.d = external_reference - external_reference["a"] = [] - save_path = os.path.join(self.get_temp_dir(), "ckpt") - with self.assertRaisesRegexp(ValueError, "modified outside the wrapper"): - model.save_weights(save_path) - - def testOverwriteCanStillSave(self): - model = training.Model() - model.d = {} - model.d["a"] = {} - model.d["a"] = {} - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - - def testIter(self): - model = training.Model() - model.d = {1: 3} - model.d[1] = 3 - self.assertEqual([1], list(model.d)) - new_dict = {} - # This update() is super tricky. If the dict wrapper subclasses dict, - # CPython will access its storage directly instead of calling any - # methods/properties on the object. So the options are either not to - # subclass dict (in which case update will call normal iter methods, but the - # object won't pass isinstance checks) or to subclass dict and keep that - # storage updated (no shadowing all its methods like ListWrapper). - new_dict.update(model.d) - self.assertEqual({1: 3}, new_dict) - def testListShallowCopy(self): root = tracking.AutoTrackable() orig_list = [[1.]] @@ -871,157 +514,13 @@ class MappingTests(test.TestCase): self.assertIs(first_trace, second_trace) -class HasTuple(training.Model): - - def __init__(self): - super(HasTuple, self).__init__() - self.layer_list = ( - core.Dense(3), core.Dense(4), - core.Dense(5, kernel_regularizer=math_ops.reduce_sum)) - self.layers_with_updates = (normalization.BatchNormalization(),) - - def call(self, x): - aggregation = 0. - for l in self.layer_list: - x = l(x) - aggregation += math_ops.reduce_sum(x) - bn, = self.layers_with_updates - return bn(x) / aggregation - - class TupleTests(test.TestCase, parameterized.TestCase): - @test_util.run_in_graph_and_eager_modes - def testTracking(self): - model = HasTuple() - output = model(array_ops.ones([32, 2])) - self.assertAllEqual([32, 5], output.shape.as_list()) - self.assertLen(model.layers, 4) - self.assertLen(model.layer_list.layers, 3) - six.assertCountEqual( - self, - model.layers, - tuple(model.layer_list.layers) + model.layers_with_updates) - self.assertEqual(3, model.layer_list.layers[0].units) - self.assertEqual(4, model.layer_list.layers[1].units) - self.assertEqual(5, model.layer_list.layers[2].units) - self.assertLen(model._checkpoint_dependencies, 2) - self.assertIs(model.layer_list, model._checkpoint_dependencies[0].ref) - self.assertIs(model.layers_with_updates, - model._checkpoint_dependencies[1].ref) - self.assertLen( - model._checkpoint_dependencies[0].ref._checkpoint_dependencies, 3) - self.evaluate([v.initializer for v in model.variables]) - self.evaluate(model.variables[0].assign([[1., 2., 3.], [4., 5., 6.]])) - save_path = os.path.join(self.get_temp_dir(), "ckpt") - model.save_weights(save_path) - self.evaluate(model.variables[0].assign(array_ops.zeros([2, 3]))) - model.load_weights(save_path) - self.assertAllEqual([[1., 2., 3.], [4., 5., 6.]], - self.evaluate(model.variables[0])) - v = variables.Variable(1.) - model.var_list = (v,) - self.assertIn(id(v), [id(obj) for obj in model.variables]) - self.assertIn(id(v), [id(obj) for obj in model.trainable_variables]) - self.assertNotIn(id(v), [id(obj) for obj in model.non_trainable_variables]) - self.assertIn(id(model.layer_list[0].trainable_weights[0]), - [id(obj) for obj in model.trainable_weights]) - - @parameterized.named_parameters( - ("Module", module.Module), - ("Model", training.Model), - ) - def testSubModelTracking(self, module_subclass): - model = module_subclass() - model.v = variables.Variable(1.) - self.assertIn(model.v, model.trainable_variables) - model2 = module_subclass() - model2.m = (model,) - self.assertIn(model.v, model2.trainable_variables) - - def testSubSequentialTracking(self): - - class _Subclassed(training.Model): - - def __init__(self, wrapped): - super(_Subclassed, self).__init__() - self._wrapped = wrapped - - def call(self, x): - return self._wrapped(x) - - model = sequential.Sequential() - layer = core.Dense(1) - model.add(layer) - model2 = _Subclassed(model) - model2(array_ops.ones([1, 2])) - model2.m = (model,) - self.assertIn(layer.kernel, model2.trainable_weights) - def testJSONSerialization(self): obj = tracking.AutoTrackable() obj.l = (1,) json.dumps(obj.l, default=serialization.get_json_type) - def testUpdatesForwarded(self): - with ops.Graph().as_default(): - model = HasTuple() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertNotEmpty(model.layers_with_updates[0].updates) - self.assertEqual(set(model.layers_with_updates[0].updates), - set(model.updates)) - - model = HasTuple() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertEmpty(model.updates) - - @test_util.run_in_graph_and_eager_modes - def testLossesForwarded(self): - model = HasTuple() - model_input = array_ops.ones([32, 2]) - model(model_input) - self.assertLen(model.losses, 1) - - def testModelContainersCompareEqual(self): - class HasEqualContainers(training.Model): - - def __init__(self): - super(HasEqualContainers, self).__init__() - self.l1 = () - self.l2 = () - - model = HasEqualContainers() - first_layer = HasEqualContainers() - model.l1 = (first_layer,) - second_layer = HasEqualContainers() - model.l2 = (second_layer,) - self.assertEqual((first_layer,), model.l1) - d = {model.l1: 1, model.l2: 2} - self.assertEqual(1, d[model.l1]) - self.assertEqual(1, d[(first_layer,)]) - self.assertEqual(2, d[model.l2]) - self.assertEqual(2, d[(second_layer,)]) - self.assertEqual([first_layer, second_layer], model.layers) - - @test_util.run_in_graph_and_eager_modes - def testTensorConversion(self): - - class TupleToTensor(training.Model): - - def __init__(self): - super(TupleToTensor, self).__init__() - self.l = (1., 2., 3.) - - self.assertAllEqual( - (1., 2., 3.), - self.evaluate(constant_op.constant(TupleToTensor().l))) - - self.assertAllEqual( - (1., 2., 3.), - self.evaluate(array_ops.pack(TupleToTensor().l))) - def testNonLayerVariables(self): v = resource_variable_ops.ResourceVariable([1.]) l = data_structures._TupleWrapper((v,)) diff --git a/tensorflow/python/training/tracking/tracking_test.py b/tensorflow/python/training/tracking/tracking_test.py index cf2da4c9afa..4dff392cf9f 100644 --- a/tensorflow/python/training/tracking/tracking_test.py +++ b/tensorflow/python/training/tracking/tracking_test.py @@ -25,14 +25,10 @@ import time import timeit import numpy as np -import six from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import training -from tensorflow.python.module import module from tensorflow.python.ops import array_ops from tensorflow.python.platform import test -from tensorflow.python.training.tracking import base from tensorflow.python.training.tracking import data_structures from tensorflow.python.training.tracking import tracking from tensorflow.python.training.tracking import util @@ -73,28 +69,6 @@ class InterfaceTests(test.TestCase): (_, dep_object), = root._checkpoint_dependencies self.assertIs(duplicate_name_dep, dep_object) - def testNoDependency(self): - root = tracking.AutoTrackable() - hasdep = tracking.AutoTrackable() - root.hasdep = hasdep - nodep = tracking.AutoTrackable() - root.nodep = data_structures.NoDependency(nodep) - self.assertEqual(1, len(root._checkpoint_dependencies)) - self.assertIs(root._checkpoint_dependencies[0].ref, root.hasdep) - self.assertIs(root.hasdep, hasdep) - self.assertIs(root.nodep, nodep) - - class NoDependencyModel(training.Model): - - @base.no_automatic_dependency_tracking - def __init__(self): - super(NoDependencyModel, self).__init__() - self.a = [] - self.b = tracking.AutoTrackable() - - nodeps = NoDependencyModel() - self.assertEqual([nodeps], util.list_objects(nodeps)) - def testRemoveDependency(self): root = tracking.AutoTrackable() root.a = tracking.AutoTrackable() @@ -183,43 +157,6 @@ class InterfaceTests(test.TestCase): with self.assertRaisesRegexp(ValueError, "A list element was replaced"): checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) - @test_util.run_in_graph_and_eager_modes - def testDictionariesBasic(self): - a = training.Model() - b = training.Model() - a.attribute = {"b": b} - c = training.Model() - a.attribute["c"] = [] - a.attribute["c"].append(c) - a_deps = util.list_objects(a) - self.assertIn(b, a_deps) - self.assertIn(c, a_deps) - self.assertIs(b, a.attribute["b"]) - six.assertCountEqual( - self, - ["b", "c"], - [dep.name for dep in a.attribute._checkpoint_dependencies]) - self.assertEqual([b, c], a.layers) - self.assertEqual([b, c], a.attribute.layers) - self.assertEqual([c], a.attribute["c"].layers) - checkpoint = util.Checkpoint(a=a) - save_path = checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) - with self.cached_session(): - checkpoint.restore(save_path).assert_consumed().initialize_or_restore() - - @test_util.run_in_graph_and_eager_modes - def testNoDepList(self): - a = training.Model() - a.l1 = data_structures.NoDependency([]) - a.l1.insert(1, 0) - self.assertTrue(isinstance(a.l1, list)) - checkpoint = util.Checkpoint(a=a) - checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) - a.l2 = [] - a.l2.insert(1, module.Module()) - with self.assertRaisesRegexp(ValueError, "A list element was replaced"): - checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) - @test_util.run_in_graph_and_eager_modes def testAssertions(self): a = tracking.AutoTrackable() diff --git a/tensorflow/python/training/tracking/util_test.py b/tensorflow/python/training/tracking/util_test.py index 7a96fedc89b..6c0b08426e7 100644 --- a/tensorflow/python/training/tracking/util_test.py +++ b/tensorflow/python/training/tracking/util_test.py @@ -16,25 +16,18 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import functools import os import weakref from absl.testing import parameterized import six -from tensorflow.python.eager import backprop from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import input_layer -from tensorflow.python.keras.engine import sequential -from tensorflow.python.keras.engine import training -from tensorflow.python.keras.layers import core -from tensorflow.python.keras.optimizer_v2 import adam from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import resource_variable_ops @@ -46,7 +39,6 @@ from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import checkpoint_management from tensorflow.python.training import saver as saver_lib -from tensorflow.python.training import training_util from tensorflow.python.training.saving import checkpoint_options from tensorflow.python.training.tracking import base from tensorflow.python.training.tracking import graph_view @@ -62,44 +54,8 @@ class NonLayerTrackable(tracking.AutoTrackable): self, name="a_variable", shape=[]) -# pylint: disable=not-callable -class MyModel(training.Model): - """A concrete Model for testing.""" - - def __init__(self): - super(MyModel, self).__init__() - self._named_dense = core.Dense(1, use_bias=True) - self._second = core.Dense(1, use_bias=False) - # We can still track Trackables which aren't Layers. - self._non_layer = NonLayerTrackable() - - def call(self, values): - ret = self._second(self._named_dense(values)) - return ret - - class InterfaceTests(test.TestCase): - def testLayerDeduplication(self): - model = training.Model() - layer_one = core.Dense(1) - layer_two = core.Dense(1) - model.other_path = [layer_one, layer_two] - model.l2 = layer_two - model.l1 = layer_one - self.assertEqual([layer_one, layer_two], model.layers) - - def testSaveWithOnlyKerasSession(self): - - with ops.Graph().as_default(): - inp = input_layer.Input([1]) - dense = core.Dense(1)(inp) - model = training.Model(inp, dense) - model.compile(optimizer="sgd", loss="mse") - model.fit([1.], [2.]) - checkpoint = trackable_utils.Checkpoint(model=model) - checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) def testAddVariable(self): obj = NonLayerTrackable() @@ -184,22 +140,6 @@ class InterfaceTests(test.TestCase): self.assertEqual(dtypes.float64, v2.dtype) self.assertAllEqual([1., 1., 1.], self.evaluate(v2)) - def testObjectMetadata(self): - with context.eager_mode(): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - dense = core.Dense(1) - checkpoint = trackable_utils.Checkpoint(dense=dense) - dense(constant_op.constant([[1.]])) - save_path = checkpoint.save(checkpoint_prefix) - - objects = trackable_utils.object_metadata(save_path) - all_variable_names = [] - for obj in objects.nodes: - for attribute in obj.attributes: - all_variable_names.append(attribute.full_name) - self.assertIn("dense/kernel", all_variable_names) - def testNotTrackable(self): class CallsFunctionalStuff( @@ -268,100 +208,6 @@ class _OwnsMirroredVariables(base.Trackable): class CheckpointingTests(parameterized.TestCase, test.TestCase): - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) - def testNamingWithOptimizer(self): - input_value = constant_op.constant([[3.]]) - model = MyModel() - # A nuisance Model using the same optimizer. Its slot variables should not - # go in the checkpoint, since it is never depended on. - other_model = MyModel() - optimizer = adam.Adam(0.001) - step = training_util.get_or_create_global_step() - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, step=step) - - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - train_op = control_flow_ops.group( - optimizer.apply_gradients(zip(gradients, variables)), - step.assign_add(1)) - - with backprop.GradientTape() as tape: - loss = other_model(input_value) - variables = other_model.trainable_variables - gradients = tape.gradient(loss, variables) - optimizer.apply_gradients(zip(gradients, variables)) - - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - named_variables, serialized_graph, _ = graph_view.ObjectGraphView( - root_trackable).serialize_object_graph() - expected_slot_keys = ( - "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", - "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", - "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", - "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", - "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", - "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v", - ) - expected_checkpoint_names = ( - # Created in the root node, so no prefix. - "step", - "model/_second/kernel", - "model/_named_dense/kernel", - "model/_named_dense/bias", - # non-Layer dependency of the model - "model/_non_layer/a_variable", - "optimizer/learning_rate", - "optimizer/beta_1", - "optimizer/beta_2", - "optimizer/iter", - "optimizer/decay", - ) + expected_slot_keys - suffix = "/.ATTRIBUTES/VARIABLE_VALUE" - expected_checkpoint_names = [ - name + suffix for name in expected_checkpoint_names] - named_variables = {v.name: v for v in named_variables} - six.assertCountEqual(self, expected_checkpoint_names, - named_variables.keys()) - # Check that we've mapped to the right variable objects (not exhaustive) - self.assertEqual( - "global_step", - named_variables["step" + suffix].full_name) - self.assertEqual( - "my_model/dense_1/kernel", - named_variables["model/_second/kernel" + suffix].full_name) - self.assertEqual( - "my_model/dense/kernel", - named_variables["model/_named_dense/kernel" + suffix].full_name) - self.assertEqual("Adam/beta_1", - named_variables["optimizer/beta_1" + suffix].full_name) - self.assertEqual("Adam/beta_2", - named_variables["optimizer/beta_2" + suffix].full_name) - # Spot check the generated protocol buffers. - self.assertEqual("optimizer", - serialized_graph.nodes[0].children[1].local_name) - optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ - 1].node_id] - children = [node.local_name for node in optimizer_node.children] - six.assertCountEqual( - self, - # hyper variable dependencies - ["beta_1", "beta_2", "iter", "decay", "learning_rate"], - children) - serialized_slot_keys = [] - for slot in optimizer_node.slot_variables: - for attribute in ( - serialized_graph.nodes[slot.slot_variable_node_id].attributes): - serialized_slot_keys.append(attribute.checkpoint_key) - six.assertCountEqual( - self, - [key + suffix for key in expected_slot_keys], - serialized_slot_keys) - @test_util.run_in_graph_and_eager_modes def testMoreComplexSaveableReturned(self): v = _OwnsMirroredVariables() @@ -432,174 +278,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): if op.type in ("SaveV2", "RestoreV2"): self.assertEqual(localhost, op.device) - @test_util.run_in_graph_and_eager_modes - def testSaveRestore(self): - model = MyModel() - optimizer = adam.Adam(0.001) - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - input_value = constant_op.constant([[3.]]) - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - train_op = optimizer.apply_gradients(zip(gradients, variables)) - self.assertFalse(root_trackable.save_counter.trainable) - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - prefix = os.path.join(self.get_temp_dir(), "ckpt") - self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.])) - m_bias_slot = optimizer.get_slot(model._named_dense.variables[1], "m") - self.evaluate(state_ops.assign(m_bias_slot, [1.5])) - save_path = root_trackable.save(file_prefix=prefix) - self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.])) - self.evaluate(state_ops.assign(root_trackable.save_counter, 3)) - optimizer_variables = self.evaluate( - sorted(optimizer.variables(), key=lambda v: v.name)) - self.evaluate(state_ops.assign(m_bias_slot, [-2.])) - # Immediate restoration - status = root_trackable.restore(save_path=save_path).assert_consumed() - status.run_restore_ops() - self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1])) - self.assertAllEqual(1, self.evaluate(root_trackable.save_counter)) - self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) - if not context.executing_eagerly(): - return # Restore-on-create is only supported when executing eagerly - on_create_model = MyModel() - on_create_optimizer = adam.Adam(0.001) - on_create_root = trackable_utils.Checkpoint( - optimizer=on_create_optimizer, model=on_create_model) - # Deferred restoration - status = on_create_root.restore(save_path=save_path) - status.assert_nontrivial_match() - status.assert_existing_objects_matched() - with self.assertRaises(AssertionError): - status.assert_consumed() - on_create_model(constant_op.constant([[3.]])) # create variables - self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) - self.assertAllEqual([42.], - self.evaluate( - on_create_model._named_dense.variables[1])) - on_create_m_bias_slot = on_create_optimizer.get_slot( - on_create_model._named_dense.variables[1], "m") - status.assert_existing_objects_matched() - if not context.executing_eagerly(): - with self.assertRaises(AssertionError): - status.assert_consumed() - # Optimizer slot variables are created when the original variable is - # restored. - self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot)) - dummy_var = resource_variable_ops.ResourceVariable([1.]) - on_create_optimizer.minimize(loss=dummy_var.read_value, - var_list=[dummy_var]) - status.assert_existing_objects_matched() - status.assert_consumed() - self.assertAllEqual( - optimizer_variables, - # Creation order is different, so .variables() needs to be re-sorted. - self.evaluate(sorted(optimizer.variables(), key=lambda v: v.name))) - - # TODO(allenl): Debug garbage created by this test in python3. - def testDeferredRestorationUsageEager(self): - """An idiomatic eager execution example.""" - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - model = MyModel() - optimizer = adam.Adam(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - root.restore(checkpoint_management.latest_checkpoint( - checkpoint_directory)) - for _ in range(num_training_steps): - # TODO(allenl): Use a Dataset and serialize/checkpoint it. - input_value = constant_op.constant([[3.]]) - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - optimizer.apply_gradients(zip(gradients, variables)) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - root.optimizer.iterations.numpy()) - - def testUsageGraph(self): - """Expected usage when graph building.""" - with context.graph_mode(): - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - with ops.Graph().as_default(): - model = MyModel() - optimizer = adam.Adam(0.001) - root = trackable_utils.CheckpointV1( - optimizer=optimizer, model=model) - input_value = constant_op.constant([[3.]]) - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - train_op = optimizer.apply_gradients(zip(gradients, variables)) - - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - with self.session(graph=ops.get_default_graph()) as session: - status = root.restore(save_path=checkpoint_path) - status.initialize_or_restore(session=session) - if checkpoint_path is None: - self.assertEqual(0, training_continuation) - with self.assertRaises(AssertionError): - status.assert_consumed() - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - else: - status.assert_consumed() - status.assert_existing_objects_matched() - for _ in range(num_training_steps): - session.run(train_op) - root.save(file_prefix=checkpoint_prefix, session=session) - self.assertEqual((training_continuation + 1) * num_training_steps, - session.run(root.optimizer.iterations)) - self.assertEqual(training_continuation + 1, - session.run(root.save_counter)) - - @test_util.run_in_graph_and_eager_modes - def testAgnosticUsage(self): - """Graph/eager agnostic usage.""" - # Does create garbage when executing eagerly due to ops.Graph() creation. - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - def _train_fn(model, input_value): - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - return optimizer.apply_gradients(zip(gradients, variables)) - for training_continuation in range(3): - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.Adam(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - manager = checkpoint_management.CheckpointManager( - root, checkpoint_directory, max_to_keep=1) - status = root.restore(save_path=manager.latest_checkpoint) - input_value = constant_op.constant([[3.]]) - train_fn = functools.partial(_train_fn, model, input_value) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - status.initialize_or_restore() - for _ in range(num_training_steps): - train_fn() - manager.save() - self.assertEqual((training_continuation + 1) * num_training_steps, - self.evaluate(root.optimizer.iterations)) - self.assertEqual(training_continuation + 1, - self.evaluate(root.save_counter)) - @test_util.run_in_graph_and_eager_modes def testFreezing(self): with test_util.use_gpu(): @@ -656,31 +334,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): self.fail("%s should have suffix %s" % (path, expected_suffix)) self.evaluate(step.assign_add(2)) - def testPartialRestoreWarningObject(self): - with context.eager_mode(): - optimizer = adam.Adam(0.0) - original_root = trackable_utils.Checkpoint(v1=variables_lib.Variable(2.), - v2=variables_lib.Variable(3.), - optimizer=optimizer) - # Create a slot variable to save - optimizer.minimize(original_root.v1.read_value, [original_root.v1]) - prefix = os.path.join(self.get_temp_dir(), "ckpt") - save_path = original_root.save(prefix) - partial_root = trackable_utils.Checkpoint(v1=variables_lib.Variable(0.)) - weak_partial_root = weakref.ref(partial_root) - weak_v1 = weakref.ref(partial_root.v1) - partial_root.restore(save_path) - self.assertEqual(2., partial_root.v1.numpy()) - with test.mock.patch.object(logging, "warning") as mock_log: - del partial_root - self.assertIsNone(weak_partial_root()) - self.assertIsNone(weak_v1()) - messages = str(mock_log.call_args_list) - self.assertIn("(root).v2'", messages) - self.assertIn("(root).optimizer's state 'm' for (root).v1", messages) - self.assertNotIn("(root).v1'", messages) - self.assertIn("expect_partial()", messages) - def testPartialRestoreWarningAttribute(self): with context.eager_mode(): original_root = trackable_utils.Checkpoint(v1=variables_lib.Variable(2.), @@ -734,49 +387,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): self.assertIsNone(weak_v1()) self.assertEmpty(mock_log.call_args_list) - # pylint: disable=cell-var-from-loop - @test_util.run_in_graph_and_eager_modes - @test_util.run_v1_only("b/120545219") - def testWithDefun(self): - num_training_steps = 2 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - with test_util.device(use_gpu=True): - model = MyModel() - # Don't actually train so we can test variable values - optimizer = adam.Adam(0.) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - status = root.restore(save_path=checkpoint_path) - def train_fn(): - @def_function.function - def _call_model(x): - return model(x) - with backprop.GradientTape() as tape: - loss = _call_model(constant_op.constant([[3.]])) - gradients = tape.gradient(loss, model.variables) - return optimizer.apply_gradients(zip(gradients, model.variables)) - if not context.executing_eagerly(): - train_fn = functools.partial( - self.evaluate, train_fn()) - status.initialize_or_restore() - for _ in range(num_training_steps): - train_fn() - if training_continuation > 0: - status.assert_consumed() - self.assertAllClose([[42.]], self.evaluate(model.variables[0])) - else: - self.evaluate(model.variables[0].assign([[42.]])) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - self.evaluate(optimizer.iterations)) - self.assertEqual(training_continuation + 1, - self.evaluate(root.save_counter)) - # pylint: enable=cell-var-from-loop - def _get_checkpoint_name(self, name): root = tracking.AutoTrackable() trackable_utils.add_variable( @@ -819,35 +429,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): self.assertEqual("..ATTRIBUTES/a/.ATTRIBUTES/VARIABLE_VALUE", named_variable.name) - def testAnonymousVarsInInit(self): - - class Model(training.Model): - - def __init__(self): - super(Model, self).__init__() - self.w = resource_variable_ops.ResourceVariable(0.0) - self.b = resource_variable_ops.ResourceVariable(0.0) - self.vars = [self.w, self.b] - - def call(self, x): - return x * self.w + self.b - - with context.eager_mode(): - model = Model() - optimizer = adam.Adam(learning_rate=0.05) - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - checkpoint = trackable_utils.Checkpoint( - model=model, optimizer=optimizer) - for _ in range(2): - checkpoint.save(checkpoint_prefix) - with backprop.GradientTape() as tape: - loss = (constant_op.constant(1.) - - model(constant_op.constant(1.))) ** 2 - grad = tape.gradient(loss, model.vars) - optimizer.apply_gradients( - [(g, v) for g, v in zip(grad, model.vars)]) - @test_util.run_in_graph_and_eager_modes def testLateDependencyTracking(self): @@ -909,72 +490,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): status.run_restore_ops() self.assertEqual(-14., self.evaluate(loaded_dep_after_var.dep.var)) - @test_util.run_in_graph_and_eager_modes - def testDeferredSlotRestoration(self): - checkpoint_directory = self.get_temp_dir() - - root = trackable_utils.Checkpoint() - root.var = trackable_utils.add_variable( - root, name="var", initializer=0.) - optimizer = adam.Adam(0.1) - variables = [root.var] - gradients = [1.] - train_op = optimizer.apply_gradients(zip(gradients, variables)) - # Note that `optimizer` has not been added as a dependency of - # `root`. Create a one-off grouping so that slot variables for `root.var` - # get initialized too. - self.evaluate(trackable_utils.gather_initializers( - trackable_utils.Checkpoint(root=root, optimizer=optimizer))) - self.evaluate(train_op) - self.evaluate(state_ops.assign(root.var, 12.)) - no_slots_path = root.save(os.path.join(checkpoint_directory, "no_slots")) - root.optimizer = optimizer - self.evaluate(state_ops.assign(root.var, 13.)) - self.evaluate(state_ops.assign( - optimizer.get_slot(slot_name="m", var=root.var), - 14.)) - slots_path = root.save(os.path.join(checkpoint_directory, "with_slots")) - new_root = trackable_utils.Checkpoint() - # Load the slot-containing checkpoint (deferred), then immediately overwrite - # the non-slot variable (also deferred). - slot_status = new_root.restore(slots_path) - no_slot_status = new_root.restore(no_slots_path) - with self.assertRaises(AssertionError): - no_slot_status.assert_consumed() - new_root.var = trackable_utils.add_variable( - new_root, name="var", shape=[]) - no_slot_status.assert_consumed() - no_slot_status.run_restore_ops() - self.assertEqual(12., self.evaluate(new_root.var)) - new_root.optimizer = adam.Adam(0.1) - slot_status.assert_existing_objects_matched() - if not context.executing_eagerly(): - with self.assertRaisesRegexp(AssertionError, "Unresolved object"): - slot_status.assert_consumed() - self.assertEqual(12., self.evaluate(new_root.var)) - if context.executing_eagerly(): - # Slot variables are only created with restoring initializers when - # executing eagerly. - self.assertEqual(14., self.evaluate( - new_root.optimizer.get_slot(slot_name="m", var=new_root.var))) - else: - # Slot variables are not created eagerly when graph building. - with self.assertRaises(KeyError): - new_root.optimizer.get_slot(slot_name="m", var=new_root.var) - variables = [new_root.var] - gradients = [1.] - train_op = new_root.optimizer.apply_gradients(zip(gradients, variables)) - # The slot variable now exists; restore() didn't create it, but we should - # now have a restore op for it. - slot_status.run_restore_ops() - if not context.executing_eagerly(): - # The train op hasn't run when graph building, so the slot variable has - # its restored value. It has run in eager, so the value will be different. - self.assertEqual(14., self.evaluate( - new_root.optimizer.get_slot(slot_name="m", var=new_root.var))) - self.evaluate(train_op) - slot_status.assert_consumed() - @test_util.run_in_graph_and_eager_modes def testOverlappingRestores(self): checkpoint_directory = self.get_temp_dir() @@ -1154,24 +669,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): status.run_restore_ops() self.assertEqual(4., self.evaluate(recreated_var1)) - def testManySavesGraph(self): - """Saves after the first should not modify the graph.""" - with context.graph_mode(): - graph = ops.Graph() - with graph.as_default(), self.session(graph): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - obj = trackable_utils.Checkpoint() - obj.var = variables_lib.Variable(0., name="v") - obj.opt = adam.Adam(0.1) - variables = [obj.var] - gradients = [1.] - obj.opt.apply_gradients(zip(gradients, variables)) - self.evaluate(trackable_utils.gather_initializers(obj)) - obj.save(checkpoint_prefix) - graph.finalize() - obj.save(checkpoint_prefix) - @test_util.run_in_graph_and_eager_modes def testCheckpointState(self): # No checkpoints are deleted by default @@ -1237,146 +734,6 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): self.assertEqual(1, self.evaluate(checkpoint.var_1)) self.assertEqual(0, self.evaluate(checkpoint.var_0)) - def testManyRestoresGraph(self): - """Restores after the first should not modify the graph.""" - with context.graph_mode(): - graph = ops.Graph() - with graph.as_default(), self.session(graph): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - obj = trackable_utils.Checkpoint() - obj.var = variables_lib.Variable(0., name="v") - obj.opt = adam.Adam(0.1) - variables = [obj.var] - gradients = [1.] - obj.opt.apply_gradients(zip(gradients, variables)) - self.evaluate(trackable_utils.gather_initializers(obj)) - save_path = obj.save(checkpoint_prefix) - obj.restore(save_path) - graph.finalize() - obj.restore(save_path) - - @test_util.run_in_graph_and_eager_modes - def test_sequential(self): - model = sequential.Sequential() - checkpoint = trackable_utils.Checkpoint(model=model) - model.add(core.Dense(4)) - second_dense = core.Dense(5) - model.add(second_dense) - model(constant_op.constant([[1.]])) - checkpoint.restore(None).initialize_or_restore() - self.evaluate(second_dense.bias.assign( - constant_op.constant([1., 2., 3., 4., 5.]))) - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - save_path = checkpoint.save(checkpoint_prefix) - self.evaluate(second_dense.bias.assign( - constant_op.constant([5., 6., 7., 8., 9.]))) - checkpoint.restore(save_path).assert_consumed().run_restore_ops() - self.assertAllEqual([1., 2., 3., 4., 5.], self.evaluate(second_dense.bias)) - - deferred_sequential = sequential.Sequential() - deferred_sequential_checkpoint = trackable_utils.Checkpoint( - model=deferred_sequential) - status = deferred_sequential_checkpoint.restore(save_path) - deferred_sequential.add(core.Dense(4)) - deferred_second_dense = core.Dense(5) - deferred_sequential.add(deferred_second_dense) - deferred_sequential(constant_op.constant([[1.]])) - status.run_restore_ops() - self.assertAllEqual([1., 2., 3., 4., 5.], - self.evaluate(deferred_second_dense.bias)) - - @test_util.run_in_graph_and_eager_modes - def test_initialize_if_not_restoring(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - optimizer_only_prefix = os.path.join(checkpoint_directory, "opt") - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.Adam(0.001) - root = trackable_utils.Checkpoint( - model=model) # Do not save the optimizer with the checkpoint. - optimizer_checkpoint = trackable_utils.Checkpoint( - optimizer=optimizer) - - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - status = root.restore(save_path=checkpoint_path) - input_value = constant_op.constant([[3.]]) - def train_fn(): - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - return optimizer.apply_gradients(zip(gradients, variables)) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - status.initialize_or_restore() - # TODO(tanzheny): Add hyper variables to .variables(), and set them with - # set_weights etc. - variables_not_in_the_variables_property = [ - obj for obj in optimizer._hyper.values() - if isinstance(obj, variables_lib.Variable)] - self.evaluate([v.initializer for v - in optimizer.variables() - + variables_not_in_the_variables_property]) - train_fn() - model_save_path = root.save(file_prefix=checkpoint_prefix) - self.evaluate(optimizer.beta_1.assign(42.)) - optimizer_save_path = optimizer_checkpoint.save(optimizer_only_prefix) - del train_fn - - # Restore into a graph with the optimizer - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.Adam(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - status = root.restore(save_path=model_save_path) - input_value = constant_op.constant([[3.]]) - def train_fn1(): - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - return optimizer.apply_gradients(zip(gradients, variables)) - if not context.executing_eagerly(): - train_fn1 = functools.partial(self.evaluate, train_fn1()) - status.initialize_or_restore() - train_fn1() - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - with self.assertRaises(AssertionError): - status.assert_consumed() - del train_fn1 - - # Make sure initialization doesn't clobber later restores - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.Adam(0.001, beta_1=1.0) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - opt_root = trackable_utils.Checkpoint( - optimizer=optimizer) - status = root.restore(save_path=model_save_path) - init_only_optimizer_status = opt_root.restore(save_path=None) - optimizer_status = opt_root.restore(save_path=optimizer_save_path) - input_value = constant_op.constant([[3.]]) - def train_fn2(): - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - return optimizer.apply_gradients(zip(gradients, variables)) - if not context.executing_eagerly(): - train_fn2 = functools.partial(self.evaluate, train_fn2()) - optimizer_status.run_restore_ops() - status.initialize_or_restore() - init_only_optimizer_status.initialize_or_restore() - train_fn2() - self.assertEqual(42., self.evaluate(optimizer.beta_1)) - @test_util.run_in_graph_and_eager_modes def test_restore_after_adding_empty_trackable_data_structure(self): model = NonLayerTrackable() @@ -1439,75 +796,8 @@ class CheckpointingTests(parameterized.TestCase, test.TestCase): self.assertAllClose(self.evaluate(load_checkpoint.b), {"a": 2, "b": 3}) -class _ManualScope(tracking.AutoTrackable): - - def __call__(self): - with variable_scope.variable_scope("ManualScope") as vs: - self.variable_scope = vs - with trackable_utils.capture_dependencies(template=self): - return self._build() - - def _build(self): - return variable_scope.get_variable(name="in_manual_scope", shape=[]) - - class TemplateTests(parameterized.TestCase, test.TestCase): - @test_util.run_in_graph_and_eager_modes - def test_trackable_save_restore(self): - - def _templated(): - v = variable_scope.get_variable( - "v", shape=[1], initializer=init_ops.zeros_initializer(), - use_resource=True) - v2 = variable_scope.get_variable( - "v2", shape=[1], initializer=init_ops.zeros_initializer(), - use_resource=True) - manual = _ManualScope() - return v, v + 1., v2, manual, manual() - - save_template = template.make_template("s1", _templated) - v1_save, _, v2_save, manual_scope, manual_scope_v = save_template() - six.assertCountEqual( - self, - [id(v1_save), id(v2_save), id(manual_scope), - id(manual_scope_v), id(save_template)], - map(id, trackable_utils.list_objects(save_template))) - manual_dep, = manual_scope._checkpoint_dependencies - self.assertEqual("in_manual_scope", manual_dep.name) - self.assertIs(manual_scope_v, manual_dep.ref) - optimizer = adam.Adam(0.0) - save_root = trackable_utils.Checkpoint( - my_template=save_template, optimizer=optimizer) - optimizer.minimize(v1_save.read_value, - var_list=[v1_save]) - self.evaluate([v.initializer for v in save_template.variables]) - optimizer_variables = optimizer.variables() + list( - optimizer._hyper.values()) - self.evaluate([v.initializer for v in optimizer_variables]) - self.evaluate(v1_save.assign([12.])) - self.evaluate(v2_save.assign([14.])) - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - save_path = save_root.save(checkpoint_prefix) - - load_template = template.make_template("s2", _templated) - load_optimizer = adam.Adam(0.0) - load_root = trackable_utils.Checkpoint( - my_template=load_template, optimizer=load_optimizer) - status = load_root.restore(save_path) - var, var_plus_one, var2, _, _ = load_template() - load_optimizer.minimize(var.read_value, var_list=[var]) - self.assertLen(load_template._checkpoint_dependencies, 3) - self.assertEqual("v", load_template._checkpoint_dependencies[0].name) - self.assertEqual("v2", load_template._checkpoint_dependencies[1].name) - self.assertEqual("ManualScope", - load_template._checkpoint_dependencies[2].name) - status.assert_consumed().run_restore_ops() - self.assertAllEqual([12.], self.evaluate(var)) - self.assertAllEqual([13.], self.evaluate(var_plus_one)) - self.assertAllEqual([14.], self.evaluate(var2)) - @test_util.run_in_graph_and_eager_modes def test_trackable_save_restore_nested(self): @@ -1554,157 +844,6 @@ class TemplateTests(parameterized.TestCase, test.TestCase): self.assertAllEqual([25.], self.evaluate(v3)) -class CheckpointCompatibilityTests(test.TestCase): - - def _initialized_model(self): - input_value = constant_op.constant([[3.]]) - model = MyModel() - optimizer = adam.Adam(0.001) - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - with backprop.GradientTape() as tape: - loss = model(input_value) - variables = model.trainable_variables - gradients = tape.gradient(loss, variables) - train_op = optimizer.apply_gradients(zip(gradients, variables)) - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - # A regular variable, a slot variable, and a non-slot Optimizer variable - # with known values to check when loading. - self.evaluate(model._named_dense.bias.assign([1.])) - self.evaluate(optimizer.get_slot( - var=model._named_dense.bias, slot_name="m").assign([2.])) - self.evaluate(optimizer.beta_1.assign(3.)) - return root_trackable - - def _set_sentinels(self, root_trackable): - self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) - self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, slot_name="m") - .assign([102.])) - self.evaluate(root_trackable.optimizer.beta_1.assign(103.)) - - def _check_sentinels(self, root_trackable): - self.assertAllEqual( - [1.], self.evaluate(root_trackable.model._named_dense.bias)) - self.assertAllEqual([2.], self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, slot_name="m"))) - self.assertAllEqual(3., - self.evaluate(root_trackable.optimizer.beta_1)) - - def _write_name_based_checkpoint(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph) as session: - root = self._initialized_model() - name_saver = saver_lib.Saver() - return name_saver.save( - sess=session, - save_path=checkpoint_prefix, - global_step=root.optimizer.iterations) - - @test_util.run_in_graph_and_eager_modes - def testLoadFromNameBasedSaver(self): - """Save a name-based checkpoint, load it using the object-based API.""" - with test_util.device(use_gpu=True): - save_path = self._write_name_based_checkpoint() - root = self._initialized_model() - self._set_sentinels(root) - with self.assertRaises(AssertionError): - self._check_sentinels(root) - object_saver = trackable_utils.TrackableSaver( - graph_view.ObjectGraphView(root)) - self._set_sentinels(root) - status = object_saver.restore(save_path) - if context.executing_eagerly(): - self._check_sentinels(root) - if context.executing_eagerly(): - status.assert_consumed() - status.assert_existing_objects_matched() - status.assert_nontrivial_match() - else: - # When graph building, we haven't read any keys, so we don't know - # whether the restore will be complete. - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_consumed() - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_existing_objects_matched() - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_nontrivial_match() - status.run_restore_ops() - self._check_sentinels(root) - self._set_sentinels(root) - status = object_saver.restore(save_path) - status.initialize_or_restore() - status.assert_nontrivial_match() - self._check_sentinels(root) - # Check that there is no error when keys are missing from the name-based - # checkpoint. - root.not_in_name_checkpoint = resource_variable_ops.ResourceVariable([1.]) - status = object_saver.restore(save_path) - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - - def testSaveGraphLoadEager(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph): - root = self._initialized_model() - save_path = root.save(file_prefix=checkpoint_prefix) - with context.eager_mode(): - root = self._initialized_model() - self._set_sentinels(root) - root.restore(save_path).assert_consumed() - self._check_sentinels(root) - - def testSaveEagerLoadGraph(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.eager_mode(): - root = self._initialized_model() - save_path = root.save(file_prefix=checkpoint_prefix) - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph): - root = self._initialized_model() - self._set_sentinels(root) - root.restore(save_path).assert_consumed().run_restore_ops() - self._check_sentinels(root) - - def testIgnoreSaveCounter(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with self.cached_session() as session: - # Create and save a model using Saver() before using a Checkpoint. This - # generates a snapshot without the Checkpoint's `save_counter`. - model = sequential.Sequential() - model.add(core.Flatten(input_shape=(1,))) - model.add(core.Dense(1)) - name_saver = saver_lib.Saver(model.trainable_variables) - save_path = name_saver.save( - sess=session, save_path=checkpoint_prefix, global_step=1) - # Checkpoint.restore must successfully load that checkpoint. - ckpt = trackable_utils.Checkpoint(model=model) - status = ckpt.restore(save_path) - status.assert_existing_objects_matched() - # It should, however, refuse to load a checkpoint where an unrelated - # `save_counter` variable is missing. - model.layers[1].var = variables_lib.Variable(0., name="save_counter") - status = ckpt.restore(save_path) - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - - if __name__ == "__main__": ops.enable_eager_execution() test.main() From a2fb617adb26cda76925b5009bd8cd861e2e4943 Mon Sep 17 00:00:00 2001 From: Edward Loper Date: Mon, 15 Jun 2020 09:42:58 -0700 Subject: [PATCH 0169/1390] Internal visibility change PiperOrigin-RevId: 316483590 Change-Id: Ia62c9b89846b6165be67c6cc2a3ce9f17e1b3eb3 --- tensorflow/BUILD | 6 ++++++ tensorflow/python/ops/structured/BUILD | 3 +-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tensorflow/BUILD b/tensorflow/BUILD index efbdf89ecea..ce759634232 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -531,6 +531,7 @@ package_group( # Packages that use composite tensors or dispatch. # TODO(b/154762408) Remove this package group once it's no longer needed. +# If this is modified, then copy.bara.sky must also be modified. package_group(name = "composite_tensor_whitelist") # Packages that use private types symbols, until they are exported. @@ -540,6 +541,11 @@ package_group( packages = ["//learning/deepmind/tensorflow/replicator/..."], ) +# Packages that use StructuredTensors. +# TODO(b/159007891) Remove this package once StructuredTensor is exported. +# If this is modified, then copy.bara.sky must also be modified. +package_group(name = "structured_tensor_whitelist") + filegroup( name = "intel_binary_blob", data = if_mkl_ml( diff --git a/tensorflow/python/ops/structured/BUILD b/tensorflow/python/ops/structured/BUILD index 64b7bd7f1d5..33834f0e914 100644 --- a/tensorflow/python/ops/structured/BUILD +++ b/tensorflow/python/ops/structured/BUILD @@ -4,9 +4,8 @@ load("//tensorflow:tensorflow.bzl", "py_test") package( default_visibility = [ - "//learning/tfx/autotfx:__subpackages__", - "//research/graph/convolutions/model/autotfx:__subpackages__", "//tensorflow:internal", + "//tensorflow:structured_tensor_whitelist", ], licenses = ["notice"], # Apache 2.0 ) From 05bec441e33a291e71f21de18d430feca28f00f9 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 09:57:26 -0700 Subject: [PATCH 0170/1390] Integrate LLVM at https://github.com/llvm/llvm-project/commit/f47a7766287a PiperOrigin-RevId: 316486707 Change-Id: I37ac1f5c8ad95e7a4c2cf5976e361fb06272b08b --- third_party/mlir/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD index 4c6717ea024..5f3f0a4b99b 100644 --- a/third_party/mlir/BUILD +++ b/third_party/mlir/BUILD @@ -2126,6 +2126,7 @@ cc_library( ":Analysis", ":ControlFlowInterfaces", ":IR", + ":LinalgOps", ":LoopLikeInterface", ":Pass", ":SCFDialect", From a208b5cd2900eacae5aaf66d06e34fb30ed2ea43 Mon Sep 17 00:00:00 2001 From: Gaurav Jain Date: Mon, 15 Jun 2020 09:57:28 -0700 Subject: [PATCH 0171/1390] Allow ZeroDivisionTest to run eagerly PiperOrigin-RevId: 316486715 Change-Id: I5a705b27562c57760ed8efeae52c67a91539ee7c --- tensorflow/python/kernel_tests/zero_division_test.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tensorflow/python/kernel_tests/zero_division_test.py b/tensorflow/python/kernel_tests/zero_division_test.py index 7f2d100f1e3..f62ac9f7f26 100644 --- a/tensorflow/python/kernel_tests/zero_division_test.py +++ b/tensorflow/python/kernel_tests/zero_division_test.py @@ -20,26 +20,25 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors_impl +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.platform import test class ZeroDivisionTest(test.TestCase): - @test_util.run_deprecated_v1 def testZeros(self): with test_util.use_gpu(): for dtype in dtypes.uint8, dtypes.int16, dtypes.int32, dtypes.int64: zero = constant_op.constant(0, dtype=dtype) one = constant_op.constant(1, dtype=dtype) - bads = [one // zero] + bads = [lambda x, y: x // y] if dtype in (dtypes.int32, dtypes.int64): - bads.append(one % zero) + bads.append(lambda x, y: x % y) for bad in bads: try: - result = self.evaluate(bad) - except errors_impl.OpError as e: + result = self.evaluate(bad(one, zero)) + except (errors.OpError, errors.InvalidArgumentError) as e: # Ideally, we'd get a nice exception. In theory, this should only # happen on CPU, but 32 bit integer GPU division is actually on # CPU due to a placer bug. From b09d410f3beb692c700e1951fb91f6235452940f Mon Sep 17 00:00:00 2001 From: Rahul Joshi Date: Mon, 15 Jun 2020 09:58:30 -0700 Subject: [PATCH 0172/1390] Add transforms to convert between functional and region based If. - Add this pair of transforms in TF -> TFLite conversion pass manager - Add inliner hook to the TF dialect to allow inlining of call's within the regions of the IfRegion ops' - Add a end-to-end test case using 2 If's that demonstrate inlining and followed by constant sinking and constant folding. PiperOrigin-RevId: 316486959 Change-Id: I467b9f3fbc3eafc2cb37e46f61ff6ccfa34afcd5 --- .../mlir/lite/tests/end2end/if_op.pbtxt | 421 ++++++++++++++++++ .../compiler/mlir/lite/tf_tfl_passes.cc | 6 + tensorflow/compiler/mlir/tensorflow/BUILD | 3 + .../compiler/mlir/tensorflow/ir/tf_ops.cc | 9 + .../functional-control-flow-to-regions.mlir | 113 +++++ .../region-control-flow-to-functional.mlir | 188 ++++++++ .../functional_control_flow_to_regions.cc | 117 +++++ .../mlir/tensorflow/transforms/passes.h | 14 +- .../region_control_flow_to_functional.cc | 340 ++++++++++++++ 9 files changed, 1209 insertions(+), 2 deletions(-) create mode 100644 tensorflow/compiler/mlir/lite/tests/end2end/if_op.pbtxt create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/functional_control_flow_to_regions.cc create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc diff --git a/tensorflow/compiler/mlir/lite/tests/end2end/if_op.pbtxt b/tensorflow/compiler/mlir/lite/tests/end2end/if_op.pbtxt new file mode 100644 index 00000000000..f482e3db6b9 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/end2end/if_op.pbtxt @@ -0,0 +1,421 @@ +# RUN: tf_tfl_translate -tf-input-arrays=a,b -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-input-shapes=4:4 -tf-output-arrays=StatefulIf,StatelessIf %s -o - --output-mlir | FileCheck %s +node { + name: "tf.Less" + op: "Less" + input: "a" + input: "b" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + } +} +node { + name: "my_equal" + op: "Equal" + input: "a" + input: "b" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + } +} +node { + name: "cst0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_FLOAT + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_FLOAT + tensor_shape { + dim { + size: 4 + } + } + float_val: 1.0 + float_val: 2.0 + float_val: 3.0 + float_val: 4.0 + } + } + } +} +node { + name: "cst1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_FLOAT + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_FLOAT + tensor_shape { + dim { + size: 4 + } + } + float_val: 5.0 + float_val: 6.0 + float_val: 7.0 + float_val: 8.0 + } + } + } +} +node { + name: "StatefulIf" + op: "If" + input: "tf.Less" + input: "a" + input: "b" + input: "cst0" + input: "cst1" + attr { + key: "Tcond" + value { + type: DT_BOOL + } + } + attr { + key: "Tin" + value { + list { + type: DT_FLOAT + type: DT_FLOAT + type: DT_FLOAT + type: DT_FLOAT + } + } + } + attr { + key: "Tout" + value { + list { + type: DT_FLOAT + } + } + } + attr { + key: "else_branch" + value { + func { + name: "cond_false" + } + } + } + attr { + key: "then_branch" + value { + func { + name: "cond_true" + } + } + } + experimental_debug_info { + } +} +node { + name: "StatelessIf" + op: "StatelessIf" + input: "my_equal" + input: "a" + input: "b" + attr { + key: "Tcond" + value { + type: DT_BOOL + } + } + attr { + key: "Tin" + value { + list { + type: DT_FLOAT + type: DT_FLOAT + } + } + } + attr { + key: "Tout" + value { + list { + type: DT_FLOAT + } + } + } + attr { + key: "else_branch" + value { + func { + name: "cond_false_1" + } + } + } + attr { + key: "then_branch" + value { + func { + name: "cond_true_1" + } + } + } + experimental_debug_info { + } +} +node { + name: "main" + op: "_Retval" + input: "StatefulIf" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + attr { + key: "index" + value { + i: 0 + } + } +} +node { + name: "main1" + op: "_Retval" + input: "StatelessIf" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + attr { + key: "index" + value { + i: 1 + } + } +} +node { + name: "a" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + } +} +node { + name: "b" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + } +} +library { + function { + signature { + name: "cond_true" + input_arg { + name: "cond_true_arg0" + type: DT_FLOAT + } + input_arg { + name: "cond_true_arg1" + type: DT_FLOAT + } + input_arg { + name: "cond_true_arg2" + type: DT_FLOAT + } + input_arg { + name: "cond_true_arg3" + type: DT_FLOAT + } + output_arg { + name: "cond_true_ret" + type: DT_FLOAT + } + } + node_def { + name: "tf.Add" + op: "Add" + input: "cond_true_arg2" + input: "cond_true_arg3" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + original_node_names: "tf.Add" + } + } + ret { + key: "cond_true_ret" + value: "tf.Add:z:0" + } + } + function { + signature { + name: "cond_false" + input_arg { + name: "cond_false_arg0" + type: DT_FLOAT + } + input_arg { + name: "cond_false_arg1" + type: DT_FLOAT + } + input_arg { + name: "cond_false_arg2" + type: DT_FLOAT + } + input_arg { + name: "cond_false_arg3" + type: DT_FLOAT + } + output_arg { + name: "cond_false_ret" + type: DT_FLOAT + } + } + node_def { + name: "tf.Mul" + op: "Mul" + input: "cond_false_arg0" + input: "cond_false_arg3" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + original_node_names: "tf.Mul" + } + } + ret { + key: "cond_false_ret" + value: "tf.Mul:z:0" + } + } + function { + signature { + name: "cond_true_1" + input_arg { + name: "cond_true_arg0" + type: DT_FLOAT + } + input_arg { + name: "cond_true_arg1" + type: DT_FLOAT + } + output_arg { + name: "cond_true_ret" + type: DT_FLOAT + } + } + node_def { + name: "tf.Sub" + op: "Sub" + input: "cond_true_arg0" + input: "cond_true_arg1" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + original_node_names: "tf.Sub" + } + } + ret { + key: "cond_true_ret" + value: "tf.Sub:z:0" + } + } + function { + signature { + name: "cond_false_1" + input_arg { + name: "cond_false_arg0" + type: DT_FLOAT + } + input_arg { + name: "cond_false_arg1" + type: DT_FLOAT + } + output_arg { + name: "cond_false_ret" + type: DT_FLOAT + } + } + node_def { + name: "tf.Div" + op: "Div" + input: "cond_false_arg0" + input: "cond_false_arg1" + attr { + key: "T" + value { + type: DT_FLOAT + } + } + experimental_debug_info { + original_node_names: "tf.Div" + } + } + ret { + key: "cond_false_ret" + value: "tf.Div:z:0" + } + } +} +versions { + producer: 115 + min_consumer: 12 +} + +# CHECK: func @StatefulIf_else +# CHECK-NEXT: constant dense<[5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]> +# CHECK-NEXT: tfl.mul +# CHECK: func @StatefulIf_then +# CHECK-NEXT: constant dense<[6.000000e+00, 8.000000e+00, 1.000000e+01, 1.200000e+01]> +# CHECK-NEXT: return +# CHECK: func @StatelessIf_else +# CHECK-NEXT: tfl.div +# CHECK: func @StatelessIf_then +# CHECK-NEXT: tfl.sub +# CHECK: "tf.If"{{.+}}else_branch = @StatelessIf_else{{.+}}then_branch = @StatelessIf_then +# CHECK: "tf.If"{{.+}}else_branch = @StatefulIf_else{{.+}}then_branch = @StatefulIf_then + diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc index 49fe0eb7100..06fe8684ce4 100644 --- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc +++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/tf_tfl_passes.h" #include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Function.h" // from @llvm-project #include "mlir/IR/Module.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "mlir/Pass/PassManager.h" // from @llvm-project @@ -74,6 +75,8 @@ void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config, pass_config.quant_specs.serialized_quant_stats)); } + pass_manager->addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions()); + // The conversion pipeline has to follow the following orders: // 1) Saved model related optimization like decompose resource ops // 2) Convert composite functions like lstm/rnns, along with proper function @@ -111,6 +114,9 @@ void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config, // Add a shape inference pass to optimize away the unnecessary casts. pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass()); } + + pass_manager->addPass(mlir::TF::CreateTFRegionControlFlowToFunctional()); + // Legalize while early to allow further constant folding. // TODO(jpienaar): This may not actually matter as we do canonicalization // after the legalize below, for now it needs to be below the above passes diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD index b6ff0f581d3..9e5688cd230 100644 --- a/tensorflow/compiler/mlir/tensorflow/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/BUILD @@ -416,6 +416,7 @@ cc_library( "transforms/fold_switch.cc", "transforms/freeze_global_tensors.cc", "transforms/functional_control_flow_to_cfg.cc", + "transforms/functional_control_flow_to_regions.cc", "transforms/fused_kernel_matcher.cc", "transforms/generated_canonicalize.inc", "transforms/generated_optimize.inc", @@ -430,6 +431,7 @@ cc_library( "transforms/parallel_execute_to_islands.cc", "transforms/promote_resources_to_args.cc", "transforms/readonly_references_to_resources.cc", + "transforms/region_control_flow_to_functional.cc", "transforms/replicate_invariant_op_hoisting.cc", "transforms/replicate_to_island.cc", "transforms/resource_device_inference.cc", @@ -487,6 +489,7 @@ cc_library( ":translate_utils", ":unroll_batch_matmul_pass", ":xla_sharding_util", + "//tensorflow/compiler/mlir:op_or_arg_name_mapper", "//tensorflow/compiler/mlir/lite:validators", "//tensorflow/compiler/xla:xla_data_proto_cc", "//tensorflow/compiler/xla:xla_proto_cc", diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc index 512011c3a0f..f4e5dc05eb0 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc @@ -4066,6 +4066,15 @@ struct TFInlinerInterface : public DialectInlinerInterface { // Analysis Hooks //===--------------------------------------------------------------------===// + // Defines the legality of inlinining 'src' region into the 'dest' region + // attached to a TF operation + bool isLegalToInline(Region *dest, Region *src, + BlockAndValueMapping &valueMapping) const final { + // Allow inlining in regions attached to region based control flow + // operations only if the src region is a single block region + return isa(dest->getParentOp()) && llvm::hasSingleElement(*src); + } + // Defines the legality of inlining TF operations. bool isLegalToInline(Operation *, Region *, BlockAndValueMapping &) const final { diff --git a/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir new file mode 100644 index 00000000000..a7e9b22d72b --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir @@ -0,0 +1,113 @@ +// RUN: tf-opt %s -tf-functional-control-flow-to-regions -split-input-file | FileCheck %s --dump-input=fail + +// CHECK: func @testIf1Then{{.+}} +// CHECK: func @testIf1Else{{.+}} +func @testIf1Then(tensor<*xf32>) -> tensor<*xf32> +func @testIf1Else(tensor<*xf32>) -> tensor<*xf32> + +// CHECK-LABEL: func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) +func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { + %0 = "tf.If"(%arg0, %arg1) { + then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false + } : (tensor, tensor<*xf32>) -> tensor<*xf32> + + // CHECK: "tf.IfRegion" + // CHECK: [[Result0:%.*]] = call @testIf1Then + // CHECK: "tf.Yield"([[Result0]]) + // CHECK: [[Result1:%.*]] = call @testIf1Else + // CHECK: "tf.Yield"([[Result1]]) + return %0 : tensor<*xf32> +} + +// ----- + +// With mismatching input types + +// CHECK: func @testIf1Then{{.+}} +// CHECK: func @testIf1Else{{.+}} +func @testIf1Then(tensor<*xf32>) -> tensor<*xf32> +func @testIf1Else(tensor<*xf32>) -> tensor<*xf32> + +// CHECK-LABEL: func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) +func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf32> { + %0 = "tf.If"(%arg0, %arg1) { + then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false + } : (tensor, tensor<2xf32>) -> tensor<2xf32> + + // CHECK: "tf.IfRegion" + // CHECK: "tf.Cast" + // CHECK: [[Result0:%.*]] = call @testIf1Then + // CHECK: "tf.Yield"([[Result0]]) + // CHECK: "tf.Cast" + // CHECK: [[Result1:%.*]] = call @testIf1Else + // CHECK: "tf.Yield"([[Result1]]) + return %0 : tensor<2xf32> +} + +// ----- + +// No inputs, some outputs +// CHECK: func @testIf1Then{{.+}} +// CHECK: func @testIf1Else{{.+}} +func @testIf1Then() -> tensor<*xf32> +func @testIf1Else() -> tensor<*xf32> + +// CHECK-LABEL: func @testIfNoInputs(%arg0: tensor) +func @testIfNoInputs(%arg0: tensor) -> tensor<2xf32> { + %0 = "tf.If"(%arg0) { + then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false + } : (tensor) -> tensor<2xf32> + + // CHECK: "tf.IfRegion" + // CHECK: [[Result0:%.*]] = call @testIf1Then + // CHECK: "tf.Yield"([[Result0]]) + // CHECK: [[Result1:%.*]] = call @testIf1Else + // CHECK: "tf.Yield"([[Result1]]) + return %0 : tensor<2xf32> +} + +// ----- + +// No outputs, some inputs +// CHECK: func @testIf1Then{{.+}} +// CHECK: func @testIf1Else{{.+}} +func @testIf1Then(tensor<*xf32>) -> () +func @testIf1Else(tensor<*xf32>) -> () + +// CHECK-LABEL: func @testIfNoResult(%arg0: tensor, %arg1: tensor<2xf32>) +func @testIfNoResult(%arg0: tensor, %arg1: tensor<2xf32>) -> () { + "tf.If"(%arg0, %arg1) { + then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false + } : (tensor, tensor<2xf32>) -> () + + // CHECK: "tf.IfRegion" + // CHECK: "tf.Cast" + // CHECK: call @testIf1Then + // CHECK: "tf.Yield"() + // CHECK: "tf.Cast" + // CHECK: call @testIf1Else + // CHECK: "tf.Yield"() + return +} + +// ----- +// No outputs, No inputs +// CHECK: func @testIf1Then{{.+}} +// CHECK: func @testIf1Else{{.+}} +func @testIf1Then() -> () +func @testIf1Else() -> () + +// CHECK-LABEL: func @testIfNoInputAndNoResult(%arg0: tensor) +func @testIfNoInputAndNoResult(%arg0: tensor) -> () { + "tf.If"(%arg0) { + then_branch = @testIf1Then, else_branch = @testIf1Else, is_stateless = false + } : (tensor) -> () + + // CHECK: "tf.IfRegion" + // CHECK: call @testIf1Then + // CHECK: "tf.Yield"() + // CHECK: call @testIf1Else + // CHECK: "tf.Yield"() + return +} + diff --git a/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir b/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir new file mode 100644 index 00000000000..5ea863852ad --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir @@ -0,0 +1,188 @@ +// RUN: tf-opt %s -tf-region-control-flow-to-functional -split-input-file +//| FileCheck %s --dump-input=fail + +// CHECK: func @tf.IfRegion_else(%arg0: tensor<*xf32>) -> tensor<*xf32> +// CHECK-NEXT: "tf.Neg" +// CHECK: func @tf.IfRegion_then(%arg0: tensor<*xf32>) -> tensor<*xf32> +// CHECK-NEXT: "tf.Abs" +func @testSimple(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { + // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + %0 = "tf.IfRegion"(%arg0) ({ + %1 = "tf.Abs"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%1) : (tensor<*xf32>) -> () + }, { + %2 = "tf.Neg"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%2) : (tensor<*xf32>) -> () + }) { is_stateless = true } : (tensor) -> tensor<*xf32> + return %0 : tensor<*xf32> +} + +// ----- + +// Use if condition inside the regions +// CHECK: func @tf.IfRegion_else(%arg0: tensor, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> tensor<2xf32> +// CHECK-NEXT: "tf.Select"(%arg0, %arg2, %arg3) +// CHECK: func @tf.IfRegion_then(%arg0: tensor, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> tensor<2xf32> +// CHECK-NEXT: "tf.Select"(%arg0, %arg1, %arg2) +func @testIfCondition(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf32> { + %0 = "tf.Add"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> + %1 = "tf.Mul"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> + %2 = "tf.Div"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> + + // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + %3 = "tf.IfRegion"(%arg0) ({ + %4 = "tf.Select"(%arg0, %0, %1) : (tensor, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> + "tf.Yield"(%4) : (tensor<2xf32>) -> () + }, { + %5 = "tf.Select"(%arg0, %1, %2): (tensor, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> + "tf.Yield"(%5) : (tensor<2xf32>) -> () + }) { is_stateless = true} : (tensor) -> tensor<2xf32> + return %3 : tensor<2xf32> +} + +// ----- + +// Constant sinking + +// CHECK: func @tf.IfRegion_else() -> tensor<2xf32> +// CHECK-NEXT: constant dense<1.0 +// CHECK: func @tf.IfRegion_then() -> tensor<2xf32> +// CHECK-NEXT: constant dense<0.0 +func @testIfConstant(%arg0: tensor) -> tensor<2xf32> { + %cst_zero = constant dense<0.0> : tensor<2xf32> + // CHECK: "tf.If"(%arg0) {else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + %0 = "tf.IfRegion"(%arg0) ({ + "tf.Yield"(%cst_zero) : (tensor<2xf32>) -> () + }, { + %cst_one = constant dense<1.0> : tensor<2xf32> + "tf.Yield"(%cst_one) : (tensor<2xf32>) -> () + }) { is_stateless = true} : (tensor) -> tensor<2xf32> + return %0 : tensor<2xf32> +} + +// ----- + +// Nested IfRegions +// CHECK: func @tf.IfRegion1_else +// CHECK-NEXT: "tf.Acos" +// CHECK-NEXT: "tf.Abs" + +// CHECK: func @tf.IfRegion1_then +// CHECK-NEXT: "tf.LogicalNot" +// CHECK-NEXT: "tf.Asin" +// CHECK-NEXT: "tf.If"({{.+}}) {else_branch = @tf.IfRegion_else, {{.+}} then_branch = @tf.IfRegion_then} + +// CHECK: func @tf.IfRegion_else +// CHECK-NEXT: "tf.Neg" +// CHECK: func @tf.IfRegion_then +// CHECK-NEXT: "tf.Abs" + +func @testNested(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { + // CHECK: "tf.If"({{.+}}) {else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then} + %0 = "tf.IfRegion"(%arg0) ({ + // Outer Then + %cond = "tf.LogicalNot"(%arg0) : (tensor) -> tensor + %asin = "tf.Asin"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + + // nested IfRegion + %1 = "tf.IfRegion"(%cond) ({ + %2 = "tf.Abs"(%asin) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%2) : (tensor<*xf32>) -> () + }, { + %2 = "tf.Neg"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%2) : (tensor<*xf32>) -> () + }) { is_stateless = true } : (tensor) -> tensor<*xf32> + + "tf.Yield"(%1) : (tensor<*xf32>) -> () + }, { + // Outer Else + %acos = "tf.Acos"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + %3 = "tf.Abs"(%acos) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%3) : (tensor<*xf32>) -> () + }) { is_stateless = true } : (tensor) -> tensor<*xf32> + return %0 : tensor<*xf32> +} + +// ----- + +// Match existing function->Region pattern (simple) +func @testIf1Then(tensor<*xf32>) -> tensor<*xf32> +func @testIf1Else(tensor<*xf32>) -> tensor<*xf32> +func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { + // CHECK: "tf.If"({{.+}}) {else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + %0 = "tf.IfRegion"(%arg0) ( { + %1 = call @testIf1Then(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%1) : (tensor<*xf32>) -> () + }, { + %1 = call @testIf1Else(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%1) : (tensor<*xf32>) -> () + }) {is_stateless = false} : (tensor) -> tensor<*xf32> + return %0 : tensor<*xf32> +} + +// ----- + +// Match existing function->Region pattern (with casts) + +func @testIf1Then(tensor<*xf32>) -> tensor<*xf32> +func @testIf1Else(tensor<*xf32>) -> tensor<*xf32> +func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf32> { + // CHECK: "tf.If"({{.+}}) {else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + %0 = "tf.IfRegion"(%arg0) ( { + %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xf32>) -> tensor<*xf32> + %2 = call @testIf1Then(%1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%2) : (tensor<*xf32>) -> () + }, { + %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xf32>) -> tensor<*xf32> + %2 = call @testIf1Else(%1) : (tensor<*xf32>) -> tensor<*xf32> + "tf.Yield"(%2) : (tensor<*xf32>) -> () + }) {is_stateless = false} : (tensor) -> tensor<2xf32> + return %0 : tensor<2xf32> +} + +// ----- + +// No inputs, some outputs +// CHECK: func @tf.IfRegion_else() -> tensor<2xf32> +// CHECK-NEXT: constant dense<1.000000e+00> +// CHECK-NEXT: "tf.Neg" +// CHECK: func @tf.IfRegion_then() -> tensor<2xf32> +// CHECK-NEXT: constant dense<0.000000e+00> +// CHECK-NEXT: "tf.Abs" +func @testSimple(%arg0: tensor) -> tensor<2xf32> { + // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + %0 = "tf.IfRegion"(%arg0) ({ + %cst_zero = constant dense<0.0> : tensor<2xf32> + %1 = "tf.Abs"(%cst_zero) : (tensor<2xf32>) -> tensor<2xf32> + "tf.Yield"(%1) : (tensor<2xf32>) -> () + }, { + %cst_one = constant dense<1.0> : tensor<2xf32> + %2 = "tf.Neg"(%cst_one) : (tensor<2xf32>) -> tensor<2xf32> + "tf.Yield"(%2) : (tensor<2xf32>) -> () + }) { is_stateless = true } : (tensor) -> tensor<2xf32> + return %0 : tensor<2xf32> +} + +// ----- + +// No outputs, some inputs +// +// CHECK: func @tf.IfRegion_else(%arg0: tensor<*xf32>) +// CHECK-NEXT: "tf.Neg" +// CHECK: func @tf.IfRegion_then(%arg0: tensor<*xf32>) +// CHECK-NEXT: "tf.Abs" +func @printer(tensor<*xf32>) -> () +func @testNoOutputs(%arg0: tensor, %arg1: tensor<*xf32>) -> () { + // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + "tf.IfRegion"(%arg0) ({ + %1 = "tf.Abs"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + call @printer(%1) : (tensor<*xf32>) -> () + "tf.Yield"() : () -> () + }, { + %2 = "tf.Neg"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> + call @printer(%2) : (tensor<*xf32>) -> () + "tf.Yield"() : () -> () + }) { is_stateless = false } : (tensor) -> () + return +} + diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/functional_control_flow_to_regions.cc b/tensorflow/compiler/mlir/tensorflow/transforms/functional_control_flow_to_regions.cc new file mode 100644 index 00000000000..5ab0eda08c6 --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/transforms/functional_control_flow_to_regions.cc @@ -0,0 +1,117 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +// This transformation pass transforms functional control flow operations in the +// TensorFlow dialect to their region based counterparts, i.e., +// tf.If -> tf.IfRegion + +#include "llvm/Support/raw_ostream.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/Function.h" // from @llvm-project +#include "mlir/IR/Operation.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/IR/Verifier.h" // from @llvm-project +#include "mlir/IR/Visitors.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Pass/PassRegistry.h" // from @llvm-project +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" +#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h" + +namespace mlir { +namespace TF { + +namespace { + +struct FunctionalControlFlowToRegions + : public PassWrapper> { + void runOnOperation() override; +}; + +// Create a call to function `fn` with arguments `args` and return the CallOp. +// The arguments are cast to the required type before the call. +CallOp CreateCall(Location loc, Operation::operand_range args, FuncOp fn, + OpBuilder* builder) { + FunctionType fn_type = fn.getType(); + llvm::SmallVector operands; + int num_operands = fn_type.getNumInputs(); + operands.reserve(num_operands); + for (const auto& ArgAndType : zip(args, fn_type.getInputs())) { + Value arg = std::get<0>(ArgAndType); + Type expected_type = std::get<1>(ArgAndType); + if (arg.getType() != expected_type) { + arg = builder->create(loc, expected_type, arg, + /*Truncate=*/builder->getBoolAttr(false)); + } + operands.push_back(arg); + } + return builder->create(loc, fn, operands); +} + +// Transform a functional IfOp to a region based IfRegionOp. +LogicalResult ConvertIfOp(IfOp if_op) { + auto if_region = OpBuilder(if_op).create( + if_op.getLoc(), if_op.getResultTypes(), if_op.cond(), + if_op.is_stateless()); + + // Insert call to the given function into the 'region'. + auto create_region_with_call = [&if_op](FlatSymbolRefAttr symbol, + Region& region) { + OpBuilder builder(region); + builder.createBlock(®ion); + auto func = if_op.getParentOfType().lookupSymbol( + symbol.getValue()); + auto call = CreateCall(if_op.getLoc(), if_op.input(), func, &builder); + builder.create(if_op.getLoc(), call.getResults()); + }; + + create_region_with_call(if_op.then_branchAttr(), if_region.then_branch()); + create_region_with_call(if_op.else_branchAttr(), if_region.else_branch()); + + if_op.replaceAllUsesWith(if_region.getResults()); + if_op.erase(); + return success(); +} + +void FunctionalControlFlowToRegions::runOnOperation() { + ModuleOp module = getOperation(); + auto result = module.walk([](Operation* op) { + if (IfOp if_op = llvm::dyn_cast(op)) { + if (failed(ConvertIfOp(if_op))) { + if_op.emitOpError() << " failed to convert to region form"; + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + if (result.wasInterrupted()) return signalPassFailure(); +} +} // namespace + +std::unique_ptr> +CreateTFFunctionalControlFlowToRegions() { + return std::make_unique(); +} + +static PassRegistration pass( + "tf-functional-control-flow-to-regions", + "Transform functional control flow Ops to Region based counterparts"); + +} // namespace TF +} // namespace mlir diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h index 8f18c904420..7158d0f6be0 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h +++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h @@ -32,10 +32,20 @@ std::unique_ptr> CreateFunctionalToExecutorDialectConversionPass(); namespace TF { -// Transforms functional control flow operations in the standard TensorFlow -// dialect to MLIR Control Flow Graph (CFG) form. +// Transforms functional control flow operations in the TensorFlow dialect to +// MLIR Control Flow Graph (CFG) form. std::unique_ptr> CreateTFFunctionalControlFlowToCFG(); +// Transforms functional control flow operations in the TensorFlow dialect to +// their region based counterparts. +std::unique_ptr> +CreateTFFunctionalControlFlowToRegions(); + +// Transforms region bases control flow operations in the TensorFlow dialect to +// their functional counterparts. +std::unique_ptr> +CreateTFRegionControlFlowToFunctional(); + // Materialize the MlirPassthroughOp by replacing it with the MLIR module // attached as an attribute. std::unique_ptr> CreateMaterializePassthroughOpPass(); diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc b/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc new file mode 100644 index 00000000000..ca0467942ca --- /dev/null +++ b/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc @@ -0,0 +1,340 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +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 + + http://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. +==============================================================================*/ + +// This transformation pass transforms region bases control flow operations in +// the TensorFlow dialect to their functional counterparts, i.e., +// tf.IfRegion -> tf.If + +#include "llvm/ADT/STLExtras.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/Function.h" // from @llvm-project +#include "mlir/IR/Module.h" // from @llvm-project +#include "mlir/IR/Operation.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/IR/Verifier.h" // from @llvm-project +#include "mlir/IR/Visitors.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Pass/PassRegistry.h" // from @llvm-project +#include "mlir/Transforms/RegionUtils.h" // from @llvm-project +#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" +#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h" + +namespace mlir { +namespace TF { + +namespace { + +struct RegionControlFlowToFunctional + : public PassWrapper> { + void runOnOperation() override; + + private: + LogicalResult ConvertIfOp(IfRegionOp if_region); + + // Get unique name by using the loc to name mapping. + std::string GetName(Operation* op, StringRef suffix); + + tensorflow::OpOrArgLocNameMapper mapper; + llvm::SmallVector worklist; +}; + +std::string RegionControlFlowToFunctional::GetName(Operation* op, + StringRef suffix) { + return (mapper.GetUniqueName(op) + suffix).str(); +} + +// Returns all the external values referenced from the given set of regions. If +// the external value is a constant, sink it into the region instead (and do not +// add it to the returned vector). +llvm::SmallVector CollectExternValues(ArrayRef regions) { + llvm::SetVector extern_values_set; + + for (auto region : regions) { + llvm::SetVector region_extern_values; + getUsedValuesDefinedAbove(*region, region_extern_values); + + // Sink down constants into the functions. + for (auto extern_value : region_extern_values) { + if (!matchPattern(extern_value, m_Constant())) { + extern_values_set.insert(extern_value); + continue; + } + // Add constant at start of region. + auto const_builder = OpBuilder::atBlockBegin(®ion->front()); + auto const_value = const_builder.clone(*extern_value.getDefiningOp()); + replaceAllUsesInRegionWith(extern_value, const_value->getResult(0), + *region); + } + } + + return {extern_values_set.begin(), extern_values_set.end()}; +} + +// Extracts the contents of a region with a single block into a new function. +// `extern_values` is the set of external values that the region refers to. +// +// Any inputs to the terminator of the region are converted to return values of +// the function. If any of these values is not exact type as the function's +// return type, appropriate cast operations will be inserted +void ExtractSingleBlockRegion(Region& region, FunctionType type, StringRef name, + llvm::SmallVectorImpl& extern_values, + llvm::SmallVectorImpl& worklist) { + ModuleOp module = region.getParentOfType(); + auto builder = OpBuilder::atBlockBegin(module.getBody()); + auto loc = region.getParentOp()->getLoc(); + + // Create new function and extract region body into the function. + auto outlined_func = + builder.create(loc, name, type, ArrayRef{}); + + outlined_func.getBody().takeBody(region); + Region& func_region = outlined_func.getBody(); + Block& first_block = func_region.front(); + + // Replace all external uses with function arguments. + for (auto it : llvm::enumerate(extern_values)) { + Value arg = first_block.addArgument(it.value().getType()); + replaceAllUsesInRegionWith(it.value(), arg, func_region); + } + + // Replace the existing terminator with a return. + Operation* terminator = outlined_func.getBody().front().getTerminator(); + builder.setInsertionPoint(terminator); + + SmallVector return_values; + return_values.reserve(terminator->getNumOperands()); + for (auto it : llvm::enumerate(type.getResults())) { + Value ret_val = terminator->getOperand(it.index()); + // Add a cast operation if types do not match. + if (ret_val.getType() != it.value()) { + ret_val = + builder.create(terminator->getLoc(), it.value(), ret_val); + } + return_values.push_back(ret_val); + } + builder.create(terminator->getLoc(), return_values); + terminator->erase(); + outlined_func.setVisibility(FuncOp::Visibility::Private); + + // Add the outlined function to the worklist in case its body has + // IfRegion ops that need to converted. + worklist.push_back(outlined_func); +} + +// Returns call for region with single call whose result feeds into the +// terminator of the region. Returns none if the region doesn't contain just +// call and non-truncting casts ops. +llvm::Optional IsSingleCallRegion(Region& region) { + if (!llvm::hasSingleElement(region)) return llvm::None; + + Block& block = region.front(); + auto it = block.rbegin(); + YieldOp yield = dyn_cast(*it++); + + if (it == block.rend()) return llvm::None; + + // Check if there is a Call before the Yield. + CallOp call = dyn_cast(*it++); + if (!call) return llvm::None; + + // There can only be non-truncating cast op's prior to the call. + for (; it != block.rend(); ++it) { + CastOp cast = dyn_cast(*it); + if (!cast || cast.Truncate()) return llvm::None; + } + + // All results of the call should feed into the yield. + if (call.getNumResults() != yield.getNumOperands()) return llvm::None; + + for (auto res_it : llvm::zip(call.getResults(), yield.getOperands())) + if (std::get<0>(res_it) != std::get<1>(res_it)) return llvm::None; + + return call; +} + +// Returns whether the arguments of the given call are same as the given list of +// arguments (after looking through cast ops). +bool MatchCallArgs(CallOp call, llvm::SmallVectorImpl& args) { + if (call.getNumOperands() != args.size()) return false; + + for (auto it : llvm::enumerate(args)) { + Value arg = call.getOperand(it.index()); + if (auto cast = dyn_cast_or_null(arg.getDefiningOp())) + arg = cast.getOperand(); + + if (arg != it.value()) return false; + } + return true; +} + +// Summary information for trivially transforming region based op's to +// functional ops. A trivial transformation can be done when the regions are +// just calls to functions, in which case no outlining is needed. +struct TrivialTransformInfo { + // Can the op be transformed trivially? + bool can_transform = false; + + // List of callee names (one for each region). + llvm::SmallVector callee_names; + + // List of arguments used in these call (each call uses the same arguments + // potentially through casts). + llvm::SmallVector call_args; +}; + +// Analyzes the given set of regions (attached to the same parent op) to check +// if the parent op be transformed to functional form trivially (i.e., reusing +// existing functions and without outlining). This is possible when all the +// regions are single call regions and the all the calls have the same +// arguments. +// +// If this trivial transformation is possible, return the relevant information +// needed for the transformation (in `TrivialTransformInfo`), else indicate that +// a trivial transformation is not possible by setting `can_transform` false. +TrivialTransformInfo AnalyzeForTrivialTransform(ArrayRef regions) { + const TrivialTransformInfo cannot_transform; + + if (regions.empty()) return cannot_transform; + + llvm::SmallVector calls; + calls.reserve(regions.size()); + + // Verify each region is a single call and collect these calls. + for (Region* region : regions) { + auto call = IsSingleCallRegion(*region); + if (!call.hasValue()) return cannot_transform; + calls.push_back(call.getValue()); + } + + llvm::SmallVector callees; + callees.reserve(regions.size()); + + CallOp call0 = calls[0]; + int num_args = call0.getNumOperands(); + + // Collect arguments of the first call. + llvm::SmallVector call0_args; + call0_args.reserve(num_args); + for (Value arg : call0.getArgOperands()) { + if (auto cast = dyn_cast_or_null(arg.getDefiningOp())) + arg = cast.getOperand(); + call0_args.push_back(arg); + } + + // Match arguments of rest of the calls with those of the first call. + for (auto call : calls) { + if (call != call0 && !MatchCallArgs(call, call0_args)) + return cannot_transform; + callees.push_back(call.getCallee()); + } + + return {true, callees, call0_args}; +} + +// Transform IfRegionOp to IfOp. +LogicalResult RegionControlFlowToFunctional::ConvertIfOp(IfRegionOp if_region) { + const TrivialTransformInfo tti = AnalyzeForTrivialTransform( + {&if_region.then_branch(), &if_region.else_branch()}); + + std::string then_name, else_name; + llvm::SmallVector extern_values; + + if (tti.can_transform) { + // We can transform to functional form trivially without outlining. + then_name = tti.callee_names[0].str(); + else_name = tti.callee_names[1].str(); + extern_values = tti.call_args; + } else { + // Collect external values that are used within the else and then bodies. + extern_values = CollectExternValues( + {&if_region.then_branch(), &if_region.else_branch()}); + + // These external values need to be added as inputs to the generated If. The + // order is determined by the order of these values the `extern_vales`. + + // Build the type for the outlined function. + llvm::SmallVector input_types; + input_types.reserve(extern_values.size()); + for (auto input : extern_values) input_types.push_back(input.getType()); + + FunctionType func_type = FunctionType::get( + input_types, if_region.getResultTypes(), if_region.getContext()); + + // Create 2 new functions with the input signature matching this order, + // and outline the `then` and `else` regions by moving the bodies of these + // regions into these functions. Replace tf.yield with a regular return. + then_name = GetName(if_region, "_then"); + ExtractSingleBlockRegion(if_region.then_branch(), func_type, then_name, + extern_values, worklist); + + else_name = GetName(if_region, "_else"); + ExtractSingleBlockRegion(if_region.else_branch(), func_type, else_name, + extern_values, worklist); + } + + // Once we have the `then` and `else` functions ready (either outlined or + // existing ones), replace the region based op with a functional control flow + // op. + OpBuilder builder(if_region); + auto if_op = builder.create( + if_region.getLoc(), if_region.getResultTypes(), if_region.cond(), + extern_values, then_name, else_name, if_region.is_stateless()); + if_region.replaceAllUsesWith(if_op.getResults()); + if_region.erase(); + return success(); +} + +void RegionControlFlowToFunctional::runOnOperation() { + ModuleOp module = getOperation(); + + // Seed worklist with all functions in the module. + worklist = llvm::to_vector<4>(module.getOps()); + + while (!worklist.empty()) { + FuncOp function = worklist.pop_back_val(); + + auto result = function.walk([&](Operation* op) { + if (IfRegionOp if_region = llvm::dyn_cast(op)) { + if (failed(ConvertIfOp(if_region))) { + if_region.emitOpError() << " failed to convert to functional form"; + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + + if (result.wasInterrupted()) return signalPassFailure(); + } +} + +} // namespace + +std::unique_ptr> +CreateTFRegionControlFlowToFunctional() { + return std::make_unique(); +} + +static PassRegistration pass( + "tf-region-control-flow-to-functional", + "Transform region bases control flow Ops to functional counterparts"); + +} // namespace TF +} // namespace mlir From cf83ab15b0b38c82bd136537de95cb44bcb2624e Mon Sep 17 00:00:00 2001 From: Stephan Herhut Date: Mon, 15 Jun 2020 09:59:35 -0700 Subject: [PATCH 0173/1390] Add support for f16 generated tanh kernel. PiperOrigin-RevId: 316487197 Change-Id: Id7daff4dc6264071c9371e9eb31c1f57ac044389 --- tensorflow/core/kernels/cubin_headers/BUILD | 1 + tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc | 4 +--- tensorflow/core/kernels/cwise_op_tanh.cc | 4 +--- .../kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc | 12 ++++++++++++ 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/cubin_headers/BUILD b/tensorflow/core/kernels/cubin_headers/BUILD index a7f810eeded..49ab1b8a911 100644 --- a/tensorflow/core/kernels/cubin_headers/BUILD +++ b/tensorflow/core/kernels/cubin_headers/BUILD @@ -32,6 +32,7 @@ gen_kernel_library( name = "tanh", tile_size = "256", types = [ + "f16", "f32", "f64", ], diff --git a/tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc index f3fb33a8635..59754a7260d 100644 --- a/tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc +++ b/tensorflow/core/kernels/cwise_op_gpu_tanh.cu.cc @@ -20,9 +20,7 @@ limitations under the License. namespace tensorflow { namespace functor { -#if MLIR_GENERATED_GPU_KERNELS_ENABLED -DEFINE_UNARY(tanh, Eigen::half); -#else +#ifndef MLIR_GENERATED_GPU_KERNELS_ENABLED DEFINE_UNARY3(tanh, Eigen::half, float, double); #endif DEFINE_SIMPLE_BINARY3(tanh_grad, Eigen::half, float, double); diff --git a/tensorflow/core/kernels/cwise_op_tanh.cc b/tensorflow/core/kernels/cwise_op_tanh.cc index 83152019608..1b6da56e537 100644 --- a/tensorflow/core/kernels/cwise_op_tanh.cc +++ b/tensorflow/core/kernels/cwise_op_tanh.cc @@ -21,9 +21,7 @@ REGISTER5(UnaryOp, CPU, "Tanh", functor::tanh, float, Eigen::half, double, complex64, complex128); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM -#if MLIR_GENERATED_GPU_KERNELS_ENABLED -REGISTER(UnaryOp, GPU, "Tanh", functor::tanh, Eigen::half); -#else +#ifndef MLIR_GENERATED_GPU_KERNELS_ENABLED REGISTER3(UnaryOp, GPU, "Tanh", functor::tanh, float, Eigen::half, double); #endif #endif diff --git a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc index ab4337c058f..40dd7c7e49e 100644 --- a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc +++ b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/kernels/cubin_headers/tanh_f16_kernel.h" #include "tensorflow/core/kernels/cubin_headers/tanh_f32_kernel.h" #include "tensorflow/core/kernels/cubin_headers/tanh_f64_kernel.h" #include "tensorflow/core/lib/core/errors.h" @@ -102,6 +103,14 @@ class MlirGenerateTanhOp : public OpKernel { std::mutex mu_; }; +class MlirGenerateTanhF16Op : public MlirGenerateTanhOp { + public: + explicit MlirGenerateTanhF16Op(OpKernelConstruction* ctx) + : MlirGenerateTanhOp(ctx) { + cubin_data_ = kTanhF16Kernel; + } +}; + class MlirGenerateTanhF32Op : public MlirGenerateTanhOp { public: explicit MlirGenerateTanhF32Op(OpKernelConstruction* ctx) @@ -119,6 +128,9 @@ class MlirGenerateTanhF64Op : public MlirGenerateTanhOp { }; } // namespace +REGISTER_KERNEL_BUILDER( + Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"), + MlirGenerateTanhF16Op); REGISTER_KERNEL_BUILDER( Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"), MlirGenerateTanhF32Op); From 4d1593ce36f4b555062a8107471c2651bd7d2d34 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 10:01:00 -0700 Subject: [PATCH 0174/1390] internal change only PiperOrigin-RevId: 316487501 Change-Id: Icdc8cd56a73e200fc4309e3103a40f7ef977d073 --- .../profiler/convert/xplane_to_op_stats.cc | 8 ++--- .../convert/xplane_to_op_stats_test.cc | 36 ------------------- 2 files changed, 2 insertions(+), 42 deletions(-) diff --git a/tensorflow/core/profiler/convert/xplane_to_op_stats.cc b/tensorflow/core/profiler/convert/xplane_to_op_stats.cc index 9d3aca9d831..eb2e13dbb4a 100644 --- a/tensorflow/core/profiler/convert/xplane_to_op_stats.cc +++ b/tensorflow/core/profiler/convert/xplane_to_op_stats.cc @@ -96,8 +96,7 @@ void SetRunEnvironment(int32 accelerator_count, RunEnvironment* env) { } void ProcessHostPlane(const XPlane* host_plane, bool use_device_step_events, - OpMetricsDb* op_metrics_db, StepEvents* step_events, - TfFunctionDb* tf_function_db) { + OpMetricsDb* op_metrics_db, StepEvents* step_events) { absl::flat_hash_map tf_ops = CollectTfOpsFromHostThreadsXPlane(*host_plane); OpMetricsDbCombiner combiner(op_metrics_db); @@ -108,8 +107,6 @@ void ProcessHostPlane(const XPlane* host_plane, bool use_device_step_events, CombineStepEvents(ConvertHostThreadsXLineToStepEvents( line, use_device_step_events, *step_events), step_events); - CombineTfFunctionDb(ConvertHostThreadsXLineToTfFunctionDb(line), - tf_function_db); }); } @@ -166,8 +163,7 @@ OpStats ConvertXSpaceToOpStats(const XSpace& space) { bool has_device = !device_planes.empty(); if (host_plane) { ProcessHostPlane(host_plane, has_device, - op_stats.mutable_host_op_metrics_db(), &step_events, - op_stats.mutable_tf_function_db()); + op_stats.mutable_host_op_metrics_db(), &step_events); } StepEvents nonoverlapped_step_events = ToNonOverlappedStepEvents(step_events); *op_stats.mutable_step_db() = diff --git a/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc b/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc index 1812a5592bc..138bcee72be 100644 --- a/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc +++ b/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc @@ -157,42 +157,6 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) { EXPECT_EQ(precision_stats.compute_32bit_ps(), 40); } -TEST(ConcertXPlaneToOpStats, TfFunctionTest) { - XSpace space; - XPlaneBuilder host_plane_builder(space.add_planes()); - host_plane_builder.SetName(kHostThreads); - host_plane_builder.ReserveLines(1); - std::string kFunctionName = "increment"; - - auto main_thread = host_plane_builder.GetOrCreateLine(0); - CreateTfFunctionCallEvent(&host_plane_builder, &main_thread, kFunctionName, - 10, 100, "traced-nonXla", 1); - CreateTfFunctionCallEvent(&host_plane_builder, &main_thread, kFunctionName, - 150, 20, "notTraced-nonXla", 1); - CreateTfFunctionCallEvent(&host_plane_builder, &main_thread, kFunctionName, - 200, 80, "traced-nonXla", 2); - - OpStats op_stats = ConvertXSpaceToOpStats(space); - const TfFunctionDb& tf_function_db = op_stats.tf_function_db(); - - EXPECT_EQ(tf_function_db.tf_functions().size(), 1); - EXPECT_EQ(tf_function_db.tf_functions().count(kFunctionName), 1); - const TfFunction& tf_function = - tf_function_db.tf_functions().at(kFunctionName); - EXPECT_EQ(tf_function.total_tracing_count(), 2); - EXPECT_EQ(tf_function.compiler(), OTHER_COMPILER); - const auto& metrics = tf_function.metrics(); - EXPECT_EQ(metrics.size(), 2); - EXPECT_EQ(metrics.count(TRACED_MODE), 1); - EXPECT_EQ(metrics.count(NOT_TRACED_MODE), 1); - const auto& traced_mode = metrics.at(TRACED_MODE); - EXPECT_EQ(traced_mode.count(), 2); - EXPECT_EQ(traced_mode.self_time_ps(), 180); - const auto& not_traced_mode = metrics.at(NOT_TRACED_MODE); - EXPECT_EQ(not_traced_mode.count(), 1); - EXPECT_EQ(not_traced_mode.self_time_ps(), 20); -} - TEST(ConvertXPlaneToOpStats, PropagateAndDedupErrors) { XSpace space; static constexpr char kError[] = "host: error"; From 8d32eb3bd10aceea68118556e500e87f5565a983 Mon Sep 17 00:00:00 2001 From: Prakalp Srivastava Date: Mon, 15 Jun 2020 10:12:43 -0700 Subject: [PATCH 0175/1390] Handle tf.Case in tf-tensor-list-ops-decomposition pass. PiperOrigin-RevId: 316490068 Change-Id: I8f9502c3b8361e767b6333428cffa68fe3d8a3ad --- .../tests/tensor_list_ops_decomposition.mlir | 62 ++++++++++++ .../tensor_list_ops_decomposition.cc | 97 +++++++++++-------- 2 files changed, 119 insertions(+), 40 deletions(-) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir index 49365db57f6..c453a3815f2 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir @@ -248,6 +248,68 @@ func @if_else(%arg0: tensor>>) -> tensor) -> () { + %elem_shape = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> + %max_size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor + // CHECK-NOT: tf.EmptyTensorList + %tl = "tf.EmptyTensorList"(%elem_shape, %max_size) : (tensor<0xi32>, tensor) -> tensor>> + %case_op = "tf.Case"(%arg0, %tl) {branches = [@branch_0, @branch_1, @branch_2]} + : (tensor, tensor>>) -> tensor>> + // CHECK: "tf.Slice" + %pop:2 = "tf.TensorListPopBack"(%case_op, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) + // CHECK-NOT: tf.TensorListPopBack + // CHECK: return + return +} +// CHECK: func @branch_0(%[[TARG0:.*]]: tensor<10xf32>, %[[TARG1:.*]]: tensor<1xi32>) -> (tensor<10xf32>, tensor<1xi32>) +func @branch_0(%arg0: tensor>>) -> tensor>> { + %elem = "tf._SomeOp"() : () -> tensor + // CHECK-NOT: "tf.TensorListPushBack" + // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" + // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[TARG1]], %[[CONST1]]) + // CHECK-NOT: "tf.TensorListPushBack" + %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> + // CHECK: return %[[UPDATE]], %[[ADD]] + return %push : tensor>> +} +// CHECK: func @branch_1(%[[EARG0:.*]]: tensor<10xf32>, %[[EARG1:.*]]: tensor<1xi32>) -> (tensor<10xf32>, tensor<1xi32>) +func @branch_1(%arg0: tensor>>) -> tensor>> { + %elem_shape = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NOT: "tf.TensorListPopBack" + // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[EARG0]]) + // CHECK: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[EARG1]], %[[CONST1_1]]) + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor + // CHECK-NOT: "tf.TensorListPopBack" + %pop:2 = "tf.TensorListPopBack"(%arg0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) + // CHECK: return %[[COPY]], %[[SUB]] + return %pop#0 : tensor>> +} +// CHECK: func @branch_2(%[[EARG0:.*]]: tensor<10xf32>, %[[EARG1:.*]]: tensor<1xi32>) -> (tensor<10xf32>, tensor<1xi32>) +func @branch_2(%arg0: tensor>>) -> tensor>> { + %elem_shape = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NOT: "tf.TensorListPopBack" + // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[EARG0]]) + // CHECK: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[EARG1]], %[[CONST1_1]]) + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor + // CHECK-NOT: "tf.TensorListPopBack" + %pop:2 = "tf.TensorListPopBack"(%arg0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) + // CHECK: return %[[COPY]], %[[SUB]] + return %pop#0 : tensor>> +} +// ----- + // Tests PartitionedCall/StatefulPartitionedCall. // CHECK-LABEL: func @main diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc index b2203c890e3..9733bfe2290 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc @@ -216,59 +216,62 @@ LogicalResult HandleWhileOp( return success(); } -LogicalResult HandleIfOp(TF::IfOp if_op, ModuleOp module, - llvm::SmallDenseMap* buffer_to_size, - llvm::StringMap* - decomposed_partitioned_call_callees) { +template +LogicalResult HandleCaseOrIfOp( + CaseOrIfOp op, ArrayRef branches, ModuleOp module, + llvm::SmallDenseMap* buffer_to_size, + llvm::StringMap* + decomposed_partitioned_call_callees) { // Rewrite the branches. - auto then_branch = module.lookupSymbol(if_op.then_branch()); - auto else_branch = module.lookupSymbol(if_op.else_branch()); - llvm::SmallDenseMap then_map; - llvm::SmallDenseMap else_map; + SmallVector, 2> branch_maps; + branch_maps.resize(branches.size()); auto find_arg_buffer_type = [&](int64_t index) -> llvm::Optional { - auto it = buffer_to_size->find(if_op.getOperand(index + 1)); + auto it = buffer_to_size->find(op.getOperand(index + 1)); if (it == buffer_to_size->end()) return llvm::None; return it->getFirst().getType(); }; auto arg_buffer_size_is_fixed = [&](int64_t index) { - return (*buffer_to_size)[if_op.getOperand(index + 1)].fixed; + return (*buffer_to_size)[op.getOperand(index + 1)].fixed; }; - OpBuilder builder(if_op); - ModifyFunctionSignature(then_branch, cutil::GetSizeType(builder), &then_map, - find_arg_buffer_type, arg_buffer_size_is_fixed); - ModifyFunctionSignature(else_branch, cutil::GetSizeType(builder), &else_map, - find_arg_buffer_type, arg_buffer_size_is_fixed); - const bool arg_no_changed = then_map.empty(); - if (failed(DecomposeTensorListOpsInternal( - &then_branch.front(), module, &then_map, - decomposed_partitioned_call_callees)) || - failed(DecomposeTensorListOpsInternal( - &else_branch.front(), module, &else_map, - decomposed_partitioned_call_callees))) { - return failure(); + OpBuilder builder(op); + for (const auto& pair : llvm::zip(branches, branch_maps)) { + FuncOp branch = std::get<0>(pair); + llvm::SmallDenseMap& branch_map = std::get<1>(pair); + ModifyFunctionSignature(branch, cutil::GetSizeType(builder), &branch_map, + find_arg_buffer_type, arg_buffer_size_is_fixed); + + if (failed(DecomposeTensorListOpsInternal( + &branch.front(), module, &branch_map, + decomposed_partitioned_call_callees))) + return failure(); } + + const bool arg_no_changed = branch_maps.front().empty(); auto output_buffer_to_size = - AddTensorListSizesToReturn(then_branch, then_map); - AddTensorListSizesToReturn(else_branch, else_map); + AddTensorListSizesToReturn(branches.front(), branch_maps.front()); + for (const auto& pair : llvm::drop_begin(llvm::zip(branches, branch_maps), 1)) + AddTensorListSizesToReturn(std::get<0>(pair), std::get<1>(pair)); + if (output_buffer_to_size.empty() && arg_no_changed) return success(); - // Recreate the If op. - auto new_if_operands = llvm::to_vector<8>(if_op.getOperands()); - for (int64_t i = 1; i < if_op.getNumOperands(); ++i) { - auto it = buffer_to_size->find(if_op.getOperand(i)); + + // Recreate the op. + auto new_operands = llvm::to_vector<8>(op.getOperands()); + for (int64_t i = 1; i < op.getNumOperands(); ++i) { + auto it = buffer_to_size->find(op.getOperand(i)); if (it == buffer_to_size->end()) continue; - new_if_operands.push_back(it->getSecond().size); + new_operands.push_back(it->getSecond().size); } - auto new_if = OpBuilder(if_op).create( - if_op.getLoc(), then_branch.getType().getResults(), new_if_operands, - if_op.getAttrs()); + FuncOp first_branch = branches.front(); + auto new_op = OpBuilder(op).create( + op.getLoc(), first_branch.getType().getResults(), new_operands, + op.getAttrs()); for (const auto& entry : output_buffer_to_size) { - (*buffer_to_size)[new_if.getResult(std::get<0>(entry))] = { - new_if.getResult(std::get<1>(entry)), std::get<2>(entry)}; + (*buffer_to_size)[new_op.getResult(std::get<0>(entry))] = { + new_op.getResult(std::get<1>(entry)), std::get<2>(entry)}; } - if_op.replaceAllUsesWith( - new_if.getResults().take_front(if_op.getNumResults())); - if_op.erase(); + op.replaceAllUsesWith(new_op.getResults().take_front(op.getNumResults())); + op.erase(); return success(); } @@ -710,8 +713,22 @@ LogicalResult DecomposeTensorListOpsInternal( return failure(); } } else if (auto if_op = llvm::dyn_cast(&op)) { - if (failed(HandleIfOp(if_op, module, buffer_to_size, - decomposed_partitioned_call_callees))) { + auto then_branch = module.lookupSymbol(if_op.then_branch()); + auto else_branch = module.lookupSymbol(if_op.else_branch()); + + if (failed(HandleCaseOrIfOp(if_op, {then_branch, else_branch}, module, + buffer_to_size, + decomposed_partitioned_call_callees))) { + return failure(); + } + } else if (auto case_op = llvm::dyn_cast(&op)) { + SmallVector branches; + for (auto branch_symbol : case_op.branches()) { + branches.push_back(module.lookupSymbol( + branch_symbol.cast())); + } + if (failed(HandleCaseOrIfOp(case_op, branches, module, buffer_to_size, + decomposed_partitioned_call_callees))) { return failure(); } } else if (auto pcall = llvm::dyn_cast(&op)) { From 54f1e52b95bb5f4c5a204b736a9a815ace108e68 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 10:12:45 -0700 Subject: [PATCH 0176/1390] fix deprecated message. PiperOrigin-RevId: 316490075 Change-Id: I68d52d80ccf9c2f77f870dd7f4f1d146e1e44a46 --- tensorflow/python/keras/saving/saving_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/keras/saving/saving_utils.py b/tensorflow/python/keras/saving/saving_utils.py index 9a407f64faa..3c9c33531bf 100644 --- a/tensorflow/python/keras/saving/saving_utils.py +++ b/tensorflow/python/keras/saving/saving_utils.py @@ -17,7 +17,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections import copy import os import six @@ -31,6 +30,7 @@ from tensorflow.python.keras.utils import generic_utils from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import nest +from tensorflow.python.util.compat import collections_abc def extract_model_metrics(model): @@ -79,7 +79,8 @@ def model_input_signature(model, keep_original_batch_size=False): return None input_specs = _enforce_names_consistency(input_specs) # Return a list with a single element as the model's input signature. - if isinstance(input_specs, collections.Sequence) and len(input_specs) == 1: + if isinstance(input_specs, + collections_abc.Sequence) and len(input_specs) == 1: # Note that the isinstance check filters out single-element dictionaries, # which should also be wrapped as a single-element list. return input_specs From e3b8a8e7207a4bc16287af77a5e9fd0b0df86db2 Mon Sep 17 00:00:00 2001 From: Andrew Audibert Date: Mon, 15 Jun 2020 10:39:49 -0700 Subject: [PATCH 0177/1390] [tf.data service] Increase default number of uncompress threads to 4. A single thread may not be able to uncompress data as quickly as it is requested. PiperOrigin-RevId: 316496226 Change-Id: I25842acc4485c509987654bcec8f38f4b4a067b7 --- .../data/experimental/ops/data_service_ops.py | 3 +-- .../kernel_tests/data_service_ops_test.py | 21 +++++++------------ 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/tensorflow/python/data/experimental/ops/data_service_ops.py b/tensorflow/python/data/experimental/ops/data_service_ops.py index dd81614fa45..39790d843ba 100644 --- a/tensorflow/python/data/experimental/ops/data_service_ops.py +++ b/tensorflow/python/data/experimental/ops/data_service_ops.py @@ -241,8 +241,7 @@ def _distribute(processing_mode, # TODO(b/157105111): Make this an autotuned parallel map when we have a way # to limit memory usage. dataset = dataset.map( - lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec), - num_parallel_calls=4) + lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec)) # Disable autosharding for shared jobs. if job_name: diff --git a/tensorflow/python/data/kernel_tests/data_service_ops_test.py b/tensorflow/python/data/kernel_tests/data_service_ops_test.py index 440a4f46a20..d316009ce0c 100644 --- a/tensorflow/python/data/kernel_tests/data_service_ops_test.py +++ b/tensorflow/python/data/kernel_tests/data_service_ops_test.py @@ -201,18 +201,13 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): self._new_worker = server_lib.WorkerServer( port=port, master_address=self._master._address, protocol=PROTOCOL) - # There may have been some elements prefetched from the first worker - # before it was stopped. - while True: - val = next(iterator).numpy() - if val == 0: - break - # The dataset starts over now that we read from the new worker. - # TODO(b/157086991): Iterate until end of sequence when we support - # detecting lost workers. - for i in range(1, num_elements // 2): + for i in range(num_elements): val = next(iterator).numpy() + if val == midpoint and i != midpoint: + # There may have been one last element prefetched from the first worker + # before it was stopped. + val = next(iterator).numpy() self.assertEqual(i, val) @combinations.generate(test_base.eager_only_combinations()) @@ -296,7 +291,7 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.eager_only_combinations()) def testSharedJobNameRepeat(self): - num_elements = 100 + num_elements = 10 num_repetitions = 3 master_address = self.create_cluster(1) ds = dataset_ops.Dataset.range(num_elements) @@ -307,9 +302,9 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase): results = [] iter1 = iter(ds1) iter2 = iter(ds2) - for _ in range(((num_elements * num_repetitions) // 3)): + for _ in range(((num_elements * num_repetitions) // 2) - 1): results.append(next(iter1).numpy()) - for _ in range(((num_elements * num_repetitions) // 3)): + for _ in range(((num_elements * num_repetitions) // 2) - 1): results.append(next(iter2).numpy()) for elem in iter1: results.append(elem.numpy()) From a2dc78458bd1267625b16469730b8b743ac1514c Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Mon, 15 Jun 2020 10:40:48 -0700 Subject: [PATCH 0178/1390] Fork tracking/util_with_v1_optimizers_test to keras/tests. PiperOrigin-RevId: 316496459 Change-Id: I9563ba20c0d3cd72bc0651e49ab9a48e0b355d34 --- tensorflow/python/keras/tests/BUILD | 31 + .../tracking_util_with_v1_optimizers_test.py | 703 ++++++++++++++++++ tensorflow/python/training/tracking/BUILD | 15 - .../tracking/util_with_v1_optimizers_test.py | 664 ----------------- 4 files changed, 734 insertions(+), 679 deletions(-) create mode 100644 tensorflow/python/keras/tests/tracking_util_with_v1_optimizers_test.py diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD index ad52d33abc6..d03b1bd1ee8 100644 --- a/tensorflow/python/keras/tests/BUILD +++ b/tensorflow/python/keras/tests/BUILD @@ -435,6 +435,37 @@ tf_py_test( ], ) +tf_py_test( + name = "tracking_util_with_v1_optimizers_test", + srcs = ["tracking_util_with_v1_optimizers_test.py"], + tags = [ + "notsan", # b/74395663 + ], + deps = [ + "//tensorflow/python:checkpoint_management", + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:extra_py_tests_deps", + "//tensorflow/python:framework_ops", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:saver", + "//tensorflow/python:state_ops", + "//tensorflow/python:training_lib", + "//tensorflow/python:training_util", + "//tensorflow/python/distribute:mirrored_strategy", + "//tensorflow/python/eager:backprop", + "//tensorflow/python/eager:context", + "//tensorflow/python/eager:def_function", + "//tensorflow/python/eager:test", + "//tensorflow/python/keras/engine", + "//tensorflow/python/keras/layers:core", + "//tensorflow/python/training/tracking", + "//tensorflow/python/training/tracking:graph_view", + "//tensorflow/python/training/tracking:util", + ], +) + py_library( name = "get_config_samples", srcs = ["get_config_samples.py"], diff --git a/tensorflow/python/keras/tests/tracking_util_with_v1_optimizers_test.py b/tensorflow/python/keras/tests/tracking_util_with_v1_optimizers_test.py new file mode 100644 index 00000000000..b6711ea6fe5 --- /dev/null +++ b/tensorflow/python/keras/tests/tracking_util_with_v1_optimizers_test.py @@ -0,0 +1,703 @@ +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# 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 +# +# http://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. +# ============================================================================== +"""Tests for object-based saving which use tf.train.* optimizers.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os + +import six + +from tensorflow.python.distribute import mirrored_strategy +from tensorflow.python.eager import backprop +from tensorflow.python.eager import context +from tensorflow.python.eager import def_function +from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.keras.engine import training +from tensorflow.python.keras.layers import core +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.training import adam +from tensorflow.python.training import checkpoint_management +from tensorflow.python.training import saver as saver_lib +from tensorflow.python.training import training_util +from tensorflow.python.training.tracking import graph_view +from tensorflow.python.training.tracking import tracking +from tensorflow.python.training.tracking import util as trackable_utils + + +class NonLayerTrackable(tracking.AutoTrackable): + + def __init__(self): + super(NonLayerTrackable, self).__init__() + self.a_variable = trackable_utils.add_variable( + self, name="a_variable", shape=[]) + + +# pylint: disable=not-callable +class MyModel(training.Model): + """A concrete Model for testing.""" + + def __init__(self): + super(MyModel, self).__init__() + self._named_dense = core.Dense(1, use_bias=True) + self._second = core.Dense(1, use_bias=False) + # We can still track Trackables which aren't Layers. + self._non_layer = NonLayerTrackable() + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +class CheckpointingTests(test.TestCase): + + @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) + def testNamingWithOptimizer(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + # A nuisance Model using the same optimizer. Its slot variables should not + # go in the checkpoint, since it is never depended on. + other_model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + if context.executing_eagerly(): + optimizer.minimize( + lambda: model(input_value), + global_step=optimizer_step) + optimizer.minimize( + lambda: other_model(input_value), + global_step=optimizer_step) + else: + train_op = optimizer.minimize( + model(input_value), global_step=optimizer_step) + optimizer.minimize( + other_model(input_value), + global_step=optimizer_step) + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + named_variables, serialized_graph, _ = graph_view.ObjectGraphView( + root_trackable).serialize_object_graph() + expected_checkpoint_names = ( + # Created in the root node, so no prefix. + "optimizer_step", + "model/_second/kernel", + "model/_named_dense/kernel", + "model/_named_dense/bias", + # non-Layer dependency of the model + "model/_non_layer/a_variable", + # The optimizer creates two non-slot variables + "optimizer/beta1_power", + "optimizer/beta2_power", + # Slot variables + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", + "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", + "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v", + ) + suffix = "/.ATTRIBUTES/VARIABLE_VALUE" + expected_checkpoint_names = [ + name + suffix for name in expected_checkpoint_names] + named_variables = {v.name: v for v in named_variables} + six.assertCountEqual(self, expected_checkpoint_names, + named_variables.keys()) + # Check that we've mapped to the right variable objects (not exhaustive) + self.assertEqual( + "global_step", + named_variables["optimizer_step" + suffix].full_name) + self.assertEqual( + "my_model/dense_1/kernel", + named_variables["model/_second/kernel" + suffix].full_name) + self.assertEqual( + "my_model/dense/kernel", + named_variables["model/_named_dense/kernel" + suffix].full_name) + self.assertEqual( + "beta1_power", + named_variables["optimizer/beta1_power" + suffix].full_name) + self.assertEqual( + "beta2_power", + named_variables["optimizer/beta2_power" + suffix].full_name) + # Spot check the generated protocol buffers. + self.assertEqual("optimizer", + serialized_graph.nodes[0].children[1].local_name) + optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ + 1].node_id] + self.assertEqual("beta1_power", + optimizer_node.children[0].local_name) + self.assertEqual("beta1_power", + serialized_graph.nodes[optimizer_node.children[0].node_id] + .attributes[0].full_name) + self.assertEqual( + "my_model/dense/kernel", + serialized_graph.nodes[optimizer_node.slot_variables[0] + .original_variable_node_id] + .attributes[0].full_name) + # We strip off the :0 suffix, as variable.name-based saving does. + self.assertEqual( + "my_model/dense/kernel/Adam", + serialized_graph.nodes[optimizer_node.slot_variables[0] + .slot_variable_node_id] + .attributes[0].full_name) + self.assertEqual( + "my_model/dense/kernel/Adam:0", + optimizer.get_slot( + var=model._named_dense.kernel, + name="m").name) + self.assertEqual( + "model/_named_dense/kernel" + suffix, + serialized_graph.nodes[ + optimizer_node.slot_variables[0] + .original_variable_node_id].attributes[0].checkpoint_key) + self.assertEqual("m", optimizer_node.slot_variables[0].slot_name) + self.assertEqual( + "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m" + suffix, + serialized_graph.nodes[ + optimizer_node.slot_variables[0] + .slot_variable_node_id].attributes[0].checkpoint_key) + + @test_util.run_in_graph_and_eager_modes + def testSaveRestore(self): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model) + input_value = constant_op.constant([[3.]]) + if context.executing_eagerly(): + optimizer.minimize( + lambda: model(input_value)) + else: + train_op = optimizer.minimize(model(input_value)) + # TODO(allenl): Make initialization more pleasant when graph building. + root_trackable.save_counter # pylint: disable=pointless-statement + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + prefix = os.path.join(self.get_temp_dir(), "ckpt") + self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.])) + m_bias_slot = optimizer.get_slot(model._named_dense.variables[1], "m") + self.evaluate(state_ops.assign(m_bias_slot, [1.5])) + save_path = root_trackable.save(file_prefix=prefix) + self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.])) + self.evaluate(state_ops.assign(root_trackable.save_counter, 3)) + optimizer_variables = self.evaluate(optimizer.variables()) + self.evaluate(state_ops.assign(m_bias_slot, [-2.])) + # Immediate restoration + status = root_trackable.restore(save_path=save_path).assert_consumed() + status.run_restore_ops() + self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1])) + self.assertAllEqual(1, self.evaluate(root_trackable.save_counter)) + self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) + if not context.executing_eagerly(): + return # Restore-on-create is only supported when executing eagerly + on_create_model = MyModel() + on_create_optimizer = adam.AdamOptimizer( + 0.001, + # Preserve beta1_power and beta2_power when applying gradients so we can + # test that they've been restored correctly. + beta1=1.0, + beta2=1.0) + on_create_root = trackable_utils.Checkpoint( + optimizer=on_create_optimizer, model=on_create_model) + # Deferred restoration + status = on_create_root.restore(save_path=save_path) + status.assert_nontrivial_match() + status.assert_existing_objects_matched() + with self.assertRaises(AssertionError): + status.assert_consumed() + on_create_model(constant_op.constant([[3.]])) # create variables + self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) + self.assertAllEqual([42.], + self.evaluate( + on_create_model._named_dense.variables[1])) + on_create_m_bias_slot = on_create_optimizer.get_slot( + on_create_model._named_dense.variables[1], "m") + status.assert_existing_objects_matched() + with self.assertRaises(AssertionError): + status.assert_consumed() + # Optimizer slot variables are created when the original variable is + # restored. + self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot)) + self.assertAllEqual(optimizer_variables[2:], + self.evaluate(on_create_optimizer.variables())) + dummy_var = resource_variable_ops.ResourceVariable([1.]) + on_create_optimizer.minimize(loss=dummy_var.read_value) + status.assert_existing_objects_matched() + status.assert_consumed() + beta1_power, beta2_power = on_create_optimizer._get_beta_accumulators() + self.assertAllEqual(optimizer_variables[0], self.evaluate(beta1_power)) + self.assertAllEqual(optimizer_variables[1], self.evaluate(beta2_power)) + + # TODO(allenl): Debug garbage created by this test in python3. + def testDeferredRestorationUsageEager(self): + """An idiomatic eager execution example.""" + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + optimizer_step=training_util.get_or_create_global_step()) + root.restore(checkpoint_management.latest_checkpoint( + checkpoint_directory)) + for _ in range(num_training_steps): + # TODO(allenl): Use a Dataset and serialize/checkpoint it. + input_value = constant_op.constant([[3.]]) + optimizer.minimize( + lambda: model(input_value), # pylint: disable=cell-var-from-loop + global_step=root.optimizer_step) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer_step.numpy()) + + def testEagerDistributionStrategy(self): + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + def _train_fn(optimizer, model): + input_value = constant_op.constant([[3.]]) + optimizer.minimize( + functools.partial(model, input_value), + global_step=root.optimizer_step) + + strategy = mirrored_strategy.MirroredStrategy() + with strategy.scope(): + for training_continuation in range(3): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, + model=model, + optimizer_step=training_util.get_or_create_global_step()) + root.restore( + checkpoint_management.latest_checkpoint(checkpoint_directory)) + + for _ in range(num_training_steps): + strategy.extended.call_for_each_replica( + functools.partial(_train_fn, optimizer, model)) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer_step.numpy()) + + def testGraphDistributionStrategy(self): + self.skipTest("b/121381184") + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + def _train_fn(optimizer, model): + input_value = constant_op.constant([[3.]]) + return optimizer.minimize( + functools.partial(model, input_value), + global_step=root.optimizer_step) + + for training_continuation in range(3): + with ops.Graph().as_default(): + strategy = mirrored_strategy.MirroredStrategy() + with strategy.scope(): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + optimizer_step=training_util.get_or_create_global_step()) + status = root.restore(checkpoint_management.latest_checkpoint( + checkpoint_directory)) + train_op = strategy.extended.call_for_each_replica( + functools.partial(_train_fn, optimizer, model)) + with self.session() as session: + if training_continuation > 0: + status.assert_consumed() + status.initialize_or_restore() + for _ in range(num_training_steps): + session.run(train_op) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer_step.numpy()) + + def testUsageGraph(self): + """Expected usage when graph building.""" + with context.graph_mode(): + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + with ops.Graph().as_default(): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.CheckpointV1( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + input_value = constant_op.constant([[3.]]) + train_op = optimizer.minimize( + model(input_value), + global_step=root.global_step) + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + with self.session(graph=ops.get_default_graph()) as session: + status = root.restore(save_path=checkpoint_path) + status.initialize_or_restore(session=session) + if checkpoint_path is None: + self.assertEqual(0, training_continuation) + with self.assertRaises(AssertionError): + status.assert_consumed() + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + else: + status.assert_consumed() + status.assert_existing_objects_matched() + for _ in range(num_training_steps): + session.run(train_op) + root.save(file_prefix=checkpoint_prefix, session=session) + self.assertEqual((training_continuation + 1) * num_training_steps, + session.run(root.global_step)) + self.assertEqual(training_continuation + 1, + session.run(root.save_counter)) + + @test_util.run_in_graph_and_eager_modes + def testAgnosticUsage(self): + """Graph/eager agnostic usage.""" + # Does create garbage when executing eagerly due to ops.Graph() creation. + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + for training_continuation in range(3): + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + manager = checkpoint_management.CheckpointManager( + root, checkpoint_directory, max_to_keep=1) + status = root.restore(save_path=manager.latest_checkpoint) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial( + optimizer.minimize, + functools.partial(model, input_value), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + for _ in range(num_training_steps): + train_fn() + manager.save() + self.assertEqual((training_continuation + 1) * num_training_steps, + self.evaluate(root.global_step)) + self.assertEqual(training_continuation + 1, + self.evaluate(root.save_counter)) + + # pylint: disable=cell-var-from-loop + @test_util.run_in_graph_and_eager_modes + def testWithDefun(self): + num_training_steps = 2 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + for training_continuation in range(3): + with test_util.device(use_gpu=True): + model = MyModel() + # Don't actually train so we can test variable values + optimizer = adam.AdamOptimizer(0.) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + def train_fn(): + @def_function.function + def _call_model(x): + return model(x) + with backprop.GradientTape() as tape: + loss = _call_model(constant_op.constant([[3.]])) + gradients = tape.gradient(loss, model.variables) + return optimizer.apply_gradients(zip(gradients, model.variables), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial( + self.evaluate, train_fn()) + status.initialize_or_restore() + for _ in range(num_training_steps): + train_fn() + if training_continuation > 0: + status.assert_consumed() + self.assertAllClose([[42.]], self.evaluate(model.variables[0])) + else: + self.evaluate(model.variables[0].assign([[42.]])) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + self.evaluate(root.global_step)) + self.assertEqual(training_continuation + 1, + self.evaluate(root.save_counter)) + # pylint: enable=cell-var-from-loop + + def _get_checkpoint_name(self, name): + root = tracking.AutoTrackable() + trackable_utils.add_variable( + root, name=name, shape=[1, 2], dtype=dtypes.float64) + (named_variable,), _, _ = trackable_utils._serialize_object_graph( + root, saveables_cache=None) + with ops.name_scope("root/" + named_variable.name): + pass # Make sure we can use this as an op name if we prefix it. + return named_variable.name + + def testAnonymousVarsInInit(self): + + class Model(training.Model): + + def __init__(self): + super(Model, self).__init__() + self.w = resource_variable_ops.ResourceVariable(0.0) + self.b = resource_variable_ops.ResourceVariable(0.0) + self.vars = [self.w, self.b] + + def call(self, x): + return x * self.w + self.b + + with context.eager_mode(): + model = Model() + optimizer = adam.AdamOptimizer(learning_rate=0.05) + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + checkpoint = trackable_utils.Checkpoint( + model=model, optimizer=optimizer) + for _ in range(2): + checkpoint.save(checkpoint_prefix) + with backprop.GradientTape() as tape: + loss = (constant_op.constant(1.) + - model(constant_op.constant(1.))) ** 2 + grad = tape.gradient(loss, model.vars) + optimizer.apply_gradients( + [(g, v) for g, v in zip(grad, model.vars)]) + + @test_util.run_in_graph_and_eager_modes + def test_initialize_if_not_restoring(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + optimizer_only_prefix = os.path.join(checkpoint_directory, "opt") + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + model=model, # Do not save the optimizer with the checkpoint. + global_step=training_util.get_or_create_global_step()) + optimizer_checkpoint = trackable_utils.Checkpoint( + optimizer=optimizer) + + checkpoint_path = checkpoint_management.latest_checkpoint( + checkpoint_directory) + status = root.restore(save_path=checkpoint_path) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial( + optimizer.minimize, + functools.partial(model, input_value), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + self.evaluate([v.initializer for v in optimizer.variables()]) + train_fn() + model_save_path = root.save(file_prefix=checkpoint_prefix) + self.evaluate(optimizer.variables()[0].assign(42.)) + optimizer_save_path = optimizer_checkpoint.save(optimizer_only_prefix) + + # Restore into a graph with the optimizer + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + status = root.restore(save_path=model_save_path) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial( + optimizer.minimize, + functools.partial(model, input_value), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + status.initialize_or_restore() + train_fn() + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + with self.assertRaises(AssertionError): + status.assert_consumed() + + # Make sure initialization doesn't clobber later restores + with test_util.device(use_gpu=True): + model = MyModel() + optimizer = adam.AdamOptimizer(0.001, beta1=1.0) + root = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, + global_step=training_util.get_or_create_global_step()) + opt_root = trackable_utils.Checkpoint( + optimizer=optimizer) + status = root.restore(save_path=model_save_path) + init_only_optimizer_status = opt_root.restore(save_path=None) + optimizer_status = opt_root.restore(save_path=optimizer_save_path) + input_value = constant_op.constant([[3.]]) + train_fn = functools.partial( + optimizer.minimize, + functools.partial(model, input_value), + global_step=root.global_step) + if not context.executing_eagerly(): + train_fn = functools.partial(self.evaluate, train_fn()) + optimizer_status.run_restore_ops() + status.initialize_or_restore() + init_only_optimizer_status.initialize_or_restore() + train_fn() + self.assertEqual(42., self.evaluate(optimizer.variables()[0])) + + +class CheckpointCompatibilityTests(test.TestCase): + + def _initialized_model(self): + input_value = constant_op.constant([[3.]]) + model = MyModel() + optimizer = adam.AdamOptimizer(0.001) + optimizer_step = training_util.get_or_create_global_step() + root_trackable = trackable_utils.Checkpoint( + optimizer=optimizer, model=model, optimizer_step=optimizer_step) + train_op = optimizer.minimize( + functools.partial(model, input_value), + global_step=optimizer_step) + self.evaluate(trackable_utils.gather_initializers( + root_trackable)) + self.evaluate(train_op) + # A regular variable, a slot variable, and a non-slot Optimizer variable + # with known values to check when loading. + self.evaluate(model._named_dense.bias.assign([1.])) + self.evaluate(optimizer.get_slot( + var=model._named_dense.bias, name="m").assign([2.])) + beta1_power, _ = optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(3.)) + return root_trackable + + def _set_sentinels(self, root_trackable): + self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) + self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, name="m") + .assign([102.])) + beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() + self.evaluate(beta1_power.assign(103.)) + + def _check_sentinels(self, root_trackable): + self.assertAllEqual( + [1.], self.evaluate(root_trackable.model._named_dense.bias)) + self.assertAllEqual([2.], self.evaluate( + root_trackable.optimizer.get_slot( + var=root_trackable.model._named_dense.bias, name="m"))) + beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() + self.assertAllEqual(3., self.evaluate(beta1_power)) + + def _write_name_based_checkpoint(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph) as session: + root = self._initialized_model() + name_saver = saver_lib.Saver() + return name_saver.save( + sess=session, save_path=checkpoint_prefix, + global_step=root.optimizer_step) + + @test_util.run_in_graph_and_eager_modes + def testLoadFromNameBasedSaver(self): + """Save a name-based checkpoint, load it using the object-based API.""" + with test_util.device(use_gpu=True): + save_path = self._write_name_based_checkpoint() + root = self._initialized_model() + self._set_sentinels(root) + with self.assertRaises(AssertionError): + self._check_sentinels(root) + object_saver = trackable_utils.TrackableSaver( + graph_view.ObjectGraphView(root)) + self._set_sentinels(root) + status = object_saver.restore(save_path) + if context.executing_eagerly(): + self._check_sentinels(root) + if context.executing_eagerly(): + status.assert_consumed() + status.assert_existing_objects_matched() + status.assert_nontrivial_match() + else: + # When graph building, we haven't read any keys, so we don't know + # whether the restore will be complete. + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_consumed() + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_existing_objects_matched() + with self.assertRaisesRegexp(AssertionError, "not restored"): + status.assert_nontrivial_match() + status.run_restore_ops() + self._check_sentinels(root) + self._set_sentinels(root) + status = object_saver.restore(save_path) + status.initialize_or_restore() + self._check_sentinels(root) + # Check that there is no error when keys are missing from the name-based + # checkpoint. + root.not_in_name_checkpoint = resource_variable_ops.ResourceVariable([1.]) + status = object_saver.restore(save_path) + with self.assertRaises(AssertionError): + status.assert_existing_objects_matched() + + def testSaveGraphLoadEager(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph): + root = self._initialized_model() + save_path = root.save(file_prefix=checkpoint_prefix) + with context.eager_mode(): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed() + self._check_sentinels(root) + + def testSaveEagerLoadGraph(self): + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + with context.eager_mode(): + root = self._initialized_model() + save_path = root.save(file_prefix=checkpoint_prefix) + with context.graph_mode(): + save_graph = ops.Graph() + with save_graph.as_default(), self.session( + graph=save_graph): + root = self._initialized_model() + self._set_sentinels(root) + root.restore(save_path).assert_consumed().run_restore_ops() + self._check_sentinels(root) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/training/tracking/BUILD b/tensorflow/python/training/tracking/BUILD index 36ca3cf4b66..88dfd8eba55 100644 --- a/tensorflow/python/training/tracking/BUILD +++ b/tensorflow/python/training/tracking/BUILD @@ -223,34 +223,19 @@ tf_py_test( "notsan", # b/74395663 ], deps = [ - ":base", - ":graph_view", ":tracking", ":util", - "//tensorflow/python:checkpoint_management", - "//tensorflow/python:constant_op", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", "//tensorflow/python:init_ops", - "//tensorflow/python:pywrap_tensorflow", "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:saver", "//tensorflow/python:session", "//tensorflow/python:state_ops", "//tensorflow/python:template", "//tensorflow/python:training", - "//tensorflow/python:training_util", "//tensorflow/python:variable_scope", - "//tensorflow/python/distribute:mirrored_strategy", - "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", - "//tensorflow/python/eager:def_function", "//tensorflow/python/eager:test", - "//tensorflow/python/keras:engine", - "//tensorflow/python/keras/layers", - "@absl_py//absl/testing:parameterized", "@six_archive//:six", ], ) diff --git a/tensorflow/python/training/tracking/util_with_v1_optimizers_test.py b/tensorflow/python/training/tracking/util_with_v1_optimizers_test.py index d4857677046..a5af8e1f876 100644 --- a/tensorflow/python/training/tracking/util_with_v1_optimizers_test.py +++ b/tensorflow/python/training/tracking/util_with_v1_optimizers_test.py @@ -17,486 +17,27 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import functools import os import six from tensorflow.python.client import session as session_lib -from tensorflow.python.distribute import mirrored_strategy -from tensorflow.python.eager import backprop from tensorflow.python.eager import context -from tensorflow.python.eager import def_function from tensorflow.python.eager import test -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import training -from tensorflow.python.keras.layers import core from tensorflow.python.ops import init_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import template from tensorflow.python.ops import variable_scope from tensorflow.python.training import adam -from tensorflow.python.training import checkpoint_management -from tensorflow.python.training import saver as saver_lib -from tensorflow.python.training import training_util -from tensorflow.python.training.tracking import graph_view from tensorflow.python.training.tracking import tracking from tensorflow.python.training.tracking import util as trackable_utils -class NonLayerTrackable(tracking.AutoTrackable): - - def __init__(self): - super(NonLayerTrackable, self).__init__() - self.a_variable = trackable_utils.add_variable( - self, name="a_variable", shape=[]) - - -# pylint: disable=not-callable -class MyModel(training.Model): - """A concrete Model for testing.""" - - def __init__(self): - super(MyModel, self).__init__() - self._named_dense = core.Dense(1, use_bias=True) - self._second = core.Dense(1, use_bias=False) - # We can still track Trackables which aren't Layers. - self._non_layer = NonLayerTrackable() - - def call(self, values): - ret = self._second(self._named_dense(values)) - return ret - - class CheckpointingTests(test.TestCase): - @test_util.run_in_graph_and_eager_modes(assert_no_eager_garbage=True) - def testNamingWithOptimizer(self): - input_value = constant_op.constant([[3.]]) - model = MyModel() - # A nuisance Model using the same optimizer. Its slot variables should not - # go in the checkpoint, since it is never depended on. - other_model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - optimizer_step = training_util.get_or_create_global_step() - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, optimizer_step=optimizer_step) - if context.executing_eagerly(): - optimizer.minimize( - lambda: model(input_value), - global_step=optimizer_step) - optimizer.minimize( - lambda: other_model(input_value), - global_step=optimizer_step) - else: - train_op = optimizer.minimize( - model(input_value), global_step=optimizer_step) - optimizer.minimize( - other_model(input_value), - global_step=optimizer_step) - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - named_variables, serialized_graph, _ = graph_view.ObjectGraphView( - root_trackable).serialize_object_graph() - expected_checkpoint_names = ( - # Created in the root node, so no prefix. - "optimizer_step", - "model/_second/kernel", - "model/_named_dense/kernel", - "model/_named_dense/bias", - # non-Layer dependency of the model - "model/_non_layer/a_variable", - # The optimizer creates two non-slot variables - "optimizer/beta1_power", - "optimizer/beta2_power", - # Slot variables - "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", - "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", - "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m", - "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/v", - "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/m", - "model/_named_dense/bias/.OPTIMIZER_SLOT/optimizer/v", - ) - suffix = "/.ATTRIBUTES/VARIABLE_VALUE" - expected_checkpoint_names = [ - name + suffix for name in expected_checkpoint_names] - named_variables = {v.name: v for v in named_variables} - six.assertCountEqual(self, expected_checkpoint_names, - named_variables.keys()) - # Check that we've mapped to the right variable objects (not exhaustive) - self.assertEqual( - "global_step", - named_variables["optimizer_step" + suffix].full_name) - self.assertEqual( - "my_model/dense_1/kernel", - named_variables["model/_second/kernel" + suffix].full_name) - self.assertEqual( - "my_model/dense/kernel", - named_variables["model/_named_dense/kernel" + suffix].full_name) - self.assertEqual( - "beta1_power", - named_variables["optimizer/beta1_power" + suffix].full_name) - self.assertEqual( - "beta2_power", - named_variables["optimizer/beta2_power" + suffix].full_name) - # Spot check the generated protocol buffers. - self.assertEqual("optimizer", - serialized_graph.nodes[0].children[1].local_name) - optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ - 1].node_id] - self.assertEqual("beta1_power", - optimizer_node.children[0].local_name) - self.assertEqual("beta1_power", - serialized_graph.nodes[optimizer_node.children[0].node_id] - .attributes[0].full_name) - self.assertEqual( - "my_model/dense/kernel", - serialized_graph.nodes[optimizer_node.slot_variables[0] - .original_variable_node_id] - .attributes[0].full_name) - # We strip off the :0 suffix, as variable.name-based saving does. - self.assertEqual( - "my_model/dense/kernel/Adam", - serialized_graph.nodes[optimizer_node.slot_variables[0] - .slot_variable_node_id] - .attributes[0].full_name) - self.assertEqual( - "my_model/dense/kernel/Adam:0", - optimizer.get_slot( - var=model._named_dense.kernel, - name="m").name) - self.assertEqual( - "model/_named_dense/kernel" + suffix, - serialized_graph.nodes[ - optimizer_node.slot_variables[0] - .original_variable_node_id].attributes[0].checkpoint_key) - self.assertEqual("m", optimizer_node.slot_variables[0].slot_name) - self.assertEqual( - "model/_named_dense/kernel/.OPTIMIZER_SLOT/optimizer/m" + suffix, - serialized_graph.nodes[ - optimizer_node.slot_variables[0] - .slot_variable_node_id].attributes[0].checkpoint_key) - - @test_util.run_in_graph_and_eager_modes - def testSaveRestore(self): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model) - input_value = constant_op.constant([[3.]]) - if context.executing_eagerly(): - optimizer.minimize( - lambda: model(input_value)) - else: - train_op = optimizer.minimize(model(input_value)) - # TODO(allenl): Make initialization more pleasant when graph building. - root_trackable.save_counter # pylint: disable=pointless-statement - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - prefix = os.path.join(self.get_temp_dir(), "ckpt") - self.evaluate(state_ops.assign(model._named_dense.variables[1], [42.])) - m_bias_slot = optimizer.get_slot(model._named_dense.variables[1], "m") - self.evaluate(state_ops.assign(m_bias_slot, [1.5])) - save_path = root_trackable.save(file_prefix=prefix) - self.evaluate(state_ops.assign(model._named_dense.variables[1], [43.])) - self.evaluate(state_ops.assign(root_trackable.save_counter, 3)) - optimizer_variables = self.evaluate(optimizer.variables()) - self.evaluate(state_ops.assign(m_bias_slot, [-2.])) - # Immediate restoration - status = root_trackable.restore(save_path=save_path).assert_consumed() - status.run_restore_ops() - self.assertAllEqual([42.], self.evaluate(model._named_dense.variables[1])) - self.assertAllEqual(1, self.evaluate(root_trackable.save_counter)) - self.assertAllEqual([1.5], self.evaluate(m_bias_slot)) - if not context.executing_eagerly(): - return # Restore-on-create is only supported when executing eagerly - on_create_model = MyModel() - on_create_optimizer = adam.AdamOptimizer( - 0.001, - # Preserve beta1_power and beta2_power when applying gradients so we can - # test that they've been restored correctly. - beta1=1.0, - beta2=1.0) - on_create_root = trackable_utils.Checkpoint( - optimizer=on_create_optimizer, model=on_create_model) - # Deferred restoration - status = on_create_root.restore(save_path=save_path) - status.assert_nontrivial_match() - status.assert_existing_objects_matched() - with self.assertRaises(AssertionError): - status.assert_consumed() - on_create_model(constant_op.constant([[3.]])) # create variables - self.assertAllEqual(1, self.evaluate(on_create_root.save_counter)) - self.assertAllEqual([42.], - self.evaluate( - on_create_model._named_dense.variables[1])) - on_create_m_bias_slot = on_create_optimizer.get_slot( - on_create_model._named_dense.variables[1], "m") - status.assert_existing_objects_matched() - with self.assertRaises(AssertionError): - status.assert_consumed() - # Optimizer slot variables are created when the original variable is - # restored. - self.assertAllEqual([1.5], self.evaluate(on_create_m_bias_slot)) - self.assertAllEqual(optimizer_variables[2:], - self.evaluate(on_create_optimizer.variables())) - dummy_var = resource_variable_ops.ResourceVariable([1.]) - on_create_optimizer.minimize(loss=dummy_var.read_value) - status.assert_existing_objects_matched() - status.assert_consumed() - beta1_power, beta2_power = on_create_optimizer._get_beta_accumulators() - self.assertAllEqual(optimizer_variables[0], self.evaluate(beta1_power)) - self.assertAllEqual(optimizer_variables[1], self.evaluate(beta2_power)) - - # TODO(allenl): Debug garbage created by this test in python3. - def testDeferredRestorationUsageEager(self): - """An idiomatic eager execution example.""" - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - optimizer_step=training_util.get_or_create_global_step()) - root.restore(checkpoint_management.latest_checkpoint( - checkpoint_directory)) - for _ in range(num_training_steps): - # TODO(allenl): Use a Dataset and serialize/checkpoint it. - input_value = constant_op.constant([[3.]]) - optimizer.minimize( - lambda: model(input_value), # pylint: disable=cell-var-from-loop - global_step=root.optimizer_step) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - root.optimizer_step.numpy()) - - def testEagerDistributionStrategy(self): - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - - def _train_fn(optimizer, model): - input_value = constant_op.constant([[3.]]) - optimizer.minimize( - functools.partial(model, input_value), - global_step=root.optimizer_step) - - strategy = mirrored_strategy.MirroredStrategy() - with strategy.scope(): - for training_continuation in range(3): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, - model=model, - optimizer_step=training_util.get_or_create_global_step()) - root.restore( - checkpoint_management.latest_checkpoint(checkpoint_directory)) - - for _ in range(num_training_steps): - strategy.extended.call_for_each_replica( - functools.partial(_train_fn, optimizer, model)) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - root.optimizer_step.numpy()) - - def testGraphDistributionStrategy(self): - self.skipTest("b/121381184") - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - - def _train_fn(optimizer, model): - input_value = constant_op.constant([[3.]]) - return optimizer.minimize( - functools.partial(model, input_value), - global_step=root.optimizer_step) - - for training_continuation in range(3): - with ops.Graph().as_default(): - strategy = mirrored_strategy.MirroredStrategy() - with strategy.scope(): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - optimizer_step=training_util.get_or_create_global_step()) - status = root.restore(checkpoint_management.latest_checkpoint( - checkpoint_directory)) - train_op = strategy.extended.call_for_each_replica( - functools.partial(_train_fn, optimizer, model)) - with self.session() as session: - if training_continuation > 0: - status.assert_consumed() - status.initialize_or_restore() - for _ in range(num_training_steps): - session.run(train_op) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - root.optimizer_step.numpy()) - - def testUsageGraph(self): - """Expected usage when graph building.""" - with context.graph_mode(): - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - with ops.Graph().as_default(): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.CheckpointV1( - optimizer=optimizer, model=model, - global_step=training_util.get_or_create_global_step()) - input_value = constant_op.constant([[3.]]) - train_op = optimizer.minimize( - model(input_value), - global_step=root.global_step) - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - with self.session(graph=ops.get_default_graph()) as session: - status = root.restore(save_path=checkpoint_path) - status.initialize_or_restore(session=session) - if checkpoint_path is None: - self.assertEqual(0, training_continuation) - with self.assertRaises(AssertionError): - status.assert_consumed() - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - else: - status.assert_consumed() - status.assert_existing_objects_matched() - for _ in range(num_training_steps): - session.run(train_op) - root.save(file_prefix=checkpoint_prefix, session=session) - self.assertEqual((training_continuation + 1) * num_training_steps, - session.run(root.global_step)) - self.assertEqual(training_continuation + 1, - session.run(root.save_counter)) - - @test_util.run_in_graph_and_eager_modes - def testAgnosticUsage(self): - """Graph/eager agnostic usage.""" - # Does create garbage when executing eagerly due to ops.Graph() creation. - num_training_steps = 10 - checkpoint_directory = self.get_temp_dir() - for training_continuation in range(3): - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - global_step=training_util.get_or_create_global_step()) - manager = checkpoint_management.CheckpointManager( - root, checkpoint_directory, max_to_keep=1) - status = root.restore(save_path=manager.latest_checkpoint) - input_value = constant_op.constant([[3.]]) - train_fn = functools.partial( - optimizer.minimize, - functools.partial(model, input_value), - global_step=root.global_step) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - status.initialize_or_restore() - for _ in range(num_training_steps): - train_fn() - manager.save() - self.assertEqual((training_continuation + 1) * num_training_steps, - self.evaluate(root.global_step)) - self.assertEqual(training_continuation + 1, - self.evaluate(root.save_counter)) - - # pylint: disable=cell-var-from-loop - @test_util.run_in_graph_and_eager_modes - def testWithDefun(self): - num_training_steps = 2 - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - for training_continuation in range(3): - with test_util.device(use_gpu=True): - model = MyModel() - # Don't actually train so we can test variable values - optimizer = adam.AdamOptimizer(0.) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - global_step=training_util.get_or_create_global_step()) - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - status = root.restore(save_path=checkpoint_path) - def train_fn(): - @def_function.function - def _call_model(x): - return model(x) - with backprop.GradientTape() as tape: - loss = _call_model(constant_op.constant([[3.]])) - gradients = tape.gradient(loss, model.variables) - return optimizer.apply_gradients(zip(gradients, model.variables), - global_step=root.global_step) - if not context.executing_eagerly(): - train_fn = functools.partial( - self.evaluate, train_fn()) - status.initialize_or_restore() - for _ in range(num_training_steps): - train_fn() - if training_continuation > 0: - status.assert_consumed() - self.assertAllClose([[42.]], self.evaluate(model.variables[0])) - else: - self.evaluate(model.variables[0].assign([[42.]])) - root.save(file_prefix=checkpoint_prefix) - self.assertEqual((training_continuation + 1) * num_training_steps, - self.evaluate(root.global_step)) - self.assertEqual(training_continuation + 1, - self.evaluate(root.save_counter)) - # pylint: enable=cell-var-from-loop - - def _get_checkpoint_name(self, name): - root = tracking.AutoTrackable() - trackable_utils.add_variable( - root, name=name, shape=[1, 2], dtype=dtypes.float64) - (named_variable,), _, _ = trackable_utils._serialize_object_graph( - root, saveables_cache=None) - with ops.name_scope("root/" + named_variable.name): - pass # Make sure we can use this as an op name if we prefix it. - return named_variable.name - - def testAnonymousVarsInInit(self): - - class Model(training.Model): - - def __init__(self): - super(Model, self).__init__() - self.w = resource_variable_ops.ResourceVariable(0.0) - self.b = resource_variable_ops.ResourceVariable(0.0) - self.vars = [self.w, self.b] - - def call(self, x): - return x * self.w + self.b - - with context.eager_mode(): - model = Model() - optimizer = adam.AdamOptimizer(learning_rate=0.05) - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - checkpoint = trackable_utils.Checkpoint( - model=model, optimizer=optimizer) - for _ in range(2): - checkpoint.save(checkpoint_prefix) - with backprop.GradientTape() as tape: - loss = (constant_op.constant(1.) - - model(constant_op.constant(1.))) ** 2 - grad = tape.gradient(loss, model.vars) - optimizer.apply_gradients( - [(g, v) for g, v in zip(grad, model.vars)]) - @test_util.run_in_graph_and_eager_modes def testDeferredSlotRestoration(self): checkpoint_directory = self.get_temp_dir() @@ -652,84 +193,6 @@ class CheckpointingTests(test.TestCase): beta1_power, _ = optimizer._get_beta_accumulators() self.assertAllEqual(3., self.evaluate(beta1_power)) - @test_util.run_in_graph_and_eager_modes - def test_initialize_if_not_restoring(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - optimizer_only_prefix = os.path.join(checkpoint_directory, "opt") - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - model=model, # Do not save the optimizer with the checkpoint. - global_step=training_util.get_or_create_global_step()) - optimizer_checkpoint = trackable_utils.Checkpoint( - optimizer=optimizer) - - checkpoint_path = checkpoint_management.latest_checkpoint( - checkpoint_directory) - status = root.restore(save_path=checkpoint_path) - input_value = constant_op.constant([[3.]]) - train_fn = functools.partial( - optimizer.minimize, - functools.partial(model, input_value), - global_step=root.global_step) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - status.initialize_or_restore() - self.evaluate([v.initializer for v in optimizer.variables()]) - train_fn() - model_save_path = root.save(file_prefix=checkpoint_prefix) - self.evaluate(optimizer.variables()[0].assign(42.)) - optimizer_save_path = optimizer_checkpoint.save(optimizer_only_prefix) - - # Restore into a graph with the optimizer - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - global_step=training_util.get_or_create_global_step()) - status = root.restore(save_path=model_save_path) - input_value = constant_op.constant([[3.]]) - train_fn = functools.partial( - optimizer.minimize, - functools.partial(model, input_value), - global_step=root.global_step) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - status.initialize_or_restore() - train_fn() - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - with self.assertRaises(AssertionError): - status.assert_consumed() - - # Make sure initialization doesn't clobber later restores - with test_util.device(use_gpu=True): - model = MyModel() - optimizer = adam.AdamOptimizer(0.001, beta1=1.0) - root = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, - global_step=training_util.get_or_create_global_step()) - opt_root = trackable_utils.Checkpoint( - optimizer=optimizer) - status = root.restore(save_path=model_save_path) - init_only_optimizer_status = opt_root.restore(save_path=None) - optimizer_status = opt_root.restore(save_path=optimizer_save_path) - input_value = constant_op.constant([[3.]]) - train_fn = functools.partial( - optimizer.minimize, - functools.partial(model, input_value), - global_step=root.global_step) - if not context.executing_eagerly(): - train_fn = functools.partial(self.evaluate, train_fn()) - optimizer_status.run_restore_ops() - status.initialize_or_restore() - init_only_optimizer_status.initialize_or_restore() - train_fn() - self.assertEqual(42., self.evaluate(optimizer.variables()[0])) - class _ManualScope(tracking.AutoTrackable): @@ -797,132 +260,5 @@ class TemplateTests(test.TestCase): self.assertAllEqual([14.], self.evaluate(var2)) -class CheckpointCompatibilityTests(test.TestCase): - - def _initialized_model(self): - input_value = constant_op.constant([[3.]]) - model = MyModel() - optimizer = adam.AdamOptimizer(0.001) - optimizer_step = training_util.get_or_create_global_step() - root_trackable = trackable_utils.Checkpoint( - optimizer=optimizer, model=model, optimizer_step=optimizer_step) - train_op = optimizer.minimize( - functools.partial(model, input_value), - global_step=optimizer_step) - self.evaluate(trackable_utils.gather_initializers( - root_trackable)) - self.evaluate(train_op) - # A regular variable, a slot variable, and a non-slot Optimizer variable - # with known values to check when loading. - self.evaluate(model._named_dense.bias.assign([1.])) - self.evaluate(optimizer.get_slot( - var=model._named_dense.bias, name="m").assign([2.])) - beta1_power, _ = optimizer._get_beta_accumulators() - self.evaluate(beta1_power.assign(3.)) - return root_trackable - - def _set_sentinels(self, root_trackable): - self.evaluate(root_trackable.model._named_dense.bias.assign([101.])) - self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, name="m") - .assign([102.])) - beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() - self.evaluate(beta1_power.assign(103.)) - - def _check_sentinels(self, root_trackable): - self.assertAllEqual( - [1.], self.evaluate(root_trackable.model._named_dense.bias)) - self.assertAllEqual([2.], self.evaluate( - root_trackable.optimizer.get_slot( - var=root_trackable.model._named_dense.bias, name="m"))) - beta1_power, _ = root_trackable.optimizer._get_beta_accumulators() - self.assertAllEqual(3., self.evaluate(beta1_power)) - - def _write_name_based_checkpoint(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph) as session: - root = self._initialized_model() - name_saver = saver_lib.Saver() - return name_saver.save( - sess=session, save_path=checkpoint_prefix, - global_step=root.optimizer_step) - - @test_util.run_in_graph_and_eager_modes - def testLoadFromNameBasedSaver(self): - """Save a name-based checkpoint, load it using the object-based API.""" - with test_util.device(use_gpu=True): - save_path = self._write_name_based_checkpoint() - root = self._initialized_model() - self._set_sentinels(root) - with self.assertRaises(AssertionError): - self._check_sentinels(root) - object_saver = trackable_utils.TrackableSaver( - graph_view.ObjectGraphView(root)) - self._set_sentinels(root) - status = object_saver.restore(save_path) - if context.executing_eagerly(): - self._check_sentinels(root) - if context.executing_eagerly(): - status.assert_consumed() - status.assert_existing_objects_matched() - status.assert_nontrivial_match() - else: - # When graph building, we haven't read any keys, so we don't know - # whether the restore will be complete. - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_consumed() - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_existing_objects_matched() - with self.assertRaisesRegexp(AssertionError, "not restored"): - status.assert_nontrivial_match() - status.run_restore_ops() - self._check_sentinels(root) - self._set_sentinels(root) - status = object_saver.restore(save_path) - status.initialize_or_restore() - self._check_sentinels(root) - # Check that there is no error when keys are missing from the name-based - # checkpoint. - root.not_in_name_checkpoint = resource_variable_ops.ResourceVariable([1.]) - status = object_saver.restore(save_path) - with self.assertRaises(AssertionError): - status.assert_existing_objects_matched() - - def testSaveGraphLoadEager(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph): - root = self._initialized_model() - save_path = root.save(file_prefix=checkpoint_prefix) - with context.eager_mode(): - root = self._initialized_model() - self._set_sentinels(root) - root.restore(save_path).assert_consumed() - self._check_sentinels(root) - - def testSaveEagerLoadGraph(self): - checkpoint_directory = self.get_temp_dir() - checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - with context.eager_mode(): - root = self._initialized_model() - save_path = root.save(file_prefix=checkpoint_prefix) - with context.graph_mode(): - save_graph = ops.Graph() - with save_graph.as_default(), self.session( - graph=save_graph): - root = self._initialized_model() - self._set_sentinels(root) - root.restore(save_path).assert_consumed().run_restore_ops() - self._check_sentinels(root) - - if __name__ == "__main__": test.main() From 8d5d9a50d3c8422d26cce81000310ec941b56030 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 10:47:19 -0700 Subject: [PATCH 0179/1390] More compatibility fixes for typing.Generic: * types.new_class is required in some distributions * avoid calling isinstance on some function objects in python 3.6 Required for #40132. PiperOrigin-RevId: 316497932 Change-Id: I3441e8b099d2b10c965d45cd362a4859c7c29bb9 --- tensorflow/python/framework/test_util.py | 2 -- tensorflow/python/util/tf_should_use.py | 18 +++++------------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py index 572e8aac987..2967bb3de84 100644 --- a/tensorflow/python/framework/test_util.py +++ b/tensorflow/python/framework/test_util.py @@ -732,8 +732,6 @@ def assert_no_new_tensors(f): """Finds existing Tensors, runs the test, checks for new Tensors.""" def _is_tensorflow_object(obj): - if not hasattr(obj, "__class__"): - return False try: return isinstance(obj, (ops.Tensor, variables.Variable, diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py index 41c3220f5ca..1671b078fa3 100644 --- a/tensorflow/python/util/tf_should_use.py +++ b/tensorflow/python/util/tf_should_use.py @@ -21,12 +21,15 @@ import copy import sys import textwrap import traceback -import types + +import six # pylint: disable=unused-import + from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.platform import tf_logging from tensorflow.python.util import tf_decorator +# pylint: enable=g-bad-import-order,g-import-not-at-top class _TFShouldUseHelper(object): @@ -151,18 +154,7 @@ def _get_wrapper(x, tf_should_use_helper): tx = copy.deepcopy(type_x) # Prefer using __orig_bases__, which preserve generic type arguments. bases = getattr(tx, '__orig_bases__', tx.__bases__) - - # Use types.new_class when available, which is preferred over plain type in - # some distributions. - if sys.version_info >= (3, 5): - def set_body(ns): - ns.update(tx.__dict__) - return ns - - copy_tx = types.new_class(tx.__name__, bases, exec_body=set_body) - else: - copy_tx = type(tx.__name__, bases, dict(tx.__dict__)) - + copy_tx = type(tx.__name__, bases, dict(tx.__dict__)) copy_tx.__init__ = _new__init__ copy_tx.__getattribute__ = _new__getattribute__ copy_tx.mark_used = _new_mark_used From 80b3b4fa9f8ab6565a91273dc38aeabda0258bb6 Mon Sep 17 00:00:00 2001 From: Alex Zinenko Date: Mon, 15 Jun 2020 10:47:31 -0700 Subject: [PATCH 0180/1390] [XLA] LHLO-to-Ploop: explicitly construct ValueRange when building scf::ParallelOp The builder API is about to change in LLVM, which would lead to ambiguity in overload resolution in these cases. Proactively fix the issue by constructing ValueRange of steps explicitly before it is passed to the builder function. PiperOrigin-RevId: 316497983 Change-Id: Ibb5dfe006071ab9513634c9613f43ce60a0efd22 --- .../mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc index 734a75a4307..b3112d49103 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc @@ -261,7 +261,7 @@ class ReduceOpConverter : public OpConversionPattern { rewriter->setInsertionPointToStart(outer.getBody()); } scf::ParallelOp inner = rewriter->create( - loc, reduce_lower, reduce_upper, reduce_step, init_value); + loc, reduce_lower, reduce_upper, reduce_step, ValueRange(init_value)); Value reduction_result = *inner.getResults().begin(); SmallVector out_indices; @@ -406,7 +406,7 @@ class ReduceWindowOpConverter rewriter->create(loc, window_dim.getSExtValue())); } auto window_loop = rewriter->create( - loc, window_lower, window_upper, window_step, init_value); + loc, window_lower, window_upper, window_step, ValueRange(init_value)); Value reduction_result = *window_loop.getResults().begin(); auto output_ivs = output_loop.getInductionVars(); From 67487368bbc3f35a87314bb54ec5e12a6e3d7f93 Mon Sep 17 00:00:00 2001 From: Yujing Zhang Date: Mon, 15 Jun 2020 10:48:56 -0700 Subject: [PATCH 0181/1390] Preserve FunctionDef.arg_attr in GrapplerFunctionItem. PiperOrigin-RevId: 316498288 Change-Id: I6c3288c725bb281cca17256146c9ec3fd8cec5f0 --- .../core/framework/graph_to_functiondef.cc | 4 ++- tensorflow/core/grappler/utils/functions.cc | 27 ++++++++++++++++--- tensorflow/core/grappler/utils/functions.h | 5 ++++ .../core/grappler/utils/functions_test.cc | 17 ++++++++++++ 4 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/framework/graph_to_functiondef.cc b/tensorflow/core/framework/graph_to_functiondef.cc index bbd70151849..e825aa722b5 100644 --- a/tensorflow/core/framework/graph_to_functiondef.cc +++ b/tensorflow/core/framework/graph_to_functiondef.cc @@ -434,9 +434,11 @@ Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name, // _Arg/Placeholder nodes. if (absl::StartsWith(attr.first, "_")) { arg_attrs.mutable_attr()->insert(attr); - } else if (attr.first == "shape") { + } else if (attr.first == "shape" && argdef->type() != DT_RESOURCE) { // Preserve known shapes by moving them to the _output_shapes list. // The _Arg shape function knows how to extract them from there. + // Don't preserve the shape of a resource arg node, which is a scalar + // resource handle. AttrValue value; *(value.mutable_list()->add_shape()) = attr.second.shape(); arg_attrs.mutable_attr()->insert({"_output_shapes", value}); diff --git a/tensorflow/core/grappler/utils/functions.cc b/tensorflow/core/grappler/utils/functions.cc index 780e3c7e3f2..a83fb824cc3 100644 --- a/tensorflow/core/grappler/utils/functions.cc +++ b/tensorflow/core/grappler/utils/functions.cc @@ -38,12 +38,14 @@ namespace grappler { GrapplerFunctionItem::GrapplerFunctionItem( string func_name, string description, AttrSlice func_attr, + std::vector arg_attr, std::vector input_args, std::vector output_args, std::vector control_outputs, const int graph_def_version, const bool is_stateful, GraphDef&& function_body) : description_(std::move(description)), func_attr_(func_attr), + arg_attr_(std::move(arg_attr)), input_args_(std::move(input_args)), output_args_(std::move(output_args)), control_outputs_(std::move(control_outputs)), @@ -108,6 +110,11 @@ const std::size_t GrapplerFunctionItem::control_output_size() const { const AttrSlice& GrapplerFunctionItem::func_attr() const { return func_attr_; } +const std::vector& +GrapplerFunctionItem::arg_attr() const { + return arg_attr_; +} + const GraphDef& GrapplerFunctionItem::function_body() const { return graph; } GraphDef& GrapplerFunctionItem::mutable_function_body() { return graph; } @@ -278,12 +285,17 @@ Status MakeGrapplerFunctionItem(const FunctionDef& func, control_outputs.push_back({control_ret.first, control_ret.second}); } + std::vector arg_attr(inputs.size(), nullptr); + for (const auto& attr : func.arg_attr()) { + arg_attr.at(attr.first) = &attr.second; + } + *item = GrapplerFunctionItem( /*func_name=*/signature.name(), /*description=*/signature.description(), - /*func_attr=*/AttrSlice(&func.attr()), std::move(inputs), - std::move(outputs), std::move(control_outputs), graph_def_version, - signature.is_stateful(), std::move(function_body)); + /*func_attr=*/AttrSlice(&func.attr()), std::move(arg_attr), + std::move(inputs), std::move(outputs), std::move(control_outputs), + graph_def_version, signature.is_stateful(), std::move(function_body)); return Status::OK(); } @@ -330,6 +342,7 @@ Status ReplaceInputWithConst(const NodeDef& input_const, int input_index, } item->input_args_.erase(item->input_args_.begin() + input_index); + item->arg_attr_.erase(item->arg_attr_.begin() + input_index); return Status::OK(); } @@ -566,6 +579,14 @@ Status MakeFunctionDef(const GrapplerFunctionItem& item, (*func->mutable_attr())[attr_name] = attr_value; } + // Copy function arg attributes. + for (int i = 0; i < item.arg_attr().size(); ++i) { + const auto* attr = item.arg_attr().at(i); + if (attr != nullptr) { + (*func->mutable_arg_attr())[i] = *attr; + } + } + // Copy function body nodes to the FunctionDef and update input format for (const NodeDef& func_node : item.function_body().node()) { // Skip original `_Arg` and `_Retval` nodes. If node was converted to some diff --git a/tensorflow/core/grappler/utils/functions.h b/tensorflow/core/grappler/utils/functions.h index b03b89af2ab..2f1fd5d2ed6 100644 --- a/tensorflow/core/grappler/utils/functions.h +++ b/tensorflow/core/grappler/utils/functions.h @@ -76,6 +76,7 @@ class GrapplerFunctionItem : public GrapplerItem { const std::size_t control_output_size() const; const AttrSlice& func_attr() const; + const std::vector& arg_attr() const; const GraphDef& function_body() const; GraphDef& mutable_function_body(); @@ -95,6 +96,7 @@ class GrapplerFunctionItem : public GrapplerItem { GrapplerFunctionItem(string func_name, string description, AttrSlice func_attr, + std::vector arg_attr, std::vector input_args, std::vector output_args, std::vector control_outputs, @@ -105,6 +107,9 @@ class GrapplerFunctionItem : public GrapplerItem { AttrSlice func_attr_; // Attributes specific to function definition that // produced this item (FuncDef.attr field). + // Attributes of function arguments + std::vector arg_attr_; + std::vector input_args_; std::vector output_args_; std::vector control_outputs_; diff --git a/tensorflow/core/grappler/utils/functions_test.cc b/tensorflow/core/grappler/utils/functions_test.cc index 8cc938ec845..66320d60f27 100644 --- a/tensorflow/core/grappler/utils/functions_test.cc +++ b/tensorflow/core/grappler/utils/functions_test.cc @@ -523,6 +523,14 @@ TEST_F(FunctionsTest, MakeFunctionDef) { {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}}, }); + // Add an attribute to _Arg 0; + const uint32 arg_index = 0; + const std::pair arg_attr_key_and_value = {"_arg_attr", "abc"}; + FunctionDef::ArgAttrs arg_attr; + (*arg_attr.mutable_attr())[arg_attr_key_and_value.first].set_s( + arg_attr_key_and_value.second); + (*func.mutable_arg_attr())[arg_index] = arg_attr; + protobuf::Map func_instantiation_attr; func_instantiation_attr["T"].set_type(DT_FLOAT); FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary()); @@ -541,6 +549,15 @@ TEST_F(FunctionsTest, MakeFunctionDef) { EXPECT_EQ("y", specialized.signature().output_arg(0).name()); EXPECT_EQ(DT_FLOAT, specialized.signature().output_arg(0).type()); + EXPECT_EQ(specialized.arg_attr().size(), 1); + EXPECT_EQ(specialized.arg_attr().at(arg_index).attr().size(), 1); + EXPECT_EQ(specialized.arg_attr() + .at(arg_index) + .attr() + .at(arg_attr_key_and_value.first) + .s(), + arg_attr_key_and_value.second); + // Function body specialized for instantiation types. int count = 0; for (const NodeDef &node : specialized.node_def()) { From b7d66ef92658ffcf7fffa7d606293bfc314d82d2 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 15 Jun 2020 10:50:51 -0700 Subject: [PATCH 0182/1390] Add unit test that runs TPU Embedding layouter with MLIR bridge. PiperOrigin-RevId: 316498711 Change-Id: I90365d953f7469464cd81a306ddaf50db580ba93 --- tensorflow/python/tpu/tpu_embedding_v2_test.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/tpu/tpu_embedding_v2_test.py b/tensorflow/python/tpu/tpu_embedding_v2_test.py index a8b21480919..0c257010f6a 100644 --- a/tensorflow/python/tpu/tpu_embedding_v2_test.py +++ b/tensorflow/python/tpu/tpu_embedding_v2_test.py @@ -34,6 +34,7 @@ from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver from tensorflow.python.eager import backprop from tensorflow.python.eager import def_function from tensorflow.python.eager import remote +from tensorflow.python.framework import config from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -719,7 +720,11 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase): self.assertAllClose(golden, weights0) - def test_enqueue_with_outside_compilation(self): + @parameterized.parameters([True, False]) + def test_enqueue_with_outside_compilation(self, use_mlir): + if use_mlir: + config.enable_mlir_bridge() + strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd') dataset = self._create_sparse_dataset(strategy) dataset_iter = iter(strategy.experimental_distribute_dataset(dataset)) @@ -749,7 +754,11 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase): self.assertAllClose(activations_oc0, activations0) - def test_enqueue_with_outside_compilation_in_control_flow(self): + @parameterized.parameters(True, False) + def test_enqueue_with_outside_compilation_in_control_flow(self, use_mlir): + if use_mlir: + config.enable_mlir_bridge() + strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd') dataset = self._create_sparse_dataset(strategy) dataset_iter = iter(strategy.experimental_distribute_dataset(dataset)) From ba658404f279df1b2c86ade146fbdccffcf68f8f Mon Sep 17 00:00:00 2001 From: George Karpenkov Date: Mon, 15 Jun 2020 11:10:33 -0700 Subject: [PATCH 0183/1390] [TF/XLA] Only force retracing for non-unique XLA context ID for TPUReplicatedContext Fixes https://github.com/tensorflow/tensorflow/issues/39872 PiperOrigin-RevId: 316503485 Change-Id: Ice63983fcdf2fdedca60a9054f3b76ac60e1ff15 --- .../python/eager/def_function_xla_jit_test.py | 19 ------------------- tensorflow/python/eager/function.py | 7 +++---- tensorflow/python/ops/control_flow_ops.py | 5 ----- tensorflow/python/tpu/tpu.py | 6 ------ 4 files changed, 3 insertions(+), 34 deletions(-) diff --git a/tensorflow/python/eager/def_function_xla_jit_test.py b/tensorflow/python/eager/def_function_xla_jit_test.py index 78d44a81b0b..b63a3b434d4 100644 --- a/tensorflow/python/eager/def_function_xla_jit_test.py +++ b/tensorflow/python/eager/def_function_xla_jit_test.py @@ -29,7 +29,6 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import math_ops -from tensorflow.python.ops import random_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import tensor_array_ops from tensorflow.python.platform import test @@ -386,24 +385,6 @@ class DefFunctionTest(test.TestCase): f64_input = constant_op.constant([1.1, 2.2, 3.3], dtype=dtypes.float64) self.assertAllClose([1.1, 3.3, 6.6], f(f64_input)) - def testNoExcessiveRetracing(self): - inner_retracings = 0 - - @def_function.function(experimental_compile=True) - def inner(a, b): - nonlocal inner_retracings - inner_retracings += 1 - return a * b + a - - def outer(a, b): - return inner(a, b) - - func_input = random_ops.random_normal([10, 10]) - for _ in range(2): - def_function.function(outer)(func_input, func_input) - - self.assertEqual(inner_retracings, 1) - if __name__ == '__main__': ops.enable_eager_execution() diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py index c02318cb814..a40eaf886b3 100644 --- a/tensorflow/python/eager/function.py +++ b/tensorflow/python/eager/function.py @@ -2981,10 +2981,9 @@ class Function(object): if not executing_eagerly: # We want to force function retracing for each different # XLAControlFlowContext, so add `xla_context_id` to the cache key. - xla_context = _enclosing_xla_context() - if xla_context is not None and \ - xla_context.RequiresUniqueFunctionRetracing(): - xla_context_id = id(xla_context) + tpu_context = _enclosing_xla_context() + if tpu_context is not None: + xla_context_id = id(tpu_context) with ops.init_scope(): # The graph, or whether we're executing eagerly, should be a part of the diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py index 748f842a9e0..3398308d42e 100644 --- a/tensorflow/python/ops/control_flow_ops.py +++ b/tensorflow/python/ops/control_flow_ops.py @@ -3682,11 +3682,6 @@ class XLAControlFlowContext(ControlFlowContext): def AddValue(self, x): return x - def RequiresUniqueFunctionRetracing(self): - """Returns whether the tf.function should be retraced if the context changes. - """ - return False - def from_control_flow_context_def(context_def, import_scope=None): """Deserializes `context_def` into the appropriate ControlFlowContext. diff --git a/tensorflow/python/tpu/tpu.py b/tensorflow/python/tpu/tpu.py index ce3aaa8a058..28eba69b7da 100644 --- a/tensorflow/python/tpu/tpu.py +++ b/tensorflow/python/tpu/tpu.py @@ -639,12 +639,6 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): def GetControlPivot(self): return self._pivot - def RequiresUniqueFunctionRetracing(self): - # More context: b/158152827. TPU stack uses the TPUReplicateContext to - # create replicated variable handles and cluster TPU computations, thus we - # always retrace a tf.function when the wrapped TPUReplicateContext changes. - return True - class OutsideCompilationV2Context(control_flow_ops.ControlFlowContext): """The context for outside compilation in Tensorflow 2.0. From d1a34523f44b853998bc2740d5b59a472a12eb86 Mon Sep 17 00:00:00 2001 From: Raman Sarokin Date: Mon, 15 Jun 2020 11:21:03 -0700 Subject: [PATCH 0184/1390] Added OpenCL versions 2.1/2.2/3.0. PiperOrigin-RevId: 316505978 Change-Id: I5e35dc8e625aef2feb8f59b10a6a60d175a08314 --- tensorflow/lite/delegates/gpu/cl/cl_device.cc | 18 +++++++++++++++++- tensorflow/lite/delegates/gpu/cl/cl_device.h | 10 +++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.cc b/tensorflow/lite/delegates/gpu/cl/cl_device.cc index 13e299b181b..aea81d5e659 100644 --- a/tensorflow/lite/delegates/gpu/cl/cl_device.cc +++ b/tensorflow/lite/delegates/gpu/cl/cl_device.cc @@ -105,8 +105,18 @@ OpenCLVersion ParseCLVersion(const std::string& version) { } else { return OpenCLVersion::CL_1_0; } + } else if (major == 2) { + if (minor == 2) { + return OpenCLVersion::CL_2_2; + } else if (minor == 1) { + return OpenCLVersion::CL_2_1; + } else { + return OpenCLVersion::CL_2_0; + } + } else if (major == 3) { + return OpenCLVersion::CL_3_0; } else { - return OpenCLVersion::CL_2_0; + return OpenCLVersion::CL_1_0; } } @@ -227,6 +237,12 @@ std::string OpenCLVersionToString(OpenCLVersion version) { return "1.2"; case OpenCLVersion::CL_2_0: return "2.0"; + case OpenCLVersion::CL_2_1: + return "2.1"; + case OpenCLVersion::CL_2_2: + return "2.2"; + case OpenCLVersion::CL_3_0: + return "3.0"; } } diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.h b/tensorflow/lite/delegates/gpu/cl/cl_device.h index 4fd683b78ff..1df16aa3bad 100644 --- a/tensorflow/lite/delegates/gpu/cl/cl_device.h +++ b/tensorflow/lite/delegates/gpu/cl/cl_device.h @@ -31,7 +31,15 @@ namespace cl { enum class Vendor { QUALCOMM, MALI, POWERVR, NVIDIA, AMD, INTEL, UNKNOWN }; std::string VendorToString(Vendor v); -enum class OpenCLVersion { CL_1_0, CL_1_1, CL_1_2, CL_2_0 }; +enum class OpenCLVersion { + CL_1_0, + CL_1_1, + CL_1_2, + CL_2_0, + CL_2_1, + CL_2_2, + CL_3_0 +}; std::string OpenCLVersionToString(OpenCLVersion version); // for use only in cl_device.cc, but putted here to make tests From 555be8943e70a775f4617a342587857904a9a7c1 Mon Sep 17 00:00:00 2001 From: George Karpenkov Date: Mon, 15 Jun 2020 11:27:37 -0700 Subject: [PATCH 0185/1390] [XLA] [NFC] Reduce duplication between Executable subclasses Factor out the code to mark owning arguments as to-be-released. PiperOrigin-RevId: 316507574 Change-Id: I22fce6e69d0933baa25db09e25bda4037beceb17 --- .../compiler/xla/service/cpu/cpu_executable.cc | 11 +---------- tensorflow/compiler/xla/service/executable.cc | 12 ++++++++++++ tensorflow/compiler/xla/service/executable.h | 9 +++++++++ .../compiler/xla/service/gpu/gpu_executable.cc | 10 +--------- .../xla/service/interpreter/executable_base.cc | 9 +-------- 5 files changed, 24 insertions(+), 27 deletions(-) diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index e0c8adcbbe1..4552d7b5ba9 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -350,16 +350,7 @@ StatusOr CpuExecutable::ExecuteAsyncOnStream( std::move(buffers)), hlo_execution_profile}); - // TODO(cheshire): Duplication with other executables. - for (ExecutionInput& argument : arguments) { - for (auto& index_buffer : *argument.MutableBuffers()) { - absl::optional maybe_owning_buffer = - index_buffer.second.Release(); - if (maybe_owning_buffer) { - result.AddToBeReleased(std::move(*maybe_owning_buffer)); - } - } - } + MarkToBeReleasedArguments(absl::MakeSpan(arguments), result); return std::move(result); } diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index 4f210442005..ebf7cc440dd 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -258,4 +258,16 @@ StatusOr Executable::ExecuteAsyncOnStreamWrapper( int64 Executable::SizeOfGeneratedCodeInBytes() const { return -1; } +void Executable::MarkToBeReleasedArguments(absl::Span arguments, + ExecutionOutput& result) { + for (ExecutionInput& argument : arguments) { + for (auto& index_buffer : *argument.MutableBuffers()) { + if (absl::optional maybe_owning_buffer = + index_buffer.second.Release()) { + result.AddToBeReleased(std::move(*maybe_owning_buffer)); + } + } + } +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h index 49614c1af00..2c979662d24 100644 --- a/tensorflow/compiler/xla/service/executable.h +++ b/tensorflow/compiler/xla/service/executable.h @@ -331,6 +331,15 @@ class Executable { bool dumping_snapshot() const { return hlo_proto_ != nullptr; } HloProto const* hlo_proto() const { return hlo_proto_.get(); } + // Gather unused but donated buffers, return them to the caller of this API. + // We don't free buffers inside this function since the caller could have + // different preferences for buffer deallocation. For example, in TensorFlow, + // buffers are mostly efficiently deallocated as soon as a program has been + // launched. However, in XRT, the buffers are expected to be deallocated after + // the program has finished since XRT doesn't support async deallocation. + void MarkToBeReleasedArguments(absl::Span arguments, + ExecutionOutput& result); + protected: // HloModule this was compiled from. BufferAssignment keeps pointers to // HloInstructions owned by the HloModule so we need to keep the HloModule diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index c8b11cab31a..520bbedbaeb 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -541,15 +541,7 @@ StatusOr GpuExecutable::ExecuteAsyncOnStream( buffer_allocations.TearDown(buffers_in_result, assignment_.get())); // Free allocations for arguments. - for (ExecutionInput& argument : arguments) { - for (auto& index_buffer : *argument.MutableBuffers()) { - if (absl::optional owning = - index_buffer.second.Release()) { - result.AddToBeReleased(std::move(*owning)); - } - } - } - + MarkToBeReleasedArguments(absl::MakeSpan(arguments), result); return std::move(result); } diff --git a/tensorflow/compiler/xla/service/interpreter/executable_base.cc b/tensorflow/compiler/xla/service/interpreter/executable_base.cc index 5850cbf005b..4b020ea2d32 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable_base.cc +++ b/tensorflow/compiler/xla/service/interpreter/executable_base.cc @@ -122,14 +122,7 @@ StatusOr InterpreterExecutableBase::ExecuteAsyncOnStream( const double nanoseconds = (end_micros - start_micros) * 1000.0; profile->set_compute_time_ns(std::max(nanoseconds, 1.0)); } - for (auto& argument : arguments) { - for (auto& index_buffer : *argument.MutableBuffers()) { - auto maybe_owning_buffer = index_buffer.second.Release(); - if (maybe_owning_buffer) { - result.AddToBeReleased(std::move(*maybe_owning_buffer)); - } - } - } + MarkToBeReleasedArguments(absl::MakeSpan(arguments), result); return std::move(result); } From 49ba207bdaf810293efeba4d0ce1b1bce56e5804 Mon Sep 17 00:00:00 2001 From: George Karpenkov Date: Mon, 15 Jun 2020 11:40:07 -0700 Subject: [PATCH 0186/1390] [TF/XLA] Remove the wrapping function XlaTensor::RefCountIsOne PiperOrigin-RevId: 316510507 Change-Id: I7bbee7208348a1ae43b11671ecd95fe1e104f280 --- tensorflow/compiler/jit/xla_tensor.cc | 4 ---- tensorflow/compiler/jit/xla_tensor.h | 2 -- 2 files changed, 6 deletions(-) diff --git a/tensorflow/compiler/jit/xla_tensor.cc b/tensorflow/compiler/jit/xla_tensor.cc index e56d2714b0a..6bad1b703b6 100644 --- a/tensorflow/compiler/jit/xla_tensor.cc +++ b/tensorflow/compiler/jit/xla_tensor.cc @@ -27,10 +27,6 @@ namespace tensorflow { return xla_tensor; } -/*static*/ bool XlaTensor::RefCountIsOne(const Tensor& tensor) { - return tensor.RefCountIsOne(); -} - /*static*/ se::DeviceMemoryBase XlaTensor::DeviceMemoryFromTensor( const Tensor& tensor) { const XlaTensor* xla_tensor = FromTensor(&tensor); diff --git a/tensorflow/compiler/jit/xla_tensor.h b/tensorflow/compiler/jit/xla_tensor.h index 7f7d97e3b3f..a6de405ec9e 100644 --- a/tensorflow/compiler/jit/xla_tensor.h +++ b/tensorflow/compiler/jit/xla_tensor.h @@ -39,8 +39,6 @@ class XlaTensor { // fails. static XlaTensor* FromTensor(const Tensor* tensor); - static bool RefCountIsOne(const Tensor& tensor); - // Create a DeviceMemoryBase from a Tensor. The Tensor can be an XlaTensor, in // which case the returned value is shaped_buffer()->root_buffer(), or a // normal Tensor in which case the returned value is From 1569e9d09748b452941fc54f3095cd2ec7c309a3 Mon Sep 17 00:00:00 2001 From: Amit Patankar Date: Mon, 15 Jun 2020 11:45:04 -0700 Subject: [PATCH 0187/1390] Re-enable Windows build after updating RBE images with estimator nightly pip fix. PiperOrigin-RevId: 316511566 Change-Id: Ie6e26911d080301439bc59dc7cd03b2e33ed3d45 --- tensorflow/core/platform/BUILD | 1 - tensorflow/python/distribute/BUILD | 1 - 2 files changed, 2 deletions(-) diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD index 30734a840d1..70bb8a89417 100644 --- a/tensorflow/core/platform/BUILD +++ b/tensorflow/core/platform/BUILD @@ -386,7 +386,6 @@ py_test( name = "ram_file_system_test", srcs = ["ram_file_system_test.py"], python_version = "PY3", - tags = ["no_windows"], # TODO(b/156428279): reenable this test once the image is updated. deps = [ "//tensorflow:tensorflow_py", ], diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD index 3d5ae4f4215..77ef98d1cb7 100644 --- a/tensorflow/python/distribute/BUILD +++ b/tensorflow/python/distribute/BUILD @@ -1695,7 +1695,6 @@ cuda_py_test( python_version = "PY3", tags = [ "multi_and_single_gpu", - "no_windows", # TODO(b/156428279): reenable this test once the image is updated. ], # b/141096229: Non-atomic AssignAdd xla_enable_strict_auto_jit = False, From 91cd70ddc1dd49f1d316cfd6ae0bba6413039d2d Mon Sep 17 00:00:00 2001 From: Meghna Natraj Date: Mon, 15 Jun 2020 11:45:34 -0700 Subject: [PATCH 0188/1390] Update post_training_quantization.md: Update TFLiteConverter signature, remove usage of deprecated tf.lite.constants. PiperOrigin-RevId: 316511691 Change-Id: Ic7cae9d887af7b8641e71c050e29e38ab36f6a60 --- .../performance/post_training_quantization.md | 25 ++++++------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/tensorflow/lite/g3doc/performance/post_training_quantization.md b/tensorflow/lite/g3doc/performance/post_training_quantization.md index af7d9dbf02d..1a579430656 100644 --- a/tensorflow/lite/g3doc/performance/post_training_quantization.md +++ b/tensorflow/lite/g3doc/performance/post_training_quantization.md @@ -15,11 +15,11 @@ summary table of the choices and the benefits they provide: | Technique | Benefits | Hardware | | -------------------- | ------------------------- | ---------------- | -| Dynamic range | 4x smaller, 2-3x speedup | CPU | +| Dynamic range | 4x smaller, 2x-3x speedup | CPU | : quantization : : : | Full integer | 4x smaller, 3x+ speedup | CPU, Edge TPU, | : quantization : : Microcontrollers : -| Float16 quantization | 2x smaller, potential GPU | CPU, GPU | +| Float16 quantization | 2x smaller, GPU | CPU, GPU | : : acceleration : : The following decision tree can help determine which post-training quantization @@ -34,7 +34,7 @@ weights from floating point to integer, which has 8-bits of precision: 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 tflite_quant_model = converter.convert()
@@ -48,16 +48,7 @@ activations based on their range to 8-bits and perform computations with 8-bit weights and activations. This optimization provides latencies close to fully fixed-point inference. However, the outputs are still stored using floating point so that the speedup with dynamic-range ops is less than a full fixed-point -computation. Dynamic-range ops are available for the most compute-intensive -operators in a network: - -* `tf.keras.layers.Dense` -* `tf.keras.layers.Conv2D` -* `tf.keras.layers.LSTM` -* `tf.nn.embedding_lookup` -* `tf.compat.v1.nn.rnn_cell.BasicRNNCell` -* `tf.compat.v1.nn.bidirectional_dynamic_rnn` -* `tf.compat.v1.nn.dynamic_rnn` +computation. ### Full integer quantization @@ -77,7 +68,7 @@ the following steps: 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 def representative_dataset_gen():
   for _ in range(num_calibration_steps):
@@ -105,7 +96,7 @@ the following steps:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 def representative_dataset_gen():
   for _ in range(num_calibration_steps):
@@ -129,9 +120,9 @@ quantization of weights, use the following steps:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
-converter.target_spec.supported_types = [tf.lite.constants.FLOAT16]
+converter.target_spec.supported_types = [tf.float16]
 tflite_quant_model = converter.convert()
 

 

From a975ef0f5c81d913d7bb1a3a2a2e84fe1f255a71 Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Mon, 15 Jun 2020 12:02:13 -0700
Subject: [PATCH 0189/1390] [XLA] Support aliasing in XLA:Python bindings

PiperOrigin-RevId: 316515355
Change-Id: I9dae67c98209188ddacbae3b83ccfabceec9ea23
---
 tensorflow/compiler/xla/pjrt/pjrt_client.cc   | 54 ++++++++++++++++---
 .../compiler/xla/python/xla_client_test.py    |  5 --
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.cc b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
index 56370fa23a9..c1b433845b2 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.cc
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
@@ -202,16 +202,58 @@ StatusOr PjRtClient::GetDefaultDeviceAssignment(
 
 StatusOr> PjRtClient::GetParametersThatMustBeDonated(
     const LocalExecutable& executable, bool tuple_inputs) const {
-  // TODO(b/149489114) support buffer donation on CPU/GPU when XLA supports it.
+  HloComputation* computation =
+      executable.executable()->module().entry_computation();
+  int number_of_parameters = [&]() -> int {
+    if (tuple_inputs) {
+      CHECK_EQ(computation->num_parameters(), 1);
+      const Shape& input_tuple_shape =
+          computation->parameter_instruction(0)->shape();
+      CHECK(input_tuple_shape.IsTuple());
+      return input_tuple_shape.tuple_shapes_size();
+    } else {
+      return computation->num_parameters();
+    }
+  }();
+  // If any buffer in a parameter is aliased we will donate the entire input
+  // parameter.
+  absl::flat_hash_set parameters_to_donate;
   const HloInputOutputAliasConfig& config =
       executable.executable()->module().input_output_alias_config();
   TF_RETURN_IF_ERROR(config.ForEachAliasWithStatus(
-      [](const ShapeIndex& output_index,
-         const HloInputOutputAliasConfig::Alias& alias) {
-        return InvalidArgument(
-            "Buffer aliasing is not supported by XLA for non-TPU backends.");
+      [&](const ShapeIndex& output_index,
+          const HloInputOutputAliasConfig::Alias& alias) {
+        if (tuple_inputs) {
+          if (alias.parameter_number != 0) {
+            return InvalidArgument(
+                "Unexpected parameter number %d in alias config with tupled "
+                "inputs",
+                alias.parameter_number);
+          }
+          const ShapeIndex& index = alias.parameter_index;
+          if (!index.empty()) {
+            int this_parameter = index.data()[0];
+            if (this_parameter >= number_of_parameters) {
+              return InvalidArgument(
+                  "Unexpected parameter index %s in alias config with tupled "
+                  "inputs and %d parameters",
+                  index.ToString(), number_of_parameters);
+            }
+            parameters_to_donate.insert(this_parameter);
+          }
+        } else {
+          int this_parameter = alias.parameter_number;
+          if (this_parameter >= number_of_parameters) {
+            return InvalidArgument(
+                "Unexpected parameter number %d in alias config without tupled "
+                "inputs and %d parameters",
+                this_parameter, number_of_parameters);
+          }
+          parameters_to_donate.insert(this_parameter);
+        }
+        return Status::OK();
       }));
-  return absl::flat_hash_set();
+  return parameters_to_donate;
 }
 
 namespace {
diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py
index 0fc0bcae954..6a316044734 100644
--- a/tensorflow/compiler/xla/python/xla_client_test.py
+++ b/tensorflow/compiler/xla/python/xla_client_test.py
@@ -1909,11 +1909,6 @@ def TestFactory(xla_backend, cloud_tpu=False):
       out = ops.Add(p1, p2)
       c.setup_alias([], 0, [])
       c = c.build(out)
-      if self.backend.platform != "tpu":
-        with self.assertRaisesRegex(
-            RuntimeError, "Buffer aliasing is not supported "
-            "by XLA for non-TPU backends"):
-          self.backend.compile(c)
 
   tests.append(AliasTest)
 

From 32c25efaa929b07054cde6a0b7cf6638cd11ba79 Mon Sep 17 00:00:00 2001
From: Tim Shen 
Date: Mon, 15 Jun 2020 12:28:00 -0700
Subject: [PATCH 0190/1390] [MLIR/GPU] Add fix_signature knob to turn off
 FixKernelFunctionSignatures.

PiperOrigin-RevId: 316520783
Change-Id: I56de509059bff1e3eb87f5a618e68362ee7c6e66
---
 .../mlir/tools/kernel_gen/cubin_creator.cc    | 10 ++++++---
 .../xla/service/mlir_gpu/kernel_lowering.cc   | 22 +++++++++----------
 .../xla/service/mlir_gpu/kernel_lowering.h    | 11 +++++++---
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
index d1e5c09cf59..30b60e8079f 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
@@ -237,9 +237,13 @@ StatusOr> tensorflow::kernel_gen::GenerateCubinForTfCode(
   mlir::OwningModuleRef module = mlir::parseSourceString(tf_code, &context);
 
   TF_RETURN_IF_ERROR(LowerTfOpToLhloWithDynamicShapes(module.get()));
-  TF_RETURN_IF_ERROR(
-      xla::mlir_gpu::LowerLHLOToGPU(module.get(), tile_sizes, unroll_factors,
-                                    /*collapseParallelLoops=*/false));
+  {
+    xla::mlir_gpu::LowerLHLOToGPUOptions options;
+    options.tile_sizes = tile_sizes;
+    options.unroll_factors = unroll_factors;
+    options.collapse_parallel_loops = false;
+    TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerLHLOToGPU(module.get(), options));
+  }
   TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerKernelBodiesToNVVM(module.get()));
   // TODO(b/156985522): Figure out why we get a segfault when generating Tanh
   // with 'same_shape' containing {0, 1}. We would also get the crash if we
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index 4645b084eb6..f1e01bba27e 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -387,10 +387,7 @@ struct ParallelLoopCollapsingToFirstDim
 };
 }  // namespace
 
-Status LowerLHLOToGPU(mlir::ModuleOp module,
-                      llvm::ArrayRef tile_sizes,
-                      llvm::ArrayRef unroll_factors,
-                      bool collapseParallelLoops) {
+Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   mlir::PassManager pm(module.getContext());
   applyPassManagerCLOptions(pm);
 
@@ -399,14 +396,15 @@ Status LowerLHLOToGPU(mlir::ModuleOp module,
   // needed.
   llvm::SmallVector tiling_for_unrolling;
   llvm::SmallVector as_int64;
-  if (!unroll_factors.empty()) {
-    tiling_for_unrolling.reserve(tile_sizes.size());
-    for (auto pair : llvm::zip(tile_sizes, unroll_factors)) {
+  if (!options.unroll_factors.empty()) {
+    tiling_for_unrolling.reserve(options.tile_sizes.size());
+    for (auto pair : llvm::zip(options.tile_sizes, options.unroll_factors)) {
       tiling_for_unrolling.push_back(std::get<0>(pair) * std::get<1>(pair));
       as_int64.push_back(std::get<1>(pair));
     }
   } else {
-    tiling_for_unrolling.append(tile_sizes.begin(), tile_sizes.end());
+    tiling_for_unrolling.append(options.tile_sizes.begin(),
+                                options.tile_sizes.end());
   }
 
   // Legalize from HLO to LHLO.
@@ -441,11 +439,11 @@ Status LowerLHLOToGPU(mlir::ModuleOp module,
   pm.addPass(absl::make_unique());
   // Remove now unused temporary buffers.
   pm.addPass(absl::make_unique());
-  if (!unroll_factors.empty()) {
+  if (!options.unroll_factors.empty()) {
     pm.addPass(::mlir::createParallelLoopTilingPass(as_int64));
   }
   // Project all loop dimensions to X if necessary.
-  if (collapseParallelLoops) {
+  if (options.collapse_parallel_loops) {
     pm.addPass(absl::make_unique());
   }
   // Some basic cleanup.
@@ -464,7 +462,9 @@ Status LowerLHLOToGPU(mlir::ModuleOp module,
   pm.addPass(::mlir::createGpuKernelOutliningPass());
   // Make sure the kernel signature resembled the original function's
   // signature
-  pm.addPass(absl::make_unique());
+  if (options.fix_signature) {
+    pm.addPass(absl::make_unique());
+  }
   if (failed(pm.run(module))) {
     return InternalError("Lowering to GPU kernels failed.");
   }
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
index ab045808477..7b5d5c35c05 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
@@ -23,10 +23,15 @@ limitations under the License.
 namespace xla {
 namespace mlir_gpu {
 
+struct LowerLHLOToGPUOptions {
+  llvm::ArrayRef tile_sizes = {16, 64};
+  llvm::ArrayRef unroll_factors = {};
+  bool collapse_parallel_loops = true;
+  bool fix_signature = true;
+};
+
 Status LowerLHLOToGPU(mlir::ModuleOp module,
-                      llvm::ArrayRef tile_sizes = {16, 64},
-                      llvm::ArrayRef unroll_factors = {},
-                      bool collapseParallelLoops = true);
+                      LowerLHLOToGPUOptions options = {});
 
 Status LowerKernelBodiesToNVVM(mlir::ModuleOp module);
 

From 09cf51da59d77ae14c44116473e247cd88347ac7 Mon Sep 17 00:00:00 2001
From: Christian Sigg 
Date: Mon, 15 Jun 2020 12:35:00 -0700
Subject: [PATCH 0191/1390] Fix switch statement for
 CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING.

PiperOrigin-RevId: 316522188
Change-Id: I6557506f90d4bf6a91e97b65abd494d58769e3a3
---
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 2dbd2c58ebd..be18c989861 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -254,10 +254,8 @@ cudnnConvolutionBwdFilterAlgo_t ToConvBackwardFilterAlgo(
     // Based on cudnn.h, the following is not implemented.
     // case CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD:
     case CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED:
+    case CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING:
       return algo;
-    // Produces incorrect results for some shapes. Disabled for now, see
-    // NVIDIA bug 2072856. TODO(csigg): Only disable for subset of shapes.
-    // case CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING:
     default:
       LOG(FATAL)
           << "Unsupported Cudnn convolution backward algorithm for filter: "

From 08cbfe4090c322267c69dc67b6e070bc718914cb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 12:37:32 -0700
Subject: [PATCH 0192/1390] Fixed error from printing out of IndexedSlices
 objects.

PiperOrigin-RevId: 316522660
Change-Id: Ib7cc467ec2e9f2d465aa51db441574b9aca4f36e
---
 tensorflow/python/framework/indexed_slices.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/framework/indexed_slices.py b/tensorflow/python/framework/indexed_slices.py
index a2746d22650..6ddf9410fd7 100644
--- a/tensorflow/python/framework/indexed_slices.py
+++ b/tensorflow/python/framework/indexed_slices.py
@@ -147,7 +147,7 @@ class IndexedSlices(internal.NativeObject, composite_tensor.CompositeTensor):
     return "IndexedSlices(indices=%s, values=%s%s)" % (
         self._indices, self._values,
         (", dense_shape=%s" %
-         self._dense_shape) if self._dense_shape is not None else "")
+         (self._dense_shape,)) if self._dense_shape is not None else "")
 
   def __neg__(self):
     return IndexedSlices(-self.values, self.indices, self.dense_shape)

From 51373058dec434e10113da90e0a620e29715b36e Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Mon, 15 Jun 2020 12:43:34 -0700
Subject: [PATCH 0193/1390] [XLA] Do not needlessly store wrapped Tensor and
 ScopedShapedBuffer inside XlaTensor on a heap

Use absl::optional instead of std::unique_ptr to store them inside the class instead.

PiperOrigin-RevId: 316523861
Change-Id: I8f54f64e5661a877b7c9807465983d8132920474
---
 tensorflow/compiler/jit/xla_tensor.h | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/jit/xla_tensor.h b/tensorflow/compiler/jit/xla_tensor.h
index a6de405ec9e..dc358760534 100644
--- a/tensorflow/compiler/jit/xla_tensor.h
+++ b/tensorflow/compiler/jit/xla_tensor.h
@@ -55,7 +55,7 @@ class XlaTensor {
   // manage the memory for these tensors a ShapedBuffer may be required.
 
   // Return true if this XlaTensor contains a ShapedBuffer.
-  bool has_shaped_buffer() const { return shaped_buffer_ != nullptr; }
+  bool has_shaped_buffer() const { return shaped_buffer_.has_value(); }
   // Return the contained ShapedBuffer.
   // REQUIRES: has_shaped_buffer()
   const xla::ShapedBuffer& shaped_buffer() const {
@@ -68,8 +68,7 @@ class XlaTensor {
   }
   // Mutates the XlaTensor to set the ShapedBuffer.
   void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) {
-    shaped_buffer_ =
-        absl::make_unique(std::move(shaped_buffer));
+    shaped_buffer_ = std::move(shaped_buffer);
   }
 
   // Some tensors on the device may have known values on the host. We use these
@@ -77,14 +76,12 @@ class XlaTensor {
   // host value already.
 
   // Return true if this XlaTensor contains a host tensor.
-  bool has_host_tensor() const { return host_tensor_ != nullptr; }
+  bool has_host_tensor() const { return host_tensor_.has_value(); }
   // Return the contained host tensor.
   // REQUIRES: has_host_tensor()
   const Tensor& host_tensor() const { return *host_tensor_; }
   // Sets the contained host tensor.
-  void set_host_tensor(const Tensor& tensor) {
-    host_tensor_.reset(new Tensor(tensor));
-  }
+  void set_host_tensor(const Tensor& tensor) { host_tensor_.emplace(tensor); }
 
   // Adds synchronization events to 'stream' that wait for this tensor to be
   // defined on 'stream'. Does nothing if the tensor is already defined on that
@@ -111,9 +108,9 @@ class XlaTensor {
 
  private:
   // The optional contained ShapedBuffer.
-  std::unique_ptr shaped_buffer_;
+  absl::optional shaped_buffer_;
   // An optional host tensor value.
-  std::unique_ptr host_tensor_;
+  absl::optional host_tensor_;
   // An optional event that is triggered when the tensor's content has been
   // defined. If this event is nullptr, it is assumed that the tensor's content
   // is always defined.

From 540852285d931813f599fd66811d8d3413b0da4c Mon Sep 17 00:00:00 2001
From: Geoffrey Martin-Noble 
Date: Mon, 15 Jun 2020 12:51:16 -0700
Subject: [PATCH 0194/1390] Remove forwarding aliases from LLVM Bazel BUILD
 file

All uses have been migrated, so these are no longer necessary.

PiperOrigin-RevId: 316525401
Change-Id: I2f5c886f8580bcc7fe936e4f4ea7c440d6c20635
---
 .../xla/service/gpu/llvm_gpu_backend/BUILD    |    2 +-
 third_party/llvm/llvm.autogenerated.BUILD     | 1012 +----------------
 2 files changed, 48 insertions(+), 966 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD
index 5a3dc91d48f..c3ef02a04f2 100644
--- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD
@@ -39,6 +39,7 @@ cc_library(
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@llvm-project//llvm:AMDGPUCodeGen",
         "@llvm-project//llvm:Analysis",
         "@llvm-project//llvm:BitReader",
         "@llvm-project//llvm:BitWriter",
@@ -52,7 +53,6 @@ cc_library(
         "@llvm-project//llvm:Scalar",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:Target",
-        "@llvm-project//llvm:amdgpu_code_gen",
     ],
 )
 
diff --git a/third_party/llvm/llvm.autogenerated.BUILD b/third_party/llvm/llvm.autogenerated.BUILD
index 2857de01ecc..88f007dff1d 100644
--- a/third_party/llvm/llvm.autogenerated.BUILD
+++ b/third_party/llvm/llvm.autogenerated.BUILD
@@ -411,10 +411,10 @@ cc_binary(
     linkopts = llvm_linkopts,
     stamp = 0,
     deps = [
+        ":Support",
+        ":TableGen",
         ":config",
-        ":support",
-        ":tablegen",
-        ":utils_tablegen",
+        ":tblgen",
     ],
 )
 
@@ -428,7 +428,7 @@ cc_binary(
     copts = llvm_copts,
     linkopts = llvm_linkopts,
     stamp = 0,
-    deps = [":support"],
+    deps = [":Support"],
 )
 
 llvm_target_list = [
@@ -606,24 +606,18 @@ gentbl(
     ]),
 )
 
-[[
-    [gentbl(
-        name = target["name"] + "CommonTableGen",
-        tbl_outs = target["tbl_outs"],
-        tblgen = ":llvm-tblgen",
-        td_file = "lib/Target/" + target["dir_name"] + "/" + target["short_name"] + ".td",
-        td_srcs = [
-            ":common_target_td_sources",
-        ] + glob([
-            "lib/Target/" + target["dir_name"] + "/*.td",
-        ]),
-        deps = target.get("tbl_deps", []),
-    )],
-    [alias(
-        name = target["lower_name"] + "_target_gen",
-        actual = target["name"] + "CommonTableGen",
-    )],
-] for target in llvm_target_list]
+[gentbl(
+    name = target["name"] + "CommonTableGen",
+    tbl_outs = target["tbl_outs"],
+    tblgen = ":llvm-tblgen",
+    td_file = "lib/Target/" + target["dir_name"] + "/" + target["short_name"] + ".td",
+    td_srcs = [
+        ":common_target_td_sources",
+    ] + glob([
+        "lib/Target/" + target["dir_name"] + "/*.td",
+    ]),
+    deps = target.get("tbl_deps", []),
+) for target in llvm_target_list]
 
 # This target is used to provide *.def files to x86_code_gen.
 # Files with '.def' extension are not allowed in 'srcs' of 'cc_library' rule.
@@ -660,7 +654,7 @@ cc_binary(
     copts = llvm_copts,
     linkopts = llvm_linkopts,
     deps = [
-        ":support",
+        ":Support",
     ],
 )
 
@@ -698,11 +692,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "aarch64_asm_parser",
-    actual = ":AArch64AsmParser",
-)
-
 cc_library(
     name = "AArch64CodeGen",
     srcs = glob([
@@ -737,11 +726,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "aarch64_code_gen",
-    actual = ":AArch64CodeGen",
-)
-
 cc_library(
     name = "AArch64Desc",
     srcs = glob([
@@ -757,12 +741,12 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/AArch64"],
     deps = [
+        ":AArch64CommonTableGen",
         ":AArch64Info",
         ":AArch64Utils",
         ":BinaryFormat",
         ":MC",
         ":Support",
-        ":aarch64_target_gen",
         ":attributes_gen",
         ":config",
         ":intrinsic_enums_gen",
@@ -770,11 +754,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "aarch64_desc",
-    actual = ":AArch64Desc",
-)
-
 cc_library(
     name = "AArch64Disassembler",
     srcs = glob([
@@ -800,11 +779,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "aarch64_disassembler",
-    actual = ":AArch64Disassembler",
-)
-
 cc_library(
     name = "AArch64Info",
     srcs = glob([
@@ -823,18 +797,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/AArch64"],
     deps = [
+        ":CodeGen",
         ":Support",
-        ":code_gen",
+        ":Target",
         ":config",
-        ":target",
     ],
 )
 
-alias(
-    name = "aarch64_info",
-    actual = ":AArch64Info",
-)
-
 cc_library(
     name = "AArch64Utils",
     srcs = glob([
@@ -851,18 +820,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/AArch64"],
     deps = [
+        ":AArch64CommonTableGen",
+        ":MC",
         ":Support",
-        ":aarch64_target_gen",
         ":config",
-        ":mc",
     ],
 )
 
-alias(
-    name = "aarch64_utils",
-    actual = ":AArch64Utils",
-)
-
 cc_library(
     name = "AMDGPUAsmParser",
     srcs = glob([
@@ -888,11 +852,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "amdgpu_asm_parser",
-    actual = ":AMDGPUAsmParser",
-)
-
 cc_library(
     name = "AMDGPUCodeGen",
     srcs = glob([
@@ -930,11 +889,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "amdgpu_code_gen",
-    actual = ":AMDGPUCodeGen",
-)
-
 cc_library(
     name = "AMDGPUDesc",
     srcs = glob([
@@ -960,11 +914,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "amdgpu_desc",
-    actual = ":AMDGPUDesc",
-)
-
 cc_library(
     name = "AMDGPUDisassembler",
     srcs = glob([
@@ -990,11 +939,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "amdgpu_disassembler",
-    actual = ":AMDGPUDisassembler",
-)
-
 cc_library(
     name = "AMDGPUInfo",
     srcs = glob([
@@ -1010,19 +954,14 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/AMDGPU"],
     deps = [
+        ":AMDGPUCommonTableGen",
+        ":Core",
         ":Support",
-        ":amdgpu_target_gen",
         ":config",
-        ":core",
         ":r600_target_gen",
     ],
 )
 
-alias(
-    name = "amdgpu_info",
-    actual = ":AMDGPUInfo",
-)
-
 cc_library(
     name = "AMDGPUUtils",
     srcs = glob([
@@ -1038,21 +977,16 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/AMDGPU"],
     deps = [
+        ":AMDGPUCommonTableGen",
         ":BinaryFormat",
         ":Core",
         ":MC",
         ":Support",
-        ":amdgpu_target_gen",
         ":config",
         ":r600_target_gen",
     ],
 )
 
-alias(
-    name = "amdgpu_utils",
-    actual = ":AMDGPUUtils",
-)
-
 cc_library(
     name = "ARCCodeGen",
     srcs = glob([
@@ -1083,11 +1017,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arc_code_gen",
-    actual = ":ARCCodeGen",
-)
-
 cc_library(
     name = "ARCDesc",
     srcs = glob([
@@ -1110,11 +1039,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arc_desc",
-    actual = ":ARCDesc",
-)
-
 cc_library(
     name = "ARCDisassembler",
     srcs = glob([
@@ -1137,11 +1061,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arc_disassembler",
-    actual = ":ARCDisassembler",
-)
-
 cc_library(
     name = "ARCInfo",
     srcs = glob([
@@ -1162,11 +1081,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arc_info",
-    actual = ":ARCInfo",
-)
-
 cc_library(
     name = "ARMAsmParser",
     srcs = glob([
@@ -1192,11 +1106,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arm_asm_parser",
-    actual = ":ARMAsmParser",
-)
-
 cc_library(
     name = "ARMCodeGen",
     srcs = glob([
@@ -1231,11 +1140,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arm_code_gen",
-    actual = ":ARMCodeGen",
-)
-
 cc_library(
     name = "ARMDesc",
     srcs = glob([
@@ -1253,13 +1157,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/ARM"],
     deps = [
+        ":ARMCommonTableGen",
         ":ARMInfo",
         ":ARMUtils",
         ":BinaryFormat",
         ":MC",
         ":MCDisassembler",
         ":Support",
-        ":arm_target_gen",
         ":attributes_gen",
         ":config",
         ":intrinsic_enums_gen",
@@ -1267,11 +1171,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arm_desc",
-    actual = ":ARMDesc",
-)
-
 cc_library(
     name = "ARMDisassembler",
     srcs = glob([
@@ -1296,11 +1195,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "arm_disassembler",
-    actual = ":ARMDisassembler",
-)
-
 cc_library(
     name = "ARMInfo",
     srcs = glob([
@@ -1317,18 +1211,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/ARM"],
     deps = [
+        ":ARMCommonTableGen",
         ":Support",
-        ":arm_target_gen",
+        ":Target",
         ":config",
-        ":target",
     ],
 )
 
-alias(
-    name = "arm_info",
-    actual = ":ARMInfo",
-)
-
 cc_library(
     name = "ARMUtils",
     srcs = glob([
@@ -1345,18 +1234,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/ARM"],
     deps = [
+        ":ARMCommonTableGen",
+        ":MC",
         ":Support",
-        ":arm_target_gen",
         ":config",
-        ":mc",
     ],
 )
 
-alias(
-    name = "arm_utils",
-    actual = ":ARMUtils",
-)
-
 cc_library(
     name = "AVRAsmParser",
     srcs = glob([
@@ -1381,11 +1265,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "avr_asm_parser",
-    actual = ":AVRAsmParser",
-)
-
 cc_library(
     name = "AVRCodeGen",
     srcs = glob([
@@ -1414,11 +1293,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "avr_code_gen",
-    actual = ":AVRCodeGen",
-)
-
 cc_library(
     name = "AVRDesc",
     srcs = glob([
@@ -1441,11 +1315,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "avr_desc",
-    actual = ":AVRDesc",
-)
-
 cc_library(
     name = "AVRDisassembler",
     srcs = glob([
@@ -1468,11 +1337,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "avr_disassembler",
-    actual = ":AVRDisassembler",
-)
-
 cc_library(
     name = "AVRInfo",
     srcs = glob([
@@ -1493,11 +1357,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "avr_info",
-    actual = ":AVRInfo",
-)
-
 cc_library(
     name = "AggressiveInstCombine",
     srcs = glob([
@@ -1521,11 +1380,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "aggressive_inst_combine",
-    actual = ":AggressiveInstCombine",
-)
-
 cc_library(
     name = "Analysis",
     srcs = glob([
@@ -1552,11 +1406,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "analysis",
-    actual = ":Analysis",
-)
-
 cc_library(
     name = "AsmParser",
     srcs = glob([
@@ -1579,11 +1428,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "asm_parser",
-    actual = ":AsmParser",
-)
-
 cc_library(
     name = "AsmPrinter",
     srcs = glob([
@@ -1616,11 +1460,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "asm_printer",
-    actual = ":AsmPrinter",
-)
-
 cc_library(
     name = "BPFAsmParser",
     srcs = glob([
@@ -1645,11 +1484,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bpf_asm_parser",
-    actual = ":BPFAsmParser",
-)
-
 cc_library(
     name = "BPFCodeGen",
     srcs = glob([
@@ -1678,11 +1512,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bpf_code_gen",
-    actual = ":BPFCodeGen",
-)
-
 cc_library(
     name = "BPFDesc",
     srcs = glob([
@@ -1705,11 +1534,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bpf_desc",
-    actual = ":BPFDesc",
-)
-
 cc_library(
     name = "BPFDisassembler",
     srcs = glob([
@@ -1732,11 +1556,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bpf_disassembler",
-    actual = ":BPFDisassembler",
-)
-
 cc_library(
     name = "BPFInfo",
     srcs = glob([
@@ -1757,11 +1576,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bpf_info",
-    actual = ":BPFInfo",
-)
-
 cc_library(
     name = "BinaryFormat",
     srcs = glob([
@@ -1784,11 +1598,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "binary_format",
-    actual = ":BinaryFormat",
-)
-
 cc_library(
     name = "BitReader",
     srcs = glob([
@@ -1812,11 +1621,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bit_reader",
-    actual = ":BitReader",
-)
-
 cc_library(
     name = "BitWriter",
     srcs = glob([
@@ -1844,11 +1648,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bit_writer",
-    actual = ":BitWriter",
-)
-
 cc_library(
     name = "BitstreamReader",
     srcs = glob([
@@ -1869,11 +1668,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "bitstream_reader",
-    actual = ":BitstreamReader",
-)
-
 cc_library(
     name = "CFGuard",
     srcs = glob([
@@ -1895,11 +1689,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "cf_guard",
-    actual = ":CFGuard",
-)
-
 cc_library(
     name = "CodeGen",
     srcs = glob([
@@ -1920,6 +1709,7 @@ cc_library(
         ":BitReader",
         ":BitWriter",
         ":Core",
+        ":Instrumentation",
         ":MC",
         ":ProfileData",
         ":Scalar",
@@ -1927,15 +1717,9 @@ cc_library(
         ":Target",
         ":TransformUtils",
         ":config",
-        ":instrumentation",
     ],
 )
 
-alias(
-    name = "code_gen",
-    actual = ":CodeGen",
-)
-
 cc_library(
     name = "Core",
     srcs = glob([
@@ -1983,11 +1767,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "core",
-    actual = ":Core",
-)
-
 cc_library(
     name = "Coroutines",
     srcs = glob([
@@ -2013,11 +1792,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "coroutines",
-    actual = ":Coroutines",
-)
-
 cc_library(
     name = "Coverage",
     srcs = glob([
@@ -2041,11 +1815,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "coverage",
-    actual = ":Coverage",
-)
-
 cc_library(
     name = "DWARFLinker",
     srcs = glob([
@@ -2071,11 +1840,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "dwarf_linker",
-    actual = ":DWARFLinker",
-)
-
 cc_library(
     name = "DebugInfoCodeView",
     srcs = glob([
@@ -2091,18 +1855,13 @@ cc_library(
     ]),
     copts = llvm_copts,
     deps = [
+        ":BinaryFormat",
         ":DebugInfoMSF",
         ":Support",
-        ":binary_format",
         ":config",
     ],
 )
 
-alias(
-    name = "debug_info_code_view",
-    actual = ":DebugInfoCodeView",
-)
-
 cc_library(
     name = "DebugInfoDWARF",
     srcs = glob([
@@ -2126,11 +1885,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "debug_info_dwarf",
-    actual = ":DebugInfoDWARF",
-)
-
 cc_library(
     name = "DebugInfoGSYM",
     srcs = glob([
@@ -2154,11 +1908,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "debug_info_gsym",
-    actual = ":DebugInfoGSYM",
-)
-
 cc_library(
     name = "DebugInfoMSF",
     srcs = glob([
@@ -2179,11 +1928,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "debug_info_msf",
-    actual = ":DebugInfoMSF",
-)
-
 cc_library(
     name = "DebugInfoPDB",
     srcs = glob([
@@ -2208,11 +1952,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "debug_info_pdb",
-    actual = ":DebugInfoPDB",
-)
-
 cc_library(
     name = "Demangle",
     srcs = glob([
@@ -2230,11 +1969,6 @@ cc_library(
     deps = [":config"],
 )
 
-alias(
-    name = "demangle",
-    actual = ":Demangle",
-)
-
 cc_library(
     name = "DlltoolDriver",
     srcs = glob([
@@ -2257,11 +1991,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "dlltool_driver",
-    actual = ":DlltoolDriver",
-)
-
 cc_library(
     name = "ExecutionEngine",
     srcs = glob([
@@ -2287,11 +2016,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "execution_engine",
-    actual = ":ExecutionEngine",
-)
-
 cc_library(
     name = "Extensions",
     srcs = glob([
@@ -2309,11 +2033,6 @@ cc_library(
     deps = [":config"],
 )
 
-alias(
-    name = "extensions",
-    actual = ":Extensions",
-)
-
 cc_library(
     name = "FrontendOpenMP",
     srcs = glob([
@@ -2336,11 +2055,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "frontend_open_mp",
-    actual = ":FrontendOpenMP",
-)
-
 cc_library(
     name = "FuzzMutate",
     srcs = glob([
@@ -2367,11 +2081,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "fuzz_mutate",
-    actual = ":FuzzMutate",
-)
-
 cc_library(
     name = "GlobalISel",
     srcs = glob([
@@ -2399,11 +2108,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "global_i_sel",
-    actual = ":GlobalISel",
-)
-
 cc_library(
     name = "HexagonAsmParser",
     srcs = glob([
@@ -2428,11 +2132,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "hexagon_asm_parser",
-    actual = ":HexagonAsmParser",
-)
-
 cc_library(
     name = "HexagonCodeGen",
     srcs = glob([
@@ -2466,11 +2165,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "hexagon_code_gen",
-    actual = ":HexagonCodeGen",
-)
-
 cc_library(
     name = "HexagonDesc",
     srcs = glob([
@@ -2493,11 +2187,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "hexagon_desc",
-    actual = ":HexagonDesc",
-)
-
 cc_library(
     name = "HexagonDisassembler",
     srcs = glob([
@@ -2522,11 +2211,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "hexagon_disassembler",
-    actual = ":HexagonDisassembler",
-)
-
 cc_library(
     name = "HexagonInfo",
     srcs = glob([
@@ -2547,11 +2231,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "hexagon_info",
-    actual = ":HexagonInfo",
-)
-
 cc_library(
     name = "IPO",
     srcs = glob([
@@ -2590,11 +2269,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ipo",
-    actual = ":IPO",
-)
-
 cc_library(
     name = "IRReader",
     srcs = glob([
@@ -2618,11 +2292,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ir_reader",
-    actual = ":IRReader",
-)
-
 cc_library(
     name = "InstCombine",
     srcs = glob([
@@ -2640,18 +2309,13 @@ cc_library(
     deps = [
         ":Analysis",
         ":Core",
+        ":InstCombineTableGen",
         ":Support",
         ":TransformUtils",
         ":config",
-        ":instcombine_transforms_gen",
     ],
 )
 
-alias(
-    name = "inst_combine",
-    actual = ":InstCombine",
-)
-
 cc_library(
     name = "Instrumentation",
     srcs = glob([
@@ -2681,11 +2345,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "instrumentation",
-    actual = ":Instrumentation",
-)
-
 cc_library(
     name = "Interpreter",
     srcs = glob([
@@ -2709,11 +2368,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "interpreter",
-    actual = ":Interpreter",
-)
-
 cc_library(
     name = "JITLink",
     srcs = glob([
@@ -2736,11 +2390,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "jit_link",
-    actual = ":JITLink",
-)
-
 cc_library(
     name = "LTO",
     srcs = glob([
@@ -2780,11 +2429,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lto",
-    actual = ":LTO",
-)
-
 cc_library(
     name = "LanaiAsmParser",
     srcs = glob([
@@ -2809,11 +2453,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lanai_asm_parser",
-    actual = ":LanaiAsmParser",
-)
-
 cc_library(
     name = "LanaiCodeGen",
     srcs = glob([
@@ -2845,11 +2484,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lanai_code_gen",
-    actual = ":LanaiCodeGen",
-)
-
 cc_library(
     name = "LanaiDesc",
     srcs = glob([
@@ -2873,11 +2507,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lanai_desc",
-    actual = ":LanaiDesc",
-)
-
 cc_library(
     name = "LanaiDisassembler",
     srcs = glob([
@@ -2902,11 +2531,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lanai_disassembler",
-    actual = ":LanaiDisassembler",
-)
-
 cc_library(
     name = "LanaiInfo",
     srcs = glob([
@@ -2927,11 +2551,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lanai_info",
-    actual = ":LanaiInfo",
-)
-
 cc_library(
     name = "LibDriver",
     srcs = glob([
@@ -2956,11 +2575,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "lib_driver",
-    actual = ":LibDriver",
-)
-
 cc_library(
     name = "LineEditor",
     srcs = glob([
@@ -2981,11 +2595,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "line_editor",
-    actual = ":LineEditor",
-)
-
 cc_library(
     name = "Linker",
     srcs = glob([
@@ -3008,11 +2617,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "linker",
-    actual = ":Linker",
-)
-
 cc_library(
     name = "MC",
     srcs = glob([
@@ -3031,17 +2635,10 @@ cc_library(
         ":BinaryFormat",
         ":DebugInfoCodeView",
         ":Support",
-        ":binary_format",
         ":config",
-        ":debug_info_code_view",
     ],
 )
 
-alias(
-    name = "mc",
-    actual = ":MC",
-)
-
 cc_library(
     name = "MCA",
     srcs = glob([
@@ -3063,11 +2660,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mca",
-    actual = ":MCA",
-)
-
 cc_library(
     name = "MCDisassembler",
     srcs = glob([
@@ -3089,11 +2681,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mc_disassembler",
-    actual = ":MCDisassembler",
-)
-
 cc_library(
     name = "MCJIT",
     srcs = glob([
@@ -3119,11 +2706,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mcjit",
-    actual = ":MCJIT",
-)
-
 cc_library(
     name = "MCParser",
     srcs = glob([
@@ -3145,11 +2727,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mc_parser",
-    actual = ":MCParser",
-)
-
 cc_library(
     name = "MIRParser",
     srcs = glob([
@@ -3176,11 +2753,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mir_parser",
-    actual = ":MIRParser",
-)
-
 cc_library(
     name = "MLPolicies",
     srcs = glob([
@@ -3202,11 +2774,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ml_policies",
-    actual = ":MLPolicies",
-)
-
 cc_library(
     name = "MSP430AsmParser",
     srcs = glob([
@@ -3231,11 +2798,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "msp430_asm_parser",
-    actual = ":MSP430AsmParser",
-)
-
 cc_library(
     name = "MSP430CodeGen",
     srcs = glob([
@@ -3264,11 +2826,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "msp430_code_gen",
-    actual = ":MSP430CodeGen",
-)
-
 cc_library(
     name = "MSP430Desc",
     srcs = glob([
@@ -3291,11 +2848,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "msp430_desc",
-    actual = ":MSP430Desc",
-)
-
 cc_library(
     name = "MSP430Disassembler",
     srcs = glob([
@@ -3318,11 +2870,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "msp430_disassembler",
-    actual = ":MSP430Disassembler",
-)
-
 cc_library(
     name = "MSP430Info",
     srcs = glob([
@@ -3343,11 +2890,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "msp430_info",
-    actual = ":MSP430Info",
-)
-
 cc_library(
     name = "MipsAsmParser",
     srcs = glob([
@@ -3372,11 +2914,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mips_asm_parser",
-    actual = ":MipsAsmParser",
-)
-
 cc_library(
     name = "MipsCodeGen",
     srcs = glob([
@@ -3407,11 +2944,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mips_code_gen",
-    actual = ":MipsCodeGen",
-)
-
 cc_library(
     name = "MipsDesc",
     srcs = glob([
@@ -3434,11 +2966,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mips_desc",
-    actual = ":MipsDesc",
-)
-
 cc_library(
     name = "MipsDisassembler",
     srcs = glob([
@@ -3461,11 +2988,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mips_disassembler",
-    actual = ":MipsDisassembler",
-)
-
 cc_library(
     name = "MipsInfo",
     srcs = glob([
@@ -3486,11 +3008,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "mips_info",
-    actual = ":MipsInfo",
-)
-
 cc_library(
     name = "NVPTXCodeGen",
     srcs = glob([
@@ -3524,11 +3041,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "nvptx_code_gen",
-    actual = ":NVPTXCodeGen",
-)
-
 cc_library(
     name = "NVPTXDesc",
     srcs = glob([
@@ -3544,19 +3056,14 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/NVPTX"],
     deps = [
-        "nvptx_target_gen",
         ":MC",
+        ":NVPTXCommonTableGen",
         ":NVPTXInfo",
         ":Support",
         ":config",
     ],
 )
 
-alias(
-    name = "nvptx_desc",
-    actual = ":NVPTXDesc",
-)
-
 cc_library(
     name = "NVPTXInfo",
     srcs = glob([
@@ -3574,20 +3081,15 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/NVPTX"],
     deps = [
-        "nvptx_target_gen",
+        ":Core",
+        ":NVPTXCommonTableGen",
         ":Support",
+        ":Target",
         ":attributes_gen",
         ":config",
-        ":core",
-        ":target",
     ],
 )
 
-alias(
-    name = "nvptx_info",
-    actual = ":NVPTXInfo",
-)
-
 cc_library(
     name = "ObjCARC",
     srcs = glob([
@@ -3612,11 +3114,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "objc_arc",
-    actual = ":ObjCARC",
-)
-
 cc_library(
     name = "Object",
     srcs = glob([
@@ -3643,11 +3140,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "object",
-    actual = ":Object",
-)
-
 cc_library(
     name = "ObjectYAML",
     srcs = glob([
@@ -3671,11 +3163,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "object_yaml",
-    actual = ":ObjectYAML",
-)
-
 cc_library(
     name = "Option",
     srcs = glob([
@@ -3696,11 +3183,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "option",
-    actual = ":Option",
-)
-
 cc_library(
     name = "OrcError",
     srcs = glob([
@@ -3721,11 +3203,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "orc_error",
-    actual = ":OrcError",
-)
-
 cc_library(
     name = "OrcJIT",
     srcs = glob([
@@ -3756,11 +3233,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "orc_jit",
-    actual = ":OrcJIT",
-)
-
 cc_library(
     name = "Passes",
     srcs = glob([
@@ -3794,11 +3266,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "passes",
-    actual = ":Passes",
-)
-
 cc_library(
     name = "PowerPCAsmParser",
     srcs = glob([
@@ -3823,11 +3290,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "powerpc_asm_parser",
-    actual = ":PowerPCAsmParser",
-)
-
 cc_library(
     name = "PowerPCCodeGen",
     srcs = glob([
@@ -3859,11 +3321,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "powerpc_code_gen",
-    actual = ":PowerPCCodeGen",
-)
-
 cc_library(
     name = "PowerPCDesc",
     srcs = glob([
@@ -3881,21 +3338,16 @@ cc_library(
     deps = [
         ":BinaryFormat",
         ":MC",
+        ":PowerPCCommonTableGen",
         ":PowerPCInfo",
         ":Support",
         ":attributes_gen",
         ":config",
         ":intrinsic_enums_gen",
         ":intrinsics_impl_gen",
-        ":powerpc_target_gen",
     ],
 )
 
-alias(
-    name = "powerpc_desc",
-    actual = ":PowerPCDesc",
-)
-
 cc_library(
     name = "PowerPCDisassembler",
     srcs = glob([
@@ -3918,11 +3370,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "powerpc_disassembler",
-    actual = ":PowerPCDisassembler",
-)
-
 cc_library(
     name = "PowerPCInfo",
     srcs = glob([
@@ -3940,20 +3387,15 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/PowerPC"],
     deps = [
+        ":Core",
+        ":PowerPCCommonTableGen",
         ":Support",
+        ":Target",
         ":attributes_gen",
         ":config",
-        ":core",
-        ":powerpc_target_gen",
-        ":target",
     ],
 )
 
-alias(
-    name = "powerpc_info",
-    actual = ":PowerPCInfo",
-)
-
 cc_library(
     name = "ProfileData",
     srcs = glob([
@@ -3975,11 +3417,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "profile_data",
-    actual = ":ProfileData",
-)
-
 cc_library(
     name = "RISCVAsmParser",
     srcs = glob([
@@ -4005,11 +3442,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_asm_parser",
-    actual = ":RISCVAsmParser",
-)
-
 cc_library(
     name = "RISCVCodeGen",
     srcs = glob([
@@ -4041,11 +3473,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_code_gen",
-    actual = ":RISCVCodeGen",
-)
-
 cc_library(
     name = "RISCVDesc",
     srcs = glob([
@@ -4069,11 +3496,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_desc",
-    actual = ":RISCVDesc",
-)
-
 cc_library(
     name = "RISCVDisassembler",
     srcs = glob([
@@ -4096,11 +3518,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_disassembler",
-    actual = ":RISCVDisassembler",
-)
-
 cc_library(
     name = "RISCVInfo",
     srcs = glob([
@@ -4121,11 +3538,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_info",
-    actual = ":RISCVInfo",
-)
-
 cc_library(
     name = "RISCVUtils",
     srcs = glob([
@@ -4146,11 +3558,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "riscv_utils",
-    actual = ":RISCVUtils",
-)
-
 cc_library(
     name = "Remarks",
     srcs = glob([
@@ -4172,11 +3579,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "remarks",
-    actual = ":Remarks",
-)
-
 cc_library(
     name = "RuntimeDyld",
     srcs = glob([
@@ -4201,18 +3603,13 @@ cc_library(
     copts = llvm_copts,
     deps = [
         ":MC",
+        ":MCDisassembler",
         ":Object",
         ":Support",
         ":config",
-        ":mc_disassembler",
     ],
 )
 
-alias(
-    name = "runtime_dyld",
-    actual = ":RuntimeDyld",
-)
-
 cc_library(
     name = "Scalar",
     srcs = glob([
@@ -4238,17 +3635,12 @@ cc_library(
         ":Core",
         ":InstCombine",
         ":Support",
+        ":Target",
         ":TransformUtils",
         ":config",
-        ":target",
     ],
 )
 
-alias(
-    name = "scalar",
-    actual = ":Scalar",
-)
-
 cc_library(
     name = "SelectionDAG",
     srcs = glob([
@@ -4275,11 +3667,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "selection_dag",
-    actual = ":SelectionDAG",
-)
-
 cc_library(
     name = "SparcAsmParser",
     srcs = glob([
@@ -4304,11 +3691,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "sparc_asm_parser",
-    actual = ":SparcAsmParser",
-)
-
 cc_library(
     name = "SparcCodeGen",
     srcs = glob([
@@ -4337,11 +3719,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "sparc_code_gen",
-    actual = ":SparcCodeGen",
-)
-
 cc_library(
     name = "SparcDesc",
     srcs = glob([
@@ -4364,11 +3741,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "sparc_desc",
-    actual = ":SparcDesc",
-)
-
 cc_library(
     name = "SparcDisassembler",
     srcs = glob([
@@ -4391,11 +3763,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "sparc_disassembler",
-    actual = ":SparcDisassembler",
-)
-
 cc_library(
     name = "SparcInfo",
     srcs = glob([
@@ -4416,11 +3783,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "sparc_info",
-    actual = ":SparcInfo",
-)
-
 cc_library(
     name = "Support",
     srcs = glob([
@@ -4452,11 +3814,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "support",
-    actual = ":Support",
-)
-
 cc_library(
     name = "Symbolize",
     srcs = glob([
@@ -4481,11 +3838,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "symbolize",
-    actual = ":Symbolize",
-)
-
 cc_library(
     name = "SystemZAsmParser",
     srcs = glob([
@@ -4510,11 +3862,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "system_z_asm_parser",
-    actual = ":SystemZAsmParser",
-)
-
 cc_library(
     name = "SystemZCodeGen",
     srcs = glob([
@@ -4545,11 +3892,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "system_z_code_gen",
-    actual = ":SystemZCodeGen",
-)
-
 cc_library(
     name = "SystemZDesc",
     srcs = glob([
@@ -4572,11 +3914,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "system_z_desc",
-    actual = ":SystemZDesc",
-)
-
 cc_library(
     name = "SystemZDisassembler",
     srcs = glob([
@@ -4601,11 +3938,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "system_z_disassembler",
-    actual = ":SystemZDisassembler",
-)
-
 cc_library(
     name = "SystemZInfo",
     srcs = glob([
@@ -4626,11 +3958,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "system_z_info",
-    actual = ":SystemZInfo",
-)
-
 cc_library(
     name = "TableGen",
     srcs = glob([
@@ -4648,17 +3975,12 @@ cc_library(
     ]),
     copts = llvm_copts,
     deps = [
+        ":MC",
         ":Support",
         ":config",
-        ":mc",
     ],
 )
 
-alias(
-    name = "tablegen",
-    actual = ":TableGen",
-)
-
 cc_library(
     name = "Target",
     srcs = glob([
@@ -4687,11 +4009,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "target",
-    actual = ":Target",
-)
-
 cc_library(
     name = "TestingSupport",
     srcs = glob([
@@ -4712,11 +4029,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "testing_support",
-    actual = ":TestingSupport",
-)
-
 cc_library(
     name = "TextAPI",
     srcs = glob([
@@ -4752,11 +4064,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "text_api",
-    actual = ":TextAPI",
-)
-
 cc_library(
     name = "TransformUtils",
     srcs = glob([
@@ -4781,11 +4088,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "transform_utils",
-    actual = ":TransformUtils",
-)
-
 cc_library(
     name = "VEAsmParser",
     srcs = glob([
@@ -4810,11 +4112,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ve_asm_parser",
-    actual = ":VEAsmParser",
-)
-
 cc_library(
     name = "VECodeGen",
     srcs = glob([
@@ -4844,11 +4141,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ve_code_gen",
-    actual = ":VECodeGen",
-)
-
 cc_library(
     name = "VEDesc",
     srcs = glob([
@@ -4871,11 +4163,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ve_desc",
-    actual = ":VEDesc",
-)
-
 cc_library(
     name = "VEDisassembler",
     srcs = glob([
@@ -4898,11 +4185,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ve_disassembler",
-    actual = ":VEDisassembler",
-)
-
 cc_library(
     name = "VEInfo",
     srcs = glob([
@@ -4923,11 +4205,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "ve_info",
-    actual = ":VEInfo",
-)
-
 cc_library(
     name = "Vectorize",
     srcs = glob([
@@ -4947,18 +4224,13 @@ cc_library(
     deps = [
         ":Analysis",
         ":Core",
+        ":Scalar",
         ":Support",
         ":TransformUtils",
         ":config",
-        ":scalar",
     ],
 )
 
-alias(
-    name = "vectorize",
-    actual = ":Vectorize",
-)
-
 cc_library(
     name = "WebAssemblyAsmParser",
     srcs = glob([
@@ -4982,11 +4254,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "web_assembly_asm_parser",
-    actual = ":WebAssemblyAsmParser",
-)
-
 cc_library(
     name = "WebAssemblyCodeGen",
     srcs = glob([
@@ -5019,11 +4286,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "web_assembly_code_gen",
-    actual = ":WebAssemblyCodeGen",
-)
-
 cc_library(
     name = "WebAssemblyDesc",
     srcs = glob([
@@ -5046,11 +4308,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "web_assembly_desc",
-    actual = ":WebAssemblyDesc",
-)
-
 cc_library(
     name = "WebAssemblyDisassembler",
     srcs = glob([
@@ -5075,11 +4332,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "web_assembly_disassembler",
-    actual = ":WebAssemblyDisassembler",
-)
-
 cc_library(
     name = "WebAssemblyInfo",
     srcs = glob([
@@ -5100,11 +4352,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "web_assembly_info",
-    actual = ":WebAssemblyInfo",
-)
-
 cc_library(
     name = "WindowsManifest",
     srcs = glob([
@@ -5125,11 +4372,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "windows_manifest",
-    actual = ":WindowsManifest",
-)
-
 cc_library(
     name = "X86AsmParser",
     srcs = glob([
@@ -5154,11 +4396,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x86_asm_parser",
-    actual = ":X86AsmParser",
-)
-
 cc_library(
     name = "X86CodeGen",
     srcs = glob([
@@ -5192,11 +4429,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x86_code_gen",
-    actual = ":X86CodeGen",
-)
-
 cc_library(
     name = "X86Desc",
     srcs = glob([
@@ -5221,11 +4453,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x86_desc",
-    actual = ":X86Desc",
-)
-
 cc_library(
     name = "X86Disassembler",
     srcs = glob([
@@ -5248,11 +4475,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x86_disassembler",
-    actual = ":X86Disassembler",
-)
-
 cc_library(
     name = "X86Info",
     srcs = glob([
@@ -5269,18 +4491,13 @@ cc_library(
     ]),
     copts = llvm_copts + ["-Iexternal/llvm-project/llvm/lib/Target/X86"],
     deps = [
+        ":MC",
         ":Support",
+        ":X86CommonTableGen",
         ":config",
-        ":mc",
-        ":x86_target_gen",
     ],
 )
 
-alias(
-    name = "x86_info",
-    actual = ":X86Info",
-)
-
 cc_library(
     name = "XCoreCodeGen",
     srcs = glob([
@@ -5311,11 +4528,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x_core_code_gen",
-    actual = ":XCoreCodeGen",
-)
-
 cc_library(
     name = "XCoreDesc",
     srcs = glob([
@@ -5338,11 +4550,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x_core_desc",
-    actual = ":XCoreDesc",
-)
-
 cc_library(
     name = "XCoreDisassembler",
     srcs = glob([
@@ -5365,11 +4572,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x_core_disassembler",
-    actual = ":XCoreDisassembler",
-)
-
 cc_library(
     name = "XCoreInfo",
     srcs = glob([
@@ -5390,11 +4592,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x_core_info",
-    actual = ":XCoreInfo",
-)
-
 cc_library(
     name = "XRay",
     srcs = glob([
@@ -5416,11 +4613,6 @@ cc_library(
     ],
 )
 
-alias(
-    name = "x_ray",
-    actual = ":XRay",
-)
-
 cc_library(
     name = "gtest",
     srcs = glob([
@@ -5460,113 +4652,3 @@ cc_library(
         ":gtest",
     ],
 )
-
-alias(
-    name = "aarch64_target",
-    actual = ":aarch64_code_gen",
-)
-
-alias(
-    name = "aarch64_target_disassembler",
-    actual = ":aarch64_disassembler",
-)
-
-alias(
-    name = "arm_target",
-    actual = ":arm_code_gen",
-)
-
-alias(
-    name = "arm_target_disassembler",
-    actual = ":arm_disassembler",
-)
-
-alias(
-    name = "codegen",
-    actual = ":code_gen",
-)
-
-alias(
-    name = "frontend_openmp",
-    actual = ":frontend_open_mp",
-)
-
-alias(
-    name = "ipo_transforms",
-    actual = ":ipo",
-)
-
-alias(
-    name = "ir",
-    actual = ":core",
-)
-
-alias(
-    name = "machine_code",
-    actual = ":mc",
-)
-
-alias(
-    name = "machine_code_disassembler",
-    actual = ":mc_disassembler",
-)
-
-alias(
-    name = "nvptx_target",
-    actual = ":nvptx_code_gen",
-)
-
-alias(
-    name = "objcarc_transforms",
-    actual = ":objc_arc",
-)
-
-alias(
-    name = "orcjit",
-    actual = ":orc_jit",
-)
-
-alias(
-    name = "powerpc_target",
-    actual = ":powerpc_code_gen",
-)
-
-alias(
-    name = "powerpc_target_disassembler",
-    actual = ":powerpc_disassembler",
-)
-
-alias(
-    name = "scalar_transforms",
-    actual = ":scalar",
-)
-
-alias(
-    name = "target_base",
-    actual = ":target",
-)
-
-alias(
-    name = "x86_target",
-    actual = ":x86_code_gen",
-)
-
-alias(
-    name = "x86_target_disassembler",
-    actual = ":x86_disassembler",
-)
-
-alias(
-    name = "all_targets",
-    actual = ":AllTargetsCodeGens",
-)
-
-alias(
-    name = "instcombine_transforms_gen",
-    actual = ":InstCombineTableGen",
-)
-
-alias(
-    name = "utils_tablegen",
-    actual = ":tblgen",
-)

From 18e0e6450bc7b87a118565156ff90b1effee2118 Mon Sep 17 00:00:00 2001
From: Francois Chollet 
Date: Mon, 15 Jun 2020 13:02:46 -0700
Subject: [PATCH 0195/1390] Prevent Keras dataset loading from affecting the
 global RNG seed.

PiperOrigin-RevId: 316527944
Change-Id: I13fc997ffafc02f25b94e45265c7aa97b6efc6c4
---
 tensorflow/python/keras/datasets/boston_housing.py | 4 ++--
 tensorflow/python/keras/datasets/reuters.py        | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/keras/datasets/boston_housing.py b/tensorflow/python/keras/datasets/boston_housing.py
index 2c0badfefba..8886634a4b7 100644
--- a/tensorflow/python/keras/datasets/boston_housing.py
+++ b/tensorflow/python/keras/datasets/boston_housing.py
@@ -67,9 +67,9 @@ def load_data(path='boston_housing.npz', test_split=0.2, seed=113):
     x = f['x']
     y = f['y']
 
-  np.random.seed(seed)
+  rng = np.random.RandomState(seed)
   indices = np.arange(len(x))
-  np.random.shuffle(indices)
+  rng.shuffle(indices)
   x = x[indices]
   y = y[indices]
 
diff --git a/tensorflow/python/keras/datasets/reuters.py b/tensorflow/python/keras/datasets/reuters.py
index 46ac9249637..b71440fd632 100644
--- a/tensorflow/python/keras/datasets/reuters.py
+++ b/tensorflow/python/keras/datasets/reuters.py
@@ -119,9 +119,9 @@ def load_data(path='reuters.npz',
   with np.load(path, allow_pickle=True) as f:
     xs, labels = f['x'], f['y']
 
-  np.random.seed(seed)
+  rng = np.random.RandomState(seed)
   indices = np.arange(len(xs))
-  np.random.shuffle(indices)
+  rng.shuffle(indices)
   xs = xs[indices]
   labels = labels[indices]
 

From 2fe65568b1475f23afc2ee9387e8f447d5eda4d8 Mon Sep 17 00:00:00 2001
From: sshiddib 
Date: Mon, 15 Jun 2020 13:14:03 -0700
Subject: [PATCH 0196/1390] [Intel MKL] Adding DNNL ops (part 1) supporting
 threadpool work

---
 .../core/kernels/mkl_conv_grad_filter_ops.cc  | 17 ++++++-
 .../core/kernels/mkl_conv_grad_input_ops.cc   | 11 ++++-
 tensorflow/core/kernels/mkl_conv_ops.cc       | 14 ++++++
 .../core/kernels/mkl_fused_batch_norm_op.cc   | 46 ++++++++++++++++++-
 .../core/kernels/mkl_pooling_ops_common.cc    | 25 +++++++++-
 tensorflow/core/kernels/mkl_quantize_op.cc    |  5 ++
 tensorflow/core/kernels/mkl_relu_op.cc        | 18 +++++++-
 tensorflow/core/kernels/mkl_slice_op.cc       |  8 +++-
 tensorflow/core/kernels/mkl_softmax_op.cc     |  8 +++-
 9 files changed, 143 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
index 4c3cea4b6ff..12581d0bfa5 100644
--- a/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_grad_filter_ops.cc
@@ -114,6 +114,21 @@ class MklConvBwdFilterPrimitive : public MklPrimitive {
   void Execute(const T* src_data, const T* diff_filter_data,
                const T* diff_bias_data, const T* diff_dst_data,
                std::shared_ptr bwd_filter_stream) {
+    // TODO: Create a common function and avoid the duplicate code
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *bwd_filter_stream);
+    context_.diff_filter_mem->set_data_handle(
+        static_cast(const_cast(diff_filter_data)),
+        *bwd_filter_stream);
+    if (diff_bias_data != nullptr) {
+      context_.diff_bias_mem->set_data_handle(
+          static_cast(const_cast(diff_bias_data)),
+          *bwd_filter_stream);
+    }
+    context_.diff_dst_mem->set_data_handle(
+        static_cast(const_cast(diff_dst_data)), *bwd_filter_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.diff_filter_mem->set_data_handle(
@@ -124,7 +139,7 @@ class MklConvBwdFilterPrimitive : public MklPrimitive {
     }
     context_.diff_dst_mem->set_data_handle(
         static_cast(const_cast(diff_dst_data)));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     execute_primitives(context_.bwd_filter_primitives, bwd_filter_stream,
                        context_.bwd_filter_primitives_args);
diff --git a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
index f9c8d11c67c..7177431029a 100644
--- a/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_grad_input_ops.cc
@@ -116,13 +116,22 @@ class MklConvBwdInputPrimitive : public MklPrimitive {
   void Execute(const T* diff_src_data, const T* filter_data,
                const T* diff_dst_data,
                std::shared_ptr bwd_input_stream) {
+    // TODO: Create a common function and avoid the duplicate code
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.diff_src_mem->set_data_handle(
+        static_cast(const_cast(diff_src_data)), *bwd_input_stream);
+    context_.filter_mem->set_data_handle(
+        static_cast(const_cast(filter_data)), *bwd_input_stream);
+    context_.diff_dst_mem->set_data_handle(
+        static_cast(const_cast(diff_dst_data)), *bwd_input_stream);
+#else
     context_.diff_src_mem->set_data_handle(
         static_cast(const_cast(diff_src_data)));
     context_.filter_mem->set_data_handle(
         static_cast(const_cast(filter_data)));
     context_.diff_dst_mem->set_data_handle(
         static_cast(const_cast(diff_dst_data)));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     execute_primitives(context_.bwd_input_primitives, bwd_input_stream,
                        context_.bwd_input_primitives_args);
diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc
index 7d0510d03ac..210044436aa 100644
--- a/tensorflow/core/kernels/mkl_conv_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_ops.cc
@@ -110,6 +110,19 @@ class MklConvFwdPrimitive : public MklPrimitive {
   void Execute(const Tinput* src_data, const Tfilter* filter_data,
                const Tbias* bias_data, const Toutput* dst_data,
                std::shared_ptr fwd_stream) {
+    // TODO: Create a common function and avoid the duplicate code
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *fwd_stream);
+    context_.filter_mem->set_data_handle(
+        static_cast(const_cast(filter_data)), *fwd_stream);
+    if (bias_data != nullptr) {
+      context_.bias_mem->set_data_handle(
+          static_cast(const_cast(bias_data)), *fwd_stream);
+    }
+    context_.dst_mem->set_data_handle(
+        static_cast(const_cast(dst_data)), *fwd_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.filter_mem->set_data_handle(
@@ -120,6 +133,7 @@ class MklConvFwdPrimitive : public MklPrimitive {
     }
     context_.dst_mem->set_data_handle(
         static_cast(const_cast(dst_data)));
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     DCHECK_EQ(context_.fwd_primitives.size(),
               context_.fwd_primitives_args.size());
diff --git a/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc b/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc
index 954ae0492df..3b2c4f84039 100644
--- a/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc
+++ b/tensorflow/core/kernels/mkl_fused_batch_norm_op.cc
@@ -94,6 +94,28 @@ class MklFusedBatchNormFwdPrimitive : public MklPrimitive {
   void Execute(const T* src_data, const U* weights_data, T* dst_data,
                U* mean_data, U* variance_data,
                std::shared_ptr fwd_stream, U* workspace_data) {
+    // TODO: Create a common function and avoid the duplicate code
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *fwd_stream);
+    context_.dst_mem->set_data_handle(static_cast(dst_data),
+                                      *fwd_stream);
+
+    if (IS_SET(use_scale_shift))
+      context_.weights_mem->set_data_handle(
+          static_cast(const_cast(weights_data)), *fwd_stream);
+
+    if ((context_.pkind == prop_kind::forward_training) ||
+        (IS_SET(use_global_stats))) {
+      context_.mean_mem->set_data_handle(static_cast(mean_data),
+                                         *fwd_stream);
+      context_.variance_mem->set_data_handle(static_cast(variance_data),
+                                             *fwd_stream);
+    }
+    if (workspace_data != nullptr) {
+      context_.ws_mem->set_data_handle(workspace_data, *fwd_stream);
+    }
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.dst_mem->set_data_handle(static_cast(dst_data));
@@ -110,6 +132,7 @@ class MklFusedBatchNormFwdPrimitive : public MklPrimitive {
     if (workspace_data != nullptr) {
       context_.ws_mem->set_data_handle(workspace_data);
     }
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     // Execute batch-normalization forward primitives.
     execute_primitives(context_.fwd_primitives, fwd_stream, context_.net_args);
@@ -503,6 +526,27 @@ class MklFusedBatchNormBwdPrimitive : public MklPrimitive {
                const T* diff_dst_data, const U* weights_data, T* diff_src_data,
                U* diff_weights_data, U* res_space_data,
                std::shared_ptr bwd_stream) {
+    // TODO: Create a common function and avoid the duplicate code
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *bwd_stream);
+    context_.mean_mem->set_data_handle(
+        static_cast(const_cast(mean_data)), *bwd_stream);
+    context_.variance_mem->set_data_handle(
+        static_cast(const_cast(variance_data)), *bwd_stream);
+    context_.diff_dst_mem->set_data_handle(
+        static_cast(const_cast(diff_dst_data)), *bwd_stream);
+
+    if (IS_SET(use_scale_shift)) {
+      context_.weights_mem->set_data_handle(
+          static_cast(const_cast(weights_data)), *bwd_stream);
+      context_.diff_weights_mem->set_data_handle(
+          static_cast(diff_weights_data), *bwd_stream);
+    }
+
+    context_.diff_src_mem->set_data_handle(static_cast(diff_src_data),
+                                           *bwd_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.mean_mem->set_data_handle(
@@ -520,7 +564,7 @@ class MklFusedBatchNormBwdPrimitive : public MklPrimitive {
     }
 
     context_.diff_src_mem->set_data_handle(static_cast(diff_src_data));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     // Execute backward batch-normalization primitives.
     DCHECK_EQ(context_.bwd_primitives.size(), context_.net_args.size());
diff --git a/tensorflow/core/kernels/mkl_pooling_ops_common.cc b/tensorflow/core/kernels/mkl_pooling_ops_common.cc
index 2dfc6db0075..5f1c9129ec3 100644
--- a/tensorflow/core/kernels/mkl_pooling_ops_common.cc
+++ b/tensorflow/core/kernels/mkl_pooling_ops_common.cc
@@ -127,6 +127,17 @@ template 
 void MklPoolingFwdPrimitive::Execute(const T* src_data, T* dst_data,
                                         void* ws_data,
                                         std::shared_ptr fwd_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+  context_.src_mem->set_data_handle(
+      static_cast(const_cast(src_data)), *fwd_stream);
+  context_.dst_mem->set_data_handle(static_cast(dst_data), *fwd_stream);
+  if (context_.alg_kind == ALGORITHM::pooling_max &&
+      context_.prop_kind ==
+          prop_kind::forward_training) {  // Max pooling must have workspace.
+    DCHECK(ws_data != nullptr);
+    context_.ws_mem->set_data_handle(ws_data, *fwd_stream);
+  }
+#else
   context_.src_mem->set_data_handle(
       static_cast(const_cast(src_data)));
   context_.dst_mem->set_data_handle(static_cast(dst_data));
@@ -136,7 +147,7 @@ void MklPoolingFwdPrimitive::Execute(const T* src_data, T* dst_data,
     DCHECK(ws_data != nullptr);
     context_.ws_mem->set_data_handle(ws_data);
   }
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
   execute_primitives(context_.fwd_primitives, fwd_stream, context_.net_args);
 #else
@@ -269,6 +280,16 @@ template 
 void MklPoolingBwdPrimitive::Execute(const T* diff_dst_data,
                                         T* diff_src_data, const void* ws_data,
                                         std::shared_ptr bwd_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+  context_.diff_dst_mem->set_data_handle(
+      static_cast(const_cast(diff_dst_data)), *bwd_stream);
+  context_.diff_src_mem->set_data_handle(static_cast(diff_src_data),
+                                         *bwd_stream);
+  if (context_.alg_kind == ALGORITHM::pooling_max) {
+    DCHECK(ws_data != nullptr);
+    context_.ws_mem->set_data_handle(const_cast(ws_data), *bwd_stream);
+  }
+#else
   context_.diff_dst_mem->set_data_handle(
       static_cast(const_cast(diff_dst_data)));
   context_.diff_src_mem->set_data_handle(static_cast(diff_src_data));
@@ -276,7 +297,7 @@ void MklPoolingBwdPrimitive::Execute(const T* diff_dst_data,
     DCHECK(ws_data != nullptr);
     context_.ws_mem->set_data_handle(const_cast(ws_data));
   }
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
   execute_primitives(context_.bwd_primitives, bwd_stream, context_.net_args);
 #else
diff --git a/tensorflow/core/kernels/mkl_quantize_op.cc b/tensorflow/core/kernels/mkl_quantize_op.cc
index 5adb9862250..177cbb43d0b 100644
--- a/tensorflow/core/kernels/mkl_quantize_op.cc
+++ b/tensorflow/core/kernels/mkl_quantize_op.cc
@@ -88,8 +88,13 @@ class MklReorderWithScalePrimitive : public MklPrimitive {
 
   void Execute(void* src_data, void* dst_data,
                std::shared_ptr reorder_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(src_data, *reorder_stream);
+    context_.dst_mem->set_data_handle(dst_data, *reorder_stream);
+#else
     context_.src_mem->set_data_handle(src_data);
     context_.dst_mem->set_data_handle(dst_data);
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifndef ENABLE_MKLDNN_V1
     reorder_stream->submit(context_.net);
 #else
diff --git a/tensorflow/core/kernels/mkl_relu_op.cc b/tensorflow/core/kernels/mkl_relu_op.cc
index 784bbc682dc..9af580de777 100644
--- a/tensorflow/core/kernels/mkl_relu_op.cc
+++ b/tensorflow/core/kernels/mkl_relu_op.cc
@@ -79,10 +79,16 @@ class MklEltwiseFwdPrimitive : public MklPrimitive {
   //   dst_data:  output data buffer of dst
   void Execute(const T* src_data, T* dst_data,
                std::shared_ptr fwd_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *fwd_stream);
+    context_.dst_mem->set_data_handle(static_cast(dst_data),
+                                      *fwd_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.dst_mem->set_data_handle(static_cast(dst_data));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     DCHECK_EQ(context_.fwd_primitives.size(),
               context_.fwd_primitives_args.size());
@@ -293,12 +299,20 @@ class MklEltwiseBwdPrimitive : public MklPrimitive {
   //   diff_src_data:  output data buffer of diff_src
   void Execute(const T* src_data, const T* diff_dst_data, T* diff_src_data,
                std::shared_ptr bwd_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *bwd_stream);
+    context_.diff_dst_mem->set_data_handle(
+        static_cast(const_cast(diff_dst_data)), *bwd_stream);
+    context_.diff_src_mem->set_data_handle(static_cast(diff_src_data),
+                                           *bwd_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.diff_dst_mem->set_data_handle(
         static_cast(const_cast(diff_dst_data)));
     context_.diff_src_mem->set_data_handle(static_cast(diff_src_data));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     DCHECK_EQ(context_.bwd_primitives.size(),
               context_.bwd_primitives_args.size());
diff --git a/tensorflow/core/kernels/mkl_slice_op.cc b/tensorflow/core/kernels/mkl_slice_op.cc
index 4115691c79d..7e293e14d98 100644
--- a/tensorflow/core/kernels/mkl_slice_op.cc
+++ b/tensorflow/core/kernels/mkl_slice_op.cc
@@ -189,9 +189,15 @@ class MklSlicePrimitive : public MklPrimitive {
 
   void Execute(const MklSliceParams& sliceParams,
                std::shared_ptr slice_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(sliceParams.from->get_data_handle(),
+                                      *slice_stream);
+    context_.dst_mem->set_data_handle(sliceParams.to->get_data_handle(),
+                                      *slice_stream);
+#else
     context_.src_mem->set_data_handle(sliceParams.from->get_data_handle());
     context_.dst_mem->set_data_handle(sliceParams.to->get_data_handle());
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     execute_primitives(context_.slice_primitives, slice_stream,
                        context_.slice_primitives_args);
diff --git a/tensorflow/core/kernels/mkl_softmax_op.cc b/tensorflow/core/kernels/mkl_softmax_op.cc
index 4d1cf90f28d..2f51573fe13 100644
--- a/tensorflow/core/kernels/mkl_softmax_op.cc
+++ b/tensorflow/core/kernels/mkl_softmax_op.cc
@@ -59,10 +59,16 @@ class MklSoftmaxPrimitive : public MklPrimitive {
   //   dst_data:  output data buffer of dst
   void Execute(const T* src_data, T* dst_data,
                std::shared_ptr fwd_cpu_stream) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    context_.src_mem->set_data_handle(
+        static_cast(const_cast(src_data)), *fwd_cpu_stream);
+    context_.dst_mem->set_data_handle(static_cast(dst_data),
+                                      *fwd_cpu_stream);
+#else
     context_.src_mem->set_data_handle(
         static_cast(const_cast(src_data)));
     context_.dst_mem->set_data_handle(static_cast(dst_data));
-
+#endif  // ENABLE_MKLDNN_THREADPOOL
 #ifdef ENABLE_MKLDNN_V1
     DCHECK_EQ(context_.fwd_primitives.size(), context_.fwd_net_args.size());
     execute_primitives(context_.fwd_primitives, fwd_cpu_stream,

From 5016da312802f8372672ccbfa7a4207b8683a8e5 Mon Sep 17 00:00:00 2001
From: sshiddib 
Date: Mon, 15 Jun 2020 13:21:21 -0700
Subject: [PATCH 0197/1390] [Intel MKL] Adding DNNL ops (part 2) supporting
 threadpool work

---
 tensorflow/core/kernels/mkl_aggregate_ops.cc |  7 +++++--
 tensorflow/core/kernels/mkl_concat_op.cc     | 17 ++++++++++++++---
 tensorflow/core/kernels/mkl_dequantize_op.cc |  7 +++++--
 tensorflow/core/kernels/mkl_lrn_op.cc        |  6 ++++--
 tensorflow/core/kernels/mkl_transpose_op.cc  |  5 +++--
 tensorflow/core/util/mkl_util.h              | 14 ++++++++++++--
 6 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/kernels/mkl_aggregate_ops.cc b/tensorflow/core/kernels/mkl_aggregate_ops.cc
index ec5f80cb3fa..90e0ea9aa95 100644
--- a/tensorflow/core/kernels/mkl_aggregate_ops.cc
+++ b/tensorflow/core/kernels/mkl_aggregate_ops.cc
@@ -178,6 +178,9 @@ class MklAddNOp : public OpKernel {
         dnn_fmt = MklTensorFormatToMklDnnDataFormat(mkl_data_format);
       }
 
+      std::shared_ptr fwd_cpu_stream;
+      fwd_cpu_stream.reset(CreateStream(ctx, cpu_engine));
+
       // Create memory descriptor for MKL-DNN.
       // If all input in Tensorflow format, create block memory descriptor,
       // else convert TF format to MKL memory descriptor
@@ -215,6 +218,7 @@ class MklAddNOp : public OpKernel {
         srcs_pd.push_back(memory::primitive_desc(md, cpu_engine));
 #endif
         src.SetUsrMem(md, &src_tensor);
+        src.SetUsrMemDataHandle(&src_tensor, fwd_cpu_stream);
         inputs.push_back(src.GetOpMem());
       }
 
@@ -240,11 +244,10 @@ class MklAddNOp : public OpKernel {
       }
       AllocateOutputSetMklShape(ctx, kOutputIdx, &dst_tensor, output_tf_shape,
                                 output_mkl_shape);
-      dst.SetUsrMemDataHandle(dst_tensor);
+      dst.SetUsrMemDataHandle(dst_tensor, fwd_cpu_stream);
 
       // Create Sum op, and submit net for execution.
       std::vector net;
-      stream* fwd_cpu_stream = CreateStream(ctx, cpu_engine);
 #ifdef ENABLE_MKLDNN_V1
       mkldnn::sum sum_op(sum_pd);
       std::unordered_map net_args = {
diff --git a/tensorflow/core/kernels/mkl_concat_op.cc b/tensorflow/core/kernels/mkl_concat_op.cc
index 976f778424e..4a5cb0a0d4f 100644
--- a/tensorflow/core/kernels/mkl_concat_op.cc
+++ b/tensorflow/core/kernels/mkl_concat_op.cc
@@ -281,11 +281,19 @@ class MklConcatFwdPrimitive : public MklPrimitive {
                std::shared_ptr fwd_stream) {
     DCHECK_EQ(in_data.size(), context_.data_mem.size());
     for (size_t i = 0; i < concat_fwd_dims.num_inputs; i++) {
+#ifdef ENABLE_MKLDNN_THREADPOOL
+      context_.data_mem_shdptr[i]->set_data_handle(
+          static_cast(in_data[i].get_data_handle()), *fwd_stream);
+    }
+    context_.dst_mem->set_data_handle(
+        static_cast(dst_data.get_data_handle()), *fwd_stream);
+#else
       context_.data_mem_shdptr[i]->set_data_handle(
           static_cast(in_data[i].get_data_handle()));
     }
     context_.dst_mem->set_data_handle(
         static_cast(dst_data.get_data_handle()));
+#endif  // ENABLE_MKLDNN_THREADPOOL
 
     for (size_t i = 0; i < concat_fwd_dims.num_inputs; i++) {
       context_.data_mem[i] = *context_.data_mem_shdptr[i];
@@ -788,11 +796,13 @@ class MklConcatOp : public OpKernel {
                                     dnn_shape_dst);
           DCHECK(dst_tensor != nullptr) << "Output tensor pointer is NULL";
 
+          std::shared_ptr fwd_cpu_stream;
+          fwd_cpu_stream.reset(CreateStream(context, cpu_engine));
+
           if (dnn_shape_dst.IsMklTensor())
             dst_md = dnn_shape_dst.GetMklLayout();
           dst.SetUsrMem(dst_md, dst_tensor);
-          std::shared_ptr fwd_cpu_stream;
-          fwd_cpu_stream.reset(CreateStream(context, cpu_engine));
+          dst.SetUsrMemDataHandle(dst_tensor, fwd_cpu_stream);
 #ifdef ENABLE_MKLDNN_V1
           auto concat_op = concat(concat_pd);
           std::unordered_map net_args = {
@@ -830,9 +840,10 @@ class MklConcatOp : public OpKernel {
 
           dst_md = dnn_shape_dst.IsMklTensor() ? dnn_shape_dst.GetMklLayout()
                                                : dst_md;
-          dst.SetUsrMem(dst_md, dst_tensor);
           std::shared_ptr fwd_cpu_stream;
           fwd_cpu_stream.reset(CreateStream(context, concat_fwd->GetEngine()));
+          dst.SetUsrMem(dst_md, dst_tensor);
+          dst.SetUsrMemDataHandle(dst_tensor, fwd_cpu_stream);
           // Execute concat
           concat_fwd->Execute(srcs_mem, dst.GetOpMem(), concat_fwd_dims,
                               fwd_cpu_stream);
diff --git a/tensorflow/core/kernels/mkl_dequantize_op.cc b/tensorflow/core/kernels/mkl_dequantize_op.cc
index 06570c1db1c..82d78250576 100644
--- a/tensorflow/core/kernels/mkl_dequantize_op.cc
+++ b/tensorflow/core/kernels/mkl_dequantize_op.cc
@@ -75,6 +75,9 @@ class MklDequantizeOp : public OpKernel {
       MklDnnData src(&cpu_engine);
       MklDnnData dst(&cpu_engine);
 
+      std::shared_ptr reorder_stream;
+      reorder_stream.reset(CreateStream(ctx, cpu_engine));
+
       // If input is in MKL layout, then simply grab input layout; otherwise,
       // construct input TF layout. For TF layout, although input shape
       // (src_dims) required is in MKL-DNN order, the layout is Tensorflow's
@@ -85,6 +88,7 @@ class MklDequantizeOp : public OpKernel {
               : memory::desc(src_dims, MklDnnType(), MEMORY_FORMAT::nhwc);
 
       src.SetUsrMem(src_md, &src_tensor);
+      src.SetUsrMemDataHandle(&src_tensor, reorder_stream);
 
       Tensor* output_tensor = nullptr;
       MklDnnShape output_mkl_shape;
@@ -129,6 +133,7 @@ class MklDequantizeOp : public OpKernel {
       AllocateOutputSetMklShape(ctx, 0, &output_tensor, output_tf_shape,
                                 output_mkl_shape);
       dst.SetUsrMem(dst_md, output_tensor);
+      dst.SetUsrMemDataHandle(output_tensor, reorder_stream);
 
       // The quantization logic here for mode SCALED is similar to the logic
       // in QuantizeAndDequantizeV2 and QuantizeAndDequantizeV3.
@@ -155,8 +160,6 @@ class MklDequantizeOp : public OpKernel {
       // Also it does not define round_nearest (enum).
       attr.set_int_output_round_mode(mkldnn::round_mode::round_nearest);
 #endif  // !ENABLE_MKLDNN_V1
-      std::shared_ptr reorder_stream;
-      reorder_stream.reset(CreateStream(ctx, cpu_engine));
       std::vector net;
 
       // Create reorder primitive and then execute.
diff --git a/tensorflow/core/kernels/mkl_lrn_op.cc b/tensorflow/core/kernels/mkl_lrn_op.cc
index a11e7ebcbf5..3e512d0792b 100644
--- a/tensorflow/core/kernels/mkl_lrn_op.cc
+++ b/tensorflow/core/kernels/mkl_lrn_op.cc
@@ -137,6 +137,7 @@ class MklLRNOp : public OpKernel {
       // that input is in NHWC layout with Channel being the last dimension.
       src_dnn_data.SetUsrMem(src_md, &src_tensor);
       src_dnn_data.SetOpMemDesc(input_dims, MEMORY_FORMAT::nhwc);
+      src_dnn_data.SetUsrMemDataHandle(&src_tensor, fwd_stream_);
 
       // dst_dnn_data has the same shape as input.
       dst_dnn_data.SetUsrMem(src_md);
@@ -157,7 +158,7 @@ class MklLRNOp : public OpKernel {
                            &output_tensor);
       OP_REQUIRES_OK(context, context->status());
       DCHECK(output_tensor != nullptr);
-      dst_dnn_data.SetUsrMemDataHandle(output_tensor);
+      dst_dnn_data.SetUsrMemDataHandle(output_tensor, fwd_stream_);
 
       // Handle workspace required for MKL-DNN.
       AllocateWorkspaceTensor(context, lrn_prim_desc, &workspace_dnn_data);
@@ -393,6 +394,7 @@ class MklLRNGradOp : public OpKernel {
           orig_input_dnn_shape.GetSizesAsMklDnnDims();
       orig_input_dnn_data.SetUsrMem(orig_input_md, &orig_input_tensor);
       orig_input_dnn_data.SetOpMemDesc(orig_input_dims, MEMORY_FORMAT::nhwc);
+      orig_input_dnn_data.SetUsrMemDataHandle(&orig_input_tensor, bwd_stream_);
 
       // output_dnn_data has the same shape as original input
       output_dnn_data.SetUsrMem(orig_input_md);
@@ -421,7 +423,7 @@ class MklLRNGradOp : public OpKernel {
                            orig_input_format, &output_tensor);
       OP_REQUIRES_OK(context, context->status());
       DCHECK(output_tensor != nullptr);
-      output_dnn_data.SetUsrMemDataHandle(output_tensor);
+      output_dnn_data.SetUsrMemDataHandle(output_tensor, bwd_stream_);
 
       // Create LRN primitive and add it to the net
       // At this point, workspace is enabled, so we don't need
diff --git a/tensorflow/core/kernels/mkl_transpose_op.cc b/tensorflow/core/kernels/mkl_transpose_op.cc
index 77a68afa752..2e5c6d2719b 100644
--- a/tensorflow/core/kernels/mkl_transpose_op.cc
+++ b/tensorflow/core/kernels/mkl_transpose_op.cc
@@ -137,6 +137,7 @@ Status MKLTransposeND(OpKernelContext* context, const Tensor& in_tensor,
     memory::dims out_strides =
         ReorderStrides(CalculateTFStrides(out_dims), perm);
 
+    std::shared_ptr transpose_stream;
     in.SetUsrMem(in_dims, in_strides, &in_tensor);
     // Output dimensions are same as input dimensions. We adjust the layout
     // using strides.
@@ -144,16 +145,16 @@ Status MKLTransposeND(OpKernelContext* context, const Tensor& in_tensor,
 
     std::vector net;
 #ifdef ENABLE_MKLDNN_V1
-    std::shared_ptr transpose_stream;
     auto* prim = FindOrCreateReorder(in.GetUsrMem(), out.GetUsrMem());
     transpose_stream.reset(CreateStream(context, prim->GetEngine()));
+    in.SetUsrMemDataHandle(&in_tensor, transpose_stream);
+    out.SetUsrMemDataHandle(out_tensor, transpose_stream);
     net.push_back(*(prim->GetPrimitive()));
     std::vector net_args;
     net_args.push_back({{MKLDNN_ARG_FROM, *in.GetUsrMem()},
                         {MKLDNN_ARG_TO, *out.GetUsrMem()}});
     execute_primitives(net, transpose_stream, net_args);
 #else
-    std::shared_ptr transpose_stream;
     transpose_stream.reset(new CPU_STREAM(cpu_engine));
     net.push_back(FindOrCreateReorder(in.GetUsrMem(), out.GetUsrMem()));
     transpose_stream->submit(net).wait();
diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index 7f6272b09c1..996984eebc0 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -1524,17 +1524,27 @@ class MklDnnData {
   }
 
   /// Set function for data buffer of user memory primitive.
-  inline void SetUsrMemDataHandle(void* data_buffer) {
+  inline void SetUsrMemDataHandle(void* data_buffer,
+                                  std::shared_ptr t_stream = nullptr) {
     CHECK_NOTNULL(user_memory_);
     CHECK_NOTNULL(data_buffer);
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    user_memory_->set_data_handle(data_buffer, *t_stream);
+#else
     user_memory_->set_data_handle(data_buffer);
+#endif  // ENABLE_MKLDNN_THREADPOOL
   }
 
   /// Set function for data buffer of user memory primitive.
-  inline void SetUsrMemDataHandle(const Tensor* tensor) {
+  inline void SetUsrMemDataHandle(const Tensor* tensor,
+                                  std::shared_ptr t_stream = nullptr) {
     CHECK_NOTNULL(user_memory_);
     CHECK_NOTNULL(tensor);
+#ifdef ENABLE_MKLDNN_THREADPOOL
+    user_memory_->set_data_handle(GetTensorBuffer(tensor), *t_stream);
+#else
     user_memory_->set_data_handle(GetTensorBuffer(tensor));
+#endif  // ENABLE_MKLDNN_THREADPOOL
   }
 
   /// allocate function for data buffer

From df4ea0c1a5ffd82ea6b9159f864fd354dc19d8f7 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Mon, 15 Jun 2020 13:18:42 -0700
Subject: [PATCH 0198/1390] Move tracking_util_xla_test to keras/tests

PiperOrigin-RevId: 316531126
Change-Id: I8d93adaca9c51bf85a6a938ce6a0dacc160a8a48
---
 tensorflow/python/keras/tests/BUILD           | 27 +++++++++++++++++++
 .../tests/tracking_util_xla_test.py}          |  0
 tensorflow/python/training/tracking/BUILD     | 26 ------------------
 3 files changed, 27 insertions(+), 26 deletions(-)
 rename tensorflow/python/{training/tracking/util_xla_test.py => keras/tests/tracking_util_xla_test.py} (100%)

diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD
index d03b1bd1ee8..36af32184e6 100644
--- a/tensorflow/python/keras/tests/BUILD
+++ b/tensorflow/python/keras/tests/BUILD
@@ -6,6 +6,7 @@ load("//tensorflow:tensorflow.bzl", "cuda_py_test")
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_py_test")
 load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_test")
+load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_test")
 
 package(
     default_visibility = [
@@ -466,6 +467,32 @@ tf_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "tracking_util_xla_test",
+    srcs = ["tracking_util_xla_test.py"],
+    python_version = "PY3",
+    tags = [
+        "no_pip",
+        "no_windows",
+        "nomac",
+        "notsan",  # b/74395663
+    ],
+    deps = [
+        "//tensorflow/compiler/tests:xla_test",
+        "//tensorflow/python:checkpoint_management",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:constant_op",
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/keras:metrics",
+        "//tensorflow/python/keras/engine",
+        "//tensorflow/python/keras/layers:core",
+        "//tensorflow/python/keras/optimizer_v2",
+        "//tensorflow/python/training/tracking",
+        "//tensorflow/python/training/tracking:util",
+    ],
+)
+
 py_library(
     name = "get_config_samples",
     srcs = ["get_config_samples.py"],
diff --git a/tensorflow/python/training/tracking/util_xla_test.py b/tensorflow/python/keras/tests/tracking_util_xla_test.py
similarity index 100%
rename from tensorflow/python/training/tracking/util_xla_test.py
rename to tensorflow/python/keras/tests/tracking_util_xla_test.py
diff --git a/tensorflow/python/training/tracking/BUILD b/tensorflow/python/training/tracking/BUILD
index 88dfd8eba55..ffc43964fb4 100644
--- a/tensorflow/python/training/tracking/BUILD
+++ b/tensorflow/python/training/tracking/BUILD
@@ -6,7 +6,6 @@ load(
     "//tensorflow/tools/test:performance.bzl",
     "tf_py_logged_benchmark",
 )
-load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_test")
 
 package(
     default_visibility = [
@@ -191,31 +190,6 @@ tf_py_test(
     ],
 )
 
-tf_xla_py_test(
-    name = "util_xla_test",
-    srcs = ["util_xla_test.py"],
-    python_version = "PY3",
-    tags = [
-        "no_pip",
-        "no_windows",
-        "nomac",
-        "notsan",  # b/74395663
-    ],
-    deps = [
-        ":tracking",
-        ":util",
-        "//tensorflow/compiler/tests:xla_test",
-        "//tensorflow/python:checkpoint_management",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:constant_op",
-        "//tensorflow/python:framework_ops",
-        "//tensorflow/python/eager:backprop",
-        "//tensorflow/python/keras:engine",
-        "//tensorflow/python/keras/layers",
-        "//tensorflow/python/keras/optimizer_v2",
-    ],
-)
-
 tf_py_test(
     name = "util_with_v1_optimizers_test",
     srcs = ["util_with_v1_optimizers_test.py"],

From 52736a6adc5bf2d30512e0bdb627f048d64a6561 Mon Sep 17 00:00:00 2001
From: Francois Chollet 
Date: Mon, 15 Jun 2020 13:30:30 -0700
Subject: [PATCH 0199/1390] Prevent Keras dataset loading from affecting the
 global RNG

PiperOrigin-RevId: 316533425
Change-Id: I6099847f9a7ead24786fb2fecd5ba488f53456e6
---
 tensorflow/python/keras/datasets/imdb.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/keras/datasets/imdb.py b/tensorflow/python/keras/datasets/imdb.py
index 61fbf92eaef..37403228edf 100644
--- a/tensorflow/python/keras/datasets/imdb.py
+++ b/tensorflow/python/keras/datasets/imdb.py
@@ -113,14 +113,14 @@ def load_data(path='imdb.npz',
     x_train, labels_train = f['x_train'], f['y_train']
     x_test, labels_test = f['x_test'], f['y_test']
 
-  np.random.seed(seed)
+  rng = np.random.RandomState(seed)
   indices = np.arange(len(x_train))
-  np.random.shuffle(indices)
+  rng.shuffle(indices)
   x_train = x_train[indices]
   labels_train = labels_train[indices]
 
   indices = np.arange(len(x_test))
-  np.random.shuffle(indices)
+  rng.shuffle(indices)
   x_test = x_test[indices]
   labels_test = labels_test[indices]
 

From a8950d70bfe0405fa405127cf5fd824a7a778aac Mon Sep 17 00:00:00 2001
From: Shanqing Cai 
Date: Mon, 15 Jun 2020 13:49:39 -0700
Subject: [PATCH 0200/1390] [tfdbg2] Add tfdbg_run_id to metadata of data dumps

- A data dump file set generated by tfdbg2 can contain
  multiple subsets when there are multiple hosts involved
  in the instrumented TensorFlow job (e.g., TPUs and Parameter Servers).
  Currently, there is no bit in those subset of files that
  indicates they belong to the same instrumented TF job.
  - This CL addresses this problem by adding a field to the
    metadata proto used by those files (`tfdbg_run_id`)
- The DebugEventsWriter code is revised, so that this new
  field is written to the metadata file of the file set on the writer's
  construction.
- Also in this CL: remove the previous 1-arg `GetDebugEventsWriter(dump_root)`
  that creates the writer object if it doesn't exist at the specified
  dump_root. Replace it with `LookUpDebugEventsWriter(dump_root)` that only
  looks up the writer object and returns a non-OK status if such an object
  hasn't been created at `dump_root`. This makes the code less error prone by
  keeping only the fully-explicit, 3-arg `GetDebugEventsWriter()`.

PiperOrigin-RevId: 316537044
Change-Id: Id5be0b771fbf37c0fc796f1514ed858a0e6d38f0
---
 tensorflow/core/kernels/debug_ops.h           |  11 +-
 tensorflow/core/ops/debug_ops.cc              |   1 +
 tensorflow/core/protobuf/debug_event.proto    |   6 +
 tensorflow/core/util/debug_events_writer.cc   |  23 +++-
 tensorflow/core/util/debug_events_writer.h    |  33 ++++--
 .../core/util/debug_events_writer_test.cc     | 104 ++++++++++--------
 .../client/debug_events_writer_wrapper.cc     |  59 +++++-----
 .../python/debug/lib/debug_events_reader.py   |   5 +
 .../python/debug/lib/debug_events_writer.py   |   6 +-
 .../debug/lib/debug_events_writer_test.py     |  39 +++++--
 .../python/debug/lib/debug_v2_ops_test.py     |   7 +-
 .../python/debug/lib/dumping_callback.py      |  15 ++-
 .../debug/lib/dumping_callback_test_lib.py    |   2 +
 .../api/golden/v1/tensorflow.raw_ops.pbtxt    |   2 +-
 .../api/golden/v2/tensorflow.raw_ops.pbtxt    |   2 +-
 15 files changed, 203 insertions(+), 112 deletions(-)

diff --git a/tensorflow/core/kernels/debug_ops.h b/tensorflow/core/kernels/debug_ops.h
index 3fef822244d..498cd6146a8 100644
--- a/tensorflow/core/kernels/debug_ops.h
+++ b/tensorflow/core/kernels/debug_ops.h
@@ -410,7 +410,8 @@ class DebugIdentityV2Op : public OpKernel {
       : OpKernel(context),
         device_name_(context->device()->name()),
         output_slot_(-1),
-        tensor_debug_mode_(0) {
+        tensor_debug_mode_(0),
+        tfdbg_run_id_() {
     std::vector debug_urls;
     OP_REQUIRES_OK(context, context->GetAttr("debug_urls", &debug_urls));
     for (const string& debug_url : debug_urls) {
@@ -435,14 +436,17 @@ class DebugIdentityV2Op : public OpKernel {
       circular_buffer_size_ =
           tfdbg::DebugEventsWriter::kDefaultCyclicBufferSize;
     }
+    if (context->HasAttr("tfdbg_run_id")) {
+      OP_REQUIRES_OK(context, context->GetAttr("tfdbg_run_id", &tfdbg_run_id_));
+    }
   }
 
   void Compute(OpKernelContext* context) override {
     const Tensor& tensor = context->input(0);
     for (const string& dump_root : dump_roots_) {
       tfdbg::DebugEventsWriter* debug_events_writer =
-          tfdbg::DebugEventsWriter::GetDebugEventsWriter(dump_root,
-                                                         circular_buffer_size_);
+          tfdbg::DebugEventsWriter::GetDebugEventsWriter(
+              dump_root, tfdbg_run_id_, circular_buffer_size_);
       OP_REQUIRES_OK(context, debug_events_writer->WriteGraphExecutionTrace(
                                   tfdbg_context_id_, device_name_, op_name_,
                                   output_slot_, tensor_debug_mode_, tensor));
@@ -458,6 +462,7 @@ class DebugIdentityV2Op : public OpKernel {
   int32 output_slot_;
   int32 tensor_debug_mode_;
   int64 circular_buffer_size_;
+  string tfdbg_run_id_;
 };
 
 typedef Eigen::ThreadPoolDevice CPUDevice;
diff --git a/tensorflow/core/ops/debug_ops.cc b/tensorflow/core/ops/debug_ops.cc
index 0ecc58a6a8f..ac67a0f75f3 100644
--- a/tensorflow/core/ops/debug_ops.cc
+++ b/tensorflow/core/ops/debug_ops.cc
@@ -91,6 +91,7 @@ REGISTER_OP("DebugIdentityV2")
     .Attr("tensor_debug_mode: int = -1")
     .Attr("debug_urls: list(string) = []")
     .Attr("circular_buffer_size: int = 1000")
+    .Attr("tfdbg_run_id: string = ''")
     .SetIsStateful()
     .SetShapeFn(shape_inference::UnchangedShape);
 
diff --git a/tensorflow/core/protobuf/debug_event.proto b/tensorflow/core/protobuf/debug_event.proto
index 005abe53194..5541c397fb8 100644
--- a/tensorflow/core/protobuf/debug_event.proto
+++ b/tensorflow/core/protobuf/debug_event.proto
@@ -115,6 +115,12 @@ message DebugMetadata {
   // Version of the DebugEvent file format.
   // Has a format of "debug.Event:", e.g., "debug.Event:1".
   string file_version = 2;
+
+  // A unique ID for the current run of tfdbg.
+  // A run of tfdbg is defined as a TensorFlow job instrumented by tfdbg.
+  // Multiple hosts in a distributed TensorFlow job instrumented by tfdbg
+  // have the same ID.
+  string tfdbg_run_id = 3;
 }
 
 // Content of a source file involved in the execution of the debugged TensorFlow
diff --git a/tensorflow/core/util/debug_events_writer.cc b/tensorflow/core/util/debug_events_writer.cc
index d9c3393ce3c..8ee42959131 100644
--- a/tensorflow/core/util/debug_events_writer.cc
+++ b/tensorflow/core/util/debug_events_writer.cc
@@ -122,23 +122,31 @@ DebugEventsWriter::~DebugEventsWriter() { Close().IgnoreError(); }
 
 // static
 DebugEventsWriter* DebugEventsWriter::GetDebugEventsWriter(
-    const string& dump_root, int64 circular_buffer_size) {
+    const string& dump_root, const string& tfdbg_run_id,
+    int64 circular_buffer_size) {
   mutex_lock l(DebugEventsWriter::factory_mu_);
   std::unordered_map>* writer_pool =
       DebugEventsWriter::GetDebugEventsWriterMap();
   if (writer_pool->find(dump_root) == writer_pool->end()) {
     std::unique_ptr writer(
-        new DebugEventsWriter(dump_root, circular_buffer_size));
+        new DebugEventsWriter(dump_root, tfdbg_run_id, circular_buffer_size));
     writer_pool->insert(std::make_pair(dump_root, std::move(writer)));
   }
   return (*writer_pool)[dump_root].get();
 }
 
 // static
-DebugEventsWriter* DebugEventsWriter::GetDebugEventsWriter(
-    const string& dump_root) {
-  return DebugEventsWriter::GetDebugEventsWriter(dump_root,
-                                                 kDefaultCyclicBufferSize);
+Status DebugEventsWriter::LookUpDebugEventsWriter(
+    const string& dump_root, DebugEventsWriter** debug_events_writer) {
+  mutex_lock l(DebugEventsWriter::factory_mu_);
+  std::unordered_map>* writer_pool =
+      DebugEventsWriter::GetDebugEventsWriterMap();
+  if (writer_pool->find(dump_root) == writer_pool->end()) {
+    return errors::FailedPrecondition(
+        "No DebugEventsWriter has been created at dump root ", dump_root);
+  }
+  *debug_events_writer = (*writer_pool)[dump_root].get();
+  return Status::OK();
 }
 
 Status DebugEventsWriter::Init() {
@@ -179,6 +187,7 @@ Status DebugEventsWriter::Init() {
   metadata->set_tensorflow_version(TF_VERSION_STRING);
   metadata->set_file_version(
       strings::Printf("%s%d", kVersionPrefix, kCurrentFormatVersion));
+  metadata->set_tfdbg_run_id(tfdbg_run_id_);
   TF_RETURN_IF_ERROR(SerializeAndWriteDebugEvent(&debug_event, METADATA));
   TF_RETURN_WITH_CONTEXT_IF_ERROR(
       metadata_writer_->Flush(), "Failed to flush debug event metadata writer");
@@ -457,9 +466,11 @@ DebugEventsWriter::GetDebugEventsWriterMap() {
 }
 
 DebugEventsWriter::DebugEventsWriter(const string& dump_root,
+                                     const string& tfdbg_run_id,
                                      int64 circular_buffer_size)
     : env_(Env::Default()),
       dump_root_(dump_root),
+      tfdbg_run_id_(tfdbg_run_id),
       is_initialized_(false),
       initialization_mu_(),
       circular_buffer_size_(circular_buffer_size),
diff --git a/tensorflow/core/util/debug_events_writer.h b/tensorflow/core/util/debug_events_writer.h
index 39835adf1a6..412f947e22d 100644
--- a/tensorflow/core/util/debug_events_writer.h
+++ b/tensorflow/core/util/debug_events_writer.h
@@ -93,18 +93,27 @@ class DebugEventsWriter {
   // sets of six. The singleton pattern avoids storing multiple sets in a single
   // folder, which might cause confusion.
   //
+  // If an instance of DebugEventsWriter has already been created at a
+  // `dump_root`, calling this method with the same `dump_root` will return
+  // the existing instance.
+  //
   // Args:
   //   dump_root: Dump root directory. If it doesn't exist, will be created.
+  //   tfdbg_run_id: Debugging run ID of the writer.
   //   circular_buffer_size: Circular buffer size (in number of DebugEvent
   //     protos). If set to a value <=0, will abolish the circular-buffer
   //     behavior.
   // Returns:
   //   A pointer to a DebugEventsWriter object: a per-dump_root singleton.
   static DebugEventsWriter* GetDebugEventsWriter(const string& dump_root,
+                                                 const string& tfdbg_run_id,
                                                  int64 circular_buffer_size);
-  // Same as the 2-arg factory method above, but uses the default circular
-  // buffer size.
-  static DebugEventsWriter* GetDebugEventsWriter(const string& dump_root);
+  // Look up existing events writer by dump_root.
+  // If no DebugEventsWriter has been created at the dump_root, a non-OK
+  // Status will be returned. Else an OK status will be returned, with
+  // the pointer to the existing instance provided by reference.
+  static Status LookUpDebugEventsWriter(
+      const string& dump_root, DebugEventsWriter** debug_events_writer);
   ~DebugEventsWriter();
 
   // Sets the debug event filenames and opens file for writing.
@@ -116,8 +125,8 @@ class DebugEventsWriter {
   // deleted by another process), this will open a new file.
   Status Init();
 
-  // The four DebugEvent fields below are written _without_ the circular buffer.
-  // Source file contents are written to the *.source_files file.
+  // The four DebugEvent fields below are written _without_ the circular
+  // buffer. Source file contents are written to the *.source_files file.
   // Takes ownership of source_file.
   Status WriteSourceFile(SourceFile* source_file);
   // Stack frames are written to the *.code_locations file.
@@ -132,9 +141,8 @@ class DebugEventsWriter {
 
   // The two DebugEvent fields below are written to the circular buffer
   // and saved to disk only at the FlushExecutionFiles() call.
-  // Execution events (eager execution of an op or a tf.function) are written to
-  // the *.execution file.
-  // Takes ownership of execution.
+  // Execution events (eager execution of an op or a tf.function) are written
+  // to the *.execution file. Takes ownership of execution.
   Status WriteExecution(Execution* execution);
   // Graph execution traces (graph-internal tensor values or their summaries)
   // are written to the *.graph_execution_traces file.
@@ -151,8 +159,9 @@ class DebugEventsWriter {
   //     which the trace concerns multiple tensors, this is an empty string.
   //   output_slot: Output slot index of the op that this trace is concerned
   //     with.
-  //   tensor_debug_mode: An integer that represents the tensor-debug mode enum.
-  //   tensor_value: The value of the tensor that describes the tensor(s)
+  //   tensor_debug_mode: An integer that represents the tensor-debug mode
+  //   enum. tensor_value: The value of the tensor that describes the
+  //   tensor(s)
   //     that this trace is concerned with. The semantics of this tensor value
   //     depends on the value of `tensor_debug_mode`.
   Status WriteGraphExecutionTrace(const string& tfdbg_context_id,
@@ -208,7 +217,8 @@ class DebugEventsWriter {
   // Guards calls to the GetDebugEventsWriter() method.
   static mutex factory_mu_;
 
-  DebugEventsWriter(const string& dump_root, int64 circular_buffer_size);
+  DebugEventsWriter(const string& dump_root, const string& tfdbg_run_id,
+                    int64 circular_buffer_size);
 
   // Get the path prefix. The same for all files, which differ only in the
   // suffix.
@@ -227,6 +237,7 @@ class DebugEventsWriter {
 
   Env* env_;
   const string dump_root_;
+  const string tfdbg_run_id_;
 
   string file_prefix_;
   bool is_initialized_ TF_GUARDED_BY(initialization_mu_);
diff --git a/tensorflow/core/util/debug_events_writer_test.cc b/tensorflow/core/util/debug_events_writer_test.cc
index bd0c731bc90..45895763673 100644
--- a/tensorflow/core/util/debug_events_writer_test.cc
+++ b/tensorflow/core/util/debug_events_writer_test.cc
@@ -71,6 +71,7 @@ class DebugEventsWriterTest : public ::testing::Test {
     dump_root_ = io::JoinPath(
         testing::TmpDir(),
         strings::Printf("%010lld", static_cast(env()->NowMicros())));
+    tfdbg_run_id_ = "test_tfdbg_run_id";
   }
 
   void TearDown() override {
@@ -85,14 +86,15 @@ class DebugEventsWriterTest : public ::testing::Test {
   }
 
   string dump_root_;
+  string tfdbg_run_id_;
 };
 
 TEST_F(DebugEventsWriterTest, GetDebugEventsWriterSameRootGivesSameObject) {
   // Test the per-dump_root_ singleton pattern.
-  DebugEventsWriter* writer_1 =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
-  DebugEventsWriter* writer_2 =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer_1 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
+  DebugEventsWriter* writer_2 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   EXPECT_EQ(writer_1, writer_2);
 }
 
@@ -103,8 +105,8 @@ TEST_F(DebugEventsWriterTest, ConcurrentGetDebugEventsWriterSameDumpRoot) {
   std::vector writers;
   mutex mu;
   auto fn = [this, &writers, &mu]() {
-    DebugEventsWriter* writer =
-        DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+    DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+        dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
     {
       mutex_lock l(mu);
       writers.push_back(writer);
@@ -131,8 +133,9 @@ TEST_F(DebugEventsWriterTest, ConcurrentGetDebugEventsWriterDiffDumpRoots) {
   auto fn = [this, &counter, &writers, &mu]() {
     const string new_dump_root =
         io::JoinPath(dump_root_, strings::Printf("%ld", counter.fetch_add(1)));
-    DebugEventsWriter* writer =
-        DebugEventsWriter::GetDebugEventsWriter(new_dump_root);
+    DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+        new_dump_root, tfdbg_run_id_,
+        DebugEventsWriter::kDefaultCyclicBufferSize);
     {
       mutex_lock l(mu);
       writers.push_back(writer);
@@ -151,17 +154,17 @@ TEST_F(DebugEventsWriterTest, ConcurrentGetDebugEventsWriterDiffDumpRoots) {
 
 TEST_F(DebugEventsWriterTest, GetDebugEventsWriterDifferentRoots) {
   // Test the DebugEventsWriters for different directories are different.
-  DebugEventsWriter* writer_1 =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer_1 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   const string dump_root_2 = io::JoinPath(dump_root_, "subdirectory");
-  DebugEventsWriter* writer_2 =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_2);
+  DebugEventsWriter* writer_2 = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_2, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   EXPECT_NE(writer_1, writer_2);
 }
 
 TEST_F(DebugEventsWriterTest, GetAndInitDebugEventsWriter) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
   TF_ASSERT_OK(writer->Close());
 
@@ -174,6 +177,8 @@ TEST_F(DebugEventsWriterTest, GetAndInitDebugEventsWriter) {
   const string file_version = actuals[0].debug_metadata().file_version();
   EXPECT_EQ(file_version.find(DebugEventsWriter::kVersionPrefix), 0);
   EXPECT_GT(file_version.size(), strlen(DebugEventsWriter::kVersionPrefix));
+  // Check the tfdbg run ID.
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
 
   // Verify that the .source_files file has been created and is empty.
   ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
@@ -182,22 +187,22 @@ TEST_F(DebugEventsWriterTest, GetAndInitDebugEventsWriter) {
 }
 
 TEST_F(DebugEventsWriterTest, CallingCloseWithoutInitIsOkay) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Close());
 }
 
 TEST_F(DebugEventsWriterTest, CallingCloseTwiceIsOkay) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Close());
   TF_ASSERT_OK(writer->Close());
 }
 
 TEST_F(DebugEventsWriterTest, ConcurrentInitCalls) {
   // Test that concurrent calls to Init() works correctly.
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
 
   thread::ThreadPool* thread_pool =
       new thread::ThreadPool(Env::Default(), "test_pool", 4);
@@ -218,6 +223,7 @@ TEST_F(DebugEventsWriterTest, ConcurrentInitCalls) {
   const string file_version = actuals[0].debug_metadata().file_version();
   EXPECT_EQ(file_version.find(DebugEventsWriter::kVersionPrefix), 0);
   EXPECT_GT(file_version.size(), strlen(DebugEventsWriter::kVersionPrefix));
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
 
   // Verify that the .source_files file has been created and is empty.
   ReadDebugEventProtos(writer, DebugEventFileType::SOURCE_FILES, &actuals);
@@ -227,14 +233,15 @@ TEST_F(DebugEventsWriterTest, ConcurrentInitCalls) {
 
 TEST_F(DebugEventsWriterTest, InitTwiceDoesNotCreateNewMetadataFile) {
   // Test that Init() is idempotent.
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   std::vector actuals;
   ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
   EXPECT_EQ(actuals.size(), 1);
   EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
   EXPECT_GE(actuals[0].debug_metadata().file_version().size(), 0);
 
   string metadata_path_1 =
@@ -248,12 +255,13 @@ TEST_F(DebugEventsWriterTest, InitTwiceDoesNotCreateNewMetadataFile) {
   ReadDebugEventProtos(writer, DebugEventFileType::METADATA, &actuals);
   EXPECT_EQ(actuals.size(), 1);
   EXPECT_GT(actuals[0].debug_metadata().tensorflow_version().length(), 0);
+  EXPECT_EQ(actuals[0].debug_metadata().tfdbg_run_id(), "test_tfdbg_run_id");
   EXPECT_GE(actuals[0].debug_metadata().file_version().size(), 0);
 }
 
 TEST_F(DebugEventsWriterTest, WriteSourceFile) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   SourceFile* source_file_1 = new SourceFile();
@@ -313,8 +321,8 @@ TEST_F(DebugEventsWriterTest, WriteSourceFile) {
 }
 
 TEST_F(DebugEventsWriterTest, WriteStackFramesFile) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   StackFrameWithId* stack_frame_1 = new StackFrameWithId();
@@ -375,8 +383,8 @@ TEST_F(DebugEventsWriterTest, WriteStackFramesFile) {
 }
 
 TEST_F(DebugEventsWriterTest, WriteGraphOpCreationAndDebuggedGraph) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   GraphOpCreation* graph_op_creation = new GraphOpCreation();
@@ -415,8 +423,8 @@ TEST_F(DebugEventsWriterTest, WriteGraphOpCreationAndDebuggedGraph) {
 
 TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheSameFile) {
   const size_t kConcurrentWrites = 100;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   thread::ThreadPool* thread_pool =
@@ -456,8 +464,8 @@ TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheSameFile) {
 
 TEST_F(DebugEventsWriterTest, ConcurrentWriteAndFlushCallsToTheSameFile) {
   const size_t kConcurrentWrites = 100;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   thread::ThreadPool* thread_pool =
@@ -498,8 +506,8 @@ TEST_F(DebugEventsWriterTest, ConcurrentWriteAndFlushCallsToTheSameFile) {
 
 TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheDifferentFiles) {
   const int32 kConcurrentWrites = 30;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   thread::ThreadPool* thread_pool =
@@ -576,8 +584,8 @@ TEST_F(DebugEventsWriterTest, ConcurrentWriteCallsToTheDifferentFiles) {
 TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferNoFlush) {
   // Verify that no writing to disk happens until the flushing method is called.
   const size_t kCyclicBufferSize = 10;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   // First, try writing and flushing more debug events than the capacity
@@ -601,8 +609,8 @@ TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferNoFlush) {
 TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferFlush) {
   // Verify that writing to disk happens when the flushing method is called.
   const size_t kCyclicBufferSize = 10;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   // First, try writing and flushing more debug events than the capacity
@@ -673,8 +681,8 @@ TEST_F(DebugEventsWriterTest, WriteExecutionWithCyclicBufferFlush) {
 TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferNoFlush) {
   // Check no writing to disk happens before the flushing method is called.
   const size_t kCyclicBufferSize = 10;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   // First, try writing and flushing more debug events than the capacity
@@ -697,8 +705,8 @@ TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferNoFlush) {
 
 TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithoutPreviousInitCall) {
   const size_t kCyclicBufferSize = -1;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   // NOTE(cais): `writer->Init()` is not called here before
   // WriteGraphExecutionTrace() is called. This test checks that this is okay
   // and the `GraphExecutionTrace` gets written correctly even without `Init()`
@@ -722,8 +730,8 @@ TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithoutPreviousInitCall) {
 
 TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferFlush) {
   const size_t kCyclicBufferSize = 10;
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   // First, try writing and flushing more debug events than the capacity
@@ -788,8 +796,8 @@ TEST_F(DebugEventsWriterTest, WriteGrahExecutionTraceWithCyclicBufferFlush) {
 }
 
 TEST_F(DebugEventsWriterTest, RegisterDeviceAndGetIdTrace) {
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, DebugEventsWriter::kDefaultCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   // Register and get some device IDs in a concurrent fashion.
@@ -833,8 +841,8 @@ TEST_F(DebugEventsWriterTest, RegisterDeviceAndGetIdTrace) {
 
 TEST_F(DebugEventsWriterTest, DisableCyclicBufferBehavior) {
   const size_t kCyclicBufferSize = 0;  // A value <= 0 disables cyclic behavior.
-  DebugEventsWriter* writer =
-      DebugEventsWriter::GetDebugEventsWriter(dump_root_, kCyclicBufferSize);
+  DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
+      dump_root_, tfdbg_run_id_, kCyclicBufferSize);
   TF_ASSERT_OK(writer->Init());
 
   const size_t kNumEvents = 20;
diff --git a/tensorflow/python/client/debug_events_writer_wrapper.cc b/tensorflow/python/client/debug_events_writer_wrapper.cc
index a786c6f2db6..15802df40fe 100644
--- a/tensorflow/python/client/debug_events_writer_wrapper.cc
+++ b/tensorflow/python/client/debug_events_writer_wrapper.cc
@@ -29,9 +29,10 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   using namespace tensorflow::tfdbg;  // NOLINT(build/namespaces)
 
   m.def("Init",
-        [](const std::string& dump_root, const int64 circular_buffer_size) {
+        [](const std::string& dump_root, const std::string& tfdbg_run_id,
+           const int64 circular_buffer_size) {
           DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
-              dump_root, circular_buffer_size);
+              dump_root, tfdbg_run_id, circular_buffer_size);
           if (!writer->Init().ok()) {
             throw py::value_error(tensorflow::strings::Printf(
                 "Failed to initialize debug events writer at: %s",
@@ -41,8 +42,9 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteSourceFile",
         [](const std::string& dump_root, const py::object obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedNonExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::SOURCE_FILES);
@@ -50,8 +52,9 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteStackFrameWithId",
         [](const std::string& dump_root, const py::object& obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedNonExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::STACK_FRAMES);
@@ -59,8 +62,9 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteGraphOpCreation",
         [](const std::string& dump_root, const py::object& obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedNonExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::GRAPHS);
@@ -68,8 +72,9 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteDebuggedGraph",
         [](const std::string& dump_root, const py::object& obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedNonExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::GRAPHS);
@@ -77,8 +82,9 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteExecution",
         [](const std::string& dump_root, const py::object& obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::EXECUTION);
@@ -86,31 +92,32 @@ PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
   m.def("WriteGraphExecutionTrace",
         [](const std::string& dump_root, const py::object& obj) {
           CheckProtoType(obj, "tensorflow.DebugEvent");
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
+          DebugEventsWriter* writer = nullptr;
+          TF_CHECK_OK(
+              DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
           writer->WriteSerializedExecutionDebugEvent(
               obj.attr("SerializeToString")().cast(),
               tfdbg::DebugEventFileType::GRAPH_EXECUTION_TRACES);
         });
-  m.def("RegisterDeviceAndGetId",
-        [](const std::string& dump_root, const std::string& device_name) {
-          DebugEventsWriter* writer =
-              DebugEventsWriter::GetDebugEventsWriter(dump_root);
-          return writer->RegisterDeviceAndGetId(device_name);
-        });
+  m.def("RegisterDeviceAndGetId", [](const std::string& dump_root,
+                                     const std::string& device_name) {
+    DebugEventsWriter* writer = nullptr;
+    TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
+    return writer->RegisterDeviceAndGetId(device_name);
+  });
   m.def("FlushNonExecutionFiles", [](const std::string& dump_root) {
-    DebugEventsWriter* writer =
-        DebugEventsWriter::GetDebugEventsWriter(dump_root);
+    DebugEventsWriter* writer = nullptr;
+    TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
     writer->FlushNonExecutionFiles();
   });
   m.def("FlushExecutionFiles", [](const std::string& dump_root) {
-    DebugEventsWriter* writer =
-        DebugEventsWriter::GetDebugEventsWriter(dump_root);
+    DebugEventsWriter* writer = nullptr;
+    TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
     writer->FlushExecutionFiles();
   });
   m.def("Close", [](const std::string& dump_root) {
-    DebugEventsWriter* writer =
-        DebugEventsWriter::GetDebugEventsWriter(dump_root);
+    DebugEventsWriter* writer = nullptr;
+    TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
     writer->Close();
   });
 };
diff --git a/tensorflow/python/debug/lib/debug_events_reader.py b/tensorflow/python/debug/lib/debug_events_reader.py
index 4adb97de25b..743cea7103a 100644
--- a/tensorflow/python/debug/lib/debug_events_reader.py
+++ b/tensorflow/python/debug/lib/debug_events_reader.py
@@ -863,6 +863,7 @@ class DebugDataReader(object):
     debug_event = next(metadata_iter).debug_event
     self._starting_wall_time = debug_event.wall_time
     self._tensorflow_version = debug_event.debug_metadata.tensorflow_version
+    self._tfdbg_run_id = debug_event.debug_metadata.tfdbg_run_id
 
   def _load_source_files(self):
     """Incrementally read the .source_files DebugEvent file."""
@@ -1071,6 +1072,10 @@ class DebugDataReader(object):
     """
     return self._tensorflow_version
 
+  def tfdbg_run_id(self):
+    """Get the debugger run ID of the debugged TensorFlow program."""
+    return self._tfdbg_run_id
+
   def outermost_graphs(self):
     """Get the number of outer most graphs read so far."""
     return [graph for graph in self._graph_by_id.values()
diff --git a/tensorflow/python/debug/lib/debug_events_writer.py b/tensorflow/python/debug/lib/debug_events_writer.py
index 3de0ab78b8a..f223abdd099 100644
--- a/tensorflow/python/debug/lib/debug_events_writer.py
+++ b/tensorflow/python/debug/lib/debug_events_writer.py
@@ -32,6 +32,7 @@ class DebugEventsWriter(object):
 
   def __init__(self,
                dump_root,
+               tfdbg_run_id,
                circular_buffer_size=DEFAULT_CIRCULAR_BUFFER_SIZE):
     """Construct a DebugEventsWriter object.
 
@@ -43,6 +44,7 @@ class DebugEventsWriter(object):
     Args:
       dump_root: The root directory for dumping debug data. If `dump_root` does
         not exist as a directory, it will be created.
+      tfdbg_run_id: Debugger Run ID.
       circular_buffer_size: Size of the circular buffer for each of the two
         execution-related debug events files: with the following suffixes: -
           .execution - .graph_execution_traces If <= 0, the circular-buffer
@@ -51,7 +53,9 @@ class DebugEventsWriter(object):
     if not dump_root:
       raise ValueError("Empty or None dump root")
     self._dump_root = dump_root
-    _pywrap_debug_events_writer.Init(self._dump_root, circular_buffer_size)
+    self._tfdbg_run_id = tfdbg_run_id
+    _pywrap_debug_events_writer.Init(self._dump_root, self._tfdbg_run_id,
+                                     circular_buffer_size)
 
   def WriteSourceFile(self, source_file):
     """Write a SourceFile proto with the writer.
diff --git a/tensorflow/python/debug/lib/debug_events_writer_test.py b/tensorflow/python/debug/lib/debug_events_writer_test.py
index 57721c1450f..7b06bf772be 100644
--- a/tensorflow/python/debug/lib/debug_events_writer_test.py
+++ b/tensorflow/python/debug/lib/debug_events_writer_test.py
@@ -41,7 +41,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
 
   def testMultiThreadedConstructorCallWorks(self):
     def init_writer():
-      debug_events_writer.DebugEventsWriter(self.dump_root)
+      debug_events_writer.DebugEventsWriter(self.dump_root, self.tfdbg_run_id)
 
     num_threads = 4
     threads = []
@@ -66,7 +66,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
     self._readAndCheckMetadataFile()
 
   def testWriteSourceFilesAndStackFrames(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
     num_protos = 10
     for i in range(num_protos):
       source_file = debug_event_pb2.SourceFile()
@@ -99,7 +100,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
         self.assertEqual(actuals[i].file_line_col.file_index, i * 10)
 
   def testWriteGraphOpCreationAndDebuggedGraphs(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
     num_op_creations = 10
     for i in range(num_op_creations):
       graph_op_creation = debug_event_pb2.GraphOpCreation()
@@ -122,7 +124,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
                      "deadbeaf")
 
   def testConcurrentWritesToNonExecutionFilesWorks(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
 
     source_file_state = {"counter": 0, "lock": threading.Lock()}
 
@@ -201,15 +204,18 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
 
   def testWriteAndReadMetadata(self):
     t0 = time.time()
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
     writer.Close()
     with debug_events_reader.DebugDataReader(self.dump_root) as reader:
       self.assertIsInstance(reader.starting_wall_time(), float)
       self.assertGreaterEqual(reader.starting_wall_time(), t0)
       self.assertEqual(reader.tensorflow_version(), versions.__version__)
+      self.assertTrue(reader.tfdbg_run_id())
 
   def testWriteExecutionEventsWithCircularBuffer(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
     num_execution_events = debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE * 2
     for i in range(num_execution_events):
       execution = debug_event_pb2.Execution()
@@ -232,7 +238,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
 
   def testWriteExecutionEventsWithoutCircularBufferBehavior(self):
     # A circular buffer size of 0 abolishes the circular buffer behavior.
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root, 0)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id, 0)
     num_execution_events = debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE * 2
     for i in range(num_execution_events):
       execution = debug_event_pb2.Execution()
@@ -248,7 +255,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
         self.assertEqual(execution.op_type, "OpType%d" % i)
 
   def testWriteGraphExecutionTraceEventsWithCircularBuffer(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
     num_execution_events = debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE * 2
     for i in range(num_execution_events):
       trace = debug_event_pb2.GraphExecutionTrace()
@@ -272,7 +280,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
 
   def testWriteGraphExecutionTraceEventsWithoutCircularBufferBehavior(self):
     # A circular buffer size of 0 abolishes the circular buffer behavior.
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root, 0)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id, 0)
     num_execution_events = debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE * 2
     for i in range(num_execution_events):
       trace = debug_event_pb2.GraphExecutionTrace()
@@ -290,6 +299,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testConcurrentWritesToExecutionFiles(self):
     circular_buffer_size = 5
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id,
                                                    circular_buffer_size)
     debugged_graph = debug_event_pb2.DebuggedGraph(graph_id="graph1",
                                                    graph_name="graph1")
@@ -345,7 +355,8 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
       self.assertLen(op_names, len(set(op_names)))
 
   def testConcurrentSourceFileRandomReads(self):
-    writer = debug_events_writer.DebugEventsWriter(self.dump_root)
+    writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id)
 
     for i in range(100):
       source_file = debug_event_pb2.SourceFile(
@@ -376,6 +387,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testConcurrentExecutionUpdateAndRandomRead(self):
     circular_buffer_size = -1
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id,
                                                    circular_buffer_size)
 
     writer_state = {"counter": 0, "done": False}
@@ -410,6 +422,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testConcurrentExecutionRandomReads(self):
     circular_buffer_size = -1
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id,
                                                    circular_buffer_size)
 
     for i in range(100):
@@ -445,6 +458,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testConcurrentGraphExecutionTraceUpdateAndRandomRead(self):
     circular_buffer_size = -1
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id,
                                                    circular_buffer_size)
     debugged_graph = debug_event_pb2.DebuggedGraph(graph_id="graph1",
                                                    graph_name="graph1")
@@ -487,6 +501,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testConcurrentGraphExecutionTraceRandomReads(self):
     circular_buffer_size = -1
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   self.tfdbg_run_id,
                                                    circular_buffer_size)
     debugged_graph = debug_event_pb2.DebuggedGraph(graph_id="graph1",
                                                    graph_name="graph1")
@@ -534,7 +549,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testRangeReadingExecutions(self, begin, end, expected_begin,
                                  expected_end):
     writer = debug_events_writer.DebugEventsWriter(
-        self.dump_root, circular_buffer_size=-1)
+        self.dump_root, self.tfdbg_run_id, circular_buffer_size=-1)
     for i in range(5):
       execution = debug_event_pb2.Execution(op_type="OpType%d" % i)
       writer.WriteExecution(execution)
@@ -559,7 +574,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
   def testRangeReadingGraphExecutionTraces(self, begin, end, expected_begin,
                                            expected_end):
     writer = debug_events_writer.DebugEventsWriter(
-        self.dump_root, circular_buffer_size=-1)
+        self.dump_root, self.tfdbg_run_id, circular_buffer_size=-1)
     debugged_graph = debug_event_pb2.DebuggedGraph(
         graph_id="graph1", graph_name="graph1")
     writer.WriteDebuggedGraph(debugged_graph)
diff --git a/tensorflow/python/debug/lib/debug_v2_ops_test.py b/tensorflow/python/debug/lib/debug_v2_ops_test.py
index 10de01f4f2e..d715869f359 100644
--- a/tensorflow/python/debug/lib/debug_v2_ops_test.py
+++ b/tensorflow/python/debug/lib/debug_v2_ops_test.py
@@ -52,8 +52,9 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
     super(DebugIdentityV2OpTest, self).setUp()
     # Testing using a small circular-buffer size.
     self.circular_buffer_size = 4
+    self.tfdbg_run_id = "test_tfdbg_run"
     self.writer = debug_events_writer.DebugEventsWriter(
-        self.dump_root, self.circular_buffer_size)
+        self.dump_root, self.tfdbg_run_id, self.circular_buffer_size)
 
   def tearDown(self):
     self.writer.Close()
@@ -192,7 +193,8 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
   def testTwoDumpRoots(self):
     another_dump_root = os.path.join(self.dump_root, "another")
     another_debug_url = "file://%s" % another_dump_root
-    another_writer = debug_events_writer.DebugEventsWriter(another_dump_root)
+    another_writer = debug_events_writer.DebugEventsWriter(
+        another_dump_root, "test_tfdbg_run")
 
     @def_function.function
     def write_debug_trace(x):
@@ -264,6 +266,7 @@ class DebugIdentityV2OpUninitializedWriterTest(
       self.assertAllClose(
           write_debug_trace(np.array([i]).astype(np.float32)), [i**2.0])
     writer = debug_events_writer.DebugEventsWriter(self.dump_root,
+                                                   "test_tfdbg_run",
                                                    circular_buffer_size)
     writer.FlushNonExecutionFiles()
     writer.FlushExecutionFiles()
diff --git a/tensorflow/python/debug/lib/dumping_callback.py b/tensorflow/python/debug/lib/dumping_callback.py
index 0f5836e0644..563b52f8f63 100644
--- a/tensorflow/python/debug/lib/dumping_callback.py
+++ b/tensorflow/python/debug/lib/dumping_callback.py
@@ -69,6 +69,10 @@ def _debug_identity_v2_grad(op, dy):
   return dy
 
 
+def _get_tfdbg_run_id():
+  return str(uuid.uuid4())[:8]
+
+
 def _get_id():
   """Get a short unique ID."""
   return str(uuid.uuid4())
@@ -88,6 +92,7 @@ class _DumpingCallback(object):
                op_regex,
                tensor_dtypes):
     self._dump_root = dump_root
+    self._tfdbg_run_id = _get_tfdbg_run_id()
     self._tensor_debug_mode = tensor_debug_mode
     self._circular_buffer_size = circular_buffer_size
     self._op_regex = op_regex
@@ -148,6 +153,10 @@ class _DumpingCallback(object):
       self._dump_root = dump_root
       self._writer = None
 
+  @property
+  def tfdbg_run_id(self):
+    return self._tfdbg_run_id
+
   @property
   def tensor_debug_mode(self):
     return self._tensor_debug_mode
@@ -161,6 +170,7 @@ class _DumpingCallback(object):
     if not self._writer:
       self._writer = debug_events_writer.DebugEventsWriter(
           self._dump_root,
+          self._tfdbg_run_id,
           circular_buffer_size=self._circular_buffer_size)
     return self._writer
 
@@ -365,6 +375,8 @@ class _DumpingCallback(object):
       if tf_compat.forward_compatible(2020, 6, 24):
         debug_identity_op_kwargs[
             "circular_buffer_size"] = self._circular_buffer_size
+      if tf_compat.forward_compatible(2020, 7, 1):
+        debug_identity_op_kwargs["tfdbg_run_id"] = self._tfdbg_run_id
       if tensor_debug_mode == debug_event_pb2.TensorDebugMode.NO_TENSOR:
         if (not self._should_dump_tensor(op_type, tensor.dtype) or
             not tensor.dtype.is_numpy_compatible):
@@ -873,7 +885,8 @@ def disable_dump_debug_info():
   """
   if hasattr(_state, "dumping_callback"):
     dump_root = _state.dumping_callback.dump_root
-    debug_events_writer.DebugEventsWriter(dump_root).Close()
+    tfdbg_run_id = _state.dumping_callback.tfdbg_run_id
+    debug_events_writer.DebugEventsWriter(dump_root, tfdbg_run_id).Close()
     op_callbacks.remove_op_callback(_state.dumping_callback.callback)
     function_lib.remove_function_callback(
         _state.dumping_callback.function_callback)
diff --git a/tensorflow/python/debug/lib/dumping_callback_test_lib.py b/tensorflow/python/debug/lib/dumping_callback_test_lib.py
index 164644c57fa..05bf3aeb6da 100644
--- a/tensorflow/python/debug/lib/dumping_callback_test_lib.py
+++ b/tensorflow/python/debug/lib/dumping_callback_test_lib.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 import os
 import shutil
 import tempfile
+import uuid
 
 from tensorflow.python.debug.lib import check_numerics_callback
 from tensorflow.python.debug.lib import debug_events_reader
@@ -35,6 +36,7 @@ class DumpingCallbackTestBase(test_util.TensorFlowTestCase):
   def setUp(self):
     super(DumpingCallbackTestBase, self).setUp()
     self.dump_root = tempfile.mkdtemp()
+    self.tfdbg_run_id = str(uuid.uuid4())
 
   def tearDown(self):
     if os.path.isdir(self.dump_root):
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
index 75ddf32cbe0..54d15b601c5 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
@@ -982,7 +982,7 @@ tf_module {
   }
   member_method {
     name: "DebugIdentityV2"
-    argspec: "args=[\'input\', \'tfdbg_context_id\', \'op_name\', \'output_slot\', \'tensor_debug_mode\', \'debug_urls\', \'circular_buffer_size\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'\', \'-1\', \'-1\', \'[]\', \'1000\', \'None\'], "
+    argspec: "args=[\'input\', \'tfdbg_context_id\', \'op_name\', \'output_slot\', \'tensor_debug_mode\', \'debug_urls\', \'circular_buffer_size\', \'tfdbg_run_id\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'\', \'-1\', \'-1\', \'[]\', \'1000\', \'\', \'None\'], "
   }
   member_method {
     name: "DebugNanCount"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
index 75ddf32cbe0..54d15b601c5 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
@@ -982,7 +982,7 @@ tf_module {
   }
   member_method {
     name: "DebugIdentityV2"
-    argspec: "args=[\'input\', \'tfdbg_context_id\', \'op_name\', \'output_slot\', \'tensor_debug_mode\', \'debug_urls\', \'circular_buffer_size\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'\', \'-1\', \'-1\', \'[]\', \'1000\', \'None\'], "
+    argspec: "args=[\'input\', \'tfdbg_context_id\', \'op_name\', \'output_slot\', \'tensor_debug_mode\', \'debug_urls\', \'circular_buffer_size\', \'tfdbg_run_id\', \'name\'], varargs=None, keywords=None, defaults=[\'\', \'\', \'-1\', \'-1\', \'[]\', \'1000\', \'\', \'None\'], "
   }
   member_method {
     name: "DebugNanCount"

From 7d76bc4b60e020bbfd1339923ea5a7c3ab007217 Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Mon, 15 Jun 2020 13:56:55 -0700
Subject: [PATCH 0201/1390] Fix sparse kerastensors to maintain dense shape
 information after converting to a placeholder.

PiperOrigin-RevId: 316538468
Change-Id: I8e53a7e96067a8b7edd3f57cd8a8a89eb912824b
---
 .../python/keras/engine/keras_tensor.py       | 19 +++++++++++++++----
 .../utils/composite_tensor_support_test.py    |  3 +--
 tensorflow/python/ops/array_ops.py            |  1 +
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/keras/engine/keras_tensor.py b/tensorflow/python/keras/engine/keras_tensor.py
index 4ea01da8db2..c5c0068c652 100644
--- a/tensorflow/python/keras/engine/keras_tensor.py
+++ b/tensorflow/python/keras/engine/keras_tensor.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import type_spec as type_spec_module
 from tensorflow.python.ops import array_ops
 from tensorflow.python.util import nest
@@ -210,10 +211,20 @@ class _KerasTensorIterator(object):
 def keras_tensor_to_placeholder(x):
   """TODO(kaftan): Docstring."""
   if isinstance(x, KerasTensor):
-    def tensor_spec_to_placeholder(tensorspec):
-      return array_ops.placeholder(tensorspec.dtype, tensorspec.shape)
-    ph = nest.map_structure(tensor_spec_to_placeholder, x.type_spec,
-                            expand_composites=True)
+    spec = x.type_spec
+    if isinstance(spec, sparse_tensor.SparseTensorSpec):
+      # nest.map_structure loses dense shape information for sparse tensors.
+      # So, we special-case sparse placeholder creation.
+      # This only preserves shape information for top-level sparse tensors;
+      # not for sparse tensors that are nested inside another composite
+      # tensor.
+      return array_ops.sparse_placeholder(dtype=spec.dtype, shape=spec.shape)
+
+    def component_to_placeholder(component):
+      return array_ops.placeholder(component.dtype, component.shape)
+
+    ph = nest.map_structure(
+        component_to_placeholder, spec, expand_composites=True)
     return ph
   else:
     return x
diff --git a/tensorflow/python/keras/utils/composite_tensor_support_test.py b/tensorflow/python/keras/utils/composite_tensor_support_test.py
index f31558ddba8..daba188414a 100644
--- a/tensorflow/python/keras/utils/composite_tensor_support_test.py
+++ b/tensorflow/python/keras/utils/composite_tensor_support_test.py
@@ -603,8 +603,7 @@ class RaggedTensorInputValidationTest(keras_parameterized.TestCase,
 
 
 @keras_parameterized.run_with_all_model_types()
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True,
-                                         skip_keras_tensors=True)
+@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
 class CompositeTensorModelPredictTest(keras_parameterized.TestCase):
 
   def _normalize_shape(self, shape):
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index c77977bf7d2..1c00b81c9ca 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -3184,6 +3184,7 @@ def sparse_placeholder(dtype, shape=None, name=None):
       # `SparseTensor`
       dense_shape_default = tensor_shape.TensorShape(
           tuple(None if dim == -1 else dim for dim in shape))
+      shape = tuple(tensor_shape.dimension_value(dim) for dim in shape)
       shape = tuple(-1 if dim is None else dim for dim in shape)
       shape = ops.convert_to_tensor(
           shape, dtype=dtypes.int64, name=default_shape_name)

From 0e40b3e0c30caff9427c1da54c40b6236608ec15 Mon Sep 17 00:00:00 2001
From: Marat Dukhan 
Date: Mon, 15 Jun 2020 14:05:30 -0700
Subject: [PATCH 0202/1390] Exclude dependencies on FP16 XNNPACK micro-kernels

PiperOrigin-RevId: 316540175
Change-Id: Id02758822c004d52181eb1d317ba13e94df77f49
---
 tensorflow/lite/delegates/xnnpack/BUILD | 2 +-
 tensorflow/workspace.bzl                | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index c7ff1f55a49..efbaf0cfc42 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -26,7 +26,7 @@ cc_library(
         "//tensorflow/lite/schema:schema_fbs",
         "//tensorflow/lite/tools/optimize/sparsity:format_converter",
         "@FP16",
-        "@XNNPACK",
+        "@XNNPACK//:xnnpack_f32",
     ],
 )
 
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index f44c8dea6a2..6b0143e397f 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -164,11 +164,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
 
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "7469a0a634bfa90395ed311d07a21b1d0003604b37b12745bad1cf17860984e1",
-        strip_prefix = "XNNPACK-a059b7da184954fb6c01db0e7959352ee805e9f3",
+        sha256 = "bd5fd63a09222cd092f0c058b576cf044fb4074f2c4ce8a6fc32fc43d155f9c7",
+        strip_prefix = "XNNPACK-ae046f5a5127084bfe41090afdf1c1d4c9874b77",
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/a059b7da184954fb6c01db0e7959352ee805e9f3.zip",
-            "https://github.com/google/XNNPACK/archive/a059b7da184954fb6c01db0e7959352ee805e9f3.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/ae046f5a5127084bfe41090afdf1c1d4c9874b77.zip",
+            "https://github.com/google/XNNPACK/archive/ae046f5a5127084bfe41090afdf1c1d4c9874b77.zip",
         ],
     )
 

From fe6201c57b7fd78e344b4e5ee7fe6c7f7151b08a Mon Sep 17 00:00:00 2001
From: Jinliang Wei 
Date: Mon, 15 Jun 2020 14:08:55 -0700
Subject: [PATCH 0203/1390] Fix a bug in data flow analysis for asynchronous
 collective-permute.

PiperOrigin-RevId: 316540956
Change-Id: Icb1fb9d1d445d5aa3cf7afa580eed06607c4ecb3
---
 .../compiler/xla/service/hlo_dataflow_analysis.cc   | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
index d0d533e0b06..f19882c9347 100644
--- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
+++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
@@ -698,7 +698,7 @@ bool HloDataflowAnalysis::UpdateCollectivePermuteDoneValueSet(
   CHECK_EQ(collective_permute_done->opcode(),
            HloOpcode::kCollectivePermuteDone);
   bool changed = false;
-  // CollectivePermuteDone forwards the operand value at {0} to its output.
+  // CollectivePermuteDone forwards the operand value at {1} to its output.
   const HloValueSet& operand_value_set =
       GetValueSet(collective_permute_done->operand(0), {1});
   HloValueSet& value_set = GetValueSet(collective_permute_done);
@@ -945,6 +945,17 @@ Status HloDataflowAnalysis::InitializeInstructionValueSets() {
           // CopyDone consumes a tuple produced by CopyStart and produces an
           // element. Its output aliases its input tuple element {0}.
           break;
+        case HloOpcode::kCollectivePermuteStart:
+          // CollectivePermuteStart produces a tuple of
+          // {aliased operand, destination buffer, U32 context, U32 context}.
+          define_value_at(/*index=*/{});
+          define_value_at(/*index=*/{1});
+          define_value_at(/*index=*/{2});
+          define_value_at(/*index=*/{3});
+          break;
+        case HloOpcode::kCollectivePermuteDone:
+          // CollectivePermuteDone's output aliases its input tuple element {1}.
+          break;
         case HloOpcode::kRecvDone:
           // RecvDone produces a two-element tuple. Element zero aliases its
           // input tuple element {0}; element one is a token.

From 7f3de617db0c4442ac0877cbdcf6261bbe734087 Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Mon, 15 Jun 2020 21:18:56 +0000
Subject: [PATCH 0204/1390] segrating changes 2

---
 tensorflow/core/framework/tensor_shape.cc     |  2 +-
 tensorflow/core/kernels/batch_kernels.cc      |  6 ++--
 .../core/kernels/data/prefetch_autotuner.cc   |  4 +--
 tensorflow/core/kernels/quantization_utils.h  |  2 +-
 tensorflow/core/lib/io/inputbuffer.cc         |  6 ++--
 tensorflow/core/lib/io/random_inputstream.cc  |  2 +-
 .../core/lib/io/snappy/snappy_inputbuffer.cc  |  2 +-
 .../core/lib/io/snappy/snappy_outputbuffer.cc |  6 ++--
 tensorflow/core/lib/io/zlib_outputbuffer.cc   |  6 ++--
 tensorflow/core/platform/env.cc               |  2 +-
 tensorflow/core/platform/file_system.cc       |  2 +-
 .../core/platform/file_system_helper.cc       |  2 +-
 tensorflow/core/platform/s3/s3_file_system.cc |  2 +-
 tensorflow/core/platform/status.cc            |  4 +--
 .../profiler/internal/parse_annotation.cc     |  2 +-
 .../core/profiler/utils/derived_timeline.cc   |  2 +-
 .../core/profiler/utils/derived_timeline.h    |  2 +-
 .../core/profiler/utils/xplane_utils.cc       |  2 +-
 tensorflow/core/util/bcast.h                  |  4 +--
 .../convert_trivial_tile_to_concat.cc         |  2 +-
 .../convert_trivial_transpose_to_reshape.cc   |  2 +-
 .../toco/graph_transformations/dequantize.cc  |  2 +-
 .../graph_transformations/drop_fake_quant.cc  |  2 +-
 ...int8_weights_safe_for_fast_int8_kernels.cc |  2 +-
 .../fuse_broadcast_into_following_binary.cc   |  2 +-
 .../group_bidirectional_sequence_ops.cc       |  4 +--
 .../graph_transformations/hardcode_min_max.cc |  2 +-
 .../identify_nearest_upsample.cc              |  2 +-
 .../merge_reshape_into_preceding_transpose.cc |  4 +--
 .../propagate_array_data_types.cc             |  2 +-
 .../propagate_fake_quant_num_bits.cc          |  2 +-
 .../propagate_fixed_sizes.cc                  | 28 +++++++++----------
 .../remove_successive_transpose.cc            | 10 +++----
 .../remove_trivial_passthrough.cc             |  2 +-
 .../reorder_elementwise_unary.cc              |  4 +--
 .../reorder_reshape_transpose.cc              | 12 ++++----
 .../resolve_batch_normalization.cc            | 10 +++----
 .../resolve_constant_concatenation.cc         |  2 +-
 .../resolve_constant_pack.cc                  |  2 +-
 .../resolve_constant_slice.cc                 |  2 +-
 .../resolve_constant_transpose.cc             |  2 +-
 .../resolve_constant_unary.cc                 |  4 +--
 .../unpartition_embedding_lookup.cc           |  4 +--
 tensorflow/lite/toco/model_cmdline_flags.cc   |  8 +++---
 tensorflow/lite/toco/toco_cmdline_flags.cc    |  2 +-
 45 files changed, 89 insertions(+), 91 deletions(-)

diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index 79d0cc0822d..f4b440e9cd1 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -187,7 +187,7 @@ void TensorShapeBase::InitDims(gtl::ArraySlice dim_sizes) {
                 "bad overflow check");
   bool large_size = false;
   for (auto s : dim_sizes) {
-    if (static_cast(s) > static_cast(kMaxSmall)) {
+    if (s > kMaxSmall) {
       large_size = true;
       break;
     }
diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc
index ee271f1a123..151f2367c95 100644
--- a/tensorflow/core/kernels/batch_kernels.cc
+++ b/tensorflow/core/kernels/batch_kernels.cc
@@ -486,18 +486,18 @@ class BatchResource : public ResourceBase {
     std::map> split_tensors;
 
     DCHECK_EQ(batch->task(0).context->num_outputs(), combined_outputs.size());
-    if (static_cast(combined_outputs.size()) != batch->task(0).context->num_outputs()) {
+    if (combined_outputs.size() != batch->task(0).context->num_outputs()) {
       return errors::Internal("Wrong number of batched output tensors");
     }
 
     // Generate 'split_tensors' and populate the context outputs.
-    for (size_t i = 0; i < combined_outputs.size(); ++i) {
+    for (int i = 0; i < combined_outputs.size(); ++i) {
       const Tensor& output_tensor = combined_outputs[i];
       if (output_tensor.shape().dims() == 0) {
         return errors::FailedPrecondition(
             "Batched output tensor has 0 dimensions");
       }
-      if (output_tensor.shape().dim_size(0) != static_cast(batch->size() + padding_size)) {
+      if (output_tensor.shape().dim_size(0) != batch->size() + padding_size) {
         return errors::FailedPrecondition(
             "Batched output tensor's 0th dimension does not equal the sum of "
             "the 0th dimension sizes of the input tensors");
diff --git a/tensorflow/core/kernels/data/prefetch_autotuner.cc b/tensorflow/core/kernels/data/prefetch_autotuner.cc
index a3fd9919d6b..a3bb1acc352 100644
--- a/tensorflow/core/kernels/data/prefetch_autotuner.cc
+++ b/tensorflow/core/kernels/data/prefetch_autotuner.cc
@@ -40,13 +40,13 @@ void PrefetchAutotuner::RecordConsumption(size_t current_buffer_size) {
     case Mode::kDisabled:
       return;
     case Mode::kUpswing:
-      if (static_cast(current_buffer_size) == buffer_limit_) {
+      if (current_buffer_size == buffer_limit_) {
         mode_ = Mode::kDownswing;
       }
       return;
     case Mode::kDownswing:
       if (current_buffer_size == 0) {
-        if (buffer_limit_ >= static_cast(kBufferLimitThreshold)) {
+        if (buffer_limit_ >= kBufferLimitThreshold) {
           buffer_limit_ += kBufferLimitThreshold;
         } else {
           buffer_limit_ *= 2;
diff --git a/tensorflow/core/kernels/quantization_utils.h b/tensorflow/core/kernels/quantization_utils.h
index 06c901967b0..315616f3fb3 100644
--- a/tensorflow/core/kernels/quantization_utils.h
+++ b/tensorflow/core/kernels/quantization_utils.h
@@ -268,7 +268,7 @@ inline void RequantizeManyInNewRangeReference(const qint32* input, int64 count,
   // that could be easily adapted for a SIMD implementation. It should also be
   // possible to perform all the calculations in 32-bit rather than 64, but
   // that's not been implemented yet.
-  for (size_t index = 0; static_cast(index) < count; ++index) {
+  for (size_t index = 0; index < count; ++index) {
     const int64 input_value = static_cast(input[index]);
     const int64 fp_value =
         ((input_value * range_scale_fp) >> 32) + input_offset_fp;
diff --git a/tensorflow/core/lib/io/inputbuffer.cc b/tensorflow/core/lib/io/inputbuffer.cc
index d005ee11d78..2b138b825e4 100644
--- a/tensorflow/core/lib/io/inputbuffer.cc
+++ b/tensorflow/core/lib/io/inputbuffer.cc
@@ -85,7 +85,7 @@ Status InputBuffer::ReadNBytes(int64 bytes_to_read, string* result) {
   result->resize(bytes_to_read);
   size_t bytes_read = 0;
   Status status = ReadNBytes(bytes_to_read, &(*result)[0], &bytes_read);
-  if (static_cast(bytes_read) < bytes_to_read) result->resize(bytes_read);
+  if (bytes_read < bytes_to_read) result->resize(bytes_read);
   return status;
 }
 
@@ -204,7 +204,7 @@ Status InputBuffer::Hint(int64 bytes_to_read) {
   }
 
   // The internal buffer is too small. Do nothing.
-  if (bytes_to_read > static_cast(size_)) {
+  if (bytes_to_read > size_) {
     return Status::OK();
   }
 
@@ -230,7 +230,7 @@ Status InputBuffer::Hint(int64 bytes_to_read) {
   limit_ += data.size();
   file_pos_ += data.size();
 
-  if (errors::IsOutOfRange(s) && data.size() == static_cast(bytes_to_read)) {
+  if (errors::IsOutOfRange(s) && data.size() == bytes_to_read) {
     return Status::OK();
   } else {
     return s;
diff --git a/tensorflow/core/lib/io/random_inputstream.cc b/tensorflow/core/lib/io/random_inputstream.cc
index bd0054ce753..10f734a5bae 100644
--- a/tensorflow/core/lib/io/random_inputstream.cc
+++ b/tensorflow/core/lib/io/random_inputstream.cc
@@ -92,7 +92,7 @@ Status RandomAccessInputStream::SkipNBytes(int64 bytes_to_skip) {
     } else {
       return s;
     }
-    if (data.size() < static_cast(bytes_to_read)) {
+    if (data.size() < bytes_to_read) {
       return errors::OutOfRange("reached end of file");
     }
     bytes_to_skip -= bytes_to_read;
diff --git a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
index 53939f2d8a3..a331d4173cf 100644
--- a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
@@ -134,7 +134,7 @@ Status SnappyInputBuffer::ReadCompressedBlockLength(uint32* length) {
     }
     size_t readable = std::min(bytes_to_read, avail_in_);
 
-    for (size_t i = 0; i < readable; i++) {
+    for (int i = 0; i < readable; i++) {
       // The "unsigned char" type cast is intentional to avoid implicit type
       // casting of the signed char to unsigned int during bitwise OR which
       // causes weird overflow errors.
diff --git a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
index fe3a53c6c25..563503a1319 100644
--- a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
@@ -76,7 +76,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
 
   // If there is sufficient free space in input_buffer_ to fit data we
   // add it there and return.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -87,7 +87,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered());
 
   // input_buffer_ should be empty at this point.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -144,7 +144,7 @@ void SnappyOutputBuffer::AddToInputBuffer(StringPiece data) {
   const int32 free_tail_bytes =
       input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (static_cast(bytes_to_write) > free_tail_bytes) {
+  if (bytes_to_write > free_tail_bytes) {
     memmove(input_buffer_.get(), next_in_, avail_in_);
     next_in_ = input_buffer_.get();
   }
diff --git a/tensorflow/core/lib/io/zlib_outputbuffer.cc b/tensorflow/core/lib/io/zlib_outputbuffer.cc
index d475d0eaa5c..5840ca60242 100644
--- a/tensorflow/core/lib/io/zlib_outputbuffer.cc
+++ b/tensorflow/core/lib/io/zlib_outputbuffer.cc
@@ -98,7 +98,7 @@ void ZlibOutputBuffer::AddToInputBuffer(StringPiece data) {
   int32 unread_bytes = z_stream_->avail_in;
   int32 free_tail_bytes = input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (static_cast(bytes_to_write) > free_tail_bytes) {
+  if (bytes_to_write > free_tail_bytes) {
     memmove(z_stream_input_.get(), z_stream_->next_in, z_stream_->avail_in);
     z_stream_->next_in = z_stream_input_.get();
   }
@@ -154,7 +154,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
 
   size_t bytes_to_write = data.size();
 
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -162,7 +162,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered(zlib_options_.flush_mode));
 
   // At this point input stream should be empty.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index 05d95ba0425..b29cad05459 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -214,7 +214,7 @@ bool Env::FilesExist(const std::vector& files,
     }
     if (fs_status) {
       result &= fs_result;
-      for (size_t i = 0; i < itr.second.size(); ++i) {
+      for (int i = 0; i < itr.second.size(); ++i) {
         per_file_status[itr.second[i]] = fs_status->at(i);
       }
     } else if (!fs_result) {
diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc
index c9657e2339f..9e96ceedbdc 100644
--- a/tensorflow/core/platform/file_system.cc
+++ b/tensorflow/core/platform/file_system.cc
@@ -308,7 +308,7 @@ StringPiece FileSystem::Basename(StringPiece path) const {
 StringPiece FileSystem::Extension(StringPiece path) const {
   StringPiece basename = this->Basename(path);
 
-  size_t pos = basename.rfind('.');
+  int pos = basename.rfind('.');
   if (pos == StringPiece::npos) {
     return StringPiece(path.data() + path.size(), 0);
   } else {
diff --git a/tensorflow/core/platform/file_system_helper.cc b/tensorflow/core/platform/file_system_helper.cc
index 909752389e1..64b175c4d17 100644
--- a/tensorflow/core/platform/file_system_helper.cc
+++ b/tensorflow/core/platform/file_system_helper.cc
@@ -103,7 +103,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
                 children_dir_status[i] = fs->IsDirectory(child_path);
               }
             });
-    for (size_t i = 0; i < children.size(); ++i) {
+    for (int i = 0; i < children.size(); ++i) {
       const string child_path = io::JoinPath(current_dir, children[i]);
       // If the IsDirectory call was cancelled we bail.
       if (children_dir_status[i].code() == tensorflow::error::CANCELLED) {
diff --git a/tensorflow/core/platform/s3/s3_file_system.cc b/tensorflow/core/platform/s3/s3_file_system.cc
index 45d648abcc0..1726c9fbc6c 100644
--- a/tensorflow/core/platform/s3/s3_file_system.cc
+++ b/tensorflow/core/platform/s3/s3_file_system.cc
@@ -906,7 +906,7 @@ Status S3FileSystem::MultiPartCopy(const Aws::String& source,
       // wait on the mutex until notify is called
       // then check the finished parts as there could be false notifications
       multi_part_copy_cv.wait(lock, [&finishedPartStates, num_parts] {
-        return static_cast(finishedPartStates.size()) == num_parts;
+        return finishedPartStates.size() == num_parts;
       });
     }
     // check if there was any error for any part
diff --git a/tensorflow/core/platform/status.cc b/tensorflow/core/platform/status.cc
index e303c18091c..756b8314148 100644
--- a/tensorflow/core/platform/status.cc
+++ b/tensorflow/core/platform/status.cc
@@ -74,9 +74,7 @@ class StatusLogSink : public TFLogSink {
 
     mutex_lock lock(mu_);
     messages_.emplace_back(entry.ToString());
-    if (messages_.size() > static_cast(num_messages_)){
-        messages_.pop_front();
-    }
+    if (messages_.size() > num_messages_) messages_.pop_front();
   }
 
  private:
diff --git a/tensorflow/core/profiler/internal/parse_annotation.cc b/tensorflow/core/profiler/internal/parse_annotation.cc
index a4cdc09739d..32c26befa3d 100644
--- a/tensorflow/core/profiler/internal/parse_annotation.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation.cc
@@ -50,7 +50,7 @@ std::vector SplitNameAndMetadata(
 std::vector SplitPairs(absl::string_view metadata) {
   std::vector key_value_pairs;
   std::stack quotes;
-  size_t start = 0, end = 0;
+  int start = 0, end = 0;
   for (; end < metadata.size(); ++end) {
     char ch = metadata[end];
     switch (ch) {
diff --git a/tensorflow/core/profiler/utils/derived_timeline.cc b/tensorflow/core/profiler/utils/derived_timeline.cc
index 3d03fc22c16..112c0977763 100644
--- a/tensorflow/core/profiler/utils/derived_timeline.cc
+++ b/tensorflow/core/profiler/utils/derived_timeline.cc
@@ -130,7 +130,7 @@ void DerivedXLineBuilder::ExpandOrAddLevelEvent(const XEvent& event,
 }
 
 void DerivedXLineBuilder::ResetLastEvents(int level) {
-  for (int i = level; i < static_cast(last_event_by_level_.size()); ++i) {
+  for (int i = level; i < last_event_by_level_.size(); ++i) {
     last_event_by_level_[i] = absl::nullopt;
   }
   if (level == 0) ResetDependentLines();
diff --git a/tensorflow/core/profiler/utils/derived_timeline.h b/tensorflow/core/profiler/utils/derived_timeline.h
index 92489399b8f..cd4da7996c5 100644
--- a/tensorflow/core/profiler/utils/derived_timeline.h
+++ b/tensorflow/core/profiler/utils/derived_timeline.h
@@ -37,7 +37,7 @@ class DerivedXLineBuilder {
                       std::vector dependent_lines);
 
   void ExpandOrAddEvents(const std::vector& event_per_level) {
-    for (size_t level = 0; level < event_per_level.size(); ++level) {
+    for (int level = 0; level < event_per_level.size(); ++level) {
       ExpandOrAddLevelEvent(event_per_level[level], level);
     }
   }
diff --git a/tensorflow/core/profiler/utils/xplane_utils.cc b/tensorflow/core/profiler/utils/xplane_utils.cc
index 1fe476ce79c..7f5221c5391 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.cc
+++ b/tensorflow/core/profiler/utils/xplane_utils.cc
@@ -266,7 +266,7 @@ void SortXSpace(XSpace* space) {
 // smaller than these value.
 void NormalizeTimestamps(XPlane* plane, uint64 start_time_ns) {
   for (XLine& line : *plane->mutable_lines()) {
-    if (line.timestamp_ns() >= static_cast(start_time_ns)) {
+    if (line.timestamp_ns() >= start_time_ns) {
       line.set_timestamp_ns(line.timestamp_ns() - start_time_ns);
     }
   }
diff --git a/tensorflow/core/util/bcast.h b/tensorflow/core/util/bcast.h
index 075de84964e..7bb8ea18ad3 100644
--- a/tensorflow/core/util/bcast.h
+++ b/tensorflow/core/util/bcast.h
@@ -139,7 +139,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N],
     if (x[i] != x[0]) {
       all_equal = false;
     }
-    if (static_cast(x[i].size()) > largest_rank) {
+    if (x[i].size() > largest_rank) {
       largest_rank = x[i].size();
     }
   }
@@ -176,7 +176,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N],
 
   // 1-extend and align all vectors.
   for (int i = 0; i < N; ++i) {
-    if (static_cast(copy[i].size()) < largest_rank) {
+    if (copy[i].size() < largest_rank) {
       copy[i].resize(largest_rank, 1);
     }
   }
diff --git a/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc b/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc
index c19ccf676c9..46288d2a1ed 100644
--- a/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc
+++ b/tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc
@@ -52,7 +52,7 @@ namespace toco {
   // It then just becomes a concat along that dimension.
   int non_one_dims = 0;
   int concat_axis = 0;
-  for (size_t i = 0; i < multiples.size(); ++i) {
+  for (int i = 0; i < multiples.size(); ++i) {
     if (multiples[i] != 1) {
       ++non_one_dims;
       concat_axis = i;
diff --git a/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc b/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc
index fa8a69a1e7a..2b5aaea2b23 100644
--- a/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc
+++ b/tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc
@@ -31,7 +31,7 @@ bool TransposeAffectsMemoryOrder(std::vector perm,
   // just the shape) then the flat buffer representation shouldn't change.
   std::vector old_major_index_ordering;
   std::vector new_major_index_ordering;
-  for (int i = 0; static_cast(i) < in_shape.size(); i++) {
+  for (int i = 0; i < in_shape.size(); i++) {
     if (in_shape[i] != 1) {
       old_major_index_ordering.push_back(i);
     }
diff --git a/tensorflow/lite/toco/graph_transformations/dequantize.cc b/tensorflow/lite/toco/graph_transformations/dequantize.cc
index c87c305a70d..cc5dddbb40e 100644
--- a/tensorflow/lite/toco/graph_transformations/dequantize.cc
+++ b/tensorflow/lite/toco/graph_transformations/dequantize.cc
@@ -35,7 +35,7 @@ void DequantizeBuffer(Array* array) {
   auto& new_data = array->GetMutableBuffer().data;
   new_data.resize(old_data.size());
   const auto& qparams = array->GetQuantizationParams();
-  for (size_t i = 0; i < old_data.size(); i++) {
+  for (int i = 0; i < old_data.size(); i++) {
     new_data[i] = qparams.scale * (old_data[i] - qparams.zero_point);
   }
 }
diff --git a/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc b/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc
index 3a0b4d0103f..bb8679bced8 100644
--- a/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc
+++ b/tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc
@@ -45,7 +45,7 @@ namespace toco {
   }
 
   // Drop min/max inputs
-  for (size_t i = 1; i < fakequant_op->inputs.size(); i++) {
+  for (int i = 1; i < fakequant_op->inputs.size(); i++) {
     if (CountOpsWithInput(*model, fakequant_op->inputs[i]) == 1) {
       model->EraseArray(fakequant_op->inputs[i]);
     }
diff --git a/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc b/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc
index ce4574cdfbf..918bb489995 100644
--- a/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc
+++ b/tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc
@@ -166,7 +166,7 @@ namespace toco {
   int index_of_previous_bad_value = 0;
   bool changed = false;
 
-  for (size_t i = 0; i < buffer_data.size(); i++) {
+  for (int i = 0; i < buffer_data.size(); i++) {
     if (buffer_data[i] == 0) {
       count_bad++;
       if (count_bad > 1) {
diff --git a/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc b/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc
index 2c5c2cbb5f1..ba3e277f676 100644
--- a/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc
+++ b/tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc
@@ -34,7 +34,7 @@ bool IsBroadcastingOp(const Model& model, Operator* op) {
   // Concatenation of identical inputs is usually a broadcast.
   if (op->type == OperatorType::kConcatenation) {
     // Verify that all inputs are the same.
-    for (size_t i = 1; i < op->inputs.size(); ++i) {
+    for (int i = 1; i < op->inputs.size(); ++i) {
       if (op->inputs[i] != op->inputs[0]) {
         return false;
       }
diff --git a/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc b/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc
index a6d95ec43b1..fa252b1a61b 100644
--- a/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc
+++ b/tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc
@@ -125,7 +125,7 @@ bool CheckTwoUnidirectionalSequenceOpsAreValid(
       return false;
 
     // Make sure the inputs datatype matches.
-    for (size_t i = 0; i < fw_sequence_op->inputs.size(); ++i) {
+    for (int i = 0; i < fw_sequence_op->inputs.size(); ++i) {
       const auto& fw_input_array_name = fw_sequence_op->inputs[i];
       const auto& bw_input_array_name = bw_sequence_op->inputs[i];
       if (model.HasArray(fw_input_array_name) &&
@@ -137,7 +137,7 @@ bool CheckTwoUnidirectionalSequenceOpsAreValid(
     }
 
     // Make sure the outputs datatype matches.
-    for (size_t i = 0; i < fw_sequence_op->outputs.size(); ++i) {
+    for (int i = 0; i < fw_sequence_op->outputs.size(); ++i) {
       const auto& fw_output_array_name = fw_sequence_op->outputs[i];
       const auto& bw_output_array_name = bw_sequence_op->outputs[i];
       if (model.HasArray(fw_output_array_name) &&
diff --git a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
index 4250668bcf5..171d522daa7 100644
--- a/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
+++ b/tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc
@@ -405,7 +405,7 @@ bool HardcodeMinMaxForPack(Model* model, Operator* op) {
   }
   const auto& first_input_minmax = first_input_array.GetMinMax();
 
-  for (size_t i = 1; i < op->inputs.size(); i++) {
+  for (int i = 1; i < op->inputs.size(); i++) {
     const auto& input_array = model->GetArray(op->inputs[i]);
     if (!input_array.minmax) {
       return false;
diff --git a/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc b/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc
index 08894c93a5b..2ab6692a3a8 100644
--- a/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc
+++ b/tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc
@@ -199,7 +199,7 @@ std::vector>::iterator FindOperator(
   shape_array.data_type = ArrayDataType::kInt32;
   auto& shape_buffer = shape_array.GetMutableBuffer();
   // This is what imagined as the original shape.
-  for (size_t i = 0; i < imagined_original_shape.size(); ++i) {
+  for (int i = 0; i < imagined_original_shape.size(); ++i) {
     shape_buffer.data.push_back(imagined_original_shape.at(i));
   }
 
diff --git a/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc b/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc
index a76ae1a0635..80170fe8bcb 100644
--- a/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc
+++ b/tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc
@@ -70,7 +70,7 @@ std::vector ReshapeToTranspose(const Model& model,
   std::vector not_one_indices;
 
   // Separate into one indices and not one indices.
-  for (size_t i = 0; i < in_shape.size(); i++) {
+  for (int i = 0; i < in_shape.size(); i++) {
     if (in_shape[i] == 1) {
       one_indices.push_back(i);
     } else {
@@ -167,7 +167,7 @@ std::vector ReshapeToTranspose(const Model& model,
 
   // Combine the permutations.
   const auto& transpose_perm = transpose_op->perm;
-  for (size_t i = 0; i < merged_perm.size(); i++) {
+  for (int i = 0; i < merged_perm.size(); i++) {
     merged_perm[i] = transpose_perm[merged_perm[i]];
   }
 
diff --git a/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc
index 2f316934311..49d59de860b 100644
--- a/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc
+++ b/tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc
@@ -170,7 +170,7 @@ void SetDataTypeForAllOutputs(Model* model, Operator* op,
       if (unsupported_op->output_data_types.size() < op->outputs.size()) {
         return ::tensorflow::Status::OK();
       }
-      for (size_t i = 0; i < op->outputs.size(); ++i) {
+      for (int i = 0; i < op->outputs.size(); ++i) {
         const string& output = op->outputs[i];
         const ArrayDataType data_type = unsupported_op->output_data_types[i];
         model->GetArray(output).data_type = data_type;
diff --git a/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc b/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc
index 94779f54af2..1ed618879c1 100644
--- a/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc
+++ b/tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc
@@ -149,7 +149,7 @@ bool RecursivelyBackwardPropagateDataType(GraphTransformation* transformation,
                                           ArrayDataType new_data_type,
                                           const MinMax& new_minmax) {
   bool did_change = false;
-  for (size_t input_index = 0; input_index < op->inputs.size(); ++input_index) {
+  for (int input_index = 0; input_index < op->inputs.size(); ++input_index) {
     const auto& input = op->inputs[input_index];
     auto& input_array = model->GetArray(input);
 
diff --git a/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index 520cd8b495a..006e624eb7a 100644
--- a/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -431,7 +431,7 @@ void ProcessTensorFlowReshapeOperator(Model* model,
   bool has_wildcard = false;
   int wildcard_index = 0;
   int product_non_wildcard_dims = 1;
-  for (size_t i = 0; i < shape_data.size(); i++) {
+  for (int i = 0; i < shape_data.size(); i++) {
     if (shape_data[i] == -1) {
       CHECK(!has_wildcard);
       has_wildcard = true;
@@ -574,7 +574,7 @@ void ProcessTensorFlowReductionOperator(Model* model, Operator* op) {
     std::set true_indices;
     const auto& reduction_indices =
         reduction_indices_array.GetBuffer().data;
-    for (size_t i = 0; i < reduction_indices.size(); ++i) {
+    for (int i = 0; i < reduction_indices.size(); ++i) {
       const int32 reduction_index = reduction_indices[i];
       if (reduction_index < -input_rank || reduction_index >= input_rank) {
         CHECK(false) << "Invalid reduction dimension " << reduction_index
@@ -627,7 +627,7 @@ void ProcessSliceOperator(Model* model, SliceOperator* op) {
   CHECK_EQ(op->begin.size(), op->size.size());
 
   std::vector output_dims;
-  for (size_t i = 0; i < op->begin.size(); ++i) {
+  for (int i = 0; i < op->begin.size(); ++i) {
     int size = op->size[i];
     if (size == -1) {
       size = input_array.shape().dims(i) - op->begin[i];
@@ -883,7 +883,7 @@ void ProcessTensorFlowSplitVOperator(Model* model,
 
   CHECK_EQ(op->outputs.size(), op->num_split);
 
-  for (size_t i = 0; i < op->outputs.size(); ++i) {
+  for (int i = 0; i < op->outputs.size(); ++i) {
     const auto& output = op->outputs[i];
     Shape output_shape = input_shape;
     (*output_shape.mutable_dims())[axis] = size_splits_vector.at(i);
@@ -1514,7 +1514,7 @@ void ProcessPadOperator(Model* model, PadOperator* op) {
   std::vector& dims = *output_shape.mutable_dims();
   CHECK_EQ(op->left_padding.size(), dims.size());
 
-  for (size_t i = 0; i < op->left_padding.size(); ++i) {
+  for (int i = 0; i < op->left_padding.size(); ++i) {
     dims[i] += op->left_padding[i] + op->right_padding[i];
   }
 
@@ -1540,7 +1540,7 @@ void ProcessPadV2Operator(Model* model, PadV2Operator* op) {
   std::vector& dims = *output_shape.mutable_dims();
   CHECK_EQ(op->left_padding.size(), dims.size());
 
-  for (size_t i = 0; i < op->left_padding.size(); ++i) {
+  for (int i = 0; i < op->left_padding.size(); ++i) {
     dims[i] += op->left_padding[i] + op->right_padding[i];
   }
 
@@ -1683,7 +1683,7 @@ void ProcessStridedSliceOperator(Model* model, StridedSliceOperator* op) {
   CHECK_LE(op->strides.size(), num_input_axes)
       << "StridedSlice op with output \"" << op->outputs[0]
       << "\", requires no more than " << num_input_axes << " strides";
-  for (size_t i = 0; i < op->strides.size(); i++) {
+  for (int i = 0; i < op->strides.size(); i++) {
     CHECK_NE(op->strides[i], 0) << "Strides must be non-zero. Axis " << i
                                 << " has stride=" << op->strides[i] << ".";
   }
@@ -1814,7 +1814,7 @@ void ProcessTransposeOperator(Model* model, TransposeOperator* op) {
       << "Transpose permutation input " << op->inputs[1]
       << " must be same length as input dimensions";
   std::vector* output_dims = output_array.mutable_shape()->mutable_dims();
-  for (size_t i = 0; i < perm.size(); i++) {
+  for (int i = 0; i < perm.size(); i++) {
     int axis = perm[i];
     CHECK_GE(axis, 0);
     CHECK_LT(axis, input_shape.dimensions_count());
@@ -1856,8 +1856,8 @@ void ProcessArgMinMaxOperator(Model* model, Op* op) {
   std::vector output_dims;
 
   output_dims.reserve(input_dims.size() - 1);
-  for (size_t i = 0; i < input_dims.size(); ++i) {
-    if ( static_cast(i) != axis) {
+  for (int i = 0; i < input_dims.size(); ++i) {
+    if (i != axis) {
       output_dims.push_back(input_dims[i]);
     }
   }
@@ -1938,7 +1938,7 @@ void ProcessTileOperator(Model* model, TensorFlowTileOperator* op) {
 
   auto* mutable_dims = output_array.mutable_shape()->mutable_dims();
   mutable_dims->resize(multiples.size());
-  for (size_t i = 0; i < mutable_dims->size(); ++i) {
+  for (int i = 0; i < mutable_dims->size(); ++i) {
     (*mutable_dims)[i] = input_shape.dims(i) * multiples[i];
   }
 }
@@ -2010,8 +2010,8 @@ void ProcessUnpackOperator(Model* model, UnpackOperator* op) {
   std::vector output_dims;
 
   output_dims.reserve(input_dims.size() - 1);
-  for (size_t i = 0; i < input_dims.size(); ++i) {
-    if ( static_cast(i) != op->axis) {
+  for (int i = 0; i < input_dims.size(); ++i) {
+    if (i != op->axis) {
       output_dims.push_back(input_dims[i]);
     }
   }
@@ -2399,7 +2399,7 @@ void ProcessScatterNdOperator(Model* model, ScatterNdOperator* op) {
       if (unsupported_op->output_shapes.size() < op->outputs.size()) {
         return ::tensorflow::Status::OK();
       }
-      for (size_t i = 0; i < op->outputs.size(); ++i) {
+      for (int i = 0; i < op->outputs.size(); ++i) {
         const string& output = op->outputs[i];
         model->GetArray(output).copy_shape(unsupported_op->output_shapes.at(i));
       }
diff --git a/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc b/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc
index 1cb3a300127..6eccda04c18 100644
--- a/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc
+++ b/tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc
@@ -31,12 +31,12 @@ bool TransformsToIdentity(std::vector const& perm1,
   // perm1 is the order of the indices after first transpose. When perm1 is
   // reordered according to perm2, if the result is simple increasing sequence
   // i.e., range(0, perm1.size()), then the two transposes cancel each other.
-  for (size_t i = 0; i < perm1.size(); ++i) {
-    if (perm1[i] < 0 || perm1[i] >= static_cast(perm1.size()) || perm2[i] < 0 ||
-        perm2[i] >= static_cast(perm1.size())) {
+  for (int i = 0; i < perm1.size(); ++i) {
+    if (perm1[i] < 0 || perm1[i] >= perm1.size() || perm2[i] < 0 ||
+        perm2[i] >= perm1.size()) {
       return false;
     }
-    if (perm1[perm2[i]] != static_cast(i)) {
+    if (perm1[perm2[i]] != i) {
       return false;
     }
   }
@@ -46,7 +46,7 @@ bool TransformsToIdentity(std::vector const& perm1,
 void ReplaceOpInputsWith(Model* model, const string& lookfor,
                          const string& replacewith) {
   for (const auto& op : model->operators) {
-    for (size_t i = 0; i < op->inputs.size(); ++i) {
+    for (int i = 0; i < op->inputs.size(); ++i) {
       if (op->inputs[i] == lookfor) {
         op->inputs[i] = replacewith;
       }
diff --git a/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc b/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc
index eeb8751bf86..bd529bd9ecd 100644
--- a/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc
+++ b/tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc
@@ -82,7 +82,7 @@ bool RemoveTrivialPassthroughOp(GraphTransformation* transformation,
     // We call 'main input' the unique nonconstant input array if there is one,
     // or else the 0-th input.
     int count_nonconstant_input_arrays = 0;
-    for (size_t i = 0; i < passthru_op->inputs.size(); i++) {
+    for (int i = 0; i < passthru_op->inputs.size(); i++) {
       if (!model->GetArray(passthru_op->inputs[i]).buffer) {
         count_nonconstant_input_arrays++;
         if (count_nonconstant_input_arrays == 1) {
diff --git a/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc b/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc
index 38edff76d55..17a5e9a1d6a 100644
--- a/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc
+++ b/tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc
@@ -127,9 +127,9 @@ bool IsMoveOperator(OperatorType optype) {
     move_op->outputs[0] = output_name;
   } else {
     // The intermediate array is now the output array.
-    for (size_t i = 0; i < model->operators.size(); i++) {
+    for (int i = 0; i < model->operators.size(); i++) {
       Operator* consumer = model->operators[i].get();
-      for (size_t j = 0; j < consumer->inputs.size(); j++) {
+      for (int j = 0; j < consumer->inputs.size(); j++) {
         if (consumer->inputs[j] == output_name) {
           consumer->inputs[j] = intermediate_name;
         }
diff --git a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc
index b2d184cdc31..0fbcf9f73b1 100644
--- a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc
+++ b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc
@@ -60,7 +60,7 @@ std::vector ComputeNewPerm(std::vector input_dims,
                                 std::vector perm) {
   // These are the major axis of the input.
   std::vector input_indices;
-  for (size_t i = 0; i < input_dims.size(); i++) {
+  for (int i = 0; i < input_dims.size(); i++) {
     if (input_dims[i] != 1) {
       input_indices.push_back(i);
     }
@@ -69,7 +69,7 @@ std::vector ComputeNewPerm(std::vector input_dims,
   // This maps which indices of the input produced the intermediate indices for
   // non-unary dimensions.
   std::unordered_map intermediate_to_input_indices_map;
-  for (size_t i = 0; i < intermediate_dims.size(); i++) {
+  for (int i = 0; i < intermediate_dims.size(); i++) {
     if (intermediate_dims[i] != 1) {
       intermediate_to_input_indices_map[i] =
           input_indices[intermediate_to_input_indices_map.size()];
@@ -80,14 +80,14 @@ std::vector ComputeNewPerm(std::vector input_dims,
   // major indices.
   std::vector new_perm;
   new_perm.reserve(input_dims.size());
-  for (size_t i = 0; i < perm.size(); i++) {
+  for (int i = 0; i < perm.size(); i++) {
     if (intermediate_dims[perm[i]] == 1) continue;
 
     new_perm.push_back(intermediate_to_input_indices_map[perm[i]]);
   }
 
   // Fill the rest of the transpose in with the ones.
-  for (size_t index = 0; index < input_dims.size(); index++) {
+  for (int index = 0; index < input_dims.size(); index++) {
     if (input_dims[index] == 1) {
       new_perm.push_back(index);
     }
@@ -193,9 +193,9 @@ std::vector ComputeNewPerm(std::vector input_dims,
     DeleteArrayIfUnused(intermediate_name, model);
   } else {
     // The intermediate array is now the output array.
-    for (size_t i = 0; i < model->operators.size(); i++) {
+    for (int i = 0; i < model->operators.size(); i++) {
       Operator* consumer = model->operators[i].get();
-      for (size_t j = 0; j < consumer->inputs.size(); j++) {
+      for (int j = 0; j < consumer->inputs.size(); j++) {
         if (consumer->inputs[j] == output_name) {
           consumer->inputs[j] = intermediate_name;
         }
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc b/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc
index 545c53fb31a..6e5815ee94d 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc
@@ -124,11 +124,11 @@ namespace toco {
   const auto& offset_float_data =
       offset_array.GetBuffer().data;
 
-  CHECK(static_cast(mul_float_data.size()) == buffer_size);
-  CHECK(static_cast(add_float_data.size()) == buffer_size);
-  CHECK(static_cast(mean_float_data.size()) == buffer_size);
-  CHECK(static_cast(multiplier_float_data.size()) == buffer_size);
-  CHECK(static_cast(offset_float_data.size()) == buffer_size);
+  CHECK(mul_float_data.size() == buffer_size);
+  CHECK(add_float_data.size() == buffer_size);
+  CHECK(mean_float_data.size() == buffer_size);
+  CHECK(multiplier_float_data.size() == buffer_size);
+  CHECK(offset_float_data.size() == buffer_size);
 
   for (int i = 0; i < buffer_size; i++) {
     mul_float_data[i] = multiplier_float_data[i];
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc
index 20e805a29e0..7c9aa025f64 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc
@@ -64,7 +64,7 @@ void CopyTensorSegments(const std::vector& input_arrays,
   // Copy the data from input_arrays to concatenated_array_buffer.
   T* dest_ptr = concatenated_array_buffer.data();
   for (int s = 0; s < total_copy_steps; s++) {
-    for (size_t i = 0; i < input_arrays.size(); i++) {
+    for (int i = 0; i < input_arrays.size(); i++) {
       std::copy(src_ptr[i], src_ptr[i] + array_copy_size[i], dest_ptr);
       src_ptr[i] += array_copy_size[i];
       dest_ptr += array_copy_size[i];
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc
index c6dc093ba00..0df35509d3d 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc
@@ -36,7 +36,7 @@ void Pack(Model* model, PackOperator const& op) {
   // Pack inputs into buffer
   CHECK_EQ(op.axis, 0) << "Packing only supported along first axis";
   int dst_offset = 0;
-  for (size_t i = 0; i < op.inputs.size(); i++) {
+  for (int i = 0; i < op.inputs.size(); i++) {
     // Append array data to output for each input array
     const auto& input_array = model->GetArray(op.inputs[i]);
     int input_size = RequiredBufferSizeForShape(input_array.shape());
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc
index 34a1a1ce899..fd71fb1873a 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc
@@ -50,7 +50,7 @@ bool Slice(SliceOperator const& op, Array const& input_array,
   CHECK_LE(size.size(), 4);
   std::vector begin = op.begin;
   std::vector end;
-  for (size_t i = 0; i < begin.size(); ++i) {
+  for (int i = 0; i < begin.size(); ++i) {
     int dim_size = size[i];
     if (dim_size == -1) {
       // -1 means the rest of the dimension.
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc
index a822f7b79e3..7ceffe6307e 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc
@@ -40,7 +40,7 @@ void Transpose(Model* model, const Array& input_array,
   CHECK(input_shape.dimensions_count() == output_shape.dimensions_count());
   const int dim = input_shape.dimensions_count();
   CHECK_LE(dim, 4);
-  CHECK(static_cast(perm.size()) >= dim);
+  CHECK(perm.size() >= dim);
   for (int i = 0; i < dim; i++) {
     CHECK(perm[i] >= 0 && perm[i] < dim);
     CHECK(input_shape.dims(perm[i]) == output_shape.dims(i));
diff --git a/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc b/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc
index 4d6cd188729..197e17eee16 100644
--- a/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc
+++ b/tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc
@@ -62,7 +62,7 @@ void ReduceGeneric(bool keep_dims, const std::vector& axes,
   }
 
   std::vector output_indices(input_shape.dimensions_count());
-  for (size_t input_offset = 0; input_offset < input.size(); ++input_offset) {
+  for (int input_offset = 0; input_offset < input.size(); ++input_offset) {
     std::vector input_indices = ReverseOffset(input_shape, input_offset);
     // Calculate the output location by squashing input indices to 0
     // in reduced axes.
@@ -319,7 +319,7 @@ bool CopyMinMaxFromFirstInput(const Operator& op, Model* model) {
   } else if (unary_op->type == OperatorType::kRelu6 ||
              unary_op->type == OperatorType::kRelu1 ||
              unary_op->type == OperatorType::kRelu) {
-    for (int i = 0; i < output_buffer_size; ++i) {
+    for (size_t i = 0; i < output_buffer_size; ++i) {
       const float value = (*input_float_data)[i];
       float new_value = 0.0f;
       switch (unary_op->type) {
diff --git a/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc b/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc
index 84d5922aae8..1f7035c21e2 100644
--- a/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc
+++ b/tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc
@@ -57,10 +57,10 @@ namespace toco {
   // Split up the DynamicStitch inputs into the indices and data.
   std::vector stitch_indices_inputs;
   std::vector stitch_data_inputs;
-  for (int i = 0; i < stitch_op->num_partitions; ++i) {
+  for (size_t i = 0; i < stitch_op->num_partitions; ++i) {
     stitch_indices_inputs.push_back(stitch_op->inputs[i]);
   }
-  for (int i = stitch_op->num_partitions; i < stitch_op->num_partitions * 2;
+  for (size_t i = stitch_op->num_partitions; i < stitch_op->num_partitions * 2;
        ++i) {
     stitch_data_inputs.push_back(stitch_op->inputs[i]);
   }
diff --git a/tensorflow/lite/toco/model_cmdline_flags.cc b/tensorflow/lite/toco/model_cmdline_flags.cc
index 351884fbf1e..2434481272f 100644
--- a/tensorflow/lite/toco/model_cmdline_flags.cc
+++ b/tensorflow/lite/toco/model_cmdline_flags.cc
@@ -263,7 +263,7 @@ void ReadModelFlagsFromCommandLineFlags(
     QCHECK(uses_multi_input_flags);
     std::vector mean_values =
         absl::StrSplit(parsed_model_flags.mean_values.value(), ',');
-    QCHECK(static_cast(mean_values.size()) == model_flags->input_arrays_size());
+    QCHECK(mean_values.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < mean_values.size(); ++i) {
       char* last = nullptr;
       model_flags->mutable_input_arrays(i)->set_mean_value(
@@ -280,7 +280,7 @@ void ReadModelFlagsFromCommandLineFlags(
     QCHECK(uses_multi_input_flags);
     std::vector std_values =
         absl::StrSplit(parsed_model_flags.std_values.value(), ',');
-    QCHECK( static_cast(std_values.size()) == model_flags->input_arrays_size());
+    QCHECK(std_values.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < std_values.size(); ++i) {
       char* last = nullptr;
       model_flags->mutable_input_arrays(i)->set_std_value(
@@ -298,7 +298,7 @@ void ReadModelFlagsFromCommandLineFlags(
     QCHECK(uses_multi_input_flags);
     std::vector input_data_types =
         absl::StrSplit(parsed_model_flags.input_data_types.value(), ',');
-    QCHECK(static_cast(input_data_types.size()) == model_flags->input_arrays_size());
+    QCHECK(input_data_types.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < input_data_types.size(); ++i) {
       IODataType type;
       QCHECK(IODataType_Parse(input_data_types[i], &type));
@@ -321,7 +321,7 @@ void ReadModelFlagsFromCommandLineFlags(
     QCHECK(uses_multi_input_flags);
     std::vector input_shapes =
         absl::StrSplit(parsed_model_flags.input_shapes.value(), ':');
-    QCHECK(static_cast(input_shapes.size()) == model_flags->input_arrays_size());
+    QCHECK(input_shapes.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < input_shapes.size(); ++i) {
       auto* shape = model_flags->mutable_input_arrays(i)->mutable_shape();
       shape->clear_dims();
diff --git a/tensorflow/lite/toco/toco_cmdline_flags.cc b/tensorflow/lite/toco/toco_cmdline_flags.cc
index 9697a1ecbbd..c133db8f2a4 100644
--- a/tensorflow/lite/toco/toco_cmdline_flags.cc
+++ b/tensorflow/lite/toco/toco_cmdline_flags.cc
@@ -320,7 +320,7 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags,
     std::vector input_types =
         absl::StrSplit(parsed_toco_flags.input_types.value(), ',');
     QCHECK(!input_types.empty());
-    for (size_t i = 1; i < input_types.size(); i++) {
+    for (int i = 1; i < input_types.size(); i++) {
       QCHECK_EQ(input_types[i], input_types[0]);
     }
     toco::IODataType input_type;

From 0e634188b335bf90c154dc1488d5b292ddddb0d0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 14:16:30 -0700
Subject: [PATCH 0205/1390] fix kokoro build :deprecated message.

PiperOrigin-RevId: 316542716
Change-Id: I62ad7d88c36cd5f8551dea8efa8193c6ac8691ec
---
 tensorflow/python/training/tracking/data_structures.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/training/tracking/data_structures.py b/tensorflow/python/training/tracking/data_structures.py
index 8dc252efbf9..a7d9b70a560 100644
--- a/tensorflow/python/training/tracking/data_structures.py
+++ b/tensorflow/python/training/tracking/data_structures.py
@@ -736,7 +736,7 @@ class _DictWrapper(TrackableDataStructure, wrapt.ObjectProxy):
     if wrapped_dict is None:
       # Allow zero-argument construction, e.g. from session.run's re-wrapping.
       wrapped_dict = {}
-    if not isinstance(wrapped_dict, collections.Mapping):
+    if not isinstance(wrapped_dict, collections_abc.Mapping):
       # Allow construction from a sequence, e.g. from nest.pack_sequence_as.
       wrapped_dict = dict(wrapped_dict)
     wrapt.ObjectProxy.__init__(self, wrapped_dict)

From a2442ea4077e61a564ab598ac983f4160d9546be Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Mon, 15 Jun 2020 21:24:02 +0000
Subject: [PATCH 0206/1390] segragrating changes

---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 3724f06ba4b..9db20363349 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 409  // Updated: 2020/5/22
+#define TF_GRAPH_DEF_VERSION 408  // Updated: 2020/5/21
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From fc531df8ce3df4fd434230dd3bd0e7fdfe5a6cf1 Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Mon, 15 Jun 2020 14:20:38 -0700
Subject: [PATCH 0207/1390] Stop skipping KerasTensors for Keras tests that now
 work with KerasTensors.

PiperOrigin-RevId: 316543523
Change-Id: Iea54fa7ed735e239cda293304c6a17207b136ab7
---
 .../python/keras/engine/training_test.py      |  5 +---
 .../preprocessing/category_encoding_test.py   | 26 ++++++++-----------
 .../python/keras/premade/linear_test.py       |  3 +--
 .../python/keras/premade/wide_deep_test.py    |  3 +--
 tensorflow/python/keras/regularizers_test.py  |  8 +++---
 5 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index bb6bfc32921..aa01463582c 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -1563,10 +1563,7 @@ class TrainingTest(keras_parameterized.TestCase):
     self.assertEqual(self.evaluate(layer.v), 1.)
 
   @keras_parameterized.run_all_keras_modes(
-      always_skip_v1=True,
-      # TODO(kaftan): this is failing with KerasTensors
-      # in a way that seems orthogonal to what the code is testing
-      skip_keras_tensors=True)
+      always_skip_v1=True)
   @parameterized.named_parameters(
       ('numpy_array', 'numpy_array'),
       ('dataset_array', 'dataset_array'),
diff --git a/tensorflow/python/keras/layers/preprocessing/category_encoding_test.py b/tensorflow/python/keras/layers/preprocessing/category_encoding_test.py
index ff1a06a3ae7..7e7f7f32be0 100644
--- a/tensorflow/python/keras/layers/preprocessing/category_encoding_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/category_encoding_test.py
@@ -31,7 +31,6 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.keras import backend
 from tensorflow.python.keras import keras_parameterized
-from tensorflow.python.keras import testing_utils
 from tensorflow.python.keras.layers import core
 from tensorflow.python.keras.layers.preprocessing import category_encoding
 from tensorflow.python.keras.layers.preprocessing import category_encoding_v1
@@ -252,24 +251,21 @@ class CategoryEncodingInputTest(keras_parameterized.TestCase,
         sparse_ops.sparse_tensor_to_dense(sp_output_dataset, default_value=0),
         output_dataset)
 
-  # TODO(b/158570051): Support KerasTensor
-  # Keras functional model doesn't support dense layer stacked with sparse out.
   def test_sparse_output_and_dense_layer(self):
-    with testing_utils.use_keras_tensors_scope(False):
-      input_array = constant_op.constant([[1, 2, 3], [3, 3, 0]])
+    input_array = constant_op.constant([[1, 2, 3], [3, 3, 0]])
 
-      max_tokens = 4
+    max_tokens = 4
 
-      input_data = keras.Input(shape=(None,), dtype=dtypes.int32)
-      encoding_layer = get_layer_class()(
-          max_tokens=max_tokens, output_mode=category_encoding.COUNT,
-          sparse=True)
-      int_data = encoding_layer(input_data)
-      dense_layer = keras.layers.Dense(units=1)
-      output_data = dense_layer(int_data)
+    input_data = keras.Input(shape=(None,), dtype=dtypes.int32)
+    encoding_layer = get_layer_class()(
+        max_tokens=max_tokens, output_mode=category_encoding.COUNT,
+        sparse=True)
+    int_data = encoding_layer(input_data)
+    dense_layer = keras.layers.Dense(units=1)
+    output_data = dense_layer(int_data)
 
-      model = keras.Model(inputs=input_data, outputs=output_data)
-      _ = model.predict(input_array, steps=1)
+    model = keras.Model(inputs=input_data, outputs=output_data)
+    _ = model.predict(input_array, steps=1)
 
 
 @keras_parameterized.run_all_keras_modes
diff --git a/tensorflow/python/keras/premade/linear_test.py b/tensorflow/python/keras/premade/linear_test.py
index 6fa1767a60a..676f29bb840 100644
--- a/tensorflow/python/keras/premade/linear_test.py
+++ b/tensorflow/python/keras/premade/linear_test.py
@@ -40,8 +40,7 @@ from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True,
-                                         skip_keras_tensors=True)
+@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
 class LinearModelTest(keras_parameterized.TestCase):
 
   def test_linear_model_with_single_input(self):
diff --git a/tensorflow/python/keras/premade/wide_deep_test.py b/tensorflow/python/keras/premade/wide_deep_test.py
index eae28c31df8..591b53e9a84 100644
--- a/tensorflow/python/keras/premade/wide_deep_test.py
+++ b/tensorflow/python/keras/premade/wide_deep_test.py
@@ -37,8 +37,7 @@ from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
 
-@keras_parameterized.run_all_keras_modes(always_skip_v1=True,
-                                         skip_keras_tensors=True)
+@keras_parameterized.run_all_keras_modes(always_skip_v1=True)
 class WideDeepModelTest(keras_parameterized.TestCase):
 
   def test_wide_deep_model(self):
diff --git a/tensorflow/python/keras/regularizers_test.py b/tensorflow/python/keras/regularizers_test.py
index c2c2e6c4a01..b10218ba114 100644
--- a/tensorflow/python/keras/regularizers_test.py
+++ b/tensorflow/python/keras/regularizers_test.py
@@ -83,7 +83,7 @@ class KerasRegularizersTest(keras_parameterized.TestCase,
     self.assertEqual(len(model.losses), 1)
     model.fit(x_train, y_train, batch_size=10, epochs=1, verbose=0)
 
-  @keras_parameterized.run_all_keras_modes(skip_keras_tensors=True)
+  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters([
       ('l1', regularizers.l1()),
       ('l2', regularizers.l2()),
@@ -126,7 +126,7 @@ class KerasRegularizersTest(keras_parameterized.TestCase,
         model.get_config(), custom_objects={'my_regularizer': my_regularizer})
     self.assertEqual(model2.layers[1].kernel_regularizer, my_regularizer)
 
-  @keras_parameterized.run_all_keras_modes(skip_keras_tensors=True)
+  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters([
       ('l1', regularizers.l1()),
       ('l2', regularizers.l2()),
@@ -144,7 +144,7 @@ class KerasRegularizersTest(keras_parameterized.TestCase,
         run_eagerly=testing_utils.should_run_eagerly())
     self.assertLen(model.losses, 5)
 
-  @keras_parameterized.run_all_keras_modes(skip_keras_tensors=True)
+  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters([
       ('l1', regularizers.l1()),
       ('l2', regularizers.l2()),
@@ -166,7 +166,7 @@ class KerasRegularizersTest(keras_parameterized.TestCase,
         run_eagerly=testing_utils.should_run_eagerly())
     self.assertLen(model.losses, 6)
 
-  @keras_parameterized.run_all_keras_modes(skip_keras_tensors=True)
+  @keras_parameterized.run_all_keras_modes
   @parameterized.named_parameters([
       ('l1', regularizers.l1()),
       ('l2', regularizers.l2()),

From 4c004feb3e9b08961d2e3e17639b30104800efd5 Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Mon, 15 Jun 2020 21:37:05 +0000
Subject: [PATCH 0208/1390] segragation attempt 4

---
 tensorflow/core/framework/tensor_shape.cc             | 2 +-
 tensorflow/core/lib/io/random_inputstream.cc          | 2 +-
 tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc   | 2 +-
 tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc  | 6 +++---
 tensorflow/core/lib/io/zlib_outputbuffer.cc           | 6 +++---
 tensorflow/core/platform/env.cc                       | 2 +-
 tensorflow/core/platform/file_system.cc               | 2 +-
 tensorflow/core/platform/file_system_helper.cc        | 2 +-
 tensorflow/core/platform/status.cc                    | 4 +++-
 tensorflow/core/profiler/internal/parse_annotation.cc | 2 +-
 tensorflow/core/public/version.h                      | 2 +-
 11 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index f4b440e9cd1..8040a316a45 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -182,7 +182,7 @@ void TensorShapeBase::InitDims(gtl::ArraySlice dim_sizes) {
 
   // Allow sizes that are under kint64max^0.25 so that 4-way multiplication
   // below cannot overflow.
-  static const uint64 kMaxSmall = 0xd744;
+  static const int64 kMaxSmall = 0xd744;
   static_assert(kMaxSmall * kMaxSmall * kMaxSmall * kMaxSmall <= kint64max,
                 "bad overflow check");
   bool large_size = false;
diff --git a/tensorflow/core/lib/io/random_inputstream.cc b/tensorflow/core/lib/io/random_inputstream.cc
index 10f734a5bae..bd0054ce753 100644
--- a/tensorflow/core/lib/io/random_inputstream.cc
+++ b/tensorflow/core/lib/io/random_inputstream.cc
@@ -92,7 +92,7 @@ Status RandomAccessInputStream::SkipNBytes(int64 bytes_to_skip) {
     } else {
       return s;
     }
-    if (data.size() < bytes_to_read) {
+    if (data.size() < static_cast(bytes_to_read)) {
       return errors::OutOfRange("reached end of file");
     }
     bytes_to_skip -= bytes_to_read;
diff --git a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
index a331d4173cf..53939f2d8a3 100644
--- a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
@@ -134,7 +134,7 @@ Status SnappyInputBuffer::ReadCompressedBlockLength(uint32* length) {
     }
     size_t readable = std::min(bytes_to_read, avail_in_);
 
-    for (int i = 0; i < readable; i++) {
+    for (size_t i = 0; i < readable; i++) {
       // The "unsigned char" type cast is intentional to avoid implicit type
       // casting of the signed char to unsigned int during bitwise OR which
       // causes weird overflow errors.
diff --git a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
index 563503a1319..fe3a53c6c25 100644
--- a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
@@ -76,7 +76,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
 
   // If there is sufficient free space in input_buffer_ to fit data we
   // add it there and return.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -87,7 +87,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered());
 
   // input_buffer_ should be empty at this point.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -144,7 +144,7 @@ void SnappyOutputBuffer::AddToInputBuffer(StringPiece data) {
   const int32 free_tail_bytes =
       input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (bytes_to_write > free_tail_bytes) {
+  if (static_cast(bytes_to_write) > free_tail_bytes) {
     memmove(input_buffer_.get(), next_in_, avail_in_);
     next_in_ = input_buffer_.get();
   }
diff --git a/tensorflow/core/lib/io/zlib_outputbuffer.cc b/tensorflow/core/lib/io/zlib_outputbuffer.cc
index 5840ca60242..d475d0eaa5c 100644
--- a/tensorflow/core/lib/io/zlib_outputbuffer.cc
+++ b/tensorflow/core/lib/io/zlib_outputbuffer.cc
@@ -98,7 +98,7 @@ void ZlibOutputBuffer::AddToInputBuffer(StringPiece data) {
   int32 unread_bytes = z_stream_->avail_in;
   int32 free_tail_bytes = input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (bytes_to_write > free_tail_bytes) {
+  if (static_cast(bytes_to_write) > free_tail_bytes) {
     memmove(z_stream_input_.get(), z_stream_->next_in, z_stream_->avail_in);
     z_stream_->next_in = z_stream_input_.get();
   }
@@ -154,7 +154,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
 
   size_t bytes_to_write = data.size();
 
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -162,7 +162,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered(zlib_options_.flush_mode));
 
   // At this point input stream should be empty.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index b29cad05459..05d95ba0425 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -214,7 +214,7 @@ bool Env::FilesExist(const std::vector& files,
     }
     if (fs_status) {
       result &= fs_result;
-      for (int i = 0; i < itr.second.size(); ++i) {
+      for (size_t i = 0; i < itr.second.size(); ++i) {
         per_file_status[itr.second[i]] = fs_status->at(i);
       }
     } else if (!fs_result) {
diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc
index 9e96ceedbdc..c9657e2339f 100644
--- a/tensorflow/core/platform/file_system.cc
+++ b/tensorflow/core/platform/file_system.cc
@@ -308,7 +308,7 @@ StringPiece FileSystem::Basename(StringPiece path) const {
 StringPiece FileSystem::Extension(StringPiece path) const {
   StringPiece basename = this->Basename(path);
 
-  int pos = basename.rfind('.');
+  size_t pos = basename.rfind('.');
   if (pos == StringPiece::npos) {
     return StringPiece(path.data() + path.size(), 0);
   } else {
diff --git a/tensorflow/core/platform/file_system_helper.cc b/tensorflow/core/platform/file_system_helper.cc
index 64b175c4d17..909752389e1 100644
--- a/tensorflow/core/platform/file_system_helper.cc
+++ b/tensorflow/core/platform/file_system_helper.cc
@@ -103,7 +103,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
                 children_dir_status[i] = fs->IsDirectory(child_path);
               }
             });
-    for (int i = 0; i < children.size(); ++i) {
+    for (size_t i = 0; i < children.size(); ++i) {
       const string child_path = io::JoinPath(current_dir, children[i]);
       // If the IsDirectory call was cancelled we bail.
       if (children_dir_status[i].code() == tensorflow::error::CANCELLED) {
diff --git a/tensorflow/core/platform/status.cc b/tensorflow/core/platform/status.cc
index 756b8314148..e303c18091c 100644
--- a/tensorflow/core/platform/status.cc
+++ b/tensorflow/core/platform/status.cc
@@ -74,7 +74,9 @@ class StatusLogSink : public TFLogSink {
 
     mutex_lock lock(mu_);
     messages_.emplace_back(entry.ToString());
-    if (messages_.size() > num_messages_) messages_.pop_front();
+    if (messages_.size() > static_cast(num_messages_)){
+        messages_.pop_front();
+    }
   }
 
  private:
diff --git a/tensorflow/core/profiler/internal/parse_annotation.cc b/tensorflow/core/profiler/internal/parse_annotation.cc
index 32c26befa3d..a4cdc09739d 100644
--- a/tensorflow/core/profiler/internal/parse_annotation.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation.cc
@@ -50,7 +50,7 @@ std::vector SplitNameAndMetadata(
 std::vector SplitPairs(absl::string_view metadata) {
   std::vector key_value_pairs;
   std::stack quotes;
-  int start = 0, end = 0;
+  size_t start = 0, end = 0;
   for (; end < metadata.size(); ++end) {
     char ch = metadata[end];
     switch (ch) {
diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 9db20363349..3724f06ba4b 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 408  // Updated: 2020/5/21
+#define TF_GRAPH_DEF_VERSION 409  // Updated: 2020/5/22
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 813bd968d0c33c7eec0f6839ce372abc160e8e24 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Mon, 15 Jun 2020 14:24:42 -0700
Subject: [PATCH 0209/1390] Choosing better convolution for Intel GPUs.

PiperOrigin-RevId: 316544335
Change-Id: Idf9d653eebf4cff8195ab9ca3b3e34b1a8959c1f
---
 .../delegates/gpu/cl/kernels/conv_powervr.cc  | 20 +++++++++++++++++++
 .../gpu/cl/selectors/convolution_selector.cc  |  2 ++
 2 files changed, 22 insertions(+)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
index 9bb52b3e9c2..184e070202a 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
@@ -802,6 +802,26 @@ ConvPowerVR::ConvParams ConvPowerVR::GuessBestParams(
     conv_params.fixed_work_group_size = false;
     conv_params.src_depth_loop_size = 1;
     conv_params.weights_upload_type = WeightsUploadType::GLOBAL_MEM;
+  } else if (device.IsIntel()) {
+    conv_params.block_size = int3(1, 1, 4);
+    conv_params.work_group_size = int3(8, 2, 1);
+    conv_params.work_group_launch_order = int3(0, 1, 2);
+    conv_params.fixed_work_group_size = true;
+    conv_params.src_depth_loop_size = 1;
+    conv_params.weights_upload_type = WeightsUploadType::LOCAL_MEM_BY_THREADS;
+    if (dst_depth % 4 == 0 || dst_depth >= 8) {
+      conv_params.block_size.z = 4;
+    } else if (dst_depth % 2 == 0 || dst_depth >= 4) {
+      conv_params.block_size.z = 2;
+    } else {
+      conv_params.block_size.z = dst_depth;
+    }
+    if (src_depth % 2 == 0) {
+      conv_params.src_depth_loop_size = 2;
+    }
+    if (src_depth % 4 == 0 && conv_params.block_size.z <= 2) {
+      conv_params.src_depth_loop_size = 4;
+    }
   } else {
     conv_params.block_size = int3(1, 1, 4);
     conv_params.work_group_size = int3(8, 2, 1);
diff --git a/tensorflow/lite/delegates/gpu/cl/selectors/convolution_selector.cc b/tensorflow/lite/delegates/gpu/cl/selectors/convolution_selector.cc
index dc34dd7faee..3841c415301 100644
--- a/tensorflow/lite/delegates/gpu/cl/selectors/convolution_selector.cc
+++ b/tensorflow/lite/delegates/gpu/cl/selectors/convolution_selector.cc
@@ -170,6 +170,7 @@ absl::Status SelectConvolution(const Convolution2DAttributes& attr,
       return SelectConvolutionAdreno(attr, dst_shape, creation_context, op_def,
                                      hints, ptr);
     case Vendor::POWERVR:
+    case Vendor::INTEL:
     case Vendor::AMD:
       return SelectConvolutionPowerVR(attr, creation_context, op_def, ptr);
     case Vendor::NVIDIA:
@@ -193,6 +194,7 @@ absl::Status SelectConvolutionForWinograd(
                                              op_def, hints, ptr);
     case Vendor::POWERVR:
     case Vendor::AMD:
+    case Vendor::INTEL:
     case Vendor::NVIDIA: {
       ConvPowerVR conv;
       RETURN_IF_ERROR(

From 92bd07fee5541e941b7ea3da9bc90469b0dcf049 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 14:26:51 -0700
Subject: [PATCH 0210/1390] Polish some comments in dot decomposer.

PiperOrigin-RevId: 316544755
Change-Id: I46ce48dcbf64119e4795b923f5b45b814e8bb8c7
---
 .../compiler/xla/service/dot_decomposer.cc       | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/xla/service/dot_decomposer.cc b/tensorflow/compiler/xla/service/dot_decomposer.cc
index 40354dec3c6..573b82a1e0f 100644
--- a/tensorflow/compiler/xla/service/dot_decomposer.cc
+++ b/tensorflow/compiler/xla/service/dot_decomposer.cc
@@ -29,10 +29,12 @@ namespace xla {
 
 namespace {
 
-// Convert a dot into a canonical form where non-contracting and contracting
-// dimensions are reshaped together and batch dimensions are the most major
-// dimensions. This requires transposing and reshapes of the lhs and rhs and
-// reshaping the output batch to the original shape.
+// Convert a dot into a canonical form;
+// * Non-contracting dimensions are reshaped together,
+// * Contracting dimensions are reshaped together,
+// * Batch dimensions are the most major dimensions.
+// This requires transposing and reshaping of the lhs and rhs, and reshaping the
+// output batch to the original shape.
 Status CanonicalizeDot(HloInstruction* original_dot) {
   auto computation = original_dot->parent();
   const auto& original_dnums = original_dot->dot_dimension_numbers();
@@ -63,7 +65,8 @@ Status CanonicalizeDot(HloInstruction* original_dot) {
     }
   }
   // The canonical form of the lhs is
-  // [BatchDims, NonContractingDims, ContractingsDims]
+  // [BatchDims, NonContractingDimsProduct, ContractingsDimsProduct]
+  // If NonContractingDimsProduct is 1, it is omitted.
   std::vector lhs_transpose;
   lhs_transpose.reserve(lhs_rank);
   lhs_transpose.insert(lhs_transpose.end(),
@@ -109,7 +112,8 @@ Status CanonicalizeDot(HloInstruction* original_dot) {
   }
 
   // The canonical form of the rhs is
-  // [BatchDims, ContractingsDims, NonContractingDims]
+  // [BatchDims, NonContractingDimsProduct, ContractingsDimsProduct]
+  // If NonContractingDimsProduct is 1, it is omitted.
   std::vector rhs_transpose;
   rhs_transpose.reserve(rhs_rank);
   rhs_transpose.insert(rhs_transpose.end(),

From 0096d0a19b5543b368a5d2426cb2810931913272 Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Mon, 15 Jun 2020 21:49:18 +0000
Subject: [PATCH 0211/1390] final segratation

---
 tensorflow/compiler/xla/window_util.cc                | 2 +-
 tensorflow/core/framework/tensor_shape.cc             | 2 +-
 tensorflow/core/lib/io/random_inputstream.cc          | 2 +-
 tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc   | 2 +-
 tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc  | 6 +++---
 tensorflow/core/lib/io/zlib_outputbuffer.cc           | 6 +++---
 tensorflow/core/platform/env.cc                       | 2 +-
 tensorflow/core/platform/file_system.cc               | 2 +-
 tensorflow/core/platform/file_system_helper.cc        | 2 +-
 tensorflow/core/platform/status.cc                    | 4 +---
 tensorflow/core/profiler/internal/parse_annotation.cc | 2 +-
 11 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/tensorflow/compiler/xla/window_util.cc b/tensorflow/compiler/xla/window_util.cc
index e33d0b6d1dc..a58179c3ee0 100644
--- a/tensorflow/compiler/xla/window_util.cc
+++ b/tensorflow/compiler/xla/window_util.cc
@@ -42,7 +42,7 @@ Window MakeWindow(absl::Span sizes,
                   absl::Span strides) {
   Window window;
   CHECK_EQ(sizes.size(), strides.size());
-  for (auto nb = 0; static_cast(nb) < sizes.size(); ++nb) {
+  for (auto nb = 0; nb < sizes.size(); ++nb) {
     auto* dimension = window.add_dimensions();
     dimension->set_size(sizes[nb]);
     dimension->set_stride(strides[nb]);
diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index 8040a316a45..f4b440e9cd1 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -182,7 +182,7 @@ void TensorShapeBase::InitDims(gtl::ArraySlice dim_sizes) {
 
   // Allow sizes that are under kint64max^0.25 so that 4-way multiplication
   // below cannot overflow.
-  static const int64 kMaxSmall = 0xd744;
+  static const uint64 kMaxSmall = 0xd744;
   static_assert(kMaxSmall * kMaxSmall * kMaxSmall * kMaxSmall <= kint64max,
                 "bad overflow check");
   bool large_size = false;
diff --git a/tensorflow/core/lib/io/random_inputstream.cc b/tensorflow/core/lib/io/random_inputstream.cc
index bd0054ce753..10f734a5bae 100644
--- a/tensorflow/core/lib/io/random_inputstream.cc
+++ b/tensorflow/core/lib/io/random_inputstream.cc
@@ -92,7 +92,7 @@ Status RandomAccessInputStream::SkipNBytes(int64 bytes_to_skip) {
     } else {
       return s;
     }
-    if (data.size() < static_cast(bytes_to_read)) {
+    if (data.size() < bytes_to_read) {
       return errors::OutOfRange("reached end of file");
     }
     bytes_to_skip -= bytes_to_read;
diff --git a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
index 53939f2d8a3..a331d4173cf 100644
--- a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
@@ -134,7 +134,7 @@ Status SnappyInputBuffer::ReadCompressedBlockLength(uint32* length) {
     }
     size_t readable = std::min(bytes_to_read, avail_in_);
 
-    for (size_t i = 0; i < readable; i++) {
+    for (int i = 0; i < readable; i++) {
       // The "unsigned char" type cast is intentional to avoid implicit type
       // casting of the signed char to unsigned int during bitwise OR which
       // causes weird overflow errors.
diff --git a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
index fe3a53c6c25..563503a1319 100644
--- a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
@@ -76,7 +76,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
 
   // If there is sufficient free space in input_buffer_ to fit data we
   // add it there and return.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -87,7 +87,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered());
 
   // input_buffer_ should be empty at this point.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -144,7 +144,7 @@ void SnappyOutputBuffer::AddToInputBuffer(StringPiece data) {
   const int32 free_tail_bytes =
       input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (static_cast(bytes_to_write) > free_tail_bytes) {
+  if (bytes_to_write > free_tail_bytes) {
     memmove(input_buffer_.get(), next_in_, avail_in_);
     next_in_ = input_buffer_.get();
   }
diff --git a/tensorflow/core/lib/io/zlib_outputbuffer.cc b/tensorflow/core/lib/io/zlib_outputbuffer.cc
index d475d0eaa5c..5840ca60242 100644
--- a/tensorflow/core/lib/io/zlib_outputbuffer.cc
+++ b/tensorflow/core/lib/io/zlib_outputbuffer.cc
@@ -98,7 +98,7 @@ void ZlibOutputBuffer::AddToInputBuffer(StringPiece data) {
   int32 unread_bytes = z_stream_->avail_in;
   int32 free_tail_bytes = input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (static_cast(bytes_to_write) > free_tail_bytes) {
+  if (bytes_to_write > free_tail_bytes) {
     memmove(z_stream_input_.get(), z_stream_->next_in, z_stream_->avail_in);
     z_stream_->next_in = z_stream_input_.get();
   }
@@ -154,7 +154,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
 
   size_t bytes_to_write = data.size();
 
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -162,7 +162,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered(zlib_options_.flush_mode));
 
   // At this point input stream should be empty.
-  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
+  if (bytes_to_write <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index 05d95ba0425..b29cad05459 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -214,7 +214,7 @@ bool Env::FilesExist(const std::vector& files,
     }
     if (fs_status) {
       result &= fs_result;
-      for (size_t i = 0; i < itr.second.size(); ++i) {
+      for (int i = 0; i < itr.second.size(); ++i) {
         per_file_status[itr.second[i]] = fs_status->at(i);
       }
     } else if (!fs_result) {
diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc
index c9657e2339f..9e96ceedbdc 100644
--- a/tensorflow/core/platform/file_system.cc
+++ b/tensorflow/core/platform/file_system.cc
@@ -308,7 +308,7 @@ StringPiece FileSystem::Basename(StringPiece path) const {
 StringPiece FileSystem::Extension(StringPiece path) const {
   StringPiece basename = this->Basename(path);
 
-  size_t pos = basename.rfind('.');
+  int pos = basename.rfind('.');
   if (pos == StringPiece::npos) {
     return StringPiece(path.data() + path.size(), 0);
   } else {
diff --git a/tensorflow/core/platform/file_system_helper.cc b/tensorflow/core/platform/file_system_helper.cc
index 909752389e1..64b175c4d17 100644
--- a/tensorflow/core/platform/file_system_helper.cc
+++ b/tensorflow/core/platform/file_system_helper.cc
@@ -103,7 +103,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
                 children_dir_status[i] = fs->IsDirectory(child_path);
               }
             });
-    for (size_t i = 0; i < children.size(); ++i) {
+    for (int i = 0; i < children.size(); ++i) {
       const string child_path = io::JoinPath(current_dir, children[i]);
       // If the IsDirectory call was cancelled we bail.
       if (children_dir_status[i].code() == tensorflow::error::CANCELLED) {
diff --git a/tensorflow/core/platform/status.cc b/tensorflow/core/platform/status.cc
index e303c18091c..756b8314148 100644
--- a/tensorflow/core/platform/status.cc
+++ b/tensorflow/core/platform/status.cc
@@ -74,9 +74,7 @@ class StatusLogSink : public TFLogSink {
 
     mutex_lock lock(mu_);
     messages_.emplace_back(entry.ToString());
-    if (messages_.size() > static_cast(num_messages_)){
-        messages_.pop_front();
-    }
+    if (messages_.size() > num_messages_) messages_.pop_front();
   }
 
  private:
diff --git a/tensorflow/core/profiler/internal/parse_annotation.cc b/tensorflow/core/profiler/internal/parse_annotation.cc
index a4cdc09739d..32c26befa3d 100644
--- a/tensorflow/core/profiler/internal/parse_annotation.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation.cc
@@ -50,7 +50,7 @@ std::vector SplitNameAndMetadata(
 std::vector SplitPairs(absl::string_view metadata) {
   std::vector key_value_pairs;
   std::stack quotes;
-  size_t start = 0, end = 0;
+  int start = 0, end = 0;
   for (; end < metadata.size(); ++end) {
     char ch = metadata[end];
     switch (ch) {

From 176ab11d0a9211fe9f8da0a7e6c75381c7ca26e8 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Mon, 15 Jun 2020 14:54:50 -0700
Subject: [PATCH 0212/1390] Added handling of Intel in choosing best storage
 types.

PiperOrigin-RevId: 316550810
Change-Id: I7e81cd0df1522be4b705df57a5658397328b5a18
---
 tensorflow/lite/delegates/gpu/cl/environment.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/lite/delegates/gpu/cl/environment.cc b/tensorflow/lite/delegates/gpu/cl/environment.cc
index 01d034fb1f7..6b6ab84f148 100644
--- a/tensorflow/lite/delegates/gpu/cl/environment.cc
+++ b/tensorflow/lite/delegates/gpu/cl/environment.cc
@@ -242,6 +242,8 @@ TensorStorageType GetFastestStorageType(const CLDevice& gpu) {
   } else if (gpu.IsAMD()) {
     return gpu.SupportsImageBuffer() ? TensorStorageType::IMAGE_BUFFER
                                      : TensorStorageType::BUFFER;
+  } else if (gpu.IsIntel()) {
+    return TensorStorageType::BUFFER;
   }
   return TensorStorageType::BUFFER;
 }
@@ -264,6 +266,8 @@ TensorStorageType GetStorageTypeWithMinimalMemoryConsumption(
   } else if (gpu.IsAMD()) {
     return gpu.SupportsImageBuffer() ? TensorStorageType::IMAGE_BUFFER
                                      : TensorStorageType::BUFFER;
+  } else if (gpu.IsIntel()) {
+    return TensorStorageType::BUFFER;
   }
   return TensorStorageType::BUFFER;
 }

From b59b2a10b16e17dcf439bf973b7bbd5da65d3c25 Mon Sep 17 00:00:00 2001
From: Lucy Fox 
Date: Mon, 15 Jun 2020 15:08:04 -0700
Subject: [PATCH 0213/1390] [NFC] Add comment to specify status of
 FusedKernelMatcher pass.

We're porting over this pass based on usage/need, so I'm updating the comments to make this clear.

PiperOrigin-RevId: 316553704
Change-Id: I5db9a4637edd3bcb27b4db6259b7113ae0f8dd0c
---
 .../tensorflow/transforms/fused_kernel_matcher.cc  | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
index 3ccdd957798..4b10550df7b 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
@@ -32,11 +32,15 @@ namespace TF {
 
 namespace {
 
-// Note: This implements fusions performed in the old Remapper Grappler pass.
-// That pass has specific cases for GPU and based on different target
-// configurations on both CPU and GPU (Intel MKL, ROCm, etc.). This MLIR pass
-// covers the general CPU case and at the moment does not account for any
-// target-specific configurations.
+// Note: This implements the fusions performed in the old Remapper Grappler
+// pass. That pass has specific cases for GPU and based on different
+// target configurations on both CPU and GPU (Intel MKL, ROCm, etc.). This MLIR
+// pass covers (some of) the general CPU case and at the moment does not account
+// for any target-specific configurations.
+
+// This pass is being ported over from the Grappler Remapper pass based on
+// need/usage. File a bug to request porting over additional fusions.
+
 // TODO(b/158265178): Support GPU-specific fusions.
 // TODO(b/158266710): Support CPU MKL configurations.
 

From 1b412edc891b0d2cf98dbd2291d871d5ee67a0fb Mon Sep 17 00:00:00 2001
From: Lu Wang 
Date: Mon, 15 Jun 2020 15:09:29 -0700
Subject: [PATCH 0214/1390] Create the script to generate TFLite Java API
 documentation.

PiperOrigin-RevId: 316553983
Change-Id: I08888ba83f5335d38ba588ed632c2471dcfa2d9f
---
 tensorflow/lite/g3doc/tools/BUILD             | 11 ++++
 .../lite/g3doc/tools/build_java_api_docs.py   | 66 +++++++++++++++++++
 2 files changed, 77 insertions(+)
 create mode 100644 tensorflow/lite/g3doc/tools/build_java_api_docs.py

diff --git a/tensorflow/lite/g3doc/tools/BUILD b/tensorflow/lite/g3doc/tools/BUILD
index f2c6d8efedc..5c891a67128 100644
--- a/tensorflow/lite/g3doc/tools/BUILD
+++ b/tensorflow/lite/g3doc/tools/BUILD
@@ -13,3 +13,14 @@ py_binary(
         "@absl_py//absl/flags",
     ],
 )
+
+py_binary(
+    name = "build_java_api_docs",
+    srcs = ["build_java_api_docs.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+    ],
+)
diff --git a/tensorflow/lite/g3doc/tools/build_java_api_docs.py b/tensorflow/lite/g3doc/tools/build_java_api_docs.py
new file mode 100644
index 00000000000..9c598ad8ec5
--- /dev/null
+++ b/tensorflow/lite/g3doc/tools/build_java_api_docs.py
@@ -0,0 +1,66 @@
+# Lint as: python3
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Generate TensorFlow Lite Java reference docs for TensorFlow.org."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pathlib
+import shutil
+import tempfile
+
+from absl import app
+from absl import flags
+
+from tensorflow_docs.api_generator import gen_java
+
+FLAGS = flags.FLAGS
+
+# These flags are required by infrastructure, not all of them are used.
+flags.DEFINE_string('output_dir', '/tmp/lite_api/',
+                    ("Use this branch as the root version and don't"
+                     ' create in version directory'))
+
+flags.DEFINE_string('site_path', 'lite/api_docs/java',
+                    'Path prefix in the _toc.yaml')
+
+flags.DEFINE_string('code_url_prefix', None,
+                    '[UNUSED] The url prefix for links to code.')
+
+flags.DEFINE_bool(
+    'search_hints', True,
+    '[UNUSED] Include metadata search hints in the generated files')
+
+# __file__ is the path to this file
+DOCS_TOOLS_DIR = pathlib.Path(__file__).resolve().parent
+TENSORFLOW_ROOT = DOCS_TOOLS_DIR.parents[3]
+SOURCE_PATH = TENSORFLOW_ROOT / 'tensorflow/lite/java/src/main/java/'
+
+
+def main(unused_argv):
+  merged_source = pathlib.Path(tempfile.mkdtemp())
+  shutil.copytree(SOURCE_PATH, merged_source / 'java')
+
+  gen_java.gen_java_docs(
+      package='org.tensorflow.lite',
+      source_path=merged_source / 'java',
+      output_dir=pathlib.Path(FLAGS.output_dir),
+      site_path=pathlib.Path(FLAGS.site_path))
+
+
+if __name__ == '__main__':
+  flags.mark_flags_as_required(['output_dir'])
+  app.run(main)

From 9136f5775e29cd0540fd594585687703849e64ab Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 15:13:45 -0700
Subject: [PATCH 0215/1390] [XLA] Refactor memory_space_assignment.cc.

Refactors IsIntervalAllowedInAlternateMemory() to a separate utils file so that it can be reused.

PiperOrigin-RevId: 316554798
Change-Id: Ibc6a4cffde6a1df233d375358164b373ea4ee7a6
---
 tensorflow/compiler/xla/service/BUILD         | 10 ++
 .../xla/service/memory_space_assignment.cc    | 82 +---------------
 .../xla/service/memory_space_assignment.h     |  4 -
 .../service/memory_space_assignment_utils.cc  | 95 +++++++++++++++++++
 .../service/memory_space_assignment_utils.h   | 34 +++++++
 5 files changed, 143 insertions(+), 82 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/memory_space_assignment_utils.cc
 create mode 100644 tensorflow/compiler/xla/service/memory_space_assignment_utils.h

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 40c237c5e6d..acd35cbc153 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -3304,6 +3304,15 @@ tf_cc_test(
     ],
 )
 
+cc_library(
+    name = "memory_space_assignment_utils",
+    srcs = ["memory_space_assignment_utils.cc"],
+    hdrs = ["memory_space_assignment_utils.h"],
+    deps = [
+        ":heap_simulator",
+    ],
+)
+
 cc_library(
     name = "memory_space_assignment",
     srcs = ["memory_space_assignment.cc"],
@@ -3311,6 +3320,7 @@ cc_library(
     deps = [
         ":heap_simulator",
         ":hlo_cost_analysis",
+        ":memory_space_assignment_utils",
         "//tensorflow/compiler/xla:debug_options_flags",
         "//tensorflow/core/lib/math:math_util",
     ],
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment.cc b/tensorflow/compiler/xla/service/memory_space_assignment.cc
index 21baaf1c7d5..388a2e18f38 100644
--- a/tensorflow/compiler/xla/service/memory_space_assignment.cc
+++ b/tensorflow/compiler/xla/service/memory_space_assignment.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/memory_space_assignment.h"
 
 #include "tensorflow/compiler/xla/debug_options_flags.h"
+#include "tensorflow/compiler/xla/service/memory_space_assignment_utils.h"
 #include "tensorflow/core/lib/math/math_util.h"
 namespace xla {
 
@@ -597,81 +598,6 @@ AlternateMemoryBestFitHeap::GetSortedColocatedIntervals(
   return colocated_intervals;
 }
 
-bool AlternateMemoryBestFitHeap::IsIntervalAllowedInAlternateMemory(
-    const BufferInterval& interval) const {
-  // If the buffer is a tuple, don't use this algorithm for now. The buffers
-  // that are pointed to by the tuple will still use this algorithm.  Because
-  // tuples are cheap to place in the alternate memory (they are just pointers)
-  // we don't need to use prefetch/evict logic.
-  if (interval.buffer->shape().IsTuple()) {
-    VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-            << " in default mem because it is a tuple.";
-    return false;
-  }
-
-  // Don't place scalars in the alternate memory.
-  if (ShapeUtil::IsEffectiveScalar(interval.buffer->shape())) {
-    VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-            << " in default mem because it is a scalar.";
-    return false;
-  }
-
-  // The semantics of TupleSelect are weird: TupleSelect doesn't define a
-  // buffer, but just forwards the buffers in the either left or right side.
-  // This means the two different inputs to TupleSelect must not alias, yet they
-  // should be allocated in the same memory space, and both buffers must be kept
-  // alive for the entire live range of TupleSelect. Instead, just don't
-  // allocate TupleSelect in the alternate memory space.
-  // TODO(berkin): Not allocating add-dependencies either since they need to be
-  // treated specially. We should revisit this later.
-  for (const HloPosition& position : interval.buffer->positions()) {
-    if (position.instruction->opcode() == HloOpcode::kTupleSelect ||
-        position.instruction->opcode() == HloOpcode::kAddDependency) {
-      VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-              << " in default mem because it has a tuple-select or "
-              << "add-dependency position.";
-      return false;
-    }
-  }
-
-  // Send and Recv HLOs return a request identifier. These should not be
-  // allocated in the alternate memory.
-  for (const HloPosition& position : interval.buffer->positions()) {
-    if ((position.instruction->opcode() == HloOpcode::kSend ||
-         position.instruction->opcode() == HloOpcode::kRecv)) {
-      // TODO(berkin): Send/recv buffers need a stable buffer allocation
-      // throughout sending/receiving. Disable memory space allocation for these
-      // for now.
-      if (position.index == ShapeIndex({0})) {
-        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-                << " in default mem because it is a send/recv buffer.";
-        return false;
-      } else if (position.index == ShapeIndex({1})) {
-        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-                << " in default mem because it is a request identifier for "
-                   "send/recv.";
-        return false;
-      }
-    }
-
-    if ((position.instruction->opcode() == HloOpcode::kCollectivePermuteStart ||
-         position.instruction->opcode() == HloOpcode::kCollectivePermuteDone)) {
-      // Disable memory space allocation for these for now.
-      if (position.index == ShapeIndex({0})) {
-        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-                << " in default mem because it is a collective-permute buffer.";
-        return false;
-      } else if (position.index == ShapeIndex({1})) {
-        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
-                << " in default mem because it is a collective-permute buffer.";
-        return false;
-      }
-    }
-  }
-
-  return true;
-}
-
 bool AlternateMemoryBestFitHeap::IsUseAllowedInAlternateMemory(
     const AllocationValue& value, const HloUse& use) const {
   const auto& instruction_schedule = hlo_live_range_.instruction_schedule();
@@ -710,8 +636,7 @@ bool AlternateMemoryBestFitHeap::IsUseAllowedInAlternateMemory(
     if (!options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy(
             shape, parameter_time, min_use_time)) {
       VLOG(4) << "While allocation not allowed in alternate memory. "
-              << "use time = " << min_use_time
-              << ", root time = " << root_time;
+              << "use time = " << min_use_time << ", root time = " << root_time;
       return false;
     }
     // Check if there is a required assignment for the while loop output.
@@ -897,7 +822,8 @@ HeapSimulator::Result AlternateMemoryBestFitHeap::Finish() {
       continue;
     }
 
-    if (!IsIntervalAllowedInAlternateMemory(interval)) {
+    if (!MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory(
+            interval)) {
       continue;
     }
 
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment.h b/tensorflow/compiler/xla/service/memory_space_assignment.h
index b8f47e73b8c..f9e5738d17e 100644
--- a/tensorflow/compiler/xla/service/memory_space_assignment.h
+++ b/tensorflow/compiler/xla/service/memory_space_assignment.h
@@ -909,10 +909,6 @@ class AlternateMemoryBestFitHeap : public GlobalDecreasingSizeBestFitHeap {
   static MemorySpaceAssignment::Allocation* GetLiveAllocationAt(
       const MemorySpaceAssignment::AllocationSequence& allocations, int64 time);
 
-  // Returns true if this buffer is allowed to be placed in the alternate
-  // memory.
-  bool IsIntervalAllowedInAlternateMemory(const BufferInterval& interval) const;
-
   // Returns true if the use is allowed in the alternate memory.
   bool IsUseAllowedInAlternateMemory(const AllocationValue& value,
                                      const HloUse& use) const;
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment_utils.cc b/tensorflow/compiler/xla/service/memory_space_assignment_utils.cc
new file mode 100644
index 00000000000..0215f007c9c
--- /dev/null
+++ b/tensorflow/compiler/xla/service/memory_space_assignment_utils.cc
@@ -0,0 +1,95 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/memory_space_assignment_utils.h"
+
+namespace xla {
+
+bool MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory(
+    const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval) {
+  // If the buffer is a tuple, don't use this algorithm for now. The buffers
+  // that are pointed to by the tuple will still use this algorithm.  Because
+  // tuples are cheap to place in the alternate memory (they are just pointers)
+  // we don't need to use prefetch/evict logic.
+  if (interval.buffer->shape().IsTuple()) {
+    VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+            << " in default mem because it is a tuple.";
+    return false;
+  }
+
+  // Don't place scalars in the alternate memory.
+  if (ShapeUtil::IsEffectiveScalar(interval.buffer->shape())) {
+    VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+            << " in default mem because it is a scalar.";
+    return false;
+  }
+
+  // The semantics of TupleSelect are weird: TupleSelect doesn't define a
+  // buffer, but just forwards the buffers in the either left or right side.
+  // This means the two different inputs to TupleSelect must not alias, yet they
+  // should be allocated in the same memory space, and both buffers must be kept
+  // alive for the entire live range of TupleSelect. Instead, just don't
+  // allocate TupleSelect in the alternate memory space.
+  // TODO(berkin): Not allocating add-dependencies either since they need to be
+  // treated specially. We should revisit this later.
+  for (const HloPosition& position : interval.buffer->positions()) {
+    if (position.instruction->opcode() == HloOpcode::kTupleSelect ||
+        position.instruction->opcode() == HloOpcode::kAddDependency) {
+      VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+              << " in default mem because it has a tuple-select or "
+              << "add-dependency position.";
+      return false;
+    }
+  }
+
+  // Send and Recv HLOs return a request identifier. These should not be
+  // allocated in the alternate memory.
+  for (const HloPosition& position : interval.buffer->positions()) {
+    if ((position.instruction->opcode() == HloOpcode::kSend ||
+         position.instruction->opcode() == HloOpcode::kRecv)) {
+      // TODO(berkin): Send/recv buffers need a stable buffer allocation
+      // throughout sending/receiving. Disable memory space allocation for these
+      // for now.
+      if (position.index == ShapeIndex({0})) {
+        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+                << " in default mem because it is a send/recv buffer.";
+        return false;
+      } else if (position.index == ShapeIndex({1})) {
+        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+                << " in default mem because it is a request identifier for "
+                   "send/recv.";
+        return false;
+      }
+    }
+
+    if ((position.instruction->opcode() == HloOpcode::kCollectivePermuteStart ||
+         position.instruction->opcode() == HloOpcode::kCollectivePermuteDone)) {
+      // Disable memory space allocation for these for now.
+      if (position.index == ShapeIndex({0})) {
+        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+                << " in default mem because it is a collective-permute buffer.";
+        return false;
+      } else if (position.index == ShapeIndex({1})) {
+        VLOG(4) << "Keeping value " << interval.buffer->ToShortString()
+                << " in default mem because it is a collective-permute buffer.";
+        return false;
+      }
+    }
+  }
+
+  return true;
+}
+
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment_utils.h b/tensorflow/compiler/xla/service/memory_space_assignment_utils.h
new file mode 100644
index 00000000000..651ac107c25
--- /dev/null
+++ b/tensorflow/compiler/xla/service/memory_space_assignment_utils.h
@@ -0,0 +1,34 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_
+
+#include "tensorflow/compiler/xla/service/heap_simulator.h"
+
+namespace xla {
+
+// Encapsulates common utility methods for memory space assignment.
+class MemorySpaceAssignmentUtils {
+ public:
+  // Returns true if this buffer is allowed to be placed in the alternate
+  // memory.
+  static bool IsIntervalAllowedInAlternateMemory(
+      const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval);
+};
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_

From 072c2f5d0dead92e2c41ef168986e927481db463 Mon Sep 17 00:00:00 2001
From: Nick Kreeger 
Date: Mon, 15 Jun 2020 15:22:07 -0700
Subject: [PATCH 0216/1390] Reduce excessive RAM in TFLM by using the existing
 flatbuffer quantization data for scales.

Currently, TFLM manually allocates a tail chunk to store "quantization" tensor data on TfLiteTensor objects. The size of these allocations vary based on the type of model - conv1d/2d models tend to be rich since quantization data is stored "per channel".

This change simply points the scale data at the existing value in the flatbuffer. The flatbuffer schema stores float values as flatbuffers::Vector and the TfLiteAffineQuantization struct can point the scale pointer at these values. Unfortunately, the zero point values are stored as flatbuffers::Vector and can not be reused. This allocation will be addressed in a future change.

Keyword Model ~2% reduction in tail allocation:
-----------------------------------------------
[RecordingMicroAllocator] Arena allocation total 21040 bytes
[RecordingMicroAllocator] Arena allocation head 672 bytes
[RecordingMicroAllocator] Arena allocation tail 20368 bytes
[RecordingMicroAllocator] 'TfLiteTensor struct' used 6048 bytes with alignment overhead (requested 6048 bytes for 54 tensors)
[RecordingMicroAllocator] 'TfLiteTensor quantization data' used 1728 bytes with alignment overhead (requested 1728 bytes for 108 allocations)
[RecordingMicroAllocator] 'TfLiteTensor variable buffer data' used 10240 bytes with alignment overhead (requested 10240 bytes for 7 allocations)
[RecordingMicroAllocator] 'NodeAndRegistration struct' used 1200 bytes with alignment overhead (requested 1200 bytes for 15 NodeAndRegistration structs)
[RecordingMicroAllocator] 'Operator runtime data' used 148 bytes with alignment overhead (requested 148 bytes for 13 OpData structs)

Test Conv Model ~10% reduction in tail allocation:
-----------------------------------------------
[RecordingMicroAllocator] Arena allocation total 11680 bytes
[RecordingMicroAllocator] Arena allocation head 7744 bytes
[RecordingMicroAllocator] Arena allocation tail 3936 bytes
[RecordingMicroAllocator] 'TfLiteTensor struct' used 1680 bytes with alignment overhead (requested 1680 bytes for 15 tensors)
[RecordingMicroAllocator] 'TfLiteTensor quantization data' used 768 bytes with alignment overhead (requested 752 bytes for 24 allocations)
[RecordingMicroAllocator] 'TfLiteTensor variable buffer data' used 0 bytes with alignment overhead (requested 0 bytes for 0 allocations)
[RecordingMicroAllocator] 'NodeAndRegistration struct' used 560 bytes with alignment overhead (requested 560 bytes for 7 NodeAndRegistration structs)
[RecordingMicroAllocator] 'Operator runtime data' used 136 bytes with alignment overhead (requested 136 bytes for 5 OpData structs)

PiperOrigin-RevId: 316556393
Change-Id: Iadadab51019d2787d11af9713b3639f087afa7bc
---
 .../lite/micro/memory_arena_threshold_test.cc | 31 +++++++++++++------
 tensorflow/lite/micro/micro_allocator.cc      | 22 ++++++-------
 .../micro/recording_micro_allocator_test.cc   |  6 ++--
 3 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/tensorflow/lite/micro/memory_arena_threshold_test.cc b/tensorflow/lite/micro/memory_arena_threshold_test.cc
index 19c3d0f1e06..58d3eff8df5 100644
--- a/tensorflow/lite/micro/memory_arena_threshold_test.cc
+++ b/tensorflow/lite/micro/memory_arena_threshold_test.cc
@@ -41,11 +41,11 @@ constexpr int kKeywordModelNodeAndRegistrationCount = 15;
 
 // NOTE: These values are measured on x86-64:
 // TODO(b/158651472): Consider auditing these values on non-64 bit systems.
-constexpr int kKeywordModelTotalSize = 21472;
+constexpr int kKeywordModelTotalSize = 21040;
 constexpr int kKeywordModelHeadSize = 672;
-constexpr int kKeywordModelTailSize = 20800;
+constexpr int kKeywordModelTailSize = 20368;
 constexpr int kKeywordModelTfLiteTensorVariableBufferDataSize = 10240;
-constexpr int kKeywordModelTfLiteTensorQuantizationDataSize = 2160;
+constexpr int kKeywordModelTfLiteTensorQuantizationDataSize = 1728;
 constexpr int kKeywordModelOpRuntimeDataSize = 148;
 
 constexpr int kTestConvModelArenaSize = 12 * 1024;
@@ -56,10 +56,10 @@ constexpr int kTestConvModelNodeAndRegistrationCount = 7;
 
 // NOTE: These values are measured on x86-64:
 // TODO(b/158651472): Consider auditing these values on non-64 bit systems.
-constexpr int kTestConvModelTotalSize = 12128;
+constexpr int kTestConvModelTotalSize = 11680;
 constexpr int kTestConvModelHeadSize = 7744;
-constexpr int kTestConvModelTailSize = 4384;
-constexpr int kTestConvModelTfLiteTensorQuantizationDataSize = 1216;
+constexpr int kTestConvModelTailSize = 3936;
+constexpr int kTestConvModelTfLiteTensorQuantizationDataSize = 768;
 constexpr int kTestConvModelOpRuntimeDataSize = 136;
 
 struct ModelAllocationThresholds {
@@ -73,11 +73,17 @@ struct ModelAllocationThresholds {
   size_t op_runtime_data_size = 0;
 };
 
-void EnsureAllocatedSizeThreshold(size_t actual, size_t expected) {
+void EnsureAllocatedSizeThreshold(const char* allocation_type, size_t actual,
+                                  size_t expected) {
   // TODO(b/158651472): Better auditing of non-64 bit systems:
   if (kIs64BitSystem) {
     // 64-bit systems should check floor and ceiling to catch memory savings:
     TF_LITE_MICRO_EXPECT_NEAR(actual, expected, kAllocationThreshold);
+    if (actual != expected) {
+      TF_LITE_REPORT_ERROR(micro_test::reporter,
+                           "%s threshold failed: %ld != %ld", allocation_type,
+                           actual, expected);
+    }
   } else {
     // Non-64 bit systems should just expect allocation does not exceed the
     // ceiling:
@@ -91,33 +97,37 @@ void ValidateModelAllocationThresholds(
   allocator.PrintAllocations();
 
   EnsureAllocatedSizeThreshold(
-      allocator.GetSimpleMemoryAllocator()->GetUsedBytes(),
+      "Total", allocator.GetSimpleMemoryAllocator()->GetUsedBytes(),
       thresholds.total_alloc_size);
   EnsureAllocatedSizeThreshold(
-      allocator.GetSimpleMemoryAllocator()->GetHeadUsedBytes(),
+      "Head", allocator.GetSimpleMemoryAllocator()->GetHeadUsedBytes(),
       thresholds.head_alloc_size);
   EnsureAllocatedSizeThreshold(
-      allocator.GetSimpleMemoryAllocator()->GetTailUsedBytes(),
+      "Tail", allocator.GetSimpleMemoryAllocator()->GetTailUsedBytes(),
       thresholds.tail_alloc_size);
   EnsureAllocatedSizeThreshold(
+      "TfLiteTensor",
       allocator
           .GetRecordedAllocation(
               tflite::RecordedAllocationType::kTfLiteTensorArray)
           .used_bytes,
       sizeof(TfLiteTensor) * thresholds.tensor_count);
   EnsureAllocatedSizeThreshold(
+      "VariableBufferData",
       allocator
           .GetRecordedAllocation(
               tflite::RecordedAllocationType::kTfLiteTensorVariableBufferData)
           .used_bytes,
       thresholds.tensor_variable_buffer_data_size);
   EnsureAllocatedSizeThreshold(
+      "QuantizationData",
       allocator
           .GetRecordedAllocation(tflite::RecordedAllocationType::
                                      kTfLiteTensorArrayQuantizationData)
           .used_bytes,
       thresholds.tensor_quantization_data_size);
   EnsureAllocatedSizeThreshold(
+      "NodeAndRegistration",
       allocator
           .GetRecordedAllocation(
               tflite::RecordedAllocationType::kNodeAndRegistrationArray)
@@ -125,6 +135,7 @@ void ValidateModelAllocationThresholds(
       sizeof(tflite::NodeAndRegistration) *
           thresholds.node_and_registration_count);
   EnsureAllocatedSizeThreshold(
+      "OpData",
       allocator.GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
           .used_bytes,
       thresholds.op_runtime_data_size);
diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc
index bfe44cab73a..f3b64bc9f39 100644
--- a/tensorflow/lite/micro/micro_allocator.cc
+++ b/tensorflow/lite/micro/micro_allocator.cc
@@ -466,6 +466,8 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
   TF_LITE_ENSURE_STATUS(BytesRequiredForTensor(
       flatbuffer_tensor, &result->bytes, &type_size, error_reporter));
 
+  // TODO(b/159043126): Cleanup endian casting by doing all endian casting in
+  // one spot:
   if (flatbuffer_tensor.shape() == nullptr) {
     // flatbuffer_tensor.shape() can return a nullptr in the case of a scalar
     // tensor.
@@ -513,6 +515,10 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
                            "Unable to allocate TfLiteAffineQuantization.\n");
       return kTfLiteError;
     }
+
+    // TODO(b/153688719): Reduce tail allocation by using a global zero-point
+    // buffer. This value can not be reused from the flatbuffer since the
+    // zero_point is stored as a int64_t.
     quantization->zero_point =
         reinterpret_cast(allocator->AllocateFromTail(
             TfLiteIntArrayGetSizeInBytes(channels), alignof(TfLiteIntArray)));
@@ -522,22 +528,14 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
       return kTfLiteError;
     }
 
-    quantization->scale = reinterpret_cast(
-        allocator->AllocateFromTail(TfLiteFloatArrayGetSizeInBytes(channels),
-                                    alignof(TfLiteFloatArray)));
-    if (quantization->scale == nullptr) {
-      TF_LITE_REPORT_ERROR(error_reporter,
-                           "Unable to allocate quantization->scale.\n");
-      return kTfLiteError;
-    }
+    // TODO(b/159043126): Check for big endian before casting flatbuffer values.
+    quantization->scale = const_cast(
+        reinterpret_cast(src_quantization->scale()));
 
     quantization->zero_point->size = channels;
-    quantization->scale->size = channels;
     int* zero_point_data = quantization->zero_point->data;
-    float* scale_data = quantization->scale->data;
     for (int i = 0; i < channels; i++) {
       zero_point_data[i] = src_quantization->zero_point()->Get(i);
-      scale_data[i] = src_quantization->scale()->Get(i);
     }
     // TODO(rocky): Need to add a micro_allocator test case that fails when
     // this is not copied:
@@ -815,8 +813,10 @@ TfLiteStatus MicroAllocator::PrepareNodeAndRegistrationDataFromFlatbuffer(
     }
 
     // Disregard const qualifier to workaround with existing API.
+    // TODO(b/159043126): Check for big endian before casting flatbuffer values.
     TfLiteIntArray* inputs_array = const_cast(
         reinterpret_cast(op->inputs()));
+    // TODO(b/159043126): Check for big endian before casting flatbuffer values.
     TfLiteIntArray* outputs_array = const_cast(
         reinterpret_cast(op->outputs()));
 
diff --git a/tensorflow/lite/micro/recording_micro_allocator_test.cc b/tensorflow/lite/micro/recording_micro_allocator_test.cc
index 9bbe0f405d4..775a2de2dfd 100644
--- a/tensorflow/lite/micro/recording_micro_allocator_test.cc
+++ b/tensorflow/lite/micro/recording_micro_allocator_test.cc
@@ -93,7 +93,6 @@ TF_LITE_MICRO_TEST(TestRecordsTensorArrayQuantizationData) {
       quantized_tensor_count++;
       size_t num_channels = quantization_params->scale()->size();
       quantized_channel_bytes += TfLiteIntArrayGetSizeInBytes(num_channels);
-      quantized_channel_bytes += TfLiteFloatArrayGetSizeInBytes(num_channels);
     }
   }
 
@@ -106,10 +105,9 @@ TF_LITE_MICRO_TEST(TestRecordsTensorArrayQuantizationData) {
       micro_allocator->GetRecordedAllocation(
           tflite::RecordedAllocationType::kTfLiteTensorArrayQuantizationData);
 
-  // Each quantized tensors has 3 mallocs (quant struct, scale dimensions, zero
-  // point dimensions):
+  // Each quantized tensors has 2 mallocs (quant struct, zero point dimensions):
   TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.count,
-                          quantized_tensor_count * 3);
+                          quantized_tensor_count * 2);
   TF_LITE_MICRO_EXPECT_EQ(recorded_allocation.requested_bytes,
                           expected_requested_bytes);
   TF_LITE_MICRO_EXPECT_GE(recorded_allocation.used_bytes,

From 5d4c6e105f775cd60c376a28eb540b2d286c4605 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Mon, 15 Jun 2020 15:23:29 -0700
Subject: [PATCH 0217/1390] Added new ways of weights uploading.

PiperOrigin-RevId: 316556676
Change-Id: I343e4f6461a26a7d921699b23ee3ccf65ecb3bee
---
 .../delegates/gpu/cl/kernels/conv_powervr.cc  | 101 +++++++++++++-----
 .../delegates/gpu/cl/kernels/conv_powervr.h   |   5 +
 2 files changed, 82 insertions(+), 24 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
index 184e070202a..6ab22bf545d 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
@@ -308,6 +308,41 @@ std::string GenerateConv(
       conv_params.weights_upload_type ==
           ConvPowerVR::WeightsUploadType::LOCAL_MEM_ASYNC_SUBGROUP;
 
+  const int local_mem_size =
+      conv_params.block_size.z * 4 * conv_params.src_depth_loop_size;
+
+  const bool use_simd_broadcast =
+      conv_params.weights_upload_type ==
+          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST ||
+      conv_params.weights_upload_type ==
+          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST ||
+      conv_params.weights_upload_type ==
+          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST ||
+      conv_params.weights_upload_type ==
+          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST ||
+      conv_params.weights_upload_type ==
+          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST;
+
+  int simd_size = 1;
+  if (conv_params.weights_upload_type ==
+      ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST) {
+    simd_size = 8;
+  } else if (conv_params.weights_upload_type ==
+             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST) {
+    simd_size = 16;
+  } else if (conv_params.weights_upload_type ==
+             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST) {
+    simd_size = 32;
+  } else if (conv_params.weights_upload_type ==
+             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST) {
+    simd_size = 64;
+  } else if (conv_params.weights_upload_type ==
+             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST) {
+    simd_size = 128;
+  }
+
+  bool late_oob_check = need_local_mem || use_simd_broadcast;
+
   const std::string weights_space =
       conv_params.weights_upload_type ==
               ConvPowerVR::WeightsUploadType::CONSTANT_MEM
@@ -328,6 +363,10 @@ std::string GenerateConv(
          std::to_string(work_group_size.y) + ", " +
          std::to_string(work_group_size.z) + ")))\n";
   }
+  if (use_simd_broadcast && device.IsIntel()) {
+    c += "__attribute__((intel_reqd_work_group_size(" +
+         std::to_string(simd_size) + ")))\n";
+  }
   c += "__kernel void main_function(\n";
   c += src_tensor.GetDeclaration(AccessType::READ) + ",\n";
   c += "    " + weights_global_ptr + " filters_buffer,    \n";
@@ -355,7 +394,7 @@ std::string GenerateConv(
   for (int y = 0; y < conv_params.block_size.y; ++y) {
     dst_y[y] = "(Y + " + std::to_string(y) + ")";
   }
-  if (!need_local_mem) {
+  if (!late_oob_check) {
     c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) {\n";
     c += "    return;\n";
     c += "  }\n";
@@ -396,13 +435,8 @@ std::string GenerateConv(
   }
   if (need_local_mem) {
     c += "  __local " + weights_data_type + " weights_cache[" +
-         std::to_string(block_size.z * 4 * conv_params.src_depth_loop_size) +
-         "];\n";
-  }
-  if (conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::GLOBAL_MEM ||
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::CONSTANT_MEM) {
+         std::to_string(local_mem_size) + "];\n";
+  } else {
     c += "    " + weights_global_ptr + " weights_cache;\n";
   }
   if (is1x1) {
@@ -521,9 +555,17 @@ std::string GenerateConv(
           for (int y = 0; y < block_size.y; ++y) {
             for (int x = 0; x < block_size.x; ++x) {
               std::string id = std::to_string(y) + std::to_string(x);
-              c += "    r" + std::to_string(z) + id + " += weights_cache[" +
-                   std::to_string(z * 4 + ch + shared_offset) + "] * src" + id +
-                   "." + channels[ch] + ";\n";
+              std::string w_val = "weights_cache[" +
+                                  std::to_string(z * 4 + ch + shared_offset) +
+                                  "]";
+              if (use_simd_broadcast) {
+                int simd_id = (z * 4 + ch + shared_offset) / simd_size;
+                int thread_id = (z * 4 + ch + shared_offset) % simd_size;
+                w_val = "sub_group_broadcast(simd_w" + std::to_string(simd_id) +
+                        ", " + std::to_string(thread_id) + "u)";
+              }
+              c += "    r" + std::to_string(z) + id + " += " + w_val +
+                   " * src" + id + "." + channels[ch] + ";\n";
             }
           }
         }
@@ -554,17 +596,30 @@ std::string GenerateConv(
       work_group_size.x * work_group_size.y * work_group_size.z;
   if (conv_params.weights_upload_type ==
       ConvPowerVR::WeightsUploadType::LOCAL_MEM_ASYNC_SUBGROUP) {
-    c +=
-        GenerateAsyncUpload("weights_cache", "filters_loc",
-                            /*global_offset_name*/ "",
-                            block_size.z * 4 * conv_params.src_depth_loop_size);
+    c += GenerateAsyncUpload("weights_cache", "filters_loc",
+                             /*global_offset_name*/ "", local_mem_size);
   } else if (conv_params.weights_upload_type ==
              ConvPowerVR::WeightsUploadType::LOCAL_MEM_BY_THREADS) {
     c += "    barrier(CLK_LOCAL_MEM_FENCE);\n";
-    c += GenerateUploadByThreads(
-        "weights_cache", "filters_loc",
-        /*global_offset_name*/ "", "lid", total_work_items,
-        block_size.z * 4 * conv_params.src_depth_loop_size);
+    c += GenerateUploadByThreads("weights_cache", "filters_loc",
+                                 /*global_offset_name*/ "", "lid",
+                                 total_work_items, local_mem_size);
+  } else if (use_simd_broadcast) {
+    int parts = local_mem_size / simd_size;
+    int reminder = local_mem_size % simd_size;
+    for (int i = 0; i < parts; ++i) {
+      c += "    FLT4 simd_w" + std::to_string(i) +
+           " = filters_loc[get_sub_group_local_id() + " +
+           std::to_string(i * simd_size) + "];\n";
+    }
+    if (reminder) {
+      c += "    FLT4 simd_w" + std::to_string(parts) + ";\n";
+      c += "    if (simd_id < " + std::to_string(reminder) + ") {\n";
+      c += "      simd_w" + std::to_string(parts) +
+           " = filters_loc[get_sub_group_local_id() + " +
+           std::to_string(parts * simd_size) + "];\n";
+      c += "    }\n";
+    }
   } else {  // GLOBAL_MEM/CONSTANT_MEM
     c += "    weights_cache = filters_loc;\n";
   }
@@ -580,9 +635,7 @@ std::string GenerateConv(
     conv_core(i * block_size.z * 4);
     c += "    s += 1;\n";
   }
-  c += "    filters_loc += " +
-       std::to_string(block_size.z * 4 * conv_params.src_depth_loop_size) +
-       ";\n";
+  c += "    filters_loc += " + std::to_string(local_mem_size) + ";\n";
   c += "  } while (s < src_size.z);\n";
   if (!is1x1) {
     c += "  };\n";
@@ -597,10 +650,10 @@ std::string GenerateConv(
     c += GenerateUploadByThreads("weights_cache", "biases", "Z", "lid",
                                  total_work_items, block_size.z);
     c += "    barrier(CLK_LOCAL_MEM_FENCE);\n";
-  } else {  // GLOBAL_MEM/CONSTANT_MEM
+  } else {
     c += "    weights_cache = biases + Z;\n";
   }
-  if (need_local_mem) {
+  if (late_oob_check) {
     c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) {\n";
     c += "    return;\n";
     c += "  }\n";
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
index 5eff4b36053..a729098bded 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
@@ -64,6 +64,11 @@ class ConvPowerVR : public GPUOperation {
     LOCAL_MEM_BY_THREADS,
     GLOBAL_MEM,
     CONSTANT_MEM,
+    PRIVATE_MEM_SIMD8_BROADCAST,
+    PRIVATE_MEM_SIMD16_BROADCAST,
+    PRIVATE_MEM_SIMD32_BROADCAST,
+    PRIVATE_MEM_SIMD64_BROADCAST,
+    PRIVATE_MEM_SIMD128_BROADCAST,
   };
 
   struct ConvParams {

From 67fb07ba9fe4587d65ed2cfe83a364797f9c97b4 Mon Sep 17 00:00:00 2001
From: Ken Franko 
Date: Mon, 15 Jun 2020 15:28:17 -0700
Subject: [PATCH 0218/1390] Add load option for loading SavedModel from
 specific io_device for distributed training.

A new class LoadOptions is created similar to the existing SavedOptions.  The option experimental_io_device is the only option added at this time and usd to set the io_device when loading a SavedModel for distributed training.

PiperOrigin-RevId: 316557681
Change-Id: If3f1eae18b09085ff11dc8a6882fabcb18f5f48e
---
 tensorflow/python/keras/engine/training.py    | 10 +++-
 tensorflow/python/keras/saving/save.py        |  6 +-
 .../python/keras/saving/saved_model/load.py   |  9 ++-
 .../python/keras/saving/saved_model/save.py   |  2 +-
 tensorflow/python/saved_model/BUILD           |  8 +++
 tensorflow/python/saved_model/load.py         | 28 ++++++---
 tensorflow/python/saved_model/load_options.py | 57 +++++++++++++++++++
 tensorflow/python/saved_model/load_test.py    |  7 +++
 .../golden/v1/tensorflow.keras.-model.pbtxt   |  2 +-
 .../v1/tensorflow.keras.-sequential.pbtxt     |  2 +-
 ...low.keras.experimental.-linear-model.pbtxt |  2 +-
 ....keras.experimental.-wide-deep-model.pbtxt |  2 +-
 .../v1/tensorflow.keras.models.-model.pbtxt   |  2 +-
 .../tensorflow.keras.models.-sequential.pbtxt |  2 +-
 .../golden/v1/tensorflow.keras.models.pbtxt   |  2 +-
 .../golden/v1/tensorflow.saved_model.pbtxt    |  2 +-
 .../golden/v2/tensorflow.keras.-model.pbtxt   |  2 +-
 .../v2/tensorflow.keras.-sequential.pbtxt     |  2 +-
 ...low.keras.experimental.-linear-model.pbtxt |  2 +-
 ....keras.experimental.-wide-deep-model.pbtxt |  2 +-
 .../v2/tensorflow.keras.models.-model.pbtxt   |  2 +-
 .../tensorflow.keras.models.-sequential.pbtxt |  2 +-
 .../golden/v2/tensorflow.keras.models.pbtxt   |  2 +-
 ...tensorflow.saved_model.-load-options.pbtxt | 13 +++++
 .../golden/v2/tensorflow.saved_model.pbtxt    |  6 +-
 25 files changed, 144 insertions(+), 32 deletions(-)
 create mode 100644 tensorflow/python/saved_model/load_options.py
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.saved_model.-load-options.pbtxt

diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index 87782adff46..b7a4795d768 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -2078,7 +2078,11 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
           save_relative_paths=True,
           all_model_checkpoint_paths=[filepath])
 
-  def load_weights(self, filepath, by_name=False, skip_mismatch=False):
+  def load_weights(self,
+                   filepath,
+                   by_name=False,
+                   skip_mismatch=False,
+                   options=None):
     """Loads all layer weights, either from a TensorFlow or an HDF5 weight file.
 
     If `by_name` is False weights are loaded based on the network's
@@ -2108,6 +2112,8 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         skip_mismatch: Boolean, whether to skip loading of layers where there is
             a mismatch in the number of weights, or a mismatch in the shape of
             the weight (only valid when `by_name=True`).
+        options: Optional `tf.train.CheckpointOptions` object that specifies
+            options for loading weights.
 
     Returns:
         When loading a weight file in TensorFlow format, returns the same status
@@ -2145,7 +2151,7 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         # The checkpoint is not readable in TensorFlow format. Try HDF5.
         save_format = 'h5'
     if save_format == 'tf':
-      status = self._trackable_saver.restore(filepath)
+      status = self._trackable_saver.restore(filepath, options)
       if by_name:
         raise NotImplementedError(
             'Weights may only be loaded based on topology into Models when '
diff --git a/tensorflow/python/keras/saving/save.py b/tensorflow/python/keras/saving/save.py
index 7f725d3978e..9c83914d380 100644
--- a/tensorflow/python/keras/saving/save.py
+++ b/tensorflow/python/keras/saving/save.py
@@ -135,7 +135,7 @@ def save_model(model,
 
 
 @keras_export('keras.models.load_model')
-def load_model(filepath, custom_objects=None, compile=True):  # pylint: disable=redefined-builtin
+def load_model(filepath, custom_objects=None, compile=True, options=None):  # pylint: disable=redefined-builtin
   """Loads a model saved via `model.save()`.
 
   Usage:
@@ -162,6 +162,8 @@ def load_model(filepath, custom_objects=None, compile=True):  # pylint: disable=
           considered during deserialization.
       compile: Boolean, whether to compile the model
           after loading.
+      options: Optional `tf.saved_model.LoadOptions` object that specifies
+        options for loading from SavedModel.
 
   Returns:
       A Keras model instance. If the original model was compiled, and saved with
@@ -182,7 +184,7 @@ def load_model(filepath, custom_objects=None, compile=True):  # pylint: disable=
     filepath = path_to_string(filepath)
     if isinstance(filepath, six.string_types):
       loader_impl.parse_saved_model(filepath)
-      return saved_model_load.load(filepath, compile)
+      return saved_model_load.load(filepath, compile, options)
 
   raise IOError(
       'Unable to load model. Filepath is not an hdf5 file (or h5py is not '
diff --git a/tensorflow/python/keras/saving/saved_model/load.py b/tensorflow/python/keras/saving/saved_model/load.py
index ca8164c9407..7e67bf6305c 100644
--- a/tensorflow/python/keras/saving/saved_model/load.py
+++ b/tensorflow/python/keras/saving/saved_model/load.py
@@ -90,7 +90,7 @@ KERAS_OBJECT_IDENTIFIERS = (
     '_tf_keras_rnn_layer')
 
 
-def load(path, compile=True):  # pylint: disable=redefined-builtin
+def load(path, compile=True, options=None):  # pylint: disable=redefined-builtin
   """Loads Keras objects from a SavedModel.
 
   Any Keras layer or model saved to the SavedModel will be loaded back
@@ -107,13 +107,18 @@ def load(path, compile=True):  # pylint: disable=redefined-builtin
   Args:
     path: Path to SavedModel.
     compile: If true, compile the model after loading it.
+    options: Optional `tf.saved_model.LoadOptions` object that specifies
+      options for loading from SavedModel.
+
 
   Returns:
     Object loaded from SavedModel.
   """
   # TODO(kathywu): Add saving/loading of optimizer, compiled losses and metrics.
   # TODO(kathywu): Add code to load from objects that contain all endpoints
-  model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
+
+  model = tf_load.load_internal(
+      path, options=options, loader_cls=KerasObjectLoader)
 
   # pylint: disable=protected-access
   if isinstance(model, training_lib.Model) and compile:
diff --git a/tensorflow/python/keras/saving/saved_model/save.py b/tensorflow/python/keras/saving/saved_model/save.py
index 9d4ca5e2c59..7d6bc120758 100644
--- a/tensorflow/python/keras/saving/saved_model/save.py
+++ b/tensorflow/python/keras/saving/saved_model/save.py
@@ -49,7 +49,7 @@ def save(model, filepath, overwrite, include_optimizer, signatures=None,
     signatures: Signatures to save with the SavedModel. Applicable to the 'tf'
       format only. Please see the `signatures` argument in `tf.saved_model.save`
       for details.
-    options: Optional`tf.saved_model.SaveOptions` object that specifies
+    options: Optional `tf.saved_model.SaveOptions` object that specifies
       options for saving to SavedModel.
 
   Raises:
diff --git a/tensorflow/python/saved_model/BUILD b/tensorflow/python/saved_model/BUILD
index 6e17b8af206..240b60f43f6 100644
--- a/tensorflow/python/saved_model/BUILD
+++ b/tensorflow/python/saved_model/BUILD
@@ -348,6 +348,7 @@ py_library(
     deps = [
         ":constants",
         ":function_deserialization",
+        ":load_options",
         ":load_v1_in_v2",
         ":loader",
         ":nested_structure_coder",
@@ -522,6 +523,13 @@ py_library(
     ],
 )
 
+py_library(
+    name = "load_options",
+    srcs = ["load_options.py"],
+    deps = [
+    ],
+)
+
 py_library(
     name = "method_name_updater",
     srcs = ["method_name_updater.py"],
diff --git a/tensorflow/python/saved_model/load.py b/tensorflow/python/saved_model/load.py
index fe2919c88dc..74b030a3797 100644
--- a/tensorflow/python/saved_model/load.py
+++ b/tensorflow/python/saved_model/load.py
@@ -37,11 +37,13 @@ from tensorflow.python.ops import lookup_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.saved_model import function_deserialization
+from tensorflow.python.saved_model import load_options
 from tensorflow.python.saved_model import load_v1_in_v2
 from tensorflow.python.saved_model import loader_impl
 from tensorflow.python.saved_model import nested_structure_coder
 from tensorflow.python.saved_model import revived_types
 from tensorflow.python.saved_model import utils_impl as saved_model_utils
+from tensorflow.python.training.saving import checkpoint_options
 from tensorflow.python.training.tracking import base
 from tensorflow.python.training.tracking import graph_view
 from tensorflow.python.training.tracking import tracking
@@ -105,7 +107,8 @@ class _WrapperFunction(function.ConcreteFunction):
 class Loader(object):
   """Helper class to load an object-based SavedModel."""
 
-  def __init__(self, object_graph_proto, saved_model_proto, export_dir):
+  def __init__(self, object_graph_proto, saved_model_proto, export_dir,
+               ckpt_options):
     meta_graph = saved_model_proto.meta_graphs[0]
     self._asset_file_def = meta_graph.asset_file_def
     self._operation_attributes = {
@@ -115,6 +118,7 @@ class Loader(object):
     self._concrete_functions = (
         function_deserialization.load_function_def_library(
             meta_graph.graph_def.library))
+    self._checkpoint_options = ckpt_options
 
     for name, concrete_function in self._concrete_functions.items():
       # Wrap all the concrete function so that they are capable of dealing with
@@ -306,9 +310,10 @@ class Loader(object):
     with ops.device("CPU"):
       saver._file_prefix_placeholder = constant_op.constant(variables_path)
     if self._expect_partial_checkpoint:
-      load_status = saver.restore(variables_path).expect_partial()
+      load_status = saver.restore(variables_path,
+                                  self._checkpoint_options).expect_partial()
     else:
-      load_status = saver.restore(variables_path)
+      load_status = saver.restore(variables_path, self._checkpoint_options)
     load_status.assert_existing_objects_matched()
     checkpoint = load_status._checkpoint
 
@@ -491,7 +496,7 @@ def _call_attribute(instance, *args, **kwargs):
 
 
 @tf_export("saved_model.load", v1=["saved_model.load_v2"])
-def load(export_dir, tags=None):
+def load(export_dir, tags=None, options=None):
   """Load a SavedModel from `export_dir`.
 
   Signatures associated with the SavedModel are available as functions:
@@ -569,6 +574,8 @@ def load(export_dir, tags=None):
     tags: A tag or sequence of tags identifying the MetaGraph to load. Optional
       if the SavedModel contains a single MetaGraph, as for those exported from
       `tf.saved_model.save`.
+    options: Optional, `tf.saved_model.LoadOptions` object that specifies
+      options for loading.
 
   Returns:
     A trackable object with a `signatures` attribute mapping from signature
@@ -579,11 +586,12 @@ def load(export_dir, tags=None):
   Raises:
     ValueError: If `tags` don't match a MetaGraph in the SavedModel.
   """
-  return load_internal(export_dir, tags)
+  return load_internal(export_dir, tags, options)
 
 
-def load_internal(export_dir, tags=None, loader_cls=Loader):
+def load_internal(export_dir, tags=None, options=None, loader_cls=Loader):
   """Loader implementation."""
+  options = options or load_options.LoadOptions()
   if tags is not None and not isinstance(tags, set):
     # Supports e.g. tags=SERVING and tags=[SERVING]. Sets aren't considered
     # sequences for nest.flatten, so we put those through as-is.
@@ -602,10 +610,12 @@ def load_internal(export_dir, tags=None, loader_cls=Loader):
            "it, pass 'None', or pass matching tags.")
           .format(export_dir, meta_graph_def.meta_info_def.tags, tags))
     object_graph_proto = meta_graph_def.object_graph_def
+
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_io_device=options.experimental_io_device)
     with ops.init_scope():
-      loader = loader_cls(object_graph_proto,
-                          saved_model_proto,
-                          export_dir)
+      loader = loader_cls(object_graph_proto, saved_model_proto, export_dir,
+                          ckpt_options)
       root = loader.get(0)
       if isinstance(loader, Loader):
         root.graph_debug_info = loader.adjust_debug_info_func_names(debug_info)
diff --git a/tensorflow/python/saved_model/load_options.py b/tensorflow/python/saved_model/load_options.py
new file mode 100644
index 00000000000..9718d8ffed9
--- /dev/null
+++ b/tensorflow/python/saved_model/load_options.py
@@ -0,0 +1,57 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Options for saving SavedModels."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("saved_model.LoadOptions", v1=[])
+class LoadOptions(object):
+  """Options for loading a SavedModel.
+
+  This function may be used in the `options` argument in functions that
+  load a SavedModel (`tf.saved_model.load`, `tf.keras.models.load_model`).
+  """
+
+  # Define object attributes in __slots__ for improved memory and performance.
+  __slots__ = ("experimental_io_device",)
+
+  def __init__(self,
+               experimental_io_device=None):
+    """Creates an object that stores options for SavedModel loading.
+
+    Args:
+      experimental_io_device: string. Applies in a distributed setting.
+        Tensorflow device to use to access the filesystem. If `None` (default)
+        then for each variable the filesystem is accessed from the CPU:0 device
+        of the host where that variable is assigned. If specified, the
+        filesystem is instead accessed from that device for all variables.
+        This is for example useful if you want to load from a local directory,
+        such as "/tmp" when running in a distributed setting. In that case
+        pass a device for the host where the "/tmp" directory is accessible.
+
+    Example:
+
+      load_options = tf.saved_model.LoadOptions(experimental_io_device=
+        '/job:localhost')
+      restoredmodel = tf.keras.models.load_model(saved_model_path,
+                                                 options=load_options)
+
+    """
+    self.experimental_io_device = experimental_io_device
diff --git a/tensorflow/python/saved_model/load_test.py b/tensorflow/python/saved_model/load_test.py
index 7bd2e87c739..5449cc1c9a2 100644
--- a/tensorflow/python/saved_model/load_test.py
+++ b/tensorflow/python/saved_model/load_test.py
@@ -56,6 +56,7 @@ from tensorflow.python.ops import variables
 from tensorflow.python.ops.ragged import ragged_factory_ops
 from tensorflow.python.ops.ragged import ragged_tensor
 from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import load_options
 from tensorflow.python.saved_model import save
 from tensorflow.python.saved_model import tag_constants
 from tensorflow.python.training import monitored_session
@@ -1788,6 +1789,12 @@ class LoadTest(test.TestCase, parameterized.TestCase):
     self.assertAllEqual(imported2.f(rt, 2), [[3, 4], [5]])
     self.assertAllEqual(imported2.f(rt, 3), [[4, 5], [6]])
 
+  def test_accepts_io_device(self, cycles):
+    options = load_options.LoadOptions()
+    self.assertIsNone(options.experimental_io_device)
+    options = load_options.LoadOptions(experimental_io_device="/job:localhost")
+    self.assertEqual("/job:localhost", options.experimental_io_device)
+
 
 class SingleCycleTests(test.TestCase, parameterized.TestCase):
 
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
index ea2945b5bf6..b62814e81cb 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
@@ -258,7 +258,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
index c1dea9335c0..7485a0b3c62 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
@@ -264,7 +264,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
index ba87c0a2a7a..bf980e5d116 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
@@ -259,7 +259,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
index 37fb2051f81..c214a5c3419 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
@@ -259,7 +259,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
index 7439fc2dd6b..86868c9d17f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
@@ -258,7 +258,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
index 24e8bf57611..05aa19a915a 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
@@ -264,7 +264,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.pbtxt
index 6f85b3c2150..ac80126aaa3 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.pbtxt
@@ -14,7 +14,7 @@ tf_module {
   }
   member_method {
     name: "load_model"
-    argspec: "args=[\'filepath\', \'custom_objects\', \'compile\'], varargs=None, keywords=None, defaults=[\'None\', \'True\'], "
+    argspec: "args=[\'filepath\', \'custom_objects\', \'compile\', \'options\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "model_from_config"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
index 8833f02b0db..2dde9c495cc 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt
@@ -182,7 +182,7 @@ tf_module {
   }
   member_method {
     name: "load_v2"
-    argspec: "args=[\'export_dir\', \'tags\'], varargs=None, keywords=None, defaults=[\'None\'], "
+    argspec: "args=[\'export_dir\', \'tags\', \'options\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
   }
   member_method {
     name: "main_op_with_restore"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
index ea2945b5bf6..b62814e81cb 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
@@ -258,7 +258,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
index c1dea9335c0..7485a0b3c62 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
@@ -264,7 +264,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
index ba87c0a2a7a..bf980e5d116 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
@@ -259,7 +259,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
index 37fb2051f81..c214a5c3419 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
@@ -259,7 +259,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
index 7439fc2dd6b..86868c9d17f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
@@ -258,7 +258,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
index 24e8bf57611..05aa19a915a 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
@@ -264,7 +264,7 @@ tf_class {
   }
   member_method {
     name: "load_weights"
-    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\'], varargs=None, keywords=None, defaults=[\'False\', \'False\'], "
+    argspec: "args=[\'self\', \'filepath\', \'by_name\', \'skip_mismatch\', \'options\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'None\'], "
   }
   member_method {
     name: "make_predict_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.pbtxt
index 6f85b3c2150..ac80126aaa3 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.pbtxt
@@ -14,7 +14,7 @@ tf_module {
   }
   member_method {
     name: "load_model"
-    argspec: "args=[\'filepath\', \'custom_objects\', \'compile\'], varargs=None, keywords=None, defaults=[\'None\', \'True\'], "
+    argspec: "args=[\'filepath\', \'custom_objects\', \'compile\', \'options\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'None\'], "
   }
   member_method {
     name: "model_from_config"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-load-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-load-options.pbtxt
new file mode 100644
index 00000000000..20216d93c3d
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-load-options.pbtxt
@@ -0,0 +1,13 @@
+path: "tensorflow.saved_model.LoadOptions"
+tf_class {
+  is_instance: ""
+  is_instance: ""
+  member {
+    name: "experimental_io_device"
+    mtype: ""
+  }
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'experimental_io_device\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
index 0a82cfd0873..0adfbd30102 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt
@@ -44,6 +44,10 @@ tf_module {
     name: "GPU"
     mtype: ""
   }
+  member {
+    name: "LoadOptions"
+    mtype: ""
+  }
   member {
     name: "PREDICT_INPUTS"
     mtype: ""
@@ -110,7 +114,7 @@ tf_module {
   }
   member_method {
     name: "load"
-    argspec: "args=[\'export_dir\', \'tags\'], varargs=None, keywords=None, defaults=[\'None\'], "
+    argspec: "args=[\'export_dir\', \'tags\', \'options\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
   }
   member_method {
     name: "save"

From 1158611838055b81804e3e208334544048a16366 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Mon, 15 Jun 2020 15:42:09 -0700
Subject: [PATCH 0219/1390] LSTM: do projection to output_state instead of
 output. Because the two arrays are only different in stride (state has no
 stride), this allows us to do the projection in a batched manner.

Copy the result to the strided output after projection.

PiperOrigin-RevId: 316560275
Change-Id: I60c544d10a64437ece1fa75eea891af4b97df231
---
 tensorflow/lite/kernels/lstm_eval.cc          | 85 +++++++------------
 .../calibration/builtin_logging_ops/lstm.cc   | 43 ++++------
 2 files changed, 46 insertions(+), 82 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 0a2c381ebf1..7fa3d85687c 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -391,40 +391,29 @@ inline void LstmStepFloat(
   const bool use_projection_weight = (projection_weights_ptr != nullptr);
   const bool use_projection_bias = (projection_bias_ptr != nullptr);
 
-  // For each batch: update the projection and output_state. Note that since
-  // the output batch rows may not be contiguous (output_batch_leading_dim !=
-  // n_output), we unroll batched operations.
+  // For each batch: update output_state.
   if (use_projection_weight) {
     if (use_projection_bias) {
-      for (int b = 0; b < n_batch; b++) {
-        std::copy_n(projection_bias_ptr, n_output,
-                    output_ptr + b * output_batch_leading_dim);
-      }
+      tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output,
+                                            n_batch, output_state_ptr);
     } else {
-      for (int b = 0; b < n_batch; b++) {
-        std::fill_n(output_ptr + b * output_batch_leading_dim, n_output, 0.0f);
-      }
+      std::fill_n(output_state_ptr, n_batch * n_output, 0.0f);
     }
-    for (int b = 0; b < n_batch; b++) {
-      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-          projection_weights_ptr, n_output, n_cell,
-          output_gate_scratch + b * n_cell,
-          /*n_batch=*/1, output_ptr + b * output_batch_leading_dim);
-      if (params->proj_clip > 0.0) {
-        tensor_utils::ClipVector(output_ptr + b * output_batch_leading_dim,
-                                 n_output, params->proj_clip,
-                                 output_ptr + b * output_batch_leading_dim);
-      }
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        projection_weights_ptr, n_output, n_cell, output_gate_scratch, n_batch,
+        output_state_ptr);
+    if (params->proj_clip > 0.0) {
+      tensor_utils::ClipVector(output_state_ptr, n_batch * n_output,
+                               params->proj_clip, output_state_ptr);
     }
   } else {
-    for (int b = 0; b < n_batch; b++) {
-      std::copy_n(output_gate_scratch + b * n_output, n_output,
-                  output_ptr + b * output_batch_leading_dim);
-    }
+    std::copy_n(output_gate_scratch, n_batch * n_output, output_state_ptr);
   }
+  // Copy output_state to the output. Note that the output batch rows may not be
+  // contiguous (output_batch_leading_dim != n_output).
   for (int b = 0; b < n_batch; b++) {
-    std::copy_n(output_ptr + b * output_batch_leading_dim, n_output,
-                output_state_ptr + b * n_output);
+    std::copy_n(output_state_ptr + b * n_output, n_output,
+                output_ptr + b * output_batch_leading_dim);
   }
 }
 // LINT.ThenChange(//tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc)
@@ -863,14 +852,10 @@ inline void LstmStepHybrid(
   // n_output), we unroll the batched operations.
   if (use_projection_weight) {
     if (use_projection_bias) {
-      for (int b = 0; b < n_batch; b++) {
-        std::copy_n(projection_bias_ptr, n_output,
-                    output_ptr + b * output_batch_leading_dim);
-      }
+      tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output,
+                                            n_batch, output_state_ptr);
     } else {
-      for (int b = 0; b < n_batch; b++) {
-        std::fill_n(output_ptr + b * output_batch_leading_dim, n_output, 0.0f);
-      }
+      std::fill_n(output_state_ptr, n_batch * n_output, 0.0f);
     }
     if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) {
       // Save quantization and matmul computation for all zero input.
@@ -881,35 +866,25 @@ inline void LstmStepHybrid(
         scaling_factors_scratch[b] =
             scaling_factors[b] * projection_weights_scale;
       }
-      for (int b = 0; b < n_batch; b++) {
-        tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-            projection_weights_ptr, n_output, n_cell,
-            quantized_cell_state_ptr + b * n_cell, &scaling_factors_scratch[b],
-            /*n_batch=*/1, output_ptr + b * output_batch_leading_dim,
-            /*per_channel_scale=*/nullptr,
-            asymmetric_quantize_inputs ? &zero_points[b] : nullptr,
-            accum_scratch_ptr, projection_weights_row_sums, compute_row_sums,
-            context);
-      }
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          projection_weights_ptr, n_output, n_cell, quantized_cell_state_ptr,
+          scaling_factors_scratch, n_batch, output_state_ptr,
+          /*per_channel_scale=*/nullptr,
+          asymmetric_quantize_inputs ? zero_points : nullptr, accum_scratch_ptr,
+          projection_weights_row_sums, compute_row_sums, context);
     }
     if (params->proj_clip > 0.0) {
-      for (int b = 0; b < n_batch; b++) {
-        tensor_utils::ClipVector(output_ptr + b * output_batch_leading_dim,
-                                 n_output, params->proj_clip,
-                                 output_ptr + b * output_batch_leading_dim);
-      }
+      tensor_utils::ClipVector(output_state_ptr, n_batch * n_output,
+                               params->proj_clip, output_state_ptr);
     }
   } else {
-    for (int b = 0; b < n_batch; b++) {
-      std::copy_n(output_gate_scratch + b * n_output, n_output,
-                  output_ptr + b * output_batch_leading_dim);
-    }
+    std::copy_n(output_gate_scratch, n_batch * n_output, output_state_ptr);
   }
   for (int b = 0; b < n_batch; b++) {
-    std::copy_n(output_ptr + b * output_batch_leading_dim, n_output,
-                output_state_ptr + b * n_output);
+    std::copy_n(output_state_ptr + b * n_output, n_output,
+                output_ptr + b * output_batch_leading_dim);
   }
-}
+}  // namespace
 
 // Fully quantized lstm kernel for 16 bit gate matmul output.
 //
diff --git a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
index 41a03f16d63..b58900c0bc6 100644
--- a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
+++ b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
@@ -249,40 +249,29 @@ inline void LstmStepWithAuxInput(
   const bool use_projection_weight = (projection_weights_ptr != nullptr);
   const bool use_projection_bias = (projection_bias_ptr != nullptr);
 
-  // For each batch: update the projection and output_state. Note that since
-  // the output batch rows may not be contiguous (output_batch_leading_dim !=
-  // n_output), we unroll batched operations.
+  // For each batch: update output_state.
   if (use_projection_weight) {
     if (use_projection_bias) {
-      for (int k = 0; k < n_batch; k++) {
-        std::copy_n(projection_bias_ptr, n_output,
-                    output_ptr + k * output_batch_leading_dim);
-      }
+      tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output,
+                                            n_batch, output_state_ptr);
     } else {
-      for (int k = 0; k < n_batch; k++) {
-        std::fill_n(output_ptr + k * output_batch_leading_dim, n_output, 0.0f);
-      }
+      std::fill_n(output_state_ptr, n_batch * n_output, 0.0f);
     }
-    for (int k = 0; k < n_batch; k++) {
-      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-          projection_weights_ptr, n_output, n_cell,
-          output_gate_scratch + k * n_cell,
-          /*n_batch=*/1, output_ptr + k * output_batch_leading_dim);
-      if (params->proj_clip > 0.0) {
-        tensor_utils::ClipVector(output_ptr + k * output_batch_leading_dim,
-                                 n_output, params->proj_clip,
-                                 output_ptr + k * output_batch_leading_dim);
-      }
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        projection_weights_ptr, n_output, n_cell, output_gate_scratch, n_batch,
+        output_state_ptr);
+    if (params->proj_clip > 0.0) {
+      tensor_utils::ClipVector(output_state_ptr, n_batch * n_output,
+                               params->proj_clip, output_state_ptr);
     }
   } else {
-    for (int k = 0; k < n_batch; k++) {
-      std::copy_n(output_gate_scratch + k * n_output, n_output,
-                  output_ptr + k * output_batch_leading_dim);
-    }
+    std::copy_n(output_gate_scratch, n_batch * n_output, output_state_ptr);
   }
-  for (int k = 0; k < n_batch; k++) {
-    std::copy_n(output_ptr + k * output_batch_leading_dim, n_output,
-                output_state_ptr + k * n_output);
+  // Copy output_state to the output. Note that the output batch rows may not be
+  // contiguous (output_batch_leading_dim != n_output).
+  for (int b = 0; b < n_batch; b++) {
+    std::copy_n(output_state_ptr + b * n_output, n_output,
+                output_ptr + b * output_batch_leading_dim);
   }
 }
 

From 7292433984f91ddaec03cb07fc3749b781199984 Mon Sep 17 00:00:00 2001
From: Jaesung Chung 
Date: Mon, 15 Jun 2020 15:57:32 -0700
Subject: [PATCH 0220/1390] Support flex ops in calibration optimization

This CL makes the tool generate a user-friendly error message as well.

In order to use the correct logger for mobile, it uses the error_reporter.

PiperOrigin-RevId: 316563081
Change-Id: Ib56f80330087750777725ed6ad3c97f54b1fa80b
---
 .../lite/tools/optimize/calibration/BUILD     |  2 +
 .../tools/optimize/calibration/calibrator.cc  |  4 +-
 .../calibration/logging_op_resolver.cc        | 33 ++++++++++-
 .../calibration/logging_op_resolver.h         |  3 +-
 .../calibration/logging_op_resolver_test.cc   | 55 +++++++++++++++++--
 5 files changed, 89 insertions(+), 8 deletions(-)

diff --git a/tensorflow/lite/tools/optimize/calibration/BUILD b/tensorflow/lite/tools/optimize/calibration/BUILD
index a394156786f..11d00efe103 100644
--- a/tensorflow/lite/tools/optimize/calibration/BUILD
+++ b/tensorflow/lite/tools/optimize/calibration/BUILD
@@ -89,6 +89,8 @@ cc_library(
     deps = [
         ":calibration_common",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite:util",
         "//tensorflow/lite/core/api",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator.cc b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
index 579fb8cd52f..fb1677fda99 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
@@ -378,8 +378,8 @@ TfLiteStatus BuildLoggingInterpreter(
   // Prepare the logging op resolver to use |LoggingEval| for kernel
   // invocations.
   auto logging_op_resolver = absl::make_unique(
-      builtin_op_and_versions, custom_op_and_versions, op_resolver,
-      LoggingEval);
+      builtin_op_and_versions, custom_op_and_versions, op_resolver, LoggingEval,
+      error_reporter);
   tflite::InterpreterBuilder(tflite_model, *logging_op_resolver,
                              error_reporter)(interpreter);
 
diff --git a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.cc b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.cc
index 634b2a76a3a..92601b2a459 100644
--- a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.cc
+++ b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.cc
@@ -15,6 +15,10 @@ limitations under the License.
 #include "tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h"
 
 #include "absl/memory/memory.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/util.h"
 
 namespace tflite {
 namespace optimize {
@@ -23,10 +27,18 @@ namespace calibration {
 LoggingOpResolver::LoggingOpResolver(
     const BuiltinOpsSet& builtin_ops_to_replace,
     const CustomOpsSet& custom_ops_to_replace, const OpResolver& base_resolver,
-    KernelEvalFuncPtr logging_eval_fn) {
+    KernelEvalFuncPtr logging_eval_fn, ErrorReporter* error_reporter) {
+  std::vector unresolved_builtin_ops;
+  std::vector unresolved_custom_ops;
+
   for (const auto& op_and_version : builtin_ops_to_replace) {
     const TfLiteRegistration* base_registration =
         base_resolver.FindOp(op_and_version.first, op_and_version.second);
+    if (!base_registration) {
+      unresolved_builtin_ops.push_back(
+          EnumNameBuiltinOperator(op_and_version.first));
+      continue;
+    }
     BuiltinOperatorKey key = op_and_version;
     builtin_op_evalfn_map_[key] = base_registration->invoke;
     auto logging_registration =
@@ -37,6 +49,11 @@ LoggingOpResolver::LoggingOpResolver(
   for (const auto& op_and_version : custom_ops_to_replace) {
     const TfLiteRegistration* base_registration = base_resolver.FindOp(
         op_and_version.first.c_str(), op_and_version.second);
+    if (!base_registration) {
+      if (!IsFlexOp(op_and_version.first.c_str()))
+        unresolved_custom_ops.push_back(op_and_version.first.c_str());
+      continue;
+    }
     CustomOperatorKey key = op_and_version;
     custom_op_evalfn_map_[key] = base_registration->invoke;
     auto logging_registration =
@@ -44,6 +61,20 @@ LoggingOpResolver::LoggingOpResolver(
     logging_registration->invoke = logging_eval_fn;
     custom_op_registration_map_[key] = std::move(logging_registration);
   }
+
+  if (!unresolved_builtin_ops.empty() || !unresolved_custom_ops.empty()) {
+    if (!error_reporter) return;
+    std::string error_message =
+        "Failed to initialize op resolver for calibration:";
+    if (!unresolved_builtin_ops.empty())
+      absl::StrAppend(&error_message, "\nThere are unresolved builtin ops: [",
+                      absl::StrJoin(unresolved_builtin_ops, ", "), "]");
+    if (!unresolved_custom_ops.empty()) {
+      absl::StrAppend(&error_message, "\nThere are unresolved custom ops: [",
+                      absl::StrJoin(unresolved_builtin_ops, ", "), "]");
+    }
+    TF_LITE_REPORT_ERROR(error_reporter, error_message.c_str());
+  }
 }
 
 const TfLiteRegistration* LoggingOpResolver::FindOp(BuiltinOperator op,
diff --git a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h
index bbdfef60d92..25138c38098 100644
--- a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h
+++ b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h
@@ -39,7 +39,8 @@ class LoggingOpResolver : public OpResolver {
   LoggingOpResolver(const BuiltinOpsSet& builtin_ops_to_replace,
                     const CustomOpsSet& custom_ops_to_replace,
                     const OpResolver& base_resolver,
-                    KernelEvalFuncPtr logging_eval_fn);
+                    KernelEvalFuncPtr logging_eval_fn,
+                    ErrorReporter* error_reporter);
 
   const TfLiteRegistration* FindOp(BuiltinOperator op,
                                    int version) const override;
diff --git a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver_test.cc b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver_test.cc
index 511e4d0288d..9513e1f144d 100644
--- a/tensorflow/lite/tools/optimize/calibration/logging_op_resolver_test.cc
+++ b/tensorflow/lite/tools/optimize/calibration/logging_op_resolver_test.cc
@@ -70,7 +70,7 @@ TEST(LoggingOpResolverTest, KernelInvokesAreReplaced) {
   };
 
   LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
-                             WrappingInvoke);
+                             WrappingInvoke, /*error_reporter=*/nullptr);
 
   auto reg = resolver.FindOp(BuiltinOperator_CONV_2D, 1);
 
@@ -104,7 +104,7 @@ TEST(LoggingOpResolverTest, OriginalKernelInvokesAreRetained) {
   };
 
   LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
-                             WrappingInvoke);
+                             WrappingInvoke, /*error_reporter=*/nullptr);
   auto kernel_invoke =
       resolver.GetWrappedKernelInvoke(BuiltinOperator_CONV_2D, 1);
   EXPECT_TRUE(kernel_invoke == ConvEval);
@@ -131,7 +131,7 @@ TEST(LoggingOpResolverTest, OnlyOpsInReplacementSetAreReplaces) {
   };
 
   LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
-                             WrappingInvoke);
+                             WrappingInvoke, /*error_reporter=*/nullptr);
   auto reg = resolver.FindOp(BuiltinOperator_CONV_2D, 1);
   EXPECT_EQ(reg->builtin_code, BuiltinOperator_CONV_2D);
   EXPECT_TRUE(reg->prepare == ConvPrepare);
@@ -155,7 +155,7 @@ TEST(LoggingOpResolverTest, CustomOps) {
   };
 
   LoggingOpResolver resolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
-                             WrappingInvoke);
+                             WrappingInvoke, /*error_reporter=*/nullptr);
 
   auto reg = resolver.FindOp(custom_op_name.c_str(), 1);
 
@@ -165,6 +165,53 @@ TEST(LoggingOpResolverTest, CustomOps) {
   EXPECT_TRUE(reg->invoke == WrappingInvoke);
 }
 
+TEST(LoggingOpResolverTest, UnresolvedCustomOps) {
+  // No custom op registration.
+  MutableOpResolver base_resolver;
+
+  std::string custom_op_name = "unresolved_custom_op";
+
+  CustomOpsSet ops_to_replace = {
+      {custom_op_name, /*version*/ 1},
+  };
+
+  // Expect no death.
+  LoggingOpResolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
+                    WrappingInvoke, /*error_reporter=*/nullptr);
+}
+
+TEST(LoggingOpResolverTest, UnresolvedBuiltinOps) {
+  // No builtin op registration.
+  MutableOpResolver base_resolver;
+
+  BuiltinOpsSet ops_to_replace = {
+      {BuiltinOperator_CONV_2D, /*version*/ 1},
+      {BuiltinOperator_ADD, /*version*/ 1},
+  };
+
+  // Expect no death.
+  LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+}
+
+TEST(LoggingOpResolverTest, FlexOps) {
+  // No flex op registration.
+  MutableOpResolver base_resolver;
+
+  std::string custom_op_name = "FlexAdd";
+
+  CustomOpsSet ops_to_replace = {
+      {custom_op_name, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+
+  auto reg = resolver.FindOp(custom_op_name.c_str(), 1);
+
+  EXPECT_TRUE(!reg);
+}
+
 }  // namespace
 }  // namespace calibration
 }  // namespace optimize

From d29d8af754f281844dfc6870da52faa571c8b948 Mon Sep 17 00:00:00 2001
From: Bruce Fontaine 
Date: Mon, 15 Jun 2020 16:01:40 -0700
Subject: [PATCH 0221/1390] Add InputOptions to
 experimental_distribute_dataset(s_from_function).

PiperOrigin-RevId: 316563848
Change-Id: I00d54d309395754a6182829725f42e1f968f14c4
---
 .../collective_all_reduce_strategy.py         |  5 +-
 .../python/distribute/distribute_lib.py       | 60 +++++++++++++---
 .../python/distribute/distribute_lib_test.py  |  3 +-
 .../python/distribute/mirrored_strategy.py    |  5 +-
 .../python/distribute/one_device_strategy.py  |  5 +-
 .../distribute/parameter_server_strategy.py   |  5 +-
 tensorflow/python/distribute/tpu_strategy.py  | 70 +++++++++----------
 .../python/distribute/tpu_strategy_test.py    | 42 +++++++++++
 ...orflow.distribute.-mirrored-strategy.pbtxt |  4 +-
 ...flow.distribute.-one-device-strategy.pbtxt |  4 +-
 .../v1/tensorflow.distribute.-strategy.pbtxt  |  4 +-
 ...perimental.-central-storage-strategy.pbtxt |  4 +-
 ...ntal.-multi-worker-mirrored-strategy.pbtxt |  4 +-
 ...erimental.-parameter-server-strategy.pbtxt |  4 +-
 ...tribute.experimental.-t-p-u-strategy.pbtxt |  4 +-
 ...tensorflow.distribute.-input-options.pbtxt | 19 +++++
 ...orflow.distribute.-mirrored-strategy.pbtxt |  4 +-
 .../v2/tensorflow.distribute.-strategy.pbtxt  |  4 +-
 ...ntal.-multi-worker-mirrored-strategy.pbtxt |  4 +-
 ...tribute.experimental.-t-p-u-strategy.pbtxt |  4 +-
 .../api/golden/v2/tensorflow.distribute.pbtxt |  4 ++
 21 files changed, 186 insertions(+), 76 deletions(-)
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.-input-options.pbtxt

diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy.py b/tensorflow/python/distribute/collective_all_reduce_strategy.py
index 40c60241ac0..23ed16c5cfd 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy.py
@@ -409,7 +409,7 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
         num_replicas_in_sync=self._num_replicas_in_sync)
     return input_context
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     input_context = self._make_input_context()
     return input_lib.get_distributed_dataset(
         dataset,
@@ -418,7 +418,8 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
         split_batch_by=self._num_replicas_in_sync,
         input_context=input_context)
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     input_context = self._make_input_context()
     return input_lib.get_distributed_datasets_from_function(
         dataset_fn=dataset_fn,
diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index 109cb03ca88..a6dc35507e9 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -602,6 +602,43 @@ class RunOptions(
                  cls).__new__(cls, experimental_enable_dynamic_batch_size,
                               experimental_bucketizing_dynamic_shape)
 
+
+@tf_export("distribute.InputOptions", v1=[])
+class InputOptions(
+    collections.namedtuple("InputOptions", [
+        "experimental_prefetch_to_device",
+    ])):
+  """Run options for `experimental_distribute_dataset(s_from_function)`.
+
+  This can be used to hold some strategy specific configs.
+
+  ```python
+  # Setup TPUStrategy
+  resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+  tf.config.experimental_connect_to_cluster(resolver)
+  tf.tpu.experimental.initialize_tpu_system(resolver)
+  strategy = tf.distribute.experimental.TPUStrategy(resolver)
+
+  dataset = tf.data.Dataset.range(16)
+  distributed_dataset_on_host = (
+      strategy.experimental_distribute_dataset(
+          dataset,
+          tf.distribute.InputOptions(
+              experimental_prefetch_to_device=False)))
+  ```
+
+  Attributes:
+    experimental_prefetch_to_device: Boolean. Currently only applies to
+      TPUStrategy. Defaults to True. If True, dataset elements will be
+      prefetched to accelerator device memory. When False, dataset elements are
+      prefetched to host device memory. Must be False when using TPUEmbedding
+      API.
+  """
+
+  def __new__(cls, experimental_prefetch_to_device=True):
+    return super(InputOptions, cls).__new__(cls,
+                                            experimental_prefetch_to_device)
+
 # ------------------------------------------------------------------------------
 # Base classes for all distribution strategies.
 
@@ -821,7 +858,7 @@ class StrategyBase(object):
       args = (input_iterator.get_next(),) if input_iterator is not None else ()
     return self.run(fn, args=args)
 
-  def experimental_distribute_dataset(self, dataset):
+  def experimental_distribute_dataset(self, dataset, options=None):
     """Distributes a tf.data.Dataset instance provided via `dataset`.
 
     The returned distributed dataset can be iterated over similar to how
@@ -910,14 +947,17 @@ class StrategyBase(object):
     Args:
       dataset: `tf.data.Dataset` that will be sharded across all replicas using
         the rules stated above.
+      options: `tf.distribute.InputOptions` used to control options on how this
+        dataset is distributed.
 
     Returns:
       A "distributed `Dataset`", which acts like a `tf.data.Dataset` except
       it produces "per-replica" values.
     """
-    return self._extended._experimental_distribute_dataset(dataset)  # pylint: disable=protected-access
+    return self._extended._experimental_distribute_dataset(dataset, options)  # pylint: disable=protected-access
 
-  def experimental_distribute_datasets_from_function(self, dataset_fn):
+  def experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                     options=None):
     """Distributes `tf.data.Dataset` instances created by calls to `dataset_fn`.
 
     `dataset_fn` will be called once for each worker in the strategy. Each
@@ -973,13 +1013,15 @@ class StrategyBase(object):
     Args:
       dataset_fn: A function taking a `tf.distribute.InputContext` instance and
         returning a `tf.data.Dataset`.
+      options: `tf.distribute.InputOptions` used to control options on how this
+        dataset is distributed.
 
     Returns:
       A "distributed `Dataset`", which acts like a `tf.data.Dataset` except
       it produces "per-replica" values.
     """
     return self._extended._experimental_distribute_datasets_from_function(  # pylint: disable=protected-access
-        dataset_fn)
+        dataset_fn, options)
 
   def run(self, fn, args=(), kwargs=None, options=None):
     """Run `fn` on each replica, with the given arguments.
@@ -1943,10 +1985,11 @@ class StrategyExtendedV2(object):
   def _make_input_fn_iterator(self, input_fn, replication_mode):
     raise NotImplementedError("must be implemented in descendants")
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     raise NotImplementedError("must be implemented in descendants")
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     raise NotImplementedError("must be implemented in descendants")
 
   def _experimental_distribute_values_from_function(self, value_fn):
@@ -2693,10 +2736,11 @@ class _DefaultDistributionExtended(StrategyExtendedV1):
   def variable_created_in_scope(self, v):
     return v._distribute_strategy is None  # pylint: disable=protected-access
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     return dataset
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     return dataset_fn(InputContext())
 
   def _experimental_distribute_values_from_function(self, value_fn):
diff --git a/tensorflow/python/distribute/distribute_lib_test.py b/tensorflow/python/distribute/distribute_lib_test.py
index 828e7a1aed9..8ea1cac6f02 100644
--- a/tensorflow/python/distribute/distribute_lib_test.py
+++ b/tensorflow/python/distribute/distribute_lib_test.py
@@ -89,7 +89,8 @@ class _TestExtended(distribute_lib.StrategyExtendedV1):
                                            [distribute_lib.InputContext()],
                                            self._container_strategy())
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     return dataset_fn(distribute_lib.InputContext())
 
   def _local_results(self, value):
diff --git a/tensorflow/python/distribute/mirrored_strategy.py b/tensorflow/python/distribute/mirrored_strategy.py
index fe565261f16..ac9045d2322 100644
--- a/tensorflow/python/distribute/mirrored_strategy.py
+++ b/tensorflow/python/distribute/mirrored_strategy.py
@@ -476,7 +476,7 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
                                            input_contexts,
                                            self._container_strategy())
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     return input_lib.get_distributed_dataset(
         dataset,
         self._input_workers,
@@ -487,7 +487,8 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
     return numpy_dataset.one_host_numpy_dataset(
         numpy_input, self._host_input_device, session)
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     input_contexts = []
     num_workers = self._input_workers.num_workers
     for i in range(num_workers):
diff --git a/tensorflow/python/distribute/one_device_strategy.py b/tensorflow/python/distribute/one_device_strategy.py
index 9a74832cd9d..e2bb28ac96f 100644
--- a/tensorflow/python/distribute/one_device_strategy.py
+++ b/tensorflow/python/distribute/one_device_strategy.py
@@ -297,13 +297,14 @@ class OneDeviceExtended(distribute_lib.StrategyExtendedV1):
     del destinations
     return tensor
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     # Note that split_batch_by argument is not passed because it is always 1 in
     # this strategy, and adding it adds unnecessary overhead to the dataset.
     return input_lib.get_distributed_dataset(dataset, self._input_workers,
                                              self._container_strategy())
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     return input_lib.get_distributed_datasets_from_function(
         dataset_fn,
         self._input_workers,
diff --git a/tensorflow/python/distribute/parameter_server_strategy.py b/tensorflow/python/distribute/parameter_server_strategy.py
index 42fc327351c..9675b7002c5 100644
--- a/tensorflow/python/distribute/parameter_server_strategy.py
+++ b/tensorflow/python/distribute/parameter_server_strategy.py
@@ -337,7 +337,7 @@ class ParameterServerStrategyExtended(distribute_lib.StrategyExtendedV1):
   def _validate_colocate_with_variable(self, colocate_with_variable):
     distribute_utils.validate_colocate(colocate_with_variable, self)
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _experimental_distribute_dataset(self, dataset, options):
     return input_lib.get_distributed_dataset(
         dataset,
         self._input_workers,
@@ -376,7 +376,8 @@ class ParameterServerStrategyExtended(distribute_lib.StrategyExtendedV1):
     return numpy_dataset.one_host_numpy_dataset(
         numpy_input, self._input_host_device, session)
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
     if self._cluster_spec:
       input_pipeline_id = multi_worker_util.id_in_cluster(
           self._cluster_spec, self._task_type, self._task_id)
diff --git a/tensorflow/python/distribute/tpu_strategy.py b/tensorflow/python/distribute/tpu_strategy.py
index c605abd9eae..9493ecce767 100644
--- a/tensorflow/python/distribute/tpu_strategy.py
+++ b/tensorflow/python/distribute/tpu_strategy.py
@@ -308,13 +308,14 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
     # device 0 for each replica.
     # TODO(cjfj): Create `InputWorkers` lazily, allowing users to place the
     # input onto a different logical device?
-    input_worker_devices = collections.OrderedDict()
+    self._device_input_worker_devices = collections.OrderedDict()
+    self._host_input_worker_devices = collections.OrderedDict()
     for tpu_device in self._tpu_devices[:, 0]:
       host_device = device_util.get_host_for_device(tpu_device)
-      input_worker_devices.setdefault(host_device, [])
-      input_worker_devices[host_device].append(tpu_device)
-    self._input_worker_devices = tuple(input_worker_devices.items())
-    self._input_workers_obj = None
+      self._device_input_worker_devices.setdefault(host_device, [])
+      self._device_input_worker_devices[host_device].append(tpu_device)
+      self._host_input_worker_devices.setdefault(host_device, [])
+      self._host_input_worker_devices[host_device].append(host_device)
 
     # TODO(sourabhbajaj): Remove this once performance of running one step
     # at a time is comparable to multiple steps.
@@ -322,7 +323,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
     self._require_static_shapes = True
 
     self.experimental_enable_get_next_as_optional = True
-    self._prefetch_on_host = False
+    self._prefetch_to_device = True
 
     self._logical_device_stack = [0]
 
@@ -339,38 +340,18 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
   # memory and b) TPU Embedding enqueue operation are CPU ops and this avoids
   # a copy back to the host for dense tensors
   def _set_prefetch_on_host(self, value):
-    if self._prefetch_on_host == value:
-      return
-    if self._input_workers_obj is not None:
-      raise RuntimeError("Unable to change prefetch on host behavior as "
-                         "InputWorkers are already created.")
-    self._prefetch_on_host = value
-    if value:
-      # To prefetch on the host, we must set all the input worker devices to the
-      # corresponding host devices.
-      self._input_worker_devices = tuple([
-          tuple([host,
-                 [device_util.get_host_for_device(d) for d in devices]])
-          for host, devices in self._input_worker_devices])
-      # Force creation of the workers.
-      workers = self._input_workers
-      del workers
-
-  @property
-  def _input_workers(self):
-    if self._input_workers_obj is None:
-      self._input_workers_obj = input_lib.InputWorkers(
-          self._input_worker_devices)
-    return self._input_workers_obj
+    self._prefetch_to_device = not value
 
   def _validate_colocate_with_variable(self, colocate_with_variable):
     distribute_utils. validate_colocate(colocate_with_variable, self)
 
   def _make_dataset_iterator(self, dataset):
     """Make iterators for each of the TPU hosts."""
+    input_workers = input_lib.InputWorkers(
+        tuple(self._device_input_worker_devices.items()))
     return input_lib.DatasetIterator(
         dataset,
-        self._input_workers,
+        input_workers,
         self._container_strategy(),
         split_batch_by=self._num_replicas_in_sync)
 
@@ -379,7 +360,9 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
       input_fn,
       replication_mode=distribute_lib.InputReplicationMode.PER_WORKER):
     input_contexts = []
-    num_workers = self._input_workers.num_workers
+    input_workers = input_lib.InputWorkers(
+        tuple(self._device_input_worker_devices.items()))
+    num_workers = input_workers.num_workers
     for i in range(num_workers):
       input_contexts.append(distribute_lib.InputContext(
           num_input_pipelines=num_workers,
@@ -387,7 +370,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
           num_replicas_in_sync=self._num_replicas_in_sync))
     return input_lib.InputFunctionIterator(
         input_fn,
-        self._input_workers,
+        input_workers,
         input_contexts,
         self._container_strategy())
 
@@ -396,16 +379,29 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
         numpy_input, numpy_dataset.SingleDevice(self._host_device),
         session)
 
-  def _experimental_distribute_dataset(self, dataset):
+  def _get_input_workers(self, options):
+    prefetch_to_device = self._prefetch_to_device
+    if options:
+      prefetch_to_device = options.experimental_prefetch_to_device
+    if prefetch_to_device:
+      return input_lib.InputWorkers(
+          tuple(self._device_input_worker_devices.items()))
+    else:
+      return input_lib.InputWorkers(
+          tuple(self._host_input_worker_devices.items()))
+
+  def _experimental_distribute_dataset(self, dataset, options):
     return input_lib.get_distributed_dataset(
         dataset,
-        self._input_workers,
+        self._get_input_workers(options),
         self._container_strategy(),
         split_batch_by=self._num_replicas_in_sync)
 
-  def _experimental_distribute_datasets_from_function(self, dataset_fn):
+  def _experimental_distribute_datasets_from_function(self, dataset_fn,
+                                                      options):
+    input_workers = self._get_input_workers(options)
     input_contexts = []
-    num_workers = self._input_workers.num_workers
+    num_workers = input_workers.num_workers
     for i in range(num_workers):
       input_contexts.append(distribute_lib.InputContext(
           num_input_pipelines=num_workers,
@@ -414,7 +410,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
 
     return input_lib.get_distributed_datasets_from_function(
         dataset_fn,
-        self._input_workers,
+        input_workers,
         input_contexts,
         self._container_strategy())
 
diff --git a/tensorflow/python/distribute/tpu_strategy_test.py b/tensorflow/python/distribute/tpu_strategy_test.py
index 70a38af95aa..6dd7de500e4 100644
--- a/tensorflow/python/distribute/tpu_strategy_test.py
+++ b/tensorflow/python/distribute/tpu_strategy_test.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 from tensorflow.python import keras
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
 from tensorflow.python.distribute import distribution_strategy_context
 from tensorflow.python.distribute import reduce_util
 from tensorflow.python.distribute import tpu_strategy as tpu_lib
@@ -30,6 +31,7 @@ from tensorflow.python.eager import remote
 from tensorflow.python.eager import test
 from tensorflow.python.framework import config
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
@@ -546,6 +548,46 @@ class TPUStrategyTest(test.TestCase):
       update_variable.get_concrete_function()
       self.assertEqual(trace_count[0], len(strategy.extended.worker_devices))
 
+  def test_prefetch_to_device_default(self):
+    strategy = get_tpu_strategy()
+    dataset = dataset_ops.Dataset.range(
+        strategy.num_replicas_in_sync * 2,
+        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
+
+    # Check default, should prefetch to TPU.
+    dataset_item = next(iter(strategy.experimental_distribute_dataset(dataset)))
+    dataset_location = tf_device.DeviceSpec.from_string(
+        dataset_item.values[0].device)
+    self.assertEqual(dataset_location.device_type, "TPU")
+
+  def test_prefetch_to_device_tpu(self):
+    strategy = get_tpu_strategy()
+    dataset = dataset_ops.Dataset.range(
+        strategy.num_replicas_in_sync * 2,
+        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
+
+    input_options = distribute_lib.InputOptions(
+        experimental_prefetch_to_device=True)
+    dataset_item = next(iter(strategy.experimental_distribute_dataset(
+        dataset, options=input_options)))
+    dataset_location = tf_device.DeviceSpec.from_string(
+        dataset_item.values[0].device)
+    self.assertEqual(dataset_location.device_type, "TPU")
+
+  def test_prefetch_to_device_cpu(self):
+    strategy = get_tpu_strategy()
+    dataset = dataset_ops.Dataset.range(
+        strategy.num_replicas_in_sync * 2,
+        output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
+
+    # Should be CPU when prefetch_to_device is False.
+    input_options = distribute_lib.InputOptions(
+        experimental_prefetch_to_device=False)
+    dataset_item = next(iter(strategy.experimental_distribute_dataset(
+        dataset, options=input_options)))
+    dataset_location = tf_device.DeviceSpec.from_string(
+        dataset_item.values[0].device)
+    self.assertEqual(dataset_location.device_type, "CPU")
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
index 0b74423ce62..36c78c406b7 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -26,11 +26,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
index 67d6923e86c..09865ab02ee 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
@@ -26,11 +26,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
index d22b42d9098..0e6c10bd533 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
@@ -25,11 +25,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
index 03c5b2476b0..fbc4c107a1a 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
@@ -26,11 +26,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
index baee19e2a50..cd67e7d27c4 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
@@ -26,11 +26,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
index d92dab8f5bf..0eff82474ff 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
@@ -26,11 +26,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
index c7c8c832764..2af9a5ad095 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
@@ -30,11 +30,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_local_results"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-input-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-input-options.pbtxt
new file mode 100644
index 00000000000..c3beabd938e
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-input-options.pbtxt
@@ -0,0 +1,19 @@
+path: "tensorflow.distribute.InputOptions"
+tf_class {
+  is_instance: ""
+  is_instance: ""
+  is_instance: ""
+  member {
+    name: "experimental_prefetch_to_device"
+    mtype: ""
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "count"
+  }
+  member_method {
+    name: "index"
+  }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
index 20dfe7fe5a6..be4c841aed7 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -30,11 +30,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_values_from_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
index 0844739c8eb..9f6a2ac32be 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
@@ -29,11 +29,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_values_from_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
index 0f722ecc8b9..500ae362e5f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
@@ -30,11 +30,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_values_from_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
index 6cefc4e7977..82a4362a597 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
@@ -30,11 +30,11 @@ tf_class {
   }
   member_method {
     name: "experimental_distribute_dataset"
-    argspec: "args=[\'self\', \'dataset\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_datasets_from_function"
-    argspec: "args=[\'self\', \'dataset_fn\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'dataset_fn\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "experimental_distribute_values_from_function"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
index 3e226fd8e70..19d83909120 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
@@ -16,6 +16,10 @@ tf_module {
     name: "InputContext"
     mtype: ""
   }
+  member {
+    name: "InputOptions"
+    mtype: ""
+  }
   member {
     name: "InputReplicationMode"
     mtype: ""

From 57eccc7bc29ddb105dcaa2f6a413163461ad9987 Mon Sep 17 00:00:00 2001
From: Dominic Jack 
Date: Tue, 16 Jun 2020 09:06:15 +1000
Subject: [PATCH 0222/1390] added test

---
 tensorflow/python/keras/engine/training_test.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index 111833ba8b5..72cd5ad88a3 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -3328,6 +3328,16 @@ class TestTrainingWithMetrics(keras_parameterized.TestCase):
     outer_model.fit(np.ones((10, 1)), np.ones((10, 1)), batch_size=10)
     self.assertEqual([m.name for m in outer_model.metrics],
                      ['loss', 'acc2', 'mean', 'mean1', 'mean2'])
+  
+  def test_subclassed_model_with_empty_list_attr(self):
+    class ModelSubclass(training_module.Model):
+      def __init__(self):
+        self.empty_list = []
+        inputs = layers_module.Input(shape=())
+        outputs = inputs + 1
+        super(ModelSubclass, self).__init__(inputs, outputs)
+    
+    ModelSubclass()  # empty_list attr assignment should not raise
 
 
 class BareUpdateLayer(layers_module.Layer):

From f4e20ec5ae4792a3b595a598479ebe8f5a1760ea Mon Sep 17 00:00:00 2001
From: Jose Baiocchi 
Date: Mon, 15 Jun 2020 16:04:21 -0700
Subject: [PATCH 0223/1390] XPlane schema cleanup

PiperOrigin-RevId: 316564434
Change-Id: Icabe3400c82cb4aff84c2bdae7ecf59424b36fa4
---
 tensorflow/core/profiler/convert/BUILD        |  2 ++
 .../convert/op_stats_to_tf_stats_test.cc      |  5 ++--
 .../convert/xplane_to_memory_profile_test.cc  |  3 +-
 .../convert/xplane_to_op_metrics_db_test.cc   | 24 ++++++---------
 .../profiler/convert/xplane_to_op_stats.cc    |  2 +-
 .../convert/xplane_to_op_stats_test.cc        | 22 +++++++-------
 .../convert/xplane_to_profile_response.cc     |  3 +-
 .../convert/xplane_to_step_events_test.cc     |  3 +-
 .../convert/xplane_to_tf_functions_test.cc    | 12 ++++----
 .../convert/xplane_to_trace_events.cc         |  2 +-
 .../core/profiler/internal/cpu/host_tracer.cc |  2 +-
 .../profiler/internal/cpu/host_tracer_test.cc |  2 +-
 .../internal/cpu/metadata_collector.cc        |  2 +-
 .../profiler/internal/gpu/device_tracer.cc    |  7 +++--
 .../core/profiler/lib/profiler_session.cc     |  5 ++--
 tensorflow/core/profiler/utils/BUILD          |  1 +
 .../core/profiler/utils/derived_timeline.cc   |  7 +++--
 .../profiler/utils/derived_timeline_test.cc   |  8 ++---
 .../core/profiler/utils/group_events.cc       | 19 +++++++-----
 .../core/profiler/utils/group_events_test.cc  | 18 ++++--------
 .../core/profiler/utils/xplane_schema.cc      |  6 ++--
 .../core/profiler/utils/xplane_schema.h       | 16 ++++++++--
 .../core/profiler/utils/xplane_test_utils.cc  | 10 +++++++
 .../core/profiler/utils/xplane_test_utils.h   |  4 +++
 .../core/profiler/utils/xplane_utils.cc       | 29 +++++++------------
 tensorflow/core/profiler/utils/xplane_utils.h | 15 ++++------
 26 files changed, 117 insertions(+), 112 deletions(-)

diff --git a/tensorflow/core/profiler/convert/BUILD b/tensorflow/core/profiler/convert/BUILD
index 2482698cdf0..5f287a14267 100644
--- a/tensorflow/core/profiler/convert/BUILD
+++ b/tensorflow/core/profiler/convert/BUILD
@@ -47,6 +47,7 @@ tf_cc_test(
         "//tensorflow/core/profiler/utils:time_utils",
         "//tensorflow/core/profiler/utils:xplane_builder",
         "//tensorflow/core/profiler/utils:xplane_schema",
+        "//tensorflow/core/profiler/utils:xplane_test_utils",
         "@com_google_absl//absl/strings",
     ],
 )
@@ -171,6 +172,7 @@ tf_cc_test(
         "//tensorflow/core/profiler/utils:time_utils",
         "//tensorflow/core/profiler/utils:xplane_builder",
         "//tensorflow/core/profiler/utils:xplane_schema",
+        "//tensorflow/core/profiler/utils:xplane_test_utils",
         "@com_google_absl//absl/strings",
     ],
 )
diff --git a/tensorflow/core/profiler/convert/op_stats_to_tf_stats_test.cc b/tensorflow/core/profiler/convert/op_stats_to_tf_stats_test.cc
index 9ca83b51a70..5a01bf3417b 100644
--- a/tensorflow/core/profiler/convert/op_stats_to_tf_stats_test.cc
+++ b/tensorflow/core/profiler/convert/op_stats_to_tf_stats_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/utils/time_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
+#include "tensorflow/core/profiler/utils/xplane_test_utils.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -59,8 +60,8 @@ TEST(OpStatsToTfStats, GpuTfStats) {
   constexpr int64 kKernel3DurationNs = 10000;
 
   XSpace space;
-  XPlaneBuilder device_plane(space.add_planes());
-  device_plane.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
+  XPlaneBuilder device_plane(
+      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
   XLineBuilder stream1 = device_plane.GetOrCreateLine(/*line_id=*/10);
   AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
                        kKernel1DurationNs, /*on_device=*/true, kKernel1,
diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc b/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
index 6766fd5f1b5..5ddcbcfc75d 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
@@ -32,9 +32,8 @@ namespace {
 // activities within one memory allocator captured in host trace.
 TEST(ConvertXPlaneToMemoryProfile, OneAllocatorMultiActivitiesTest) {
   XSpace space;
-  XPlane* host_plane = space.add_planes();
+  XPlane* host_plane = GetOrCreateHostXPlane(&space);
   XPlaneBuilder host_plane_builder(host_plane);
-  host_plane_builder.SetName(kHostThreads);
   host_plane_builder.ReserveLines(1);
 
   auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
diff --git a/tensorflow/core/profiler/convert/xplane_to_op_metrics_db_test.cc b/tensorflow/core/profiler/convert/xplane_to_op_metrics_db_test.cc
index 8bd0443b8f6..bdac1129c81 100644
--- a/tensorflow/core/profiler/convert/xplane_to_op_metrics_db_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_op_metrics_db_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/utils/time_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
+#include "tensorflow/core/profiler/utils/xplane_test_utils.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -43,12 +44,6 @@ void AddTensorFlowOpEvent(absl::string_view tf_op_fullname,
                              tf_op_fullname);
 }
 
-void SetXPlaneNameAndId(absl::string_view name, int64 id,
-                        XPlaneBuilder* plane) {
-  plane->SetName(name);
-  plane->SetId(id);
-}
-
 TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
   static constexpr char kTfOp1[] = "TfOp1";
   static constexpr char kTfOp2[] = "TfOp2";
@@ -57,9 +52,9 @@ TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
   constexpr int64 kTfOp2StartNs = 110000;
   constexpr int64 kTfOp2DurationNs = 10000;
 
-  XPlane xplane;
-  XPlaneBuilder host_plane(&xplane);
-  SetXPlaneNameAndId(kHostThreads, /*id=*/0, &host_plane);
+  XSpace xspace;
+  XPlane* xplane = GetOrCreateHostXPlane(&xspace);
+  XPlaneBuilder host_plane(xplane);
   XLineBuilder thread1 = host_plane.GetOrCreateLine(/*line_id=*/10);
   AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kTfOp1StartNs,
                        kTfOp1DurationNs, /*on_device=*/false,
@@ -72,7 +67,7 @@ TEST(ConvertXPlaneToOpMetricsDb, HostOpMetricsDb) {
                        kTfOp2DurationNs, /*on_device=*/false,
                        /*kernel_name=*/"", &host_plane, &thread2);
 
-  OpMetricsDb op_metrics = ConvertHostThreadsXPlaneToOpMetricsDb(xplane);
+  OpMetricsDb op_metrics = ConvertHostThreadsXPlaneToOpMetricsDb(*xplane);
   // Op1, Op2, Idle.
   EXPECT_EQ(3, op_metrics.metrics_db_size());
   uint64 total_op_duration =
@@ -115,10 +110,9 @@ TEST(ConvertXPlaneToOpMetricsDb, DeviceOpMetricsDb) {
   constexpr int64 kKernel3StartNs = 120000;
   constexpr int64 kKernel3DurationNs = 10000;
 
-  XPlane xplane;
-  XPlaneBuilder device_plane(&xplane);
-  SetXPlaneNameAndId(absl::StrCat(kGpuPlanePrefix, ":0"), /*id=*/1,
-                     &device_plane);
+  XSpace xspace;
+  XPlane* xplane = GetOrCreateGpuXPlane(&xspace, /*device_ordinal=*/0);
+  XPlaneBuilder device_plane(xplane);
   XLineBuilder stream1 = device_plane.GetOrCreateLine(/*line_id=*/10);
   AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
                        kKernel1DurationNs, /*on_device=*/true, kKernel1,
@@ -138,7 +132,7 @@ TEST(ConvertXPlaneToOpMetricsDb, DeviceOpMetricsDb) {
                        &device_plane, &stream2);
 
   OpMetricsDb op_metrics = ConvertDeviceTraceXPlaneToOpMetricsDb(
-      xplane, /*peak_tera_flops_per_second=*/0,
+      *xplane, /*peak_tera_flops_per_second=*/0,
       /*peak_hbm_bw_giga_bytes_per_second=*/0);
 
   // kernel1, kernel2, kernel3, Idle.
diff --git a/tensorflow/core/profiler/convert/xplane_to_op_stats.cc b/tensorflow/core/profiler/convert/xplane_to_op_stats.cc
index eb2e13dbb4a..df050d16ede 100644
--- a/tensorflow/core/profiler/convert/xplane_to_op_stats.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_op_stats.cc
@@ -129,7 +129,7 @@ void PropagateXSpaceDiagnosticsToOpStats(const XSpace& space,
 }
 
 OpStats ConvertXSpaceToOpStats(const XSpace& space) {
-  const XPlane* host_plane = FindPlaneWithName(space, kHostThreads);
+  const XPlane* host_plane = FindPlaneWithName(space, kHostThreadsPlaneName);
   std::vector device_planes =
       FindPlanesWithPrefix(space, kGpuPlanePrefix);
   OpStats op_stats;
diff --git a/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc b/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc
index 138bcee72be..68bb8205f5e 100644
--- a/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_op_stats_test.cc
@@ -43,8 +43,8 @@ TEST(ConvertXPlaneToOpStats, PerfEnv) {
   constexpr int kComputeCapMajor = 7;
   constexpr int kComputeCapMinor = 0;
 
-  XPlaneBuilder device_plane(space.add_planes());
-  device_plane.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
+  XPlaneBuilder device_plane(
+      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
   device_plane.ParseAndAddStatValue(
       *device_plane.GetOrCreateStatMetadata("clock_rate"),
       absl::StrCat(kClockRateKHz));
@@ -71,10 +71,10 @@ TEST(ConvertXPlaneToOpStats, PerfEnv) {
 
 TEST(ConvertXPlaneToOpStats, RunEnvironment) {
   XSpace space;
-  XPlaneBuilder device_plane1(space.add_planes());
-  device_plane1.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
-  XPlaneBuilder device_plane2(space.add_planes());
-  device_plane2.SetName(absl::StrCat(kGpuPlanePrefix, ":1"));
+  XPlaneBuilder device_plane1(
+      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
+  XPlaneBuilder device_plane2(
+      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/1));
 
   GroupTfEvents(&space, /*event_group_name_map=*/nullptr);
   OpStats op_stats = ConvertXSpaceToOpStats(space);
@@ -91,8 +91,7 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
   constexpr int64 kStepId = 0;
 
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
@@ -120,8 +119,7 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
   constexpr int64 kCorrelationId = 100;
 
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
@@ -137,8 +135,8 @@ TEST(ConvertXPlaneToOpStats, GpuStepDbTest) {
   CreateXEvent(&host_plane_builder, &tf_executor_thread, "matmul", 30, 10,
                {{StatType::kCorrelationId, kCorrelationId}});
 
-  XPlaneBuilder device_plane_builder(space.add_planes());
-  device_plane_builder.SetName(absl::StrCat(kGpuPlanePrefix, ":0"));
+  XPlaneBuilder device_plane_builder(
+      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
   device_plane_builder.ReserveLines(1);
 
   auto stream = device_plane_builder.GetOrCreateLine(0);
diff --git a/tensorflow/core/profiler/convert/xplane_to_profile_response.cc b/tensorflow/core/profiler/convert/xplane_to_profile_response.cc
index 70a07171310..22af46c4380 100644
--- a/tensorflow/core/profiler/convert/xplane_to_profile_response.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_profile_response.cc
@@ -139,7 +139,8 @@ Status ConvertXSpaceToProfileResponse(const XSpace& xspace,
     AddToolData(ToolName(kKernelStats), op_stats.kernel_stats_db(), response);
   }
   if (tools.contains(kMemoryProfile)) {
-    if (const XPlane* host_plane = FindPlaneWithName(xspace, kHostThreads)) {
+    if (const XPlane* host_plane =
+            FindPlaneWithName(xspace, kHostThreadsPlaneName)) {
       MemoryProfile memory_profile = ConvertXPlaneToMemoryProfile(*host_plane);
       std::string json_output;
       TF_RETURN_IF_ERROR(ConvertProtoToJson(memory_profile, &json_output));
diff --git a/tensorflow/core/profiler/convert/xplane_to_step_events_test.cc b/tensorflow/core/profiler/convert/xplane_to_step_events_test.cc
index ff68f1817ed..1c6dfee7cc7 100644
--- a/tensorflow/core/profiler/convert/xplane_to_step_events_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_step_events_test.cc
@@ -45,9 +45,8 @@ TEST(ConvertXPlaneToOpStats, CpuOnlyStepDbTest) {
   constexpr int64 kSecondCorrelationId = 200;
 
   XSpace space;
-  XPlane* host_plane = space.add_planes();
+  XPlane* host_plane = GetOrCreateHostXPlane(&space);
   XPlaneBuilder host_plane_builder(host_plane);
-  host_plane_builder.SetName(kHostThreads);
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
diff --git a/tensorflow/core/profiler/convert/xplane_to_tf_functions_test.cc b/tensorflow/core/profiler/convert/xplane_to_tf_functions_test.cc
index 12287217e04..a310313ff86 100644
--- a/tensorflow/core/profiler/convert/xplane_to_tf_functions_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_tf_functions_test.cc
@@ -43,11 +43,12 @@ constexpr double kMaxError = 0.001;
 
 TfFunctionDb ConvertXSpaceToTfFunctionDb(const XSpace& space) {
   TfFunctionDb result;
-  const XPlane* host_plane = FindPlaneWithName(space, kHostThreads);
+  const XPlane* host_plane = FindPlaneWithName(space, kHostThreadsPlaneName);
   if (host_plane) {
     XPlaneVisitor plane = CreateTfXPlaneVisitor(host_plane);
     plane.ForEachLine([&result](const XLineVisitor& line) {
-      CombineTfFunctionDb(ConvertHostThreadsXLineToTfFunctionDb(line), &result);
+      TfFunctionDb tf_function_db = ConvertHostThreadsXLineToTfFunctionDb(line);
+      CombineTfFunctionDb(tf_function_db, &result);
     });
   }
   return result;
@@ -56,7 +57,7 @@ TfFunctionDb ConvertXSpaceToTfFunctionDb(const XSpace& space) {
 TEST(ConvertXPlaneToTfFunctions, CombineTwoThreads) {
   XSpace space;
   XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  host_plane_builder.SetName(kHostThreadsPlaneName);
   host_plane_builder.ReserveLines(2);
   std::string kFunctionName = "decrement";
 
@@ -100,7 +101,7 @@ TEST(ConvertXPlaneToTfFunctions, CombineTwoThreads) {
 TEST(ConvertXPlaneToTfFunctions, NestedFunctions) {
   XSpace space;
   XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  host_plane_builder.SetName(kHostThreadsPlaneName);
   host_plane_builder.ReserveLines(1);
   std::string kOuterFunctionName = "outer";
   std::string kInnerFunctionName = "inner";
@@ -140,8 +141,7 @@ TEST(ConvertXPlaneToTfFunctions, NestedFunctions) {
 
 TEST(ConvertXPlaneToTfFunctions, EagerPlusConcrete) {
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(2);
   std::string kEagerFunctionName = "i_am_eager";
   std::string kConcreteFunctionName = "i_am_concrete";
diff --git a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
index ffdad034337..f4a0145d8f6 100644
--- a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
@@ -41,7 +41,7 @@ Device BuildDeviceAndResource(const XPlaneVisitor& plane) {
   device.set_name(std::string(plane.Name()));
   device.set_device_id(plane.Id());
 
-  bool sort_by_ordinal = plane.Name() == kHostThreads;
+  bool sort_by_ordinal = (plane.Name() == kHostThreadsPlaneName);
   int ordinal = 0;
   plane.ForEachLine([&](const XLineVisitor& line) {
     Resource resource;
diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer.cc b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
index be1a7a2777b..37f7baca1d3 100644
--- a/tensorflow/core/profiler/internal/cpu/host_tracer.cc
+++ b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
@@ -146,7 +146,7 @@ Status HostTracer::CollectData(XSpace* space) {
     return errors::Internal("TraceMeRecorder not stopped");
   }
   MakeCompleteEvents(&events_);
-  XPlane* plane = GetOrCreatePlane(space, kHostThreads);
+  XPlane* plane = FindOrAddMutablePlaneWithName(space, kHostThreadsPlaneName);
   plane->set_id(kHostPlaneId);
   ConvertCompleteEventsToXPlane(start_timestamp_ns_, events_, plane);
   events_.clear();
diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc b/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
index 499b7b6b564..0e4c3dd7a9b 100644
--- a/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
+++ b/tensorflow/core/profiler/internal/cpu/host_tracer_test.cc
@@ -150,7 +150,7 @@ TEST(HostTracerTest, CollectsTraceMeEventsAsXSpace) {
   ASSERT_EQ(space.planes_size(), 1);
   const auto& plane = space.planes(0);
   XPlaneVisitor xplane(&plane);
-  ASSERT_EQ(plane.name(), kHostThreads);
+  ASSERT_EQ(plane.name(), kHostThreadsPlaneName);
   ASSERT_EQ(plane.lines_size(), 1);
   ASSERT_EQ(plane.event_metadata_size(), 7);
   ASSERT_EQ(plane.stat_metadata_size(), 2);
diff --git a/tensorflow/core/profiler/internal/cpu/metadata_collector.cc b/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
index da922d4b18b..58e6385a7ec 100644
--- a/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
+++ b/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
@@ -65,7 +65,7 @@ class MetadataCollector : public ProfilerInterface {
 
   Status CollectData(XSpace* space) override {
     if (!debug_info_.empty()) {
-      XPlane* plane = GetOrCreatePlane(space, kMetadataPlane);
+      XPlane* plane = FindOrAddMutablePlaneWithName(space, kMetadataPlaneName);
       plane->set_id(kMetadataPlaneId);
       XPlaneBuilder xplane(plane);
       const XStatMetadata& hlo_proto_stat =
diff --git a/tensorflow/core/profiler/internal/gpu/device_tracer.cc b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
index 5ddee687333..bc9952302e8 100644
--- a/tensorflow/core/profiler/internal/gpu/device_tracer.cc
+++ b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
@@ -223,11 +223,12 @@ class CuptiTraceCollectorImpl : public CuptiTraceCollector {
               << " callback api events and " << num_activity_events_
               << " activity events. " << ReportDroppedEvents();
     uint64 end_gpu_ns = CuptiTracer::GetTimestamp();
-    XPlaneBuilder host_plane(GetOrCreatePlane(space, kCuptiDriverApiPlaneName));
+    XPlaneBuilder host_plane(
+        FindOrAddMutablePlaneWithName(space, kCuptiDriverApiPlaneName));
     host_plane.SetId(kCuptiDriverApiPlaneId);
     for (int device_ordinal = 0; device_ordinal < num_gpus_; ++device_ordinal) {
-      std::string name = absl::StrCat(kGpuPlanePrefix, device_ordinal);
-      XPlaneBuilder device_plane(GetOrCreatePlane(space, name));
+      std::string name = GpuPlaneName(device_ordinal);
+      XPlaneBuilder device_plane(FindOrAddMutablePlaneWithName(space, name));
       device_plane.SetId(kGpuPlaneBaseId + device_ordinal);
       per_device_collector_[device_ordinal].Flush(start_gpu_ns_, end_gpu_ns,
                                                   &device_plane, &host_plane);
diff --git a/tensorflow/core/profiler/lib/profiler_session.cc b/tensorflow/core/profiler/lib/profiler_session.cc
index 885c5e0ca4f..f7d97711da0 100644
--- a/tensorflow/core/profiler/lib/profiler_session.cc
+++ b/tensorflow/core/profiler/lib/profiler_session.cc
@@ -98,11 +98,10 @@ Status ProfilerSession::CollectData(profiler::XSpace* space) {
   const profiler::XPlane* cupti_driver_api_plane =
       profiler::FindPlaneWithName(*space, profiler::kCuptiDriverApiPlaneName);
   if (cupti_driver_api_plane) {
-    profiler::XPlane* host_plane =
-        profiler::GetOrCreatePlane(space, profiler::kHostThreads);
+    profiler::XPlane* host_plane = profiler::FindOrAddMutablePlaneWithName(
+        space, profiler::kHostThreadsPlaneName);
     profiler::MergePlanes(*cupti_driver_api_plane, host_plane);
     profiler::SortXLinesBy(host_plane, profiler::XLinesComparatorByName());
-    // This might invalidate host_plane pointer.
     profiler::RemovePlaneWithName(space, profiler::kCuptiDriverApiPlaneName);
   }
   // 2. Normalize all timestamps by shifting timeline to profiling start time.
diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index e0eaa5968c1..6942f3ea306 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -233,6 +233,7 @@ cc_library(
     deps = [
         ":xplane_builder",
         ":xplane_schema",
+        ":xplane_utils",
         "//tensorflow/core/platform:types",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/container:flat_hash_map",
diff --git a/tensorflow/core/profiler/utils/derived_timeline.cc b/tensorflow/core/profiler/utils/derived_timeline.cc
index 112c0977763..f63a8e5c2d9 100644
--- a/tensorflow/core/profiler/utils/derived_timeline.cc
+++ b/tensorflow/core/profiler/utils/derived_timeline.cc
@@ -339,9 +339,10 @@ void GenerateDerivedTimeLines(const EventGroupNameMap& event_group_name_map,
                               XSpace* space, bool step_info_only) {
   for (XPlane& plane : *space->mutable_planes()) {
     // Derived timelines only generated for device traces.
-    if (plane.id() == kHostPlaneId) continue;
-    DeriveEventsFromAnnotations(DummySymbolResolver, event_group_name_map,
-                                &plane, step_info_only);
+    if (IsGpuPlaneName(plane.name())) {
+      DeriveEventsFromAnnotations(DummySymbolResolver, event_group_name_map,
+                                  &plane, step_info_only);
+    }
   }
 }
 
diff --git a/tensorflow/core/profiler/utils/derived_timeline_test.cc b/tensorflow/core/profiler/utils/derived_timeline_test.cc
index 4ae558eb446..a75ba8ea085 100644
--- a/tensorflow/core/profiler/utils/derived_timeline_test.cc
+++ b/tensorflow/core/profiler/utils/derived_timeline_test.cc
@@ -44,7 +44,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map;
-  XPlane* plane = space.add_planes();
+  XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
   XPlaneBuilder plane_builder(plane);
   auto line_builder = plane_builder.GetOrCreateLine(0);
   CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@@ -74,7 +74,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map;
-  XPlane* plane = space.add_planes();
+  XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
   XPlaneBuilder plane_builder(plane);
   auto line_builder = plane_builder.GetOrCreateLine(0);
   CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@@ -109,7 +109,7 @@ TEST(DerivedTimelineTest, DependencyTest) {
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map({{0, "train 0"}, {1, "train 1"}});
-  XPlane* plane = space.add_planes();
+  XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
   XPlaneBuilder plane_builder(plane);
   auto line_builder = plane_builder.GetOrCreateLine(0);
   CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
@@ -138,7 +138,7 @@ TEST(DerivedTimelineTest, TfOpNameScopeTest) {
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map;
-  XPlane* plane = space.add_planes();
+  XPlane* plane = GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0);
   XPlaneBuilder plane_builder(plane);
   auto line_builder = plane_builder.GetOrCreateLine(0);
   CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100,
diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index 2f4e3c8b3f1..be8dd506b0c 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -54,13 +54,13 @@ void CreateStatMetadata(XPlane* plane) {
 }
 
 // Returns event type if it is a KernelLaunch or KernelExecute event.
-absl::optional GetKernelEventType(const XPlaneVisitor& visitor,
+absl::optional GetKernelEventType(bool is_host_plane,
+                                         const XPlaneVisitor& visitor,
                                          const XEvent& event) {
   for (const auto& stat : event.stats()) {
     if (visitor.GetStatType(stat) == StatType::kCorrelationId) {
-      // TODO(b/149095099): avoid string comparison.
-      return visitor.Name() == kHostThreads ? HostEventType::kKernelLaunch
-                                            : HostEventType::kKernelExecute;
+      return is_host_plane ? HostEventType::kKernelLaunch
+                           : HostEventType::kKernelExecute;
     }
   }
   return absl::nullopt;
@@ -72,14 +72,15 @@ bool IsTfOpEvent(const XPlaneVisitor& visitor, const XEvent& event) {
   return tf_op.category == Category::kTensorFlow;
 }
 
-int64 GetEventType(const XPlaneVisitor& visitor, const XEvent& event) {
+int64 GetEventType(bool is_host_plane, const XPlaneVisitor& visitor,
+                   const XEvent& event) {
   if (absl::optional event_type = visitor.GetEventType(event)) {
     return *event_type;
   } else if (absl::optional kernel_event_type =
-                 GetKernelEventType(visitor, event)) {
+                 GetKernelEventType(is_host_plane, visitor, event)) {
     // KernelLaunch and KernelExecute event types are not supported by
     // XPlaneVisitor and should be checked separately.
-    // TODO(148346217): Make XPlaneVisitor support KernelLaunch and
+    // TODO(b/148346217): Make XPlaneVisitor support KernelLaunch and
     // KernelExecute event types.
     return *kernel_event_type;
   } else if (IsTfOpEvent(visitor, event)) {
@@ -396,6 +397,8 @@ bool EventNode::StartsBefore(const EventNode& other) const {
 void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor,
                                      XPlane* plane,
                                      ContextGroupMap* context_groups) {
+  // TODO(b/149095099): avoid string comparison.
+  bool is_host_plane = (visitor.Name() == kHostThreadsPlaneName);
   for (auto& line : *plane->mutable_lines()) {
     std::vector parent_nodes;
     for (auto& event : *line.mutable_events()) {
@@ -418,7 +421,7 @@ void EventForest::ConnectIntraThread(const XPlaneVisitor& visitor,
       }
       parent_nodes.push_back(cur_node.get());
       // event_node_map_ keeps cur_node alive.
-      event_node_map_[GetEventType(visitor, event)].push_back(
+      event_node_map_[GetEventType(is_host_plane, visitor, event)].push_back(
           std::move(cur_node));
     }
   }
diff --git a/tensorflow/core/profiler/utils/group_events_test.cc b/tensorflow/core/profiler/utils/group_events_test.cc
index eab9527de09..e9f5d58f8d5 100644
--- a/tensorflow/core/profiler/utils/group_events_test.cc
+++ b/tensorflow/core/profiler/utils/group_events_test.cc
@@ -36,8 +36,7 @@ TEST(GroupEventsTest, GroupGpuTraceTest) {
   constexpr int64 kCorrelationId = 100;
 
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
@@ -78,8 +77,7 @@ TEST(GroupEventsTest, GroupTensorFlowLoopTest) {
   constexpr int64 kCorrelationId = 100;
 
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(1);
 
   auto tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
@@ -125,8 +123,7 @@ TEST(GroupEventsTest, GroupMultipleTensorFlowLoopsTest) {
   constexpr int64 kSecondIterNumStart = 0;
 
   XSpace space;
-  XPlaneBuilder host_plane_builder(space.add_planes());
-  host_plane_builder.SetName(kHostThreads);
+  XPlaneBuilder host_plane_builder(GetOrCreateHostXPlane(&space));
   host_plane_builder.ReserveLines(2);
 
   auto first_tf_executor_thread = host_plane_builder.GetOrCreateLine(0);
@@ -163,9 +160,8 @@ TEST(GroupEventsTest, GroupFunctionalOp) {
   constexpr int64 kFunctionStepId = 1;
 
   XSpace space;
-  XPlane* host_plane = space.add_planes();
+  XPlane* host_plane = GetOrCreateHostXPlane(&space);
   XPlaneBuilder host_plane_builder(host_plane);
-  host_plane_builder.SetName(kHostThreads);
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
@@ -209,9 +205,8 @@ TEST(GroupEventsTest, EagerOpTest) {
   constexpr int64 kCorrelationId = 100;
 
   XSpace space;
-  XPlane* host_plane = space.add_planes();
+  XPlane* host_plane = GetOrCreateHostXPlane(&space);
   XPlaneBuilder host_plane_builder(host_plane);
-  host_plane_builder.SetName(kHostThreads);
   host_plane_builder.ReserveLines(1);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
@@ -255,9 +250,8 @@ TEST(GroupEventsTest, FunctionOpTest) {
   constexpr int64 kCorrelationId = 100;
 
   XSpace space;
-  XPlane* host_plane = space.add_planes();
+  XPlane* host_plane = GetOrCreateHostXPlane(&space);
   XPlaneBuilder host_plane_builder(host_plane);
-  host_plane_builder.SetName(kHostThreads);
   host_plane_builder.ReserveLines(2);
 
   auto main_thread = host_plane_builder.GetOrCreateLine(0);
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index c7a2cb6e37e..2c79df7980f 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -26,11 +26,11 @@ limitations under the License.
 namespace tensorflow {
 namespace profiler {
 
-const absl::string_view kHostThreads = "/host:CPU";
+const absl::string_view kHostThreadsPlaneName = "/host:CPU";
 const absl::string_view kGpuPlanePrefix = "/device:GPU:";
 const absl::string_view kCuptiDriverApiPlaneName = "/host:CUPTI";
-const absl::string_view kMetadataPlane = "/host:metadata";
-const absl::string_view kTFStreamzPlane = "/host:tfstreamz";
+const absl::string_view kMetadataPlaneName = "/host:metadata";
+const absl::string_view kTFStreamzPlaneName = "/host:tfstreamz";
 
 const absl::string_view kStepLineName = "Steps";
 const absl::string_view kTensorFlowNameScopeLineName = "TensorFlow Name Scope";
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 84b374e488d..a045e20d8de 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
 #define TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
 
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 #include "tensorflow/core/platform/logging.h"
@@ -25,15 +27,15 @@ namespace tensorflow {
 namespace profiler {
 
 // Name of XPlane that contains TraceMe events.
-ABSL_CONST_INIT extern const absl::string_view kHostThreads;
+ABSL_CONST_INIT extern const absl::string_view kHostThreadsPlaneName;
 // Name prefix of XPlane that contains GPU events.
 ABSL_CONST_INIT extern const absl::string_view kGpuPlanePrefix;
 // Name of XPlane that contains CUPTI driver API generated events.
 ABSL_CONST_INIT extern const absl::string_view kCuptiDriverApiPlaneName;
 // Name of XPlane that contains profile metadata such as XLA debug info.
-ABSL_CONST_INIT extern const absl::string_view kMetadataPlane;
+ABSL_CONST_INIT extern const absl::string_view kMetadataPlaneName;
 // Name of XPlane that contains kpi related metrics.
-ABSL_CONST_INIT extern const absl::string_view kTFStreamzPlane;
+ABSL_CONST_INIT extern const absl::string_view kTFStreamzPlaneName;
 
 // Names of XLines that contain ML-level events.
 ABSL_CONST_INIT extern const absl::string_view kStepLineName;
@@ -184,6 +186,14 @@ enum StatType {
   kLastStatType = kDevCapComputeCapMinor,
 };
 
+inline std::string GpuPlaneName(int32 device_ordinal) {
+  return absl::StrCat(kGpuPlanePrefix, device_ordinal);
+}
+
+inline bool IsGpuPlaneName(absl::string_view plane_name) {
+  return absl::StartsWith(plane_name, kGpuPlanePrefix);
+}
+
 absl::string_view GetHostEventTypeStr(HostEventType event_type);
 
 bool IsHostEventType(HostEventType event_type, absl::string_view event_name);
diff --git a/tensorflow/core/profiler/utils/xplane_test_utils.cc b/tensorflow/core/profiler/utils/xplane_test_utils.cc
index cd8821f05a8..a389d3619bd 100644
--- a/tensorflow/core/profiler/utils/xplane_test_utils.cc
+++ b/tensorflow/core/profiler/utils/xplane_test_utils.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
+#include "tensorflow/core/profiler/utils/xplane_utils.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -44,6 +45,15 @@ class XStatValueVisitor {
 
 }  // namespace
 
+XPlane* GetOrCreateHostXPlane(XSpace* space) {
+  return FindOrAddMutablePlaneWithName(space, kHostThreadsPlaneName);
+}
+
+XPlane* GetOrCreateGpuXPlane(XSpace* space, int32 device_ordinal) {
+  std::string name = GpuPlaneName(device_ordinal);
+  return FindOrAddMutablePlaneWithName(space, name);
+}
+
 void CreateXEvent(
     XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
     absl::string_view event_name, int64 offset_ps, int64 duration_ps,
diff --git a/tensorflow/core/profiler/utils/xplane_test_utils.h b/tensorflow/core/profiler/utils/xplane_test_utils.h
index 9abf09fc695..89fda765771 100644
--- a/tensorflow/core/profiler/utils/xplane_test_utils.h
+++ b/tensorflow/core/profiler/utils/xplane_test_utils.h
@@ -28,6 +28,10 @@ namespace profiler {
 
 using XStatValue = absl::variant;
 
+XPlane* GetOrCreateHostXPlane(XSpace* space);
+
+XPlane* GetOrCreateGpuXPlane(XSpace* space, int32 device_ordinal);
+
 void CreateXEvent(
     XPlaneBuilder* plane_builder, XLineBuilder* line_builder,
     absl::string_view event_name, int64 offset_ps, int64 duration_ps,
diff --git a/tensorflow/core/profiler/utils/xplane_utils.cc b/tensorflow/core/profiler/utils/xplane_utils.cc
index 10d40eee3b4..3fa421c3459 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.cc
+++ b/tensorflow/core/profiler/utils/xplane_utils.cc
@@ -65,12 +65,19 @@ std::vector FindPlanesWithPrefix(const XSpace& space,
   return result;
 }
 
-XPlane* GetOrCreatePlane(XSpace* space, absl::string_view name) {
+XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name) {
   for (XPlane& plane : *space->mutable_planes()) {
     if (plane.name() == name) return &plane;
   }
-  XPlane* plane = space->add_planes();
-  plane->set_name(std::string(name));
+  return nullptr;
+}
+
+XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name) {
+  XPlane* plane = FindMutablePlaneWithName(space, name);
+  if (plane == nullptr) {
+    plane = space->add_planes();
+    plane->set_name(name.data(), name.size());
+  }
   return plane;
 }
 
@@ -128,22 +135,6 @@ void RemoveEmptyLines(XPlane* plane) {
                lines->end());
 }
 
-XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name) {
-  for (XPlane& plane : *space->mutable_planes()) {
-    if (plane.name() == name) return &plane;
-  }
-  return nullptr;
-}
-
-XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name) {
-  XPlane* plane = FindMutablePlaneWithName(space, name);
-  if (plane == nullptr) {
-    plane = space->add_planes();
-    plane->set_name(std::string(name));
-  }
-  return plane;
-}
-
 void SortXPlane(XPlane* plane) {
   for (XLine& line : *plane->mutable_lines()) {
     auto& events = *line.mutable_events();
diff --git a/tensorflow/core/profiler/utils/xplane_utils.h b/tensorflow/core/profiler/utils/xplane_utils.h
index 11a2c28f719..7575244e7bd 100644
--- a/tensorflow/core/profiler/utils/xplane_utils.h
+++ b/tensorflow/core/profiler/utils/xplane_utils.h
@@ -32,8 +32,12 @@ const XPlane* FindPlaneWithName(const XSpace& space, absl::string_view name);
 std::vector FindPlanesWithPrefix(const XSpace& space,
                                                 absl::string_view prefix);
 
-// Returns the plane with the given name, create it if necessary.
-XPlane* GetOrCreatePlane(XSpace* space, absl::string_view name);
+// Returns the plane with the given name in the container or null if not found.
+XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name);
+
+// Returns the plane with the given name in the container. If necessary, adds a
+// new plane to the container.
+XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name);
 
 // Returns true if event is nested by parent.
 bool IsNested(const tensorflow::profiler::XEvent& event,
@@ -49,13 +53,6 @@ void RemovePlaneWithName(XSpace* space, absl::string_view name);
 void RemoveEmptyPlanes(XSpace* space);
 void RemoveEmptyLines(XPlane* plane);
 
-// Returns the plane with the given name in the container or null if not found.
-XPlane* FindMutablePlaneWithName(XSpace* space, absl::string_view name);
-
-// Returns the plane with the given name in the container. If necessary, adds a
-// new plane to the container.
-XPlane* FindOrAddMutablePlaneWithName(XSpace* space, absl::string_view name);
-
 // Sort lines in plane with a provided comparator.
 template 
 void SortXLinesBy(XPlane* plane, Compare comp) {

From a428ec179e3037c164886df692236561114830c0 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Mon, 15 Jun 2020 16:06:18 -0700
Subject: [PATCH 0224/1390] Added new compiler option.

PiperOrigin-RevId: 316564867
Change-Id: I5e0b54b99395d322041207294cbdda5efc62ff76
---
 tensorflow/lite/delegates/gpu/cl/cl_program.cc | 2 ++
 tensorflow/lite/delegates/gpu/cl/cl_program.h  | 3 ++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/cl_program.cc b/tensorflow/lite/delegates/gpu/cl/cl_program.cc
index 690bc598777..3b821dc3a5d 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_program.cc
+++ b/tensorflow/lite/delegates/gpu/cl/cl_program.cc
@@ -93,6 +93,8 @@ std::string CompilerOptionToString(const CLDevice& device,
       return "-cl-fast-relaxed-math";
     case CompilerOptions::CL_OPT_DISABLE:
       return "-cl-opt-disable";
+    case CompilerOptions::CL_2_0:
+      return "-cl-std=CL2.0";
   }
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/cl_program.h b/tensorflow/lite/delegates/gpu/cl/cl_program.h
index fb2a7edb9c1..138b7d9fbd0 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_program.h
+++ b/tensorflow/lite/delegates/gpu/cl/cl_program.h
@@ -40,7 +40,8 @@ enum class CompilerOptions {
   ADRENO_FULL_SIMD_LINE,
   ADRENO_MORE_WAVES,
   POWERVR_FP16,
-  CL_OPT_DISABLE
+  CL_OPT_DISABLE,
+  CL_2_0
 };
 
 std::string CompilerOptionsToString(

From 0c68775ed3682849849113db4c957a0a16e84000 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 16:13:08 -0700
Subject: [PATCH 0225/1390] Introduces a new experimental package that: -
 Defines a schema for configuring delegates - Defines a C++ plugin mechanism
 using the schema, so that code can support   configuring arbitrary delegates
 without a build-time dependency

PiperOrigin-RevId: 316566081
Change-Id: I4d36b4e155dd30fbdf57d60ef4b546304c033b1a
---
 .../acceleration/configuration/BUILD          | 160 --------------
 .../configuration/configuration.proto         | 208 ------------------
 .../configuration/delegate_registry.cc        |  60 -----
 .../configuration/delegate_registry.h         |  95 --------
 .../acceleration/configuration/gpu_plugin.cc  |  62 ------
 .../configuration/hexagon_plugin.cc           |  73 ------
 .../configuration/nnapi_plugin.cc             |  93 --------
 .../configuration/nnapi_plugin_test.cc        | 175 ---------------
 .../configuration/proto_to_flatbuffer.cc      |  58 -----
 .../configuration/proto_to_flatbuffer.h       |  32 ---
 10 files changed, 1016 deletions(-)
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/BUILD
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/configuration.proto
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
 delete mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h

diff --git a/tensorflow/lite/experimental/acceleration/configuration/BUILD b/tensorflow/lite/experimental/acceleration/configuration/BUILD
deleted file mode 100644
index 38d28d5cc2e..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/BUILD
+++ /dev/null
@@ -1,160 +0,0 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-
-load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_java_library", "flatc_path")
-
-package(
-    default_visibility = [
-        "//visibility:public",
-    ],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-genrule(
-    name = "configuration_schema",
-    srcs = ["configuration.proto"],
-    outs = ["configuration.fbs"],
-    # We rename the namespace since otherwise the proto classes and flatbuffer
-    # classes would have the same names.
-    cmd = """
-    $(location {}) --proto -o $(@D) $(location :configuration.proto)
-    perl -p -i -e 's/tflite.proto/tflite/' $(@D)/configuration.fbs
-    """.format(flatc_path),
-    tools = [
-        flatc_path,
-    ],
-)
-
-genrule(
-    name = "configuration_fbs_contents_cc",
-    srcs = ["configuration.fbs"],
-    outs = ["configuration_fbs_contents-inl.h"],
-    cmd = """
-      echo 'constexpr char configuration_fbs_contents[] = R"Delimiter(' > $(@)
-      cat < $(<) >> $(@)
-      echo ')Delimiter";' >> $(@)
-    """,
-)
-
-proto_library(
-    name = "configuration_proto",
-    srcs = [
-        "configuration.proto",
-    ],
-)
-
-cc_proto_library(
-    name = "configuration_cc_proto",
-    deps = [":configuration_proto"],
-)
-
-java_lite_proto_library(
-    name = "configuration_java_proto_lite",
-    deps = [":configuration_proto"],
-)
-
-flatbuffer_cc_library(
-    name = "configuration_fbs",
-    srcs = [":configuration.fbs"],
-)
-
-flatbuffer_java_library(
-    name = "configuration_fbs_java",
-    srcs = [":configuration.fbs"],
-)
-
-cc_library(
-    name = "proto_to_flatbuffer",
-    srcs = [
-        "configuration_fbs_contents-inl.h",
-        "proto_to_flatbuffer.cc",
-    ],
-    hdrs = ["proto_to_flatbuffer.h"],
-    deps = [
-        ":configuration_cc_proto",
-        ":configuration_fbs",
-        "//tensorflow/core/platform:protobuf",
-        "//tensorflow/lite:minimal_logging",
-        "@flatbuffers",
-    ],
-)
-
-cc_library(
-    name = "delegate_registry",
-    srcs = ["delegate_registry.cc"],
-    hdrs = ["delegate_registry.h"],
-    deps = [
-        ":configuration_fbs",
-        "//tensorflow/lite/c:common",
-        "@com_google_absl//absl/synchronization",
-    ],
-)
-
-cc_library(
-    name = "nnapi_plugin",
-    srcs = ["nnapi_plugin.cc"],
-    deps = [
-        ":configuration_fbs",
-        ":delegate_registry",
-        "//tensorflow/lite/delegates/nnapi:nnapi_delegate",
-        "@com_google_absl//absl/memory",
-    ],
-    alwayslink = 1,  # For registration to always run.
-)
-
-cc_test(
-    name = "nnapi_plugin_test",
-    srcs = ["nnapi_plugin_test.cc"],
-    deps = [
-        ":configuration_fbs",
-        ":delegate_registry",
-        ":nnapi_plugin",
-        "//tensorflow/lite:framework",
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/delegates/nnapi:nnapi_delegate",
-        "//tensorflow/lite/delegates/nnapi:nnapi_delegate_mock_test",
-        "//tensorflow/lite/kernels:test_util",
-        "@com_google_googletest//:gtest_main",
-        "@flatbuffers",
-    ],
-)
-
-cc_library(
-    name = "hexagon_plugin",
-    srcs = ["hexagon_plugin.cc"],
-    deps = [
-        ":configuration_fbs",
-        ":delegate_registry",
-        "@com_google_absl//absl/memory",
-    ] + select({
-        "//tensorflow:android": [
-            "//tensorflow/lite/delegates/hexagon:hexagon_delegate",
-        ],
-        "//conditions:default": [],
-    }),
-    alwayslink = 1,  # For registration to always run.
-)
-
-cc_library(
-    name = "gpu_plugin",
-    srcs = ["gpu_plugin.cc"],
-    deps = [
-        ":configuration_fbs",
-        ":delegate_registry",
-        "//tensorflow/lite/delegates/gpu:delegate",
-        "@com_google_absl//absl/memory",
-    ],
-    alwayslink = 1,  # For registration to always run.
-)
diff --git a/tensorflow/lite/experimental/acceleration/configuration/configuration.proto b/tensorflow/lite/experimental/acceleration/configuration/configuration.proto
deleted file mode 100644
index e1c49f02856..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/configuration.proto
+++ /dev/null
@@ -1,208 +0,0 @@
-// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-//
-// 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
-//
-//     http://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.
-
-// This schema defines how to configure TFLite for delegation. These
-// definitions can be used in multiple ways: as output of a compatibility list,
-// in benchmarking tools and to decouple delegate instantiation from code.
-//
-// The schema is work-in-progress, covering the most broadly used delegates and
-// options.
-
-syntax = "proto2";
-
-package tflite.proto;
-
-// ExecutionPreference is used to match accelerators against the preferences of
-// the current application or usecase. Some of the values here can appear both
-// in the compatibility list and as input, some only as input.
-//
-// These are separate from NNAPIExecutionPreference - the compatibility list
-// design doesn't assume a one-to-one mapping between which usecases
-// compatibility list entries have been developed for and what settings are used
-// for NNAPI.
-enum ExecutionPreference {
-  // Match any selected preference. Whitelist (semantically - value is same as
-  // on input).
-  ANY = 0;
-  // Match low latency preference. Both compatibility list and input.
-  LOW_LATENCY = 1;
-  // Math low power preference. Both compatibility list and input.
-  LOW_POWER = 2;
-  // Never accelerate. Can be used for input to compatibility list or for
-  // standalone Acceleration configuration.
-  FORCE_CPU = 3;
-}
-
-// TFLite delegate to use.
-enum Delegate {
-  NONE = 0;
-  NNAPI = 1;
-  GPU = 2;
-  HEXAGON = 3;
-  XNNPACK = 4;
-  // TODO(b/157893534): Support exposing edgetpu tflite delegate creation
-  // options.
-  EDGETPU = 5;
-}
-
-enum NNAPIExecutionPreference {
-  // Undefined.
-  UNDEFINED = 0;
-  // Prefer executing in a way that minimizes battery drain.
-  NNAPI_LOW_POWER = 1;
-  // Prefer returning a single answer as fast as possible, even if this causes
-  // more power consumption.
-  NNAPI_FAST_SINGLE_ANSWER = 2;
-  // Prefer maximizing the throughput of successive frames, for example when
-  // processing successive frames coming from the camera.
-  NNAPI_SUSTAINED_SPEED = 3;
-}
-
-// One possible acceleration configuration.
-message ComputeSettings {
-  // Which preference to use this accelerator for.
-  optional ExecutionPreference preference = 1;
-  // How to configure TFLite
-  optional TFLiteSettings tflite_settings = 2;
-  // Identifiers to use for instrumentation and telemetry.
-  optional string model_namespace_for_statistics = 3;
-  optional string model_identifier_for_statistics = 4;
-}
-
-// NNAPI delegate settings.
-message NNAPISettings {
-  // Which instance (NNAPI accelerator) to use. One driver may provide several
-  // accelerators (though a driver may also hide several back-ends behind one
-  // name, at the choice of the driver vendor).
-  // Note that driver introspection is only available in Android Q and later.
-  optional string accelerator_name = 1;
-
-  // NNAPI model compilation caching settings to be passed to
-  // tflite::StatefulNnApiDelegate
-  optional string cache_directory = 2;
-  optional string model_token = 3;
-
-  // NNAPI execution preference to pass. See
-  // https://developer.android.com/ndk/reference/group/neural-networks.html
-  optional NNAPIExecutionPreference execution_preference = 4;
-
-  // Number of instances to cache for the same model (for input size
-  // changes). This is mandatory for getting reasonable performance in that
-  // case.
-  optional int32 no_of_nnapi_instances_to_cache = 5;
-
-  // Whether to automatically fall back to TFLite CPU path.
-  optional FallbackSettings fallback_settings = 6;
-
-  // Whether to allow use of NNAPI CPU (nnapi-reference accelerator) on Android
-  // 10+ when an accelerator name is not specified. The NNAPI CPU typically
-  // performs less well than the TfLite built-in kernels; but allowing allows a
-  // model to be partially accelerated which may be a win.
-  optional bool allow_nnapi_cpu_on_android_10_plus = 7;
-}
-
-// Which GPU backend to select. Default behaviour on Android is to try OpenCL
-// and if it's not available fall back to OpenGL.
-enum GPUBackend {
-  UNSET = 0;
-  OPENCL = 1;
-  OPENGL = 2;
-  // Not yet supported.
-  // VULKAN = 3;
-  // METAL = 4;
-}
-
-// GPU Delegate settings.
-//
-// See
-// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/gpu/delegate.h
-message GPUSettings {
-  optional bool is_precision_loss_allowed = 1;
-  optional bool enable_quantized_inference = 2 [default = true];
-  optional GPUBackend force_backend = 3;
-  // TODO(b/152019007): add remaining options.
-}
-
-// Hexagon Delegate settings.
-//
-// See
-// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/hexagon/hexagon_delegate.h
-message HexagonSettings {
-  optional int32 debug_level = 1;
-  optional int32 powersave_level = 2;
-  optional bool print_graph_profile = 3;
-  optional bool print_graph_debug = 4;
-}
-
-// XNNPack Delegate settings.
-//
-// See
-// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
-message XNNPackSettings {
-  optional int32 num_threads = 1;
-}
-
-message CPUSettings {
-  optional int32 num_threads = 1;
-}
-
-// How to configure TFLite.
-message TFLiteSettings {
-  // Which delegate to use.
-  optional Delegate delegate = 1;
-
-  // How to configure the chosen delegate.
-  // (In principle we would like to use 'oneof', but flatc turns that into an
-  // nested anonymous table rather than a union. See
-  // https://github.com/google/flatbuffers/issues/4628).
-  optional NNAPISettings nnapi_settings = 2;
-  optional GPUSettings gpu_settings = 3;
-  optional HexagonSettings hexagon_settings = 4;
-  optional XNNPackSettings xnnpack_settings = 5;
-
-  // How to configure CPU execution.
-  optional CPUSettings cpu_settings = 6;
-
-  // Shared delegation settings.
-  optional int32 max_delegated_partitions = 7;
-}
-
-// Whether to automatically fallback to TFLite CPU path on delegation errors.
-//
-// Typically fallback is enabled in production use but disabled in tests and
-// benchmarks to ensure they test the intended path.
-message FallbackSettings {
-  // Whether to allow automatically falling back to TfLite CPU path on
-  // compilation failure. Default is not allowing automatic fallback.
-  //
-  // This is useful in naive production usecases where the caller would prefer
-  // for the model to run even if it's not accelerated. More advanced users will
-  // implement fallback themselves; e.g., by using a different model on CPU.
-  //
-  // Note that compilation errors may occur either at initial
-  // ModifyGraphWithDelegate() time, or when calling AllocateTensors() after
-  // resizing.
-  optional bool allow_automatic_fallback_on_compilation_error = 7;
-  // Whether to allow automatically falling back to TfLite CPU path on
-  // execution error. Default is not allowing automatic fallback.
-  //
-  // Experimental, use with care (only when you have complete control over the
-  // client code).
-  //
-  // The caveat above for compilation error holds.  Additionally, execution-time
-  // errors are harder to handle automatically as they require invalidating the
-  // TfLite interpreter which most client code has not been designed to deal
-  // with.
-  optional bool allow_automatic_fallback_on_execution_error = 8;
-}
diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
deleted file mode 100644
index b8d80342d5f..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
+++ /dev/null
@@ -1,60 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
-
-#include "absl/synchronization/mutex.h"
-
-namespace tflite {
-namespace delegates {
-
-void DelegatePluginRegistry::RegisterImpl(
-    const std::string& name,
-    std::function<
-        std::unique_ptr(const TFLiteSettings&)>
-        creator_function) {
-  absl::MutexLock lock(&mutex_);
-  factories_[name] = creator_function;
-}
-
-std::unique_ptr DelegatePluginRegistry::CreateImpl(
-    const std::string& name, const TFLiteSettings& settings) {
-  absl::MutexLock lock(&mutex_);
-  auto it = factories_.find(name);
-  if (it != factories_.end()) {
-    return it->second(settings);
-  } else {
-    return nullptr;
-  }
-}
-
-DelegatePluginRegistry* DelegatePluginRegistry::GetSingleton() {
-  static auto* instance = new DelegatePluginRegistry();
-  return instance;
-}
-
-std::unique_ptr DelegatePluginRegistry::CreateByName(
-    const std::string& name, const TFLiteSettings& settings) {
-  auto* const instance = DelegatePluginRegistry::GetSingleton();
-  return instance->CreateImpl(name, settings);
-}
-
-DelegatePluginRegistry::Register::Register(const std::string& name,
-                                           CreatorFunction creator_function) {
-  auto* const instance = DelegatePluginRegistry::GetSingleton();
-  instance->RegisterImpl(name, creator_function);
-}
-
-}  // namespace delegates
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
deleted file mode 100644
index c86759dcc3f..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
+++ /dev/null
@@ -1,95 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
-#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
-
-#include 
-#include 
-
-#include "absl/synchronization/mutex.h"
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-
-// Defines an interface for TFLite delegate plugins.
-//
-// The acceleration library aims to support all TFLite delegates based on
-// configuration expressed as data (flatbuffers). However, consumers tend to
-// care about size and also use a subset of delegates. Hence we don't want to
-// statically build against all delegates.
-//
-// This interface allows plugins to handle specific delegates.
-//
-// Goal of this interface is not to abstract away all the differences between
-// delegates. The goal is only to avoid static linking.
-//
-// Note to implementers: this interface may change if new delegates don't fit
-// into the same design.
-namespace tflite {
-namespace delegates {
-
-// Same w/ Interpreter::TfLiteDelegatePtr to avoid pulling
-// tensorflow/lite/interpreter.h dependency
-using TfLiteDelegatePtr =
-    std::unique_ptr;
-
-class DelegatePluginInterface {
- public:
-  virtual TfLiteDelegatePtr Create() = 0;
-  virtual int GetDelegateErrno(TfLiteDelegate* from_delegate) = 0;
-  virtual ~DelegatePluginInterface() = default;
-};
-
-// A stripped-down registry that allows delegate plugins to be created by name.
-//
-// Limitations:
-// - Doesn't allow deregistration.
-// - Doesn't check for duplication registration.
-//
-class DelegatePluginRegistry {
- public:
-  typedef std::function(
-      const TFLiteSettings&)>
-      CreatorFunction;
-  // Returns a DelegatePluginInterface registered with `name` or nullptr if no
-  // matching plugin found.
-  // TFLiteSettings is per-plugin, so that the corresponding delegate options
-  // data lifetime is maintained.
-  static std::unique_ptr CreateByName(
-      const std::string& name, const TFLiteSettings& settings);
-
-  // Struct to be statically allocated for registration.
-  struct Register {
-    Register(const std::string& name, CreatorFunction creator_function);
-  };
-
- private:
-  void RegisterImpl(const std::string& name, CreatorFunction creator_function);
-  std::unique_ptr CreateImpl(
-      const std::string& name, const TFLiteSettings& settings);
-  static DelegatePluginRegistry* GetSingleton();
-  std::unordered_map factories_;
-  absl::Mutex mutex_;
-};
-
-}  // namespace delegates
-}  // namespace tflite
-
-#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f) \
-  static auto* g_delegate_plugin_##name##_ =                     \
-      new DelegatePluginRegistry::Register(#name, f);
-#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(name, f) \
-  TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f);
-
-#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
diff --git a/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
deleted file mode 100644
index 25b8171c5ea..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
+++ /dev/null
@@ -1,62 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include 
-
-#include "absl/memory/memory.h"
-#include "tensorflow/lite/delegates/gpu/delegate.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
-
-namespace tflite {
-namespace delegates {
-class GpuPlugin : public DelegatePluginInterface {
- public:
-  TfLiteDelegatePtr Create() override {
-    return TfLiteDelegatePtr(TfLiteGpuDelegateV2Create(&options_),
-                             TfLiteGpuDelegateV2Delete);
-  }
-  int GetDelegateErrno(TfLiteDelegate* from_delegate) override { return 0; }
-  static std::unique_ptr New(
-      const TFLiteSettings& acceleration) {
-    return absl::make_unique(acceleration);
-  }
-  explicit GpuPlugin(const TFLiteSettings& tflite_settings)
-      : options_(TfLiteGpuDelegateOptionsV2Default()) {
-    const auto* gpu_settings = tflite_settings.gpu_settings();
-    if (gpu_settings) {
-      options_.inference_priority1 =
-          gpu_settings->is_precision_loss_allowed()
-              ? TFLITE_GPU_INFERENCE_PRIORITY_MIN_LATENCY
-              : TFLITE_GPU_INFERENCE_PRIORITY_MAX_PRECISION;
-      if (gpu_settings->enable_quantized_inference()) {
-        options_.experimental_flags |=
-            TFLITE_GPU_EXPERIMENTAL_FLAGS_ENABLE_QUANT;
-      }
-      if (gpu_settings->force_backend() == GPUBackend_OPENCL) {
-        options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_CL_ONLY;
-      } else if (gpu_settings->force_backend() == GPUBackend_OPENGL) {
-        options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_GL_ONLY;
-      }
-    }
-  }
-
- private:
-  TfLiteGpuDelegateOptionsV2 options_;
-};
-
-TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(GpuPlugin, GpuPlugin::New);
-
-}  // namespace delegates
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
deleted file mode 100644
index 7f2674604b0..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
+++ /dev/null
@@ -1,73 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include 
-
-#include "absl/memory/memory.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
-
-#if defined(__ARM_ARCH)
-#include "tensorflow/lite/delegates/hexagon/hexagon_delegate.h"
-#endif
-
-namespace tflite {
-namespace delegates {
-class HexagonPlugin : public DelegatePluginInterface {
- public:
-  TfLiteDelegatePtr Create() override {
-#if defined(__ARM_ARCH)
-    TfLiteHexagonInit();
-    auto* delegate_ptr = TfLiteHexagonDelegateCreate(&options_);
-    TfLiteDelegatePtr delegate(delegate_ptr, [](TfLiteDelegate* delegate) {
-      TfLiteHexagonDelegateDelete(delegate);
-      TfLiteHexagonTearDown();
-    });
-    return delegate;
-#else   // !defined(__ARM_ARCH)
-    return TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
-#endif  // defined(__ARM_ARCH)
-  }
-  int GetDelegateErrno(TfLiteDelegate* /* from_delegate */) override {
-    return 0;
-  }
-  static std::unique_ptr New(
-      const TFLiteSettings& tflite_settings) {
-    return absl::make_unique(tflite_settings);
-  }
-  explicit HexagonPlugin(const TFLiteSettings& tflite_settings) {
-    const HexagonSettings* settings = tflite_settings.hexagon_settings();
-#if defined(__ARM_ARCH)
-    options_ = TfLiteHexagonDelegateOptions({0});
-    if (settings) {
-      options_.debug_level = settings->debug_level();
-      options_.powersave_level = settings->powersave_level();
-      options_.print_graph_profile = settings->print_graph_profile();
-      options_.print_graph_debug = settings->print_graph_debug();
-    }
-#else
-    (void)settings;
-#endif
-  }
-
- private:
-#if defined(__ARM_ARCH)
-  TfLiteHexagonDelegateOptions options_;
-#endif
-};
-
-TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(HexagonPlugin, HexagonPlugin::New);
-
-}  // namespace delegates
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
deleted file mode 100644
index 7301983a815..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
+++ /dev/null
@@ -1,93 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include 
-
-#include "absl/memory/memory.h"
-#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
-
-namespace tflite {
-namespace delegates {
-
-inline tflite::StatefulNnApiDelegate::Options::ExecutionPreference
-ConvertExecutionPrefence(
-    NNAPIExecutionPreference from_compatibility_preference) {
-  using TflitePreference =
-      tflite::StatefulNnApiDelegate::Options::ExecutionPreference;
-  switch (from_compatibility_preference) {
-    case NNAPIExecutionPreference_NNAPI_LOW_POWER:
-      return TflitePreference::kLowPower;
-    case NNAPIExecutionPreference_NNAPI_FAST_SINGLE_ANSWER:
-      return TflitePreference::kFastSingleAnswer;
-    case NNAPIExecutionPreference_NNAPI_SUSTAINED_SPEED:
-      return TflitePreference::kSustainedSpeed;
-    default:
-      return TflitePreference::kUndefined;
-  }
-}
-
-class NnapiPlugin : public DelegatePluginInterface {
- public:
-  TfLiteDelegatePtr Create() override {
-    auto nnapi_delegate =
-        absl::make_unique(options_);
-    return TfLiteDelegatePtr(
-        nnapi_delegate.release(), [](TfLiteDelegate* delegate) {
-          delete reinterpret_cast(delegate);
-        });
-  }
-  int GetDelegateErrno(TfLiteDelegate* from_delegate) override {
-    auto nnapi_delegate =
-        reinterpret_cast(from_delegate);
-    return nnapi_delegate->GetNnApiErrno();
-  }
-  static std::unique_ptr New(
-      const TFLiteSettings& tflite_settings) {
-    return absl::make_unique(tflite_settings);
-  }
-  explicit NnapiPlugin(const TFLiteSettings& tflite_settings) {
-    const NNAPISettings* nnapi_settings = tflite_settings.nnapi_settings();
-    if (!nnapi_settings) return;
-    if (nnapi_settings->accelerator_name() &&
-        nnapi_settings->accelerator_name()->Length() != 0) {
-      accelerator_ = nnapi_settings->accelerator_name()->str();
-      options_.accelerator_name = accelerator_.c_str();
-    }
-    if (nnapi_settings->cache_directory() &&
-        nnapi_settings->cache_directory()->Length() != 0) {
-      cache_dir_ = nnapi_settings->cache_directory()->str();
-      options_.cache_dir = cache_dir_.c_str();
-    }
-    if (nnapi_settings->model_token() &&
-        nnapi_settings->model_token()->Length() != 0) {
-      model_token_ = nnapi_settings->model_token()->str();
-      options_.model_token = model_token_.c_str();
-    }
-    options_.execution_preference =
-        ConvertExecutionPrefence(nnapi_settings->execution_preference());
-    options_.disallow_nnapi_cpu =
-        !nnapi_settings->allow_nnapi_cpu_on_android_10_plus();
-  }
-
- private:
-  std::string accelerator_, cache_dir_, model_token_;
-  tflite::StatefulNnApiDelegate::Options options_;
-};
-
-TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(NnapiPlugin, NnapiPlugin::New);
-
-}  // namespace delegates
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
deleted file mode 100644
index 4f9f5dd08c1..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
+++ /dev/null
@@ -1,175 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include 
-
-#include 
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
-#include "tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/test_util.h"
-
-// Tests for checking that the NNAPI Delegate plugin correctly handles all the
-// options from the flatbuffer.
-//
-// Checking done at NNAPI call level, as that is where we have a mockable
-// layer.
-namespace tflite {
-namespace {
-
-using delegate::nnapi::NnApiMock;
-
-class SingleAddOpModel : tflite::SingleOpModel {
- public:
-  void Build() {
-    int input = AddInput({tflite::TensorType_FLOAT32, {1, 2, 2}});
-    int constant = AddConstInput({tflite::TensorType_FLOAT32, {1, 2, 2}},
-                                 {1.0f, 1.0f, 1.0f, 1.0f});
-    AddOutput({tflite::TensorType_FLOAT32, {}});
-
-    SetBuiltinOp(tflite::BuiltinOperator_ADD, tflite::BuiltinOptions_AddOptions,
-                 tflite::CreateAddOptions(builder_).Union());
-    BuildInterpreter({GetShape(input), GetShape(constant)});
-  }
-
-  tflite::Interpreter* Interpreter() const { return interpreter_.get(); }
-};
-
-class NNAPIPluginTest : public ::testing::Test {
- protected:
-  NNAPIPluginTest() : delegate_(nullptr, [](TfLiteDelegate*) {}) {}
-  void SetUp() override {
-    nnapi_ = const_cast(NnApiImplementation());
-    nnapi_mock_ = absl::make_unique(nnapi_);
-    nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
-        [](const ANeuralNetworksModel* model,
-           const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
-           bool* supportedOps) -> int {
-      supportedOps[0] = true;
-      return 0;
-    };
-    model_.Build();
-  }
-  template 
-  void CheckExecutionPreference() {
-    // Note - this uses a template since the NNAPI functions are C function
-    // pointers rather than lambdas so can't capture variables.
-    nnapi_->ANeuralNetworksCompilation_setPreference =
-        [](ANeuralNetworksCompilation* compilation, int32_t preference) {
-          return preference - output;
-        };
-    CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0, input));
-    // Since delegation succeeds, the model becomes immutable and hence can't
-    // reuse it.
-    SingleAddOpModel model;
-    model.Build();
-    EXPECT_EQ(model.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-              kTfLiteOk)
-        << " given input: " << input << " expected output: " << output;
-  }
-
-  void CreateDelegate(flatbuffers::Offset settings) {
-    settings_ = flatbuffers::GetTemporaryPointer(
-        fbb_, CreateTFLiteSettings(fbb_, tflite::Delegate_NNAPI, settings));
-
-    plugin_ = delegates::DelegatePluginRegistry::CreateByName("NnapiPlugin",
-                                                              *settings_);
-    delegate_ = plugin_->Create();
-  }
-
-  NnApi* nnapi_;
-  std::unique_ptr nnapi_mock_;
-  SingleAddOpModel model_;
-  flatbuffers::FlatBufferBuilder fbb_;
-  const TFLiteSettings* settings_ = nullptr;
-  delegates::TfLiteDelegatePtr delegate_;
-  std::unique_ptr plugin_;
-};
-
-TEST_F(NNAPIPluginTest, PassesAcceleratorName) {
-  // Fails with non-existent "foo".
-  CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("foo")));
-  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-            kTfLiteDelegateError);
-
-  // Succeeds with "test-device" supported by the mock.
-  CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("test-device")));
-  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-            kTfLiteOk);
-}
-
-TEST_F(NNAPIPluginTest, PassesExecutionPreference) {
-  CheckExecutionPreference();
-  CheckExecutionPreference();
-  CheckExecutionPreference();
-  CheckExecutionPreference();
-}
-
-TEST_F(NNAPIPluginTest, PassesCachingParameters) {
-  nnapi_->ANeuralNetworksCompilation_setCaching =
-      [](ANeuralNetworksCompilation* compilation, const char* cacheDir,
-         const uint8_t* token) -> int {
-    if (std::string(cacheDir) != "d") return 1;
-    // Token is hashed with other bits, just check that it's not empty.
-    if (std::string(reinterpret_cast(token)).empty()) return 2;
-    return 0;
-  };
-  CreateDelegate(CreateNNAPISettings(fbb_, 0, fbb_.CreateString("d"),
-                                     fbb_.CreateString("t")));
-  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-            kTfLiteOk);
-}
-
-TEST_F(NNAPIPluginTest, PassesFalseNNAPICpuFlag) {
-  CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0,
-                                     NNAPIExecutionPreference_UNDEFINED, 0, 0,
-                                     /* allow CPU */ false));
-  nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
-      [](const ANeuralNetworksModel* model,
-         const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
-         bool* supportedOps) -> int {
-    supportedOps[0] = true;
-    // Since no CPU, should only pass one device.
-    return numDevices - 1;
-  };
-  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-            kTfLiteOk);
-}
-
-TEST_F(NNAPIPluginTest, PassesTrueNNAPICpuFlag) {
-  CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0,
-                                     NNAPIExecutionPreference_UNDEFINED, 0, 0,
-                                     /* allow CPU */ true));
-  nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
-      [](const ANeuralNetworksModel* model,
-         const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
-         bool* supportedOps) -> int {
-    supportedOps[0] = true;
-    // With CPU allowed, should pass two devices.
-    return numDevices - 2;
-  };
-  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
-            kTfLiteOk);
-}
-
-}  // namespace
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
deleted file mode 100644
index 709bb70ca70..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
+++ /dev/null
@@ -1,58 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include "tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h"
-
-#include 
-
-#include "flatbuffers/idl.h"  // from @flatbuffers
-#include "flatbuffers/util.h"  // from @flatbuffers
-#include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/lite/minimal_logging.h"
-
-namespace tflite {
-
-namespace {
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_fbs_contents-inl.h"
-}
-
-const ComputeSettings* ConvertFromProto(
-    flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings) {
-  std::string json;
-  tensorflow::protobuf::util::JsonPrintOptions options;
-  options.preserve_proto_field_names = true;
-  options.always_print_primitive_fields = true;  // For catching problems.
-  auto status = tensorflow::protobuf::util::MessageToJsonString(proto_settings,
-                                                                &json, options);
-  if (!status.ok()) {
-    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to convert to Json: %s",
-                    status.ToString().c_str());
-    return nullptr;
-  }
-  if (!parser->Parse(configuration_fbs_contents)) {
-    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse schema: %s",
-                    parser->error_.c_str());
-    return nullptr;
-  }
-  parser->SetRootType("tflite.ComputeSettings");
-  if (!parser->Parse(json.c_str())) {
-    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse json: %s",
-                    parser->error_.c_str());
-    return nullptr;
-  }
-  return flatbuffers::GetRoot(
-      parser->builder_.GetBufferPointer());
-}
-
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h
deleted file mode 100644
index 3b69e8465a5..00000000000
--- a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h
+++ /dev/null
@@ -1,32 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_
-#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_
-
-#include "flatbuffers/idl.h"  // from @flatbuffers
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration.pb.h"
-#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
-
-namespace tflite {
-
-// Converts the protobuf version ComputeSettings to the flatbuffer version, via
-// json. The parser is used for state - the returned pointer is valid only as
-// long as the parser is kept alive and unmutated.
-const ComputeSettings* ConvertFromProto(
-    flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings);
-
-}  // namespace tflite
-
-#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_

From f8410051b094aa686640f4325ebb134ceb13ca2d Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Mon, 15 Jun 2020 16:14:45 -0700
Subject: [PATCH 0226/1390] Remove a test of Keras implementation details that
 are subject to change. Specifically, the test checks that nested layers in a
 functional model have their _keras_history updated when that model is used to
 construct a bigger functional model.

In the future this incidental internal-implementation-specific-behavior will break.

PiperOrigin-RevId: 316566350
Change-Id: Ifb1e5ab04aafbbaa45e420ed34238059b16bcbdc
---
 tensorflow/python/keras/engine/functional_test.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/tensorflow/python/keras/engine/functional_test.py b/tensorflow/python/keras/engine/functional_test.py
index 25b433ce582..a7e314d4a49 100644
--- a/tensorflow/python/keras/engine/functional_test.py
+++ b/tensorflow/python/keras/engine/functional_test.py
@@ -2073,18 +2073,6 @@ class CacheCorrectnessTest(keras_parameterized.TestCase):
     # `None` value passed during construction is overridden.
     self.assertAllEqual(network(x, training=False), x * 0.0)
 
-  def test_keras_history_propagation_(self):
-    for input_shape in [(1,), (1, 1)]:
-      sub_in = input_layer_lib.Input((1,))
-      relu_layer = layers.ReLU()
-      sub_out = relu_layer(sub_in)
-      submodel = functional.Functional(sub_in, sub_out)
-      self.assertLen(relu_layer._inbound_nodes, 1)
-
-      inp = input_layer_lib.Input(input_shape)
-      submodel(inp)
-      self.assertLen(relu_layer._inbound_nodes, 2)
-
 
 if __name__ == '__main__':
   test.main()

From 2b7baa3ba36f0e346e60086232631a41db83a224 Mon Sep 17 00:00:00 2001
From: Rahul Joshi 
Date: Mon, 15 Jun 2020 16:21:27 -0700
Subject: [PATCH 0227/1390] Change HLO importer to set visibility when
 importing.

PiperOrigin-RevId: 316567573
Change-Id: I3e04c34ce022563f03f52a3bc5d24d20c4c90b0d
---
 .../mlir/xla/hlo_function_importer.cc         |  3 +
 .../mlir/xla/tests/translate/import.hlotxt    | 87 ++++++++++---------
 2 files changed, 47 insertions(+), 43 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/hlo_function_importer.cc b/tensorflow/compiler/mlir/xla/hlo_function_importer.cc
index a3b6222a8af..0627f2587de 100644
--- a/tensorflow/compiler/mlir/xla/hlo_function_importer.cc
+++ b/tensorflow/compiler/mlir/xla/hlo_function_importer.cc
@@ -115,6 +115,9 @@ StatusOr HloFunctionImporter::ImportAsFunc(
   llvm::ArrayRef attrs;
   auto function = mlir::FuncOp::create(mlir::UnknownLoc::get(context_),
                                        computation_name, func_type, attrs);
+  auto visibility = computation_name == "main" ? FuncOp::Visibility::Public
+                                               : FuncOp::Visibility::Private;
+  function.setVisibility(visibility);
   module_.push_back(function);
 
   // Add to the map right away for function calls.
diff --git a/tensorflow/compiler/mlir/xla/tests/translate/import.hlotxt b/tensorflow/compiler/mlir/xla/tests/translate/import.hlotxt
index 4565e1e4938..6336d6ed688 100644
--- a/tensorflow/compiler/mlir/xla/tests/translate/import.hlotxt
+++ b/tensorflow/compiler/mlir/xla/tests/translate/import.hlotxt
@@ -1,4 +1,4 @@
-// RUN: tf-mlir-translate -hlo-text-to-mlir-hlo %s -o - | FileCheck %s
+// RUN: tf-mlir-translate -hlo-text-to-mlir-hlo %s -o - | FileCheck %s -DPRIVATE="attributes {sym_visibility = \"private\"}"
 
 HloModule main
 
@@ -8,6 +8,7 @@ ENTRY %dummy_main (Arg_0.1: f32[]) -> f32[] {
 }
 
 // CHECK-LABEL:  func @test_simple
+// CHECK-SAME: [[PRIVATE]]
 %test_simple (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[] {
   %Arg_0.1 = f32[4]{0} parameter(0)
   %Arg_1.2 = f32[4]{0} parameter(1)
@@ -21,7 +22,7 @@ ENTRY %dummy_main (Arg_0.1: f32[]) -> f32[] {
 }
 
 // CHECK-LABEL:  func @test_after_all
-// CHECK-SAME:  ([[VAL_0:%.*]]: !xla_hlo.token, [[VAL_1:%.*]]: !xla_hlo.token) -> !xla_hlo.token
+// CHECK-SAME:  ([[VAL_0:%.*]]: !xla_hlo.token, [[VAL_1:%.*]]: !xla_hlo.token) -> !xla_hlo.token [[PRIVATE]]
 %test_after_all (token0: token[], token1: token[] ) -> token[] {
   token0 = token[] parameter(0)
   token1 = token[] parameter(1)
@@ -95,7 +96,7 @@ add {
   ROOT %batch-norm-grad = (f32[2,2,2,2], f32[2], f32[2]) batch-norm-grad(f32[2,2,2,2] %input, f32[2] %scale, f32[2] %mean, f32[2] %variance, f32[2,2,2,2] %grad_output), epsilon=0.001, feature_index=1
 }
 
-// CHECK-LABEL:  func @call(%arg0: tensor) -> tensor {
+// CHECK-LABEL:  func @call(%arg0: tensor) -> tensor
 %call (arg_1: s64[]) -> s64[] {
   %arg_1 = s64[] parameter(0), metadata={op_name="HLO_Args"}
   ROOT %compare.2 = s64[] add(%arg_1, %arg_1), metadata={op_type="Less" op_name="Less"}
@@ -136,7 +137,7 @@ add {
 }
 
 
-// CHECK-LABEL:  func @test_compare(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>, %arg2: tensor<3xf32>) -> tensor<3xi1> {
+// CHECK-LABEL:  func @test_compare(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>, %arg2: tensor<3xf32>) -> tensor<3xi1>
 %test_compare (Arg_0.1: f32[3], Arg_1.2: f32[3], Arg_2.3: f32[3]) -> pred[3] {
   %Arg_0.1 = f32[3] parameter(0)
   %Arg_1.2 = f32[3] parameter(1)
@@ -162,7 +163,7 @@ add {
   ROOT %complex.3 = c64[4] complex(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_concat(%arg0: tensor<4x1xf32>, %arg1: tensor<4x2xf32>) -> tensor<4x3xf32> {
+// CHECK-LABEL:  func @test_concat(%arg0: tensor<4x1xf32>, %arg1: tensor<4x2xf32>) -> tensor<4x3xf32>
 %test_concat (Arg_0.1: f32[4, 1], Arg_1.2: f32[4, 2]) -> f32[4, 3] {
   %Arg_0.1 = f32[4, 1] parameter(0)
   %Arg_1.2 = f32[4, 2] parameter(1)
@@ -201,7 +202,7 @@ add {
 
 // TODO(b/129422361) Potentially update when copy, reshape, and conv have actual
 // implementations with attributes, etc.
-// CHECK-LABEL:  func @test_conv(%arg0: tensor<256x32x32x6xf32>) -> tuple> {
+// CHECK-LABEL:  func @test_conv(%arg0: tensor<256x32x32x6xf32>) -> tuple>
 %test_conv {
   %arg0.1 = f32[256,32,32,6]{3,2,1,0} parameter(0), metadata={op_name="HLO_Args"}
 
@@ -257,7 +258,7 @@ add {
   ROOT %convolution = f32[1,5,1] convolution(f32[1,2,1] %input, f32[1,1,1] %filter), feature_group_count=1, dim_labels=b0f_0io->b0f, window={pad=1_2 size=1}
 }
 
-// CHECK-LABEL:  func @test_convert(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf64> {
+// CHECK-LABEL:  func @test_convert(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf64>
 %test_convert (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f64[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -272,7 +273,7 @@ add {
   ROOT %add.5 = f64[4] add(f64[4] %convert.3, f64[4] %convert.4)
 }
 
-// CHECK-LABEL:  func @test_cosine(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32> {
+// CHECK-LABEL:  func @test_cosine(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32>
 %test_cosine (arg0.1: f32[1,16,16,3]) -> f32[1,16,16,3] {
   %arg0.1 = f32[1,16,16,3]{3,2,1,0} parameter(0), metadata={op_name="HLO_Args"}
 
@@ -289,7 +290,7 @@ add {
   ROOT %custom-call = f32[1,2,3]{0,2,1} custom-call(f32[2,3] %arg1, f32[5,5] %arg2), custom_call_target="foo", backend_config="bar", custom_call_has_side_effect=true
 }
 
-// CHECK-LABEL:  func @test_div(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_div(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
 %test_div (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -298,7 +299,7 @@ add {
   ROOT %divide.3 = f32[4] divide(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_dot(%arg0: tensor<1x4xf32>, %arg1: tensor<4x1xf32>) -> tensor {
+// CHECK-LABEL:  func @test_dot(%arg0: tensor<1x4xf32>, %arg1: tensor<4x1xf32>) -> tensor
 %test_dot (Arg_0.1: f32[1, 4], Arg_1.2: f32[4, 1]) -> f32[] {
   %Arg_0.1 = f32[1, 4] parameter(0)
   %Arg_1.2 = f32[4, 1] parameter(1)
@@ -350,7 +351,7 @@ add {
   ROOT %dynamic-slice = s32[1,1,32] dynamic-slice(s32[2,2,258] %operand, s32[] %start_idx_1, s32[] %start_idx_2, s32[] %start_idx_3), dynamic_slice_sizes={1,1,32}
 }
 
-// CHECK-LABEL:  func @test_dynamic_update_slice_1(%arg0: tensor<4x4xf32>, %arg1: tensor<1x4xf32>, %arg2: tensor, %arg3: tensor) -> tensor<4x4xf32> {
+// CHECK-LABEL:  func @test_dynamic_update_slice_1(%arg0: tensor<4x4xf32>, %arg1: tensor<1x4xf32>, %arg2: tensor, %arg3: tensor) -> tensor<4x4xf32>
 %test_dynamic_update_slice_1 (Arg_0.1: f32[4, 4], Arg_1.2: f32[1, 4], Arg_2.3: f32[], Arg_3.4: f32[]) -> f32[4, 4] {
   %Arg_0.1 = f32[4, 4] parameter(0)
   %Arg_1.2 = f32[1, 4] parameter(1)
@@ -371,7 +372,7 @@ add {
   ROOT %dynamic-update-slice.5 = f32[4] dynamic-update-slice(%Arg_0.1, %Arg_1.2, %Arg_2.3)
 }
 
-// CHECK-LABEL:  func @test_exponential(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-LABEL:  func @test_exponential(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_exponential (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -379,7 +380,7 @@ add {
   ROOT %exp.2 = f32[16] exponential(f32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_expm1(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-LABEL:  func @test_expm1(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_expm1 (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -387,7 +388,7 @@ add {
   ROOT %expm1.2 = f32[16] exponential-minus-one(f32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_fft(%arg0: tensor<3x9xf32>) -> tensor<3x5xcomplex> {
+// CHECK-LABEL:  func @test_fft(%arg0: tensor<3x9xf32>) -> tensor<3x5xcomplex>
 %test_fft {
   %arg0.1 = f32[3,9]{1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
   // CHECK:  "xla_hlo.fft"(%arg0) {fft_length = dense<9> : tensor<1xi64>, fft_type = "RFFT"
@@ -395,7 +396,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_floor(
-// CHECK-SAME: [[A0:%.+]]: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-SAME: [[A0:%.+]]: tensor<16xf32>) -> tensor<16xf32>
 %test_floor (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -404,7 +405,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_gather(
-// CHECK-SAME:  [[ARG0:%.+]]: tensor<200x100x300xf32>, [[ARG1:%.+]]: tensor<10x2xi32>) -> tensor<10x300xf32> {
+// CHECK-SAME: [[ARG0:%.+]]: tensor<200x100x300xf32>, [[ARG1:%.+]]: tensor<10x2xi32>) -> tensor<10x300xf32>
 %test_gather (arg.0: f32[200,100,300], arg.1: s32[10,2]) -> f32[10,300] {
   %arg.0 = f32[200,100,300] parameter(0)
   %arg.1 = s32[10,2] parameter(1)
@@ -442,7 +443,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_infeed
-// CHECK-SAME: ([[TOKEN:%.*]]: !xla_hlo.token) -> tuple, !xla_hlo.token> {
+// CHECK-SAME: ([[TOKEN:%.*]]: !xla_hlo.token) -> tuple, !xla_hlo.token>
 %test_infeed (token0: token[]) -> (s32[3], token[]) {
   %token0 = token[] parameter(0)
   // CHECK-NEXT:  "xla_hlo.infeed"([[TOKEN]])
@@ -451,19 +452,19 @@ add {
 }
 
 
-// CHECK-LABEL:  func @test_iota_1() -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_iota_1() -> tensor<4xf32>
 %test_iota_1 () -> f32[4] {
   // CHECK-NEXT:  "xla_hlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xf32>
   ROOT %iota.0 = f32[4] iota(), iota_dimension=0
 }
 
-// CHECK-LABEL:  func @test_iota_2() -> tensor<4x5xf32> {
+// CHECK-LABEL:  func @test_iota_2() -> tensor<4x5xf32>
 %test_iota_2 () -> f32[4, 5] {
   // CHECK-NEXT:  "xla_hlo.iota"() {iota_dimension = 1 : i64} : () -> tensor<4x5xf32>
   ROOT %iota.0 = f32[4, 5] iota(), iota_dimension=1
 }
 
-// CHECK-LABEL:  func @test_log(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-LABEL:  func @test_log(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_log (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -471,7 +472,7 @@ add {
   ROOT %log.2 = f32[16] log(f32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_log1p(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-LABEL:  func @test_log1p(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_log1p (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -501,7 +502,7 @@ add {
 
 
 
-// CHECK-LABEL:  func @test_maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
 %test_maximum (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -510,7 +511,7 @@ add {
   ROOT %maximum.3 = f32[4] maximum(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
 %test_minimum (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -519,7 +520,7 @@ add {
   ROOT %minimum.3 = f32[4] minimum(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_multiply(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_multiply(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
 %test_multiply (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -528,7 +529,7 @@ add {
   ROOT %multiply.3 = f32[4] multiply(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_negate(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-LABEL:  func @test_negate(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_negate (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -536,7 +537,7 @@ add {
   ROOT %negate.2 = f32[16] negate(f32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_not(%arg0: tensor<16xi1>) -> tensor<16xi1> {
+// CHECK-LABEL:  func @test_not(%arg0: tensor<16xi1>) -> tensor<16xi1>
 %test_not (arg0.1: pred[16]) -> pred[16] {
   %arg0.1 = pred[16] parameter(0)
 
@@ -554,7 +555,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_outfeed
-// CHECK-SAME: ([[DATA:%.*]]: tensor<3xi32>, [[TOKEN:%.*]]: !xla_hlo.token) -> !xla_hlo.token {
+// CHECK-SAME: ([[DATA:%.*]]: tensor<3xi32>, [[TOKEN:%.*]]: !xla_hlo.token) -> !xla_hlo.token
 %test_outfeed (Arg_0.1: s32[3], Arg_1.2: token[]) -> token[] {
   %Arg_0.1 = s32[3] parameter(0)
   %Arg_1.2 = token[] parameter(1)
@@ -563,7 +564,7 @@ add {
   ROOT %outfeed.3 = token[] outfeed(s32[3] %Arg_0.1, token[] %Arg_1.2), outfeed_config="foobar"
 }
 
-// CHECK-LABEL:  func @test_pad(%arg0: tensor<4xf32>, %arg1: tensor) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_pad(%arg0: tensor<4xf32>, %arg1: tensor) -> tensor<4xf32>
 %test_pad (Arg_0.1: f32[4], Arg_1.2: f32[]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[] parameter(1)
@@ -572,7 +573,7 @@ add {
   ROOT %pad.3 = f32[4] pad(%Arg_0.1, %Arg_1.2), padding=0_0_0
 }
 
-// CHECK-LABEL:  func @test_pad_edge(%arg0: tensor<4x4x4xf32>, %arg1: tensor) -> tensor<7x11x15xf32> {
+// CHECK-LABEL:  func @test_pad_edge(%arg0: tensor<4x4x4xf32>, %arg1: tensor) -> tensor<7x11x15xf32>
 %test_pad_edge (Arg_0.1: f32[4, 4, 4], Arg_1.2: f32[]) -> f32[7, 11, 15] {
   %Arg_0.1 = f32[4, 4, 4] parameter(0)
   %Arg_1.2 = f32[] parameter(1)
@@ -581,7 +582,7 @@ add {
   ROOT %pad.3 = f32[7, 11, 15] pad(%Arg_0.1, %Arg_1.2), padding=1_2x3_4x5_6
 }
 
-// CHECK-LABEL:  func @test_pad_interior(%arg0: tensor<4xf32>, %arg1: tensor) -> tensor<10xf32> {
+// CHECK-LABEL:  func @test_pad_interior(%arg0: tensor<4xf32>, %arg1: tensor) -> tensor<10xf32>
 %test_pad_interior (Arg_0.1: f32[4], Arg_1.2: f32[]) -> f32[10] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[] parameter(1)
@@ -590,7 +591,7 @@ add {
   ROOT %pad.3 = f32[10] pad(%Arg_0.1, %Arg_1.2), padding=0_0_2
 }
 
-// CHECK-LABEL:  func @test_popcnt(%arg0: tensor<16xi32>) -> tensor<16xi32> {
+// CHECK-LABEL:  func @test_popcnt(%arg0: tensor<16xi32>) -> tensor<16xi32>
 %test_popcnt (arg0.1: s32[16]) -> s32[16] {
   %arg0.1 = s32[16] parameter(0)
 
@@ -598,7 +599,7 @@ add {
   ROOT %popcnt.2 = s32[16] popcnt(s32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_pow(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_pow(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
 %test_pow (Arg_0.1: f32[4], Arg_1.2: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -659,7 +660,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_reduce
-// CHECK-SAME: ([[ARG0:%.*]]: tensor<4x4xf32>, [[ARG1:%.*]]: tensor<4xf32>, [[ARG2:%.*]]: tensor) -> tuple, tensor>, tensor> {
+// CHECK-SAME: ([[ARG0:%.*]]: tensor<4x4xf32>, [[ARG1:%.*]]: tensor<4xf32>, [[ARG2:%.*]]: tensor) -> tuple, tensor>, tensor>
 %test_reduce (Arg_0.1: f32[4, 4], Arg_1.2: f32[4], Arg_2.3: f32[]) -> ((f32[], f32[]), f32[]) {
   %Arg_0.1 = f32[4, 4] parameter(0)
   %Arg_1.2 = f32[4] parameter(1)
@@ -719,7 +720,7 @@ add {
   ROOT %remainder.3 = f32[4] remainder(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_reverse_1d(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+// CHECK-LABEL:  func @test_reverse_1d(%arg0: tensor<4xf32>) -> tensor<4xf32>
 %test_reverse_1d (Arg_0.1: f32[4]) -> f32[4] {
   %Arg_0.1 = f32[4] parameter(0)
 
@@ -727,7 +728,7 @@ add {
   ROOT reverse.2 = f32[4] reverse(%Arg_0.1), dimensions={0}
 }
 
-// CHECK-LABEL:  func @test_reverse_2d(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+// CHECK-LABEL:  func @test_reverse_2d(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32
 %test_reverse_2d (Arg_0.1: f32[4, 4]) -> f32[4, 4] {
   %Arg_0.1 = f32[4, 4] parameter(0)
 
@@ -736,7 +737,7 @@ add {
 }
 
 // CHECK-LABEL:  func @test_rsqrt(
-// CHECK-SAME: [[ARG0:%.+]]: tensor<16xf32>) -> tensor<16xf32> {
+// CHECK-SAME: [[ARG0:%.+]]: tensor<16xf32>) -> tensor<16xf32>
 %test_rsqrt (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
 
@@ -744,7 +745,7 @@ add {
   ROOT %rsqrt.2 = f32[16] rsqrt(f32[16] %arg0.1)
 }
 
-// CHECK-LABEL:  func @test_scalar(%arg0: tensor) -> tensor {
+// CHECK-LABEL:  func @test_scalar(%arg0: tensor) -> tensor
 %test_scalar (Arg_0.1: f32[]) -> f32[] {
   // CHECK-NEXT:  return %arg0 : tensor
   ROOT %Arg_0.1 = f32[] parameter(0)
@@ -781,7 +782,7 @@ add {
 // CHECK-SAME:  unique_indices = false
 
 
-// CHECK-LABEL:  func @test_select(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2: tensor<2x3xi32>) -> tensor<2x3xi32> {
+// CHECK-LABEL:  func @test_select(%arg0: tensor<2x3xi1>, %arg1: tensor<2x3xi32>, %arg2: tensor<2x3xi32>) -> tensor<2x3xi32>
 %test_select {
   %Arg_0.1 = pred[2,3] parameter(0)
   %Arg_1.2 = s32[2,3] parameter(1)
@@ -838,7 +839,7 @@ add {
   ROOT %set-dimension-size.2 = f32[4,<=4] set-dimension-size(f32[4,4] %Arg_0.1, s32[] %Arg_1.2), dimensions={1}
 }
 
-// CHECK-LABEL:  func @test_sine(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32> {
+// CHECK-LABEL:  func @test_sine(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32>
 %test_sine (arg0.1: f32[1,16,16,3]) -> f32[1,16,16,3] {
   %arg0.1 = f32[1,16,16,3]{3,2,1,0} parameter(0), metadata={op_name="HLO_Args"}
 
@@ -874,7 +875,7 @@ add {
   ROOT %subtract.3 = f32[4] subtract(f32[4] %Arg_0.1, f32[4] %Arg_1.2)
 }
 
-// CHECK-LABEL:  func @test_tanh(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32> {
+// CHECK-LABEL:  func @test_tanh(%arg0: tensor<1x16x16x3xf32>) -> tensor<1x16x16x3xf32>
 %test_tanh (arg0.1: f32[1,16,16,3]) -> f32[1,16,16,3] {
   %arg0.1 = f32[1,16,16,3]{3,2,1,0} parameter(0), metadata={op_name="HLO_Args"}
 
@@ -882,7 +883,7 @@ add {
   ROOT %tanh.3 = f32[1,16,16,3]{3,2,1,0} tanh(f32[1,16,16,3]{3,2,1,0} %arg0.1), metadata={op_type="Tanh" op_name="embedded_inference/tanh_model/Tanh"}
 }
 
-// CHECK-LABEL:  func @test_transpose(%arg0: tensor<1x2x3x4xi32>) -> tensor<2x1x4x3xi32> {
+// CHECK-LABEL:  func @test_transpose(%arg0: tensor<1x2x3x4xi32>) -> tensor<2x1x4x3xi32>
 %test_transpose {
   %Arg_0.1 = s32[1,2,3,4] parameter(0)
 
@@ -903,7 +904,7 @@ add {
   ROOT %triangular-solve.3 = f32[4,3] triangular-solve(f32[4,4] %Arg_0.1, f32[4,3] %Arg_1.2), left_side=true, lower=true, transpose_a=NO_TRANSPOSE, unit_diagonal=true
 }
 
-// CHECK-LABEL:  func @test_tuple(%arg0: tensor<1xi32>, %arg1: tensor<1x2xf32>) -> tuple, tensor<1x2xf32>> {
+// CHECK-LABEL:  func @test_tuple(%arg0: tensor<1xi32>, %arg1: tensor<1x2xf32>) -> tuple, tensor<1x2xf32>>
 %test_tuple(Arg_0.1: s32[1], Arg_1.2: f32[1, 2]) -> (s32[1], f32[1,2]) {
   %Arg_0.1 = s32[1] parameter(0)
   %Arg_1.2 = f32[1, 2] parameter(1)
@@ -928,7 +929,7 @@ add {
   ROOT %compare.2 = s64[] add(%arg_1, %arg_1), metadata={op_type="Less" op_name="Less"}
 }
 
-// CHECK-LABEL:  func @test_while(%arg0: tensor) -> tensor {
+// CHECK-LABEL:  func @test_while(%arg0: tensor) -> tensor
 %test_while (arg0.1: s64[]) -> s64[] {
   %arg0.1 = s64[] parameter(0), metadata={op_name="HLO_Args"}
   // CHECK-NEXT:  "xla_hlo.while"(%arg0) ( {

From caf465347e8e9520f1b810809f88f0ea229ba835 Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Mon, 15 Jun 2020 23:39:29 +0000
Subject: [PATCH 0228/1390] segregation attempt 5

---
 tensorflow/core/framework/tensor_shape.cc             | 2 +-
 tensorflow/core/lib/io/inputbuffer.cc                 | 6 +++---
 tensorflow/core/lib/io/random_inputstream.cc          | 2 +-
 tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc   | 2 +-
 tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc  | 6 +++---
 tensorflow/core/lib/io/zlib_outputbuffer.cc           | 6 +++---
 tensorflow/core/platform/env.cc                       | 2 +-
 tensorflow/core/platform/file_system.cc               | 2 +-
 tensorflow/core/platform/file_system_helper.cc        | 2 +-
 tensorflow/core/platform/status.cc                    | 4 +++-
 tensorflow/core/profiler/internal/parse_annotation.cc | 2 +-
 11 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index f4b440e9cd1..79d0cc0822d 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -187,7 +187,7 @@ void TensorShapeBase::InitDims(gtl::ArraySlice dim_sizes) {
                 "bad overflow check");
   bool large_size = false;
   for (auto s : dim_sizes) {
-    if (s > kMaxSmall) {
+    if (static_cast(s) > static_cast(kMaxSmall)) {
       large_size = true;
       break;
     }
diff --git a/tensorflow/core/lib/io/inputbuffer.cc b/tensorflow/core/lib/io/inputbuffer.cc
index 2b138b825e4..d005ee11d78 100644
--- a/tensorflow/core/lib/io/inputbuffer.cc
+++ b/tensorflow/core/lib/io/inputbuffer.cc
@@ -85,7 +85,7 @@ Status InputBuffer::ReadNBytes(int64 bytes_to_read, string* result) {
   result->resize(bytes_to_read);
   size_t bytes_read = 0;
   Status status = ReadNBytes(bytes_to_read, &(*result)[0], &bytes_read);
-  if (bytes_read < bytes_to_read) result->resize(bytes_read);
+  if (static_cast(bytes_read) < bytes_to_read) result->resize(bytes_read);
   return status;
 }
 
@@ -204,7 +204,7 @@ Status InputBuffer::Hint(int64 bytes_to_read) {
   }
 
   // The internal buffer is too small. Do nothing.
-  if (bytes_to_read > size_) {
+  if (bytes_to_read > static_cast(size_)) {
     return Status::OK();
   }
 
@@ -230,7 +230,7 @@ Status InputBuffer::Hint(int64 bytes_to_read) {
   limit_ += data.size();
   file_pos_ += data.size();
 
-  if (errors::IsOutOfRange(s) && data.size() == bytes_to_read) {
+  if (errors::IsOutOfRange(s) && data.size() == static_cast(bytes_to_read)) {
     return Status::OK();
   } else {
     return s;
diff --git a/tensorflow/core/lib/io/random_inputstream.cc b/tensorflow/core/lib/io/random_inputstream.cc
index 10f734a5bae..bd0054ce753 100644
--- a/tensorflow/core/lib/io/random_inputstream.cc
+++ b/tensorflow/core/lib/io/random_inputstream.cc
@@ -92,7 +92,7 @@ Status RandomAccessInputStream::SkipNBytes(int64 bytes_to_skip) {
     } else {
       return s;
     }
-    if (data.size() < bytes_to_read) {
+    if (data.size() < static_cast(bytes_to_read)) {
       return errors::OutOfRange("reached end of file");
     }
     bytes_to_skip -= bytes_to_read;
diff --git a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
index a331d4173cf..53939f2d8a3 100644
--- a/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc
@@ -134,7 +134,7 @@ Status SnappyInputBuffer::ReadCompressedBlockLength(uint32* length) {
     }
     size_t readable = std::min(bytes_to_read, avail_in_);
 
-    for (int i = 0; i < readable; i++) {
+    for (size_t i = 0; i < readable; i++) {
       // The "unsigned char" type cast is intentional to avoid implicit type
       // casting of the signed char to unsigned int during bitwise OR which
       // causes weird overflow errors.
diff --git a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
index 563503a1319..fe3a53c6c25 100644
--- a/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
+++ b/tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc
@@ -76,7 +76,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
 
   // If there is sufficient free space in input_buffer_ to fit data we
   // add it there and return.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -87,7 +87,7 @@ Status SnappyOutputBuffer::Write(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered());
 
   // input_buffer_ should be empty at this point.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -144,7 +144,7 @@ void SnappyOutputBuffer::AddToInputBuffer(StringPiece data) {
   const int32 free_tail_bytes =
       input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (bytes_to_write > free_tail_bytes) {
+  if (static_cast(bytes_to_write) > free_tail_bytes) {
     memmove(input_buffer_.get(), next_in_, avail_in_);
     next_in_ = input_buffer_.get();
   }
diff --git a/tensorflow/core/lib/io/zlib_outputbuffer.cc b/tensorflow/core/lib/io/zlib_outputbuffer.cc
index 5840ca60242..d475d0eaa5c 100644
--- a/tensorflow/core/lib/io/zlib_outputbuffer.cc
+++ b/tensorflow/core/lib/io/zlib_outputbuffer.cc
@@ -98,7 +98,7 @@ void ZlibOutputBuffer::AddToInputBuffer(StringPiece data) {
   int32 unread_bytes = z_stream_->avail_in;
   int32 free_tail_bytes = input_buffer_capacity_ - (read_bytes + unread_bytes);
 
-  if (bytes_to_write > free_tail_bytes) {
+  if (static_cast(bytes_to_write) > free_tail_bytes) {
     memmove(z_stream_input_.get(), z_stream_->next_in, z_stream_->avail_in);
     z_stream_->next_in = z_stream_input_.get();
   }
@@ -154,7 +154,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
 
   size_t bytes_to_write = data.size();
 
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
@@ -162,7 +162,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
   TF_RETURN_IF_ERROR(DeflateBuffered(zlib_options_.flush_mode));
 
   // At this point input stream should be empty.
-  if (bytes_to_write <= AvailableInputSpace()) {
+  if (static_cast(bytes_to_write) <= AvailableInputSpace()) {
     AddToInputBuffer(data);
     return Status::OK();
   }
diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index b29cad05459..05d95ba0425 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -214,7 +214,7 @@ bool Env::FilesExist(const std::vector& files,
     }
     if (fs_status) {
       result &= fs_result;
-      for (int i = 0; i < itr.second.size(); ++i) {
+      for (size_t i = 0; i < itr.second.size(); ++i) {
         per_file_status[itr.second[i]] = fs_status->at(i);
       }
     } else if (!fs_result) {
diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc
index 9e96ceedbdc..c9657e2339f 100644
--- a/tensorflow/core/platform/file_system.cc
+++ b/tensorflow/core/platform/file_system.cc
@@ -308,7 +308,7 @@ StringPiece FileSystem::Basename(StringPiece path) const {
 StringPiece FileSystem::Extension(StringPiece path) const {
   StringPiece basename = this->Basename(path);
 
-  int pos = basename.rfind('.');
+  size_t pos = basename.rfind('.');
   if (pos == StringPiece::npos) {
     return StringPiece(path.data() + path.size(), 0);
   } else {
diff --git a/tensorflow/core/platform/file_system_helper.cc b/tensorflow/core/platform/file_system_helper.cc
index 64b175c4d17..909752389e1 100644
--- a/tensorflow/core/platform/file_system_helper.cc
+++ b/tensorflow/core/platform/file_system_helper.cc
@@ -103,7 +103,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
                 children_dir_status[i] = fs->IsDirectory(child_path);
               }
             });
-    for (int i = 0; i < children.size(); ++i) {
+    for (size_t i = 0; i < children.size(); ++i) {
       const string child_path = io::JoinPath(current_dir, children[i]);
       // If the IsDirectory call was cancelled we bail.
       if (children_dir_status[i].code() == tensorflow::error::CANCELLED) {
diff --git a/tensorflow/core/platform/status.cc b/tensorflow/core/platform/status.cc
index 756b8314148..e303c18091c 100644
--- a/tensorflow/core/platform/status.cc
+++ b/tensorflow/core/platform/status.cc
@@ -74,7 +74,9 @@ class StatusLogSink : public TFLogSink {
 
     mutex_lock lock(mu_);
     messages_.emplace_back(entry.ToString());
-    if (messages_.size() > num_messages_) messages_.pop_front();
+    if (messages_.size() > static_cast(num_messages_)){
+        messages_.pop_front();
+    }
   }
 
  private:
diff --git a/tensorflow/core/profiler/internal/parse_annotation.cc b/tensorflow/core/profiler/internal/parse_annotation.cc
index 32c26befa3d..a4cdc09739d 100644
--- a/tensorflow/core/profiler/internal/parse_annotation.cc
+++ b/tensorflow/core/profiler/internal/parse_annotation.cc
@@ -50,7 +50,7 @@ std::vector SplitNameAndMetadata(
 std::vector SplitPairs(absl::string_view metadata) {
   std::vector key_value_pairs;
   std::stack quotes;
-  int start = 0, end = 0;
+  size_t start = 0, end = 0;
   for (; end < metadata.size(); ++end) {
     char ch = metadata[end];
     switch (ch) {

From 67c32abcff8b5bb873a7ce02e5fe5a6653e55d6c Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Mon, 15 Jun 2020 16:22:20 -0700
Subject: [PATCH 0229/1390] More compatibility fixes for typing.Generic:  *
 types.new_class is required in some distributions  * avoid calling isinstance
 on some function objects in python 3.6 Required for #40132.

PiperOrigin-RevId: 316567732
Change-Id: I6ef39c8de6dbdf2878d30a3a5860a1047a54a55d
---
 tensorflow/python/framework/test_util.py |  7 +++++++
 tensorflow/python/util/tf_should_use.py  | 18 +++++++++++++-----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index 2967bb3de84..950e17d0d8c 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -33,6 +33,7 @@ import tempfile
 import threading
 import time
 import unittest
+import weakref
 
 from absl.testing import parameterized
 import numpy as np
@@ -732,6 +733,12 @@ def assert_no_new_tensors(f):
     """Finds existing Tensors, runs the test, checks for new Tensors."""
 
     def _is_tensorflow_object(obj):
+      if isinstance(obj, weakref.ReferenceType):
+        obj = obj()
+        if obj is None:
+          return False
+      if not hasattr(obj, "__class__"):
+        return False
       try:
         return isinstance(obj,
                           (ops.Tensor, variables.Variable,
diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py
index 1671b078fa3..41c3220f5ca 100644
--- a/tensorflow/python/util/tf_should_use.py
+++ b/tensorflow/python/util/tf_should_use.py
@@ -21,15 +21,12 @@ import copy
 import sys
 import textwrap
 import traceback
-
-import six  # pylint: disable=unused-import
-
+import types
 
 from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util import tf_decorator
-# pylint: enable=g-bad-import-order,g-import-not-at-top
 
 
 class _TFShouldUseHelper(object):
@@ -154,7 +151,18 @@ def _get_wrapper(x, tf_should_use_helper):
   tx = copy.deepcopy(type_x)
   # Prefer using __orig_bases__, which preserve generic type arguments.
   bases = getattr(tx, '__orig_bases__', tx.__bases__)
-  copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
+
+  # Use types.new_class when available, which is preferred over plain type in
+  # some distributions.
+  if sys.version_info >= (3, 5):
+    def set_body(ns):
+      ns.update(tx.__dict__)
+      return ns
+
+    copy_tx = types.new_class(tx.__name__, bases, exec_body=set_body)
+  else:
+    copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
+
   copy_tx.__init__ = _new__init__
   copy_tx.__getattribute__ = _new__getattribute__
   copy_tx.mark_used = _new_mark_used

From 609a60b44bbf934b31a1dce4f0aa84e731b83c35 Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Tue, 16 Jun 2020 00:41:09 +0100
Subject: [PATCH 0230/1390] Refactor AutoCastVariable tests to rely on
 strategy_combinations

---
 .../keras/mixed_precision/experimental/BUILD  |   3 +-
 .../experimental/autocast_variable_test.py    | 130 +++++++++---------
 2 files changed, 64 insertions(+), 69 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/BUILD b/tensorflow/python/keras/mixed_precision/experimental/BUILD
index 024b093c469..4060e455f84 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/BUILD
+++ b/tensorflow/python/keras/mixed_precision/experimental/BUILD
@@ -144,9 +144,10 @@ py_test(
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:framework",
         "//tensorflow/python:platform_test",
+        "//tensorflow/python/distribute:combinations",
         "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/distribute:strategy_combinations",
         "//tensorflow/python/eager:context",
-        "//tensorflow/python/keras:combinations",
         "//tensorflow/python/keras/optimizer_v2",
         "@absl_py//absl/testing:parameterized",
     ],
diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 78041973cc1..95957f5634e 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -17,20 +17,23 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import contextlib
 import os
 
 from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.python import tf2
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribution_strategy_context as ds_context
 from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
-from tensorflow.python.keras import combinations
 from tensorflow.python.keras.mixed_precision.experimental import autocast_variable
 from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_v2
 from tensorflow.python.ops import array_ops
@@ -40,30 +43,17 @@ from tensorflow.python.platform import test
 from tensorflow.python.training import gradient_descent as gradient_descent_v1
 from tensorflow.python.training.tracking import util as trackable_utils
 
-TESTCASES = ({
-    'testcase_name': 'base',
-    'distribute': False
-}, {
-    'testcase_name': 'distribute',
-    'distribute': True
-})
-
-
-def get_distribute_scope(distribute):
-
-  class DummyContextManager(object):
-
-    def __enter__(self):
-      pass
-
-    def __exit__(self, *args):
-      pass
-
-  if distribute:
-    return mirrored_strategy.MirroredStrategy(['cpu:0']).scope()
-  else:
-    return DummyContextManager()
+class DummyStrategy(object):
+  @contextlib.contextmanager
+  def scope(self):
+    yield
 
+maybe_distribute = combinations.combine(
+    distribution=[
+        combinations.NamedDistribution(
+            "Dummy", lambda: DummyStrategy(), required_gpus=None),
+        strategy_combinations.mirrored_strategy_with_cpu_1_and_2
+    ])
 
 def get_var(val, dtype, name=None):
   return variables.VariableV1(val, use_resource=True, dtype=dtype, name=name)
@@ -71,10 +61,13 @@ def get_var(val, dtype, name=None):
 
 @combinations.generate(combinations.combine(mode=['graph', 'eager']))
 class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
+  def setUp(self):
+    strategy_combinations.set_virtual_cpus_to_at_least(3)
+    super(AutoCastVariableTest, self).setUp()
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_read(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_read(self, distribution):
+    with distribution.scope():
       x = get_var(1., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.evaluate(x.initializer)
@@ -116,9 +109,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
       self.assertEqual(x.sparse_read([0]).dtype, dtypes.float16)
       self.assertEqual(x.gather_nd([0]).dtype, dtypes.float16)
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_read_nested_scopes(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_read_nested_scopes(self, distribution):
+    with distribution.scope():
       x = get_var(1., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.evaluate(x.initializer)
@@ -136,9 +129,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertEqual(x.dtype, dtypes.float16)
         self.assertEqual(x.read_value().dtype, dtypes.float16)
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_dtype_is_not_string(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_dtype_is_not_string(self, distribution):
+    with distribution.scope():
       x = get_var(1., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.assertEqual(x.dtype, dtypes.float32)
@@ -153,13 +146,13 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertEqual(x.true_dtype, dtypes.float32)
         self.assertIsInstance(x.true_dtype, dtypes.DType)
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_method_delegations(self, distribute):
+  @combinations.generate(maybe_distribute)
+  def test_method_delegations(self, distribution):
     # Test AutoCastVariable correctly delegates Variable methods to the
     # underlying variable.
-    with self.test_session(), get_distribute_scope(distribute):
+    with self.test_session(), distribution.scope():
       for read_dtype in (dtypes.float32, dtypes.float16):
-        if distribute:
+        if ds_context.has_strategy():
           # MirroredVariable.assign will (incorrectly) return a Mirrored value
           # instead of a MirroredVariable. So we cannot properly wrap it in an
           # AutoCastVariable.
@@ -183,14 +176,14 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
           self.assertEqual(x.aggregation, x._variable.aggregation)
           self.assertEqual(self.evaluate(x.initialized_value()), 7)
           if not context.executing_eagerly():
-            if not distribute:
+            if not ds_context.has_strategy():
               # These functions are not supported for DistributedVariables
               x.load(9)
               self.assertEqual(x.eval(), 9)
             self.assertEqual(self.evaluate(x.initial_value), 7)
             self.assertEqual(x.op, x._variable.op)
             self.assertEqual(x.graph, x._variable.graph)
-          if not distribute:
+          if not ds_context.has_strategy():
             # These attributes are not supported for DistributedVariables
             self.assertIsNone(x.constraint)
             self.assertEqual(x.initializer, x._variable.initializer)
@@ -202,7 +195,7 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
           self.assertEqual(x.shape, ())
           self.assertEqual(x.get_shape(), ())
 
-        if not distribute:
+        if not ds_context.has_strategy():
           # Test scatter_* methods. These are not supported for
           # DistributedVariables
           x = get_var([7, 8], dtypes.float32)
@@ -233,9 +226,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
             self.assertAllEqual(
                 evaluate(x.scatter_nd_update([[0], [1]], [1., 2.])), [1, 2])
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_operator_overloads(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_operator_overloads(self, distribution):
+    with distribution.scope():
       for read_dtype in (dtypes.float32, dtypes.float16):
         x = get_var(7., dtypes.float32)
         x = autocast_variable.create_autocast_variable(x)
@@ -280,9 +273,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
             self.assertAllEqual(x == [7., 8., 10.], [True, True, False])
             self.assertAllEqual(x != [7., 8., 10.], [False, False, True])
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_assign(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_assign(self, distribution):
+    with distribution.scope():
       x = get_var(0., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.evaluate(x.initializer)
@@ -318,18 +311,19 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertAllClose(3., self.evaluate(x.assign_sub(3.)))
 
         # Assign multiple times
-        assign = x.assign(1.)
-        self.assertAllClose(1., self.evaluate(assign))
-        self.assertAllClose(0., self.evaluate(assign.assign(0.)))
-        assign_add = x.assign_add(3.)
-        self.assertAllClose(3., self.evaluate(assign_add))
-        self.assertAllClose(3. * 3,
-                            self.evaluate(x.assign_add(3.).assign_add(3.)))
-        self.assertAllClose(3. * 3, x)
-        assign_sub = x.assign_sub(3.)
-        self.assertAllClose(3. * 2, self.evaluate(assign_sub))
-        self.assertAllClose(0.,
-                            self.evaluate(x.assign_sub(3.).assign_sub(3.)))
+        if not ds_context.has_strategy():
+          assign = x.assign(1.)
+          self.assertAllClose(1., self.evaluate(assign))
+          self.assertAllClose(0., self.evaluate(assign.assign(0.)))
+          assign_add = x.assign_add(3.)
+          self.assertAllClose(3., self.evaluate(assign_add))
+          self.assertAllClose(3. * 3,
+                              self.evaluate(x.assign_add(3.).assign_add(3.)))
+          self.assertAllClose(3. * 3, x)
+          assign_sub = x.assign_sub(3.)
+          self.assertAllClose(3. * 2, self.evaluate(assign_sub))
+          self.assertAllClose(0.,
+                              self.evaluate(x.assign_sub(3.).assign_sub(3.)))
 
         # Assign with read_value=False
         self.assertIsNone(self.evaluate(x.assign(1., read_value=False)))
@@ -355,9 +349,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         # assign still expect float32 value even if in float16 scope
         run_and_check()
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_assign_stays_in_true_dtype(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_assign_stays_in_true_dtype(self, distribution):
+    with distribution.scope():
       x = get_var(1., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.evaluate(x.initializer)
@@ -382,10 +376,10 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertEqual(1., self.evaluate(x.value()))
       self.assertEqual(1. + small_val, self.evaluate(x.value()))
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_checkpoint(self, distribute):
+  @combinations.generate(maybe_distribute)
+  def test_checkpoint(self, distribution):
     with self.test_session():
-      with get_distribute_scope(distribute):
+      with distribution.scope():
         x = get_var(1., dtypes.float32)
         x = autocast_variable.create_autocast_variable(x)
       self.evaluate(x.initializer)
@@ -398,9 +392,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
       checkpoint.restore(save_path).assert_consumed().run_restore_ops()
       self.assertEqual(self.evaluate(x), 123.)
 
-  @parameterized.named_parameters(*TESTCASES)
-  def test_invalid_wrapped_variable(self, distribute):
-    with get_distribute_scope(distribute):
+  @combinations.generate(maybe_distribute)
+  def test_invalid_wrapped_variable(self, distribution):
+    with distribution.scope():
       # Wrap a non-variable
       with self.assertRaisesRegexp(ValueError, 'variable must be of type'):
         x = constant_op.constant([1.], dtype=dtypes.float32)
@@ -443,7 +437,7 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         )
 
   def test_repr_distributed(self):
-    with get_distribute_scope(distribute=True):
+    with mirrored_strategy.MirroredStrategy(["/cpu:1", "/cpu:2"]).scope():
       x = get_var(1., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)
       self.assertRegexpMatches(

From 0e3a53269fe3d02d114a54430471116a400f4207 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Mon, 15 Jun 2020 16:25:46 -0700
Subject: [PATCH 0231/1390] Rename SVDF's kInputActivationStateTensor to
 kStateTensor, as there is no other state in SVDF, and making variable names
 consistent with the corresponding parameter names of
 reference_ops::Eval..Svdf functions.

PiperOrigin-RevId: 316568280
Change-Id: I61dacf43e7a97437df531372b91ccba0876663ff
---
 tensorflow/lite/kernels/svdf.cc | 35 +++++++++++++++------------------
 1 file changed, 16 insertions(+), 19 deletions(-)

diff --git a/tensorflow/lite/kernels/svdf.cc b/tensorflow/lite/kernels/svdf.cc
index 82b7b7e4ee5..57eedb6b204 100644
--- a/tensorflow/lite/kernels/svdf.cc
+++ b/tensorflow/lite/kernels/svdf.cc
@@ -54,7 +54,7 @@ constexpr int kWeightsFeatureTensor = 1;
 constexpr int kWeightsTimeTensor = 2;
 constexpr int kBiasTensor = 3;
 // This is a variable tensor, and will be modified by this op.
-constexpr int kInputActivationStateTensor = 4;
+constexpr int kStateTensor = 4;
 
 // Output tensor.
 constexpr int kOutputTensor = 0;
@@ -107,14 +107,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_EQ(context, bias->dims->data[0], num_units);
   }
 
-  const TfLiteTensor* activation_state =
-      GetInput(context, node, kInputActivationStateTensor);
+  const TfLiteTensor* state = GetInput(context, node, kStateTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   // Check the shape of input state tensors.
-  TF_LITE_ENSURE_EQ(context, NumDimensions(activation_state), 2);
-  TF_LITE_ENSURE_EQ(context, SizeOfDimension(activation_state, 0), batch_size);
-  TF_LITE_ENSURE_EQ(context, SizeOfDimension(activation_state, 1),
+  TF_LITE_ENSURE_EQ(context, NumDimensions(state), 2);
+  TF_LITE_ENSURE_EQ(context, SizeOfDimension(state, 0), batch_size);
+  TF_LITE_ENSURE_EQ(context, SizeOfDimension(state, 1),
                     memory_size * num_filters);
 
   // Resize output.
@@ -184,7 +183,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     }
 
     // Used to store dequantized weights_time matrix for hybrid computation of
-    // matmul(activation_state, weights_time), which occurs in floating point.
+    // matmul(state, weights_time), which occurs in floating point.
     node->temporaries->data[3] = scratch_tensor_index + 3;
     TfLiteTensor* float_weights_time = GetTemporary(context, node, /*index=*/3);
     float_weights_time->type = kTfLiteFloat32;
@@ -239,8 +238,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
         reinterpret_cast(input->quantization.params);
     auto* weights_feature_params = reinterpret_cast(
         weights_feature->quantization.params);
-    auto* state_params = reinterpret_cast(
-        activation_state->quantization.params);
+    auto* state_params =
+        reinterpret_cast(state->quantization.params);
     auto* weight_time_params = reinterpret_cast(
         weights_time->quantization.params);
     auto* output_params = reinterpret_cast(
@@ -272,15 +271,14 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   TfLiteTensor* scratch = GetTemporary(context, node, /*index=*/0);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
+  TfLiteTensor* state = GetVariableInput(context, node, kStateTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   switch (weights_feature->type) {
     case kTfLiteFloat32: {
       reference_ops::EvalFloatSVDF(context, node, input, weights_feature,
-                                   weights_time, bias, params, scratch,
-                                   activation_state, output);
+                                   weights_time, bias, params, scratch, state,
+                                   output);
       return kTfLiteOk;
       break;
     }
@@ -315,8 +313,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
         reference_ops::EvalHybridSVDF(
             context, node, input, weights_feature, float_weights_time, bias,
-            params, scratch, scaling_factors, input_quantized, activation_state,
-            output, zero_points, row_sums, &op_data->compute_row_sums);
+            params, scratch, scaling_factors, input_quantized, state, output,
+            zero_points, row_sums, &op_data->compute_row_sums);
         return kTfLiteOk;
       } else {
         auto* input_params = reinterpret_cast(
@@ -330,10 +328,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
         TF_LITE_ENSURE_EQ(context, params->activation, kTfLiteActRelu);
         reference_ops::EvalIntegerSVDF(
             context, node, input, weights_feature, weights_time, bias, params,
-            activation_state, output, scratch, output_temp,
-            op_data->effective_scale_1_a, op_data->effective_scale_1_b,
-            op_data->effective_scale_2_a, op_data->effective_scale_2_b,
-            input_params->zero_point->data[0],
+            state, output, scratch, output_temp, op_data->effective_scale_1_a,
+            op_data->effective_scale_1_b, op_data->effective_scale_2_a,
+            op_data->effective_scale_2_b, input_params->zero_point->data[0],
             output_params->zero_point->data[0]);
         return kTfLiteOk;
       }

From e2e974a383c06288664b8276eab1bd25659ab0c0 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Mon, 15 Jun 2020 16:35:58 -0700
Subject: [PATCH 0232/1390] Arguments declaration moved inside
 TransformToCLCode.

PiperOrigin-RevId: 316570056
Change-Id: Ia8096e251ac78ad68d2439b8ea91c03bc1dc35bf
---
 tensorflow/lite/delegates/gpu/cl/arguments.cc |  2 ++
 .../lite/delegates/gpu/cl/kernels/BUILD       |  1 -
 .../lite/delegates/gpu/cl/kernels/softmax.cc  | 24 +++++----------
 .../delegates/gpu/cl/kernels/transpose.cc     |  5 +---
 .../lite/delegates/gpu/cl/kernels/winograd.cc | 29 +++++--------------
 5 files changed, 19 insertions(+), 42 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/arguments.cc b/tensorflow/lite/delegates/gpu/cl/arguments.cc
index afb58ba46ad..53303eab079 100644
--- a/tensorflow/lite/delegates/gpu/cl/arguments.cc
+++ b/tensorflow/lite/delegates/gpu/cl/arguments.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/str_split.h"
+#include "absl/strings/substitute.h"
 #include "tensorflow/lite/delegates/gpu/cl/tensor_type.h"
 #include "tensorflow/lite/delegates/gpu/common/status.h"
 
@@ -445,6 +446,7 @@ absl::Status Arguments::TransformToCLCode(
   RETURN_IF_ERROR(AddObjectArgs());
   RETURN_IF_ERROR(ResolveSelectorsPass(linkables, code));
   ResolveArgsPass(device_info, code);
+  *code = absl::Substitute(*code, GetListOfArgs());
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/BUILD b/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
index d1f6d4014fe..24a9a962296 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
@@ -1130,7 +1130,6 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/cl:tensor",
         "//tensorflow/lite/delegates/gpu/common:status",
         "//tensorflow/lite/delegates/gpu/common:types",
-        "@com_google_absl//absl/strings",
     ],
 )
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/softmax.cc b/tensorflow/lite/delegates/gpu/cl/kernels/softmax.cc
index e5f0933401a..fda7dbba6dd 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/softmax.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/softmax.cc
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include 
 
-#include "absl/strings/substitute.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
 #include "tensorflow/lite/delegates/gpu/common/status.h"
@@ -29,7 +28,6 @@ namespace {
 
 std::string GetSoftmaxKernelCode(
     const OperationDef& op_def,
-    const std::vector& linked_operations,
     Arguments* args) {
   auto src_desc = absl::make_unique(op_def.src_tensors[0]);
   if (op_def.IsBatchSupported()) {
@@ -43,12 +41,7 @@ std::string GetSoftmaxKernelCode(
   args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
 
   std::string c = GetCommonDefines(op_def.precision);
-  std::string linked_args = GetArgsDeclaration(linked_operations);
-  if (linked_args[0] == ',') {
-    linked_args[0] = ' ';
-  }
   c += "__kernel void main_function(\n";
-  c += linked_args;
   c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
   c += "  int Y = get_global_id(1);\n";
@@ -66,7 +59,6 @@ std::string GetSoftmaxKernelCode(
   c += "    float4 t = args.src_tensor.Read(X, Y, d);\n";
   c += "    t = exp(t) / sum;\n";
   c += "    FLT4 result = TO_FLT4(t);\n";
-  c += PostProcess(linked_operations, {"result", "X", "Y", "d"});
   c += "    args.dst_tensor.Write(result, X, Y, d);\n";
   c += "  }\n";
   c += "}\n";
@@ -89,11 +81,13 @@ Softmax& Softmax::operator=(Softmax&& kernel) {
 }
 
 absl::Status Softmax::Compile(const CreationContext& creation_context) {
-  std::string code =
-      GetSoftmaxKernelCode(definition_, linked_operations_, &args_);
+  std::string code = GetSoftmaxKernelCode(definition_, &args_);
+  std::string element_wise_code;
   RETURN_IF_ERROR(
-      args_.TransformToCLCode(creation_context.device->GetInfo(), {}, &code));
-  code = absl::Substitute(code, args_.GetListOfArgs());
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
@@ -102,10 +96,8 @@ absl::Status Softmax::Compile(const CreationContext& creation_context) {
 absl::Status Softmax::BindArguments() {
   RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
   RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(args_.Bind(kernel_.kernel(), kernel_.GetBindingCounter()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Softmax::GetGridSize() const {
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc b/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
index 8c5af6e9785..e12c44566b7 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
@@ -18,7 +18,6 @@ limitations under the License.
 #include 
 
 #include "absl/strings/str_cat.h"
-#include "absl/strings/substitute.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
 
@@ -127,7 +126,6 @@ absl::Status Transpose::Compile(const CreationContext& creation_context) {
   RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
                                           {{"dst_tensor", element_wise_code}},
                                           &code));
-  code = absl::Substitute(code, args_.GetListOfArgs());
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
@@ -137,8 +135,7 @@ absl::Status Transpose::BindArguments() {
   RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
   RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
   RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
-  RETURN_IF_ERROR(args_.Bind(kernel_.kernel()));
-  return absl::OkStatus();
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Transpose::GetGridSize() const {
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
index aa47e3a1c24..d71513e4de4 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include 
 
 #include "absl/strings/str_format.h"
-#include "absl/strings/substitute.h"
 #include "tensorflow/lite/delegates/gpu/cl/cl_device.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
@@ -36,7 +35,6 @@ namespace {
 
 std::string GetWinograd4x4To36Code(
     const OperationDef& op_def,
-    const std::vector& linked_operations,
     Arguments* args) {
   std::string c = GetCommonDefines(op_def.precision);
 
@@ -82,12 +80,7 @@ std::string GetWinograd4x4To36Code(
   args->AddInt("tiles_total");
   args->AddInt("tiles_x");
 
-  std::string linked_args = GetArgsDeclaration(linked_operations);
-  if (linked_args[0] == ',') {
-    linked_args[0] = ' ';
-  }
   c += "__kernel void main_function(\n";
-  c += linked_args;
   c += "$0) {\n";
   c += "  int DST_X = get_global_id(0);\n";
   c += "  int DST_Y = get_global_id(1);\n";
@@ -181,14 +174,12 @@ std::string GetWinograd4x4To36Code(
   const LinkingContext context{"r0", "DST_X", "DST_Y", "DST_Z"};
   c += "  {\n";
   c += "    FLT4 r0 = TO_FLT4(I0 + Bt[2] * I2 + Bt[4] * I4);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
   c += "  {\n";
   c += "    FLT4 r0 = TO_FLT4(Bt[7] * I1 + Bt[8] * I2 + Bt[9] * I3 + Bt[10] * "
        "I4);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
@@ -196,7 +187,6 @@ std::string GetWinograd4x4To36Code(
   c += "    FLT4 r0 = TO_FLT4(Bt[13] * I1 + Bt[14] * I2 + Bt[15] * I3 + Bt[16] "
        "* "
        "I4);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
@@ -204,7 +194,6 @@ std::string GetWinograd4x4To36Code(
   c += "    FLT4 r0 = TO_FLT4(Bt[19] * I1 + Bt[20] * I2 + Bt[21] * I3 + Bt[22] "
        "* "
        "I4);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
@@ -212,13 +201,11 @@ std::string GetWinograd4x4To36Code(
   c += "    FLT4 r0 = TO_FLT4(Bt[25] * I1 + Bt[26] * I2 + Bt[27] * I3 + Bt[28] "
        "* "
        "I4);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
   c += "  {\n";
   c += "    FLT4 r0 = TO_FLT4(Bt[31] * I1 + Bt[33] * I3 + I5);\n";
-  c += PostProcess(linked_operations, context);
   c += "    args.dst_tensor.Write(r0, DST_X, DST_Y, DST_Z);\n";
   c += "    DST_Y++;\n";
   c += "  }\n";
@@ -389,11 +376,13 @@ absl::Status Winograd4x4To36::Compile(const CreationContext& creation_context) {
     options.push_back(CompilerOptions::POWERVR_FP16);
   }
   RETURN_IF_ERROR(UploadBt(creation_context.context));
-  std::string code =
-      GetWinograd4x4To36Code(definition_, linked_operations_, &args_);
+  std::string code = GetWinograd4x4To36Code(definition_, &args_);
+  std::string element_wise_code;
   RETURN_IF_ERROR(
-      args_.TransformToCLCode(creation_context.device->GetInfo(), {}, &code));
-  code = absl::Substitute(code, args_.GetListOfArgs());
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   RETURN_IF_ERROR(creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", options, *creation_context.context,
       *creation_context.device, &kernel_));
@@ -445,10 +434,8 @@ absl::Status Winograd4x4To36::BindArguments() {
   RETURN_IF_ERROR(args_.SetInt("padding_y", -padding_.prepended.h));
   RETURN_IF_ERROR(args_.SetInt("tiles_total", tiles_total));
   RETURN_IF_ERROR(args_.SetInt("tiles_x", tiles_x));
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(args_.Bind(kernel_.kernel(), kernel_.GetBindingCounter()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Winograd4x4To36::GetGridSize() const {

From 4aef940d519fb79f4b4c0b44fa4adf268b396183 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar 
Date: Mon, 15 Jun 2020 16:38:27 -0700
Subject: [PATCH 0233/1390] Select default implementation via DeviceIndex

For DeviceIndex that feeds into Case, simply select the default implementation for now.

PiperOrigin-RevId: 316570448
Change-Id: Ic088cfe5a2da6953526a46e4f51e69ec42c47e0f
---
 tensorflow/compiler/mlir/lite/BUILD           |  1 +
 .../lite/tests/tf_device_index_selector.mlir  | 25 ++++++
 .../compiler/mlir/lite/tf_tfl_passes.cc       |  1 +
 .../lite/transforms/device_index_selector.cc  | 85 +++++++++++++++++++
 .../compiler/mlir/lite/transforms/passes.h    |  3 +
 5 files changed, 115 insertions(+)
 create mode 100644 tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir
 create mode 100644 tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 8d4efeb3d60..8e9d615053c 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -314,6 +314,7 @@ tf_cc_test(
 cc_library(
     name = "tensorflow_lite_legalize_tf",
     srcs = [
+        "transforms/device_index_selector.cc",
         "transforms/dilated_conv.cc",
         "transforms/generated_legalize_tf.inc",
         "transforms/generated_lower_static_tensor_list.inc",
diff --git a/tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir b/tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir
new file mode 100644
index 00000000000..1ac7f30d644
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir
@@ -0,0 +1,25 @@
+// Test DeviceIndex selector.
+
+// RUN: tf-opt --tfl-device-index-selector %s | FileCheck %s
+
+// CHECK-LABEL: func @select
+func @select(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) {
+  // CHECK:  %[[first:.*]] = "tf.DeviceIndex"
+  // CHECK: constant dense<2>
+  // CHECK:  return %[[first]],
+  %0 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+  %1 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+  %4 = "tf.Case"(%1, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf.shape<>]} : (tensor, tensor, tensor) -> tensor
+
+  return %0, %4 : tensor, tensor
+}
+
+func @add(%i: tensor, %arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
+  %0 = "tf.Add"(%arg0, %arg1): (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
+
+func @sub(%i: tensor, %arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
+  %0 = "tf.Sub"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
index 06fe8684ce4..589515d6246 100644
--- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
+++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
@@ -63,6 +63,7 @@ void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config,
   standard_pipeline_options.enable_inliner = false;
   standard_pipeline_options.form_clusters = pass_config.form_clusters;
   mlir::TF::CreateTFStandardPipeline(*pass_manager, standard_pipeline_options);
+  pass_manager->addPass(mlir::TFL::CreateDeviceIndexSelectorPass());
 
   if (pass_config.shape_inference) {
     pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
diff --git a/tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc b/tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc
new file mode 100644
index 00000000000..d4aed750dc8
--- /dev/null
+++ b/tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc
@@ -0,0 +1,85 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+// Converts DeviceIndex to constant device.
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+// Folds the DeviceIndex op to a constant value. The DeviceIndex return the
+// index of the device the op should run on. The user can use this to provide
+// different op specializations. E.g.,
+//
+// ```mlir
+//  %1 = "tf.DeviceIndex"()
+//          {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+//  %4 = "tf.Case"(%1, %arg0, %arg1)
+//          {branches = [@foo, @baz], output_shapes = [#tf.shape<>]} :
+//            (tensor, tensor, tensor) -> tensor
+// ```
+//
+// Shows an example where there are 2 different functions which could be
+// executed to produce the same values but with different functions optimized
+// for CPU or GPU.
+struct DeviceIndexSelector
+    : public PassWrapper> {
+  void runOnOperation() override;
+};
+
+}  // namespace
+
+void DeviceIndexSelector::runOnOperation() {
+  FuncOp func = getOperation();
+  // Convert all the DeviceIndex ops to constant values.
+  func.getBody().walk([](TF::DeviceIndexOp op) {
+    // This just selects the default in all cases where DeviceIndex feeds into
+    // tf.Case. This could be enhanced based on explicit TFLite specification or
+    // TAC in future.
+    OpBuilder b(op);
+    RankedTensorType type = RankedTensorType::get({}, b.getIntegerType(32));
+    int index = op.device_names().size();
+    for (auto use : op.getOperation()->getUsers()) {
+      // Skip if it doesn't feed into case. Alternatively this could always
+      // return the CPU device index if it exists.
+      if (!isa(use)) return;
+    }
+    DenseElementsAttr attr =
+        DenseElementsAttr::get(type, b.getI32IntegerAttr(index));
+    auto constant = b.create(op.getLoc(), type, attr);
+    op.replaceAllUsesWith(constant.getOperation());
+    op.erase();
+  });
+}
+
+// Creates an instance of the TensorFlow DeviceIndex selector pass.
+std::unique_ptr> CreateDeviceIndexSelectorPass() {
+  return std::make_unique();
+}
+
+static PassRegistration pass(
+    "tfl-device-index-selector", "Fold tf.DeviceIndex to constant");
+
+}  // namespace TFL
+}  // namespace mlir
diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h
index 105c9394fb4..01e5eb1cb68 100644
--- a/tensorflow/compiler/mlir/lite/transforms/passes.h
+++ b/tensorflow/compiler/mlir/lite/transforms/passes.h
@@ -91,6 +91,9 @@ std::unique_ptr> CreateWhileOutlinePass();
 // Verifies runtime constraints.
 std::unique_ptr> CreateRuntimeVerifyPass();
 
+// Creates function pass to select device index/fold tf.DeviceIndex.
+std::unique_ptr> CreateDeviceIndexSelectorPass();
+
 }  // namespace TFL
 
 }  // namespace mlir

From bfbec5fb039001aaf9ae7bff05d3dee3bbf54d82 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 16:39:39 -0700
Subject: [PATCH 0234/1390] Add a configuration option to BaseLayer/V1 that
 will throw an exception if a user tries to reconstitute a model with these
 layer objects in it but does not provide the base class.

This CL also changes SavedModel serialization code to use the Keras serialization name (which is different from the class name when using the register_keras_serializable decorator). This should fully enable custom layers in SavedModels.

PiperOrigin-RevId: 316570643
Change-Id: I5755a59cb74db8bbe6c3f607804abfb67c9ea8e0
---
 tensorflow/python/keras/engine/base_layer.py  | 29 +++++++++++++---
 .../saving/saved_model/layer_serialization.py | 10 ++++--
 .../python/keras/saving/saved_model/load.py   | 14 +++++++-
 .../saving/saved_model/saved_model_test.py    | 34 +++++++++++++++++++
 4 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index 87f306b2879..a0ee25417c0 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -286,16 +286,34 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
       module.Module._TF_MODULE_IGNORED_PROPERTIES
   ))
 
+  # When loading from a SavedModel, Layers typically can be revived into a
+  # generic Layer wrapper. Sometimes, however, layers may implement methods
+  # that go beyond this wrapper, as in the case of PreprocessingLayers'
+  # `adapt` method. When this is the case, layer implementers can override
+  # must_restore_from_config to return True; layers with this property must
+  # be restored into their actual objects (and will fail if the object is
+  # not available to the restoration code).
+  _must_restore_from_config = False
+
   @trackable.no_automatic_dependency_tracking
-  def __init__(self, trainable=True, name=None, dtype=None, dynamic=False,
+  def __init__(self,
+               trainable=True,
+               name=None,
+               dtype=None,
+               dynamic=False,
                **kwargs):
     # These properties should be set by the user via keyword arguments.
     # note that 'dtype', 'input_shape' and 'batch_input_shape'
     # are only applicable to input layers: do not pass these keywords
     # to non-input layers.
     allowed_kwargs = {
-        'input_dim', 'input_shape', 'batch_input_shape', 'batch_size',
-        'weights', 'activity_regularizer', 'autocast'
+        'input_dim',
+        'input_shape',
+        'batch_input_shape',
+        'batch_size',
+        'weights',
+        'activity_regularizer',
+        'autocast',
     }
     # Validate optional keyword arguments.
     generic_utils.validate_kwargs(kwargs, allowed_kwargs)
@@ -637,7 +655,10 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
         Python dictionary.
     """
     all_args = tf_inspect.getfullargspec(self.__init__).args
-    config = {'name': self.name, 'trainable': self.trainable}
+    config = {
+        'name': self.name,
+        'trainable': self.trainable,
+    }
     if hasattr(self, '_batch_input_shape'):
       config['batch_input_shape'] = self._batch_input_shape
     config['dtype'] = policy.serialize(self._dtype_policy)
diff --git a/tensorflow/python/keras/saving/saved_model/layer_serialization.py b/tensorflow/python/keras/saving/saved_model/layer_serialization.py
index 0de38577dc0..559b6158d87 100644
--- a/tensorflow/python/keras/saving/saved_model/layer_serialization.py
+++ b/tensorflow/python/keras/saving/saved_model/layer_serialization.py
@@ -46,13 +46,15 @@ class LayerSavedModelSaver(base_serialization.SavedModelSaver):
     # TODO(kathywu): Synchronize with the keras spec (go/keras-json-spec) once
     # the python config serialization has caught up.
     metadata = dict(
-        class_name=type(self.obj).__name__,
+        class_name=generic_utils.get_registered_name(type(self.obj)),
         name=self.obj.name,
         trainable=self.obj.trainable,
         expects_training_arg=self.obj._expects_training_arg,  # pylint: disable=protected-access
         dtype=policy.serialize(self.obj._dtype_policy),  # pylint: disable=protected-access
         batch_input_shape=getattr(self.obj, '_batch_input_shape', None),
-        stateful=self.obj.stateful)
+        stateful=self.obj.stateful,
+        must_restore_from_config=self.obj._must_restore_from_config,  # pylint: disable=protected-access
+    )
 
     metadata.update(get_config(self.obj))
     if self.obj.input_spec is not None:
@@ -85,7 +87,8 @@ class LayerSavedModelSaver(base_serialization.SavedModelSaver):
     serialized_attr = keras_cache[self.obj] = (
         serialized_attributes.SerializedAttributes.new(self.obj))
 
-    if save_impl.should_skip_serialization(self.obj):
+    if (save_impl.should_skip_serialization(self.obj) or
+        self.obj._must_restore_from_config):  # pylint: disable=protected-access
       return serialized_attr
 
     object_dict, function_dict = self._get_serialized_attributes_internal(
@@ -128,6 +131,7 @@ class InputLayerSavedModelSaver(base_serialization.SavedModelSaver):
 
   @property
   def python_properties(self):
+
     return dict(
         class_name=type(self.obj).__name__,
         name=self.obj.name,
diff --git a/tensorflow/python/keras/saving/saved_model/load.py b/tensorflow/python/keras/saving/saved_model/load.py
index 7e67bf6305c..313eea4342e 100644
--- a/tensorflow/python/keras/saving/saved_model/load.py
+++ b/tensorflow/python/keras/saving/saved_model/load.py
@@ -407,6 +407,7 @@ class KerasObjectLoader(tf_load.Loader):
     #       found.
     class_name = metadata.get('class_name')
     config = metadata.get('config')
+    must_restore_from_config = metadata.get('must_restore_from_config')
     if not generic_utils.validate_config(config):
       return None
 
@@ -414,7 +415,18 @@ class KerasObjectLoader(tf_load.Loader):
       obj = layers_module.deserialize(
           generic_utils.serialize_keras_class_and_config(class_name, config))
     except ValueError:
-      return None
+      if must_restore_from_config:
+        raise RuntimeError(
+            'Unable to restore a layer of class {cls}. Layers of '
+            'class {cls} require that the class be provided to '
+            'the model loading code, either by registering the '
+            'class using @keras.utils.register_keras_serializable '
+            'on the class def and including that file in your '
+            'program, or by passing the class in a '
+            'keras.utils.CustomObjectScope that wraps this load '
+            'call.'.format(cls=class_name))
+      else:
+        return None
 
     # Use the dtype, name, and trainable status. Often times these are not
     # specified in custom configs, so retrieve their values from the metadata.
diff --git a/tensorflow/python/keras/saving/saved_model/saved_model_test.py b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
index c208805686d..7eaa75b78e2 100644
--- a/tensorflow/python/keras/saving/saved_model/saved_model_test.py
+++ b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
@@ -108,6 +108,11 @@ class LayerWithUpdate(keras.layers.Layer):
     return inputs * 2.
 
 
+@generic_utils.register_keras_serializable('Testing')
+class GlobalLayerThatShouldFailIfNotAdded(keras.layers.Layer):
+  _must_restore_from_config = True
+
+
 @keras_parameterized.run_all_keras_modes
 class TestModelSavingAndLoadingV2(keras_parameterized.TestCase):
 
@@ -331,6 +336,35 @@ class TestModelSavingAndLoadingV2(keras_parameterized.TestCase):
     self.assertAllEqual([None, 2, 3], loaded.input_spec['b'].shape)
     self.assertEqual('float16', loaded.input_spec['b'].dtype)
 
+  def test_must_restore_from_config_fails_if_layer_is_not_in_scope(self):
+
+    class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
+      _must_restore_from_config = True
+
+    layer = LayerThatShouldFailIfNotAdded()
+    saved_model_dir = self._save_model_dir()
+    tf_save.save(layer, saved_model_dir)
+    with self.assertRaisesRegex(RuntimeError, 'Unable to restore a layer of'):
+      _ = keras_load.load(saved_model_dir)
+
+  def test_must_restore_from_config_custom_object_scope(self):
+
+    class LayerThatShouldFailIfNotAdded(keras.layers.Layer):
+      _must_restore_from_config = True
+
+    layer = LayerThatShouldFailIfNotAdded()
+    saved_model_dir = self._save_model_dir()
+    tf_save.save(layer, saved_model_dir)
+    with generic_utils.CustomObjectScope(
+        {'LayerThatShouldFailIfNotAdded': LayerThatShouldFailIfNotAdded}):
+      _ = keras_load.load(saved_model_dir)
+
+  def test_must_restore_from_config_registration(self):
+    layer = GlobalLayerThatShouldFailIfNotAdded()
+    saved_model_dir = self._save_model_dir()
+    tf_save.save(layer, saved_model_dir)
+    _ = keras_load.load(saved_model_dir)
+
   def test_multi_input_model(self):
     input_1 = keras.layers.Input(shape=(3,))
     input_2 = keras.layers.Input(shape=(5,))

From 7a43680c0e69032e0c8025af45d23de4e5f6cc44 Mon Sep 17 00:00:00 2001
From: Jay Shi 
Date: Mon, 15 Jun 2020 16:44:00 -0700
Subject: [PATCH 0235/1390] [tf.data] Fix some small typo in
 `parallel_interleave_dataset_op`.

PiperOrigin-RevId: 316571471
Change-Id: Ie4b11e5256b88086a7efb3c87077936889a3fe10
---
 .../core/kernels/data/parallel_interleave_dataset_op.cc   | 6 ++++--
 .../kernels/data/parallel_interleave_dataset_op_test.cc   | 8 ++++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
index 8aca5005789..ac82666866d 100644
--- a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
@@ -1544,11 +1544,13 @@ void ParallelInterleaveDatasetOp::MakeDataset(OpKernelContext* ctx,
   int64 buffer_output_elements = model::kAutotune;
   int64 prefetch_input_elements = model::kAutotune;
   if (op_version_ >= 4) {
+    OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kBufferOutputElements,
+                                            &buffer_output_elements));
     OP_REQUIRES(ctx,
                 buffer_output_elements == model::kAutotune ||
-                    buffer_output_elements >= 0,
+                    buffer_output_elements > 0,
                 errors::InvalidArgument("`buffer_output_elements` must be ",
-                                        model::kAutotune, " or >= 0 but is ",
+                                        model::kAutotune, " or > 0 but is ",
                                         buffer_output_elements));
 
     OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kPrefetchInputElements,
diff --git a/tensorflow/core/kernels/data/parallel_interleave_dataset_op_test.cc b/tensorflow/core/kernels/data/parallel_interleave_dataset_op_test.cc
index 95a0d121f92..489a6a2a4b6 100644
--- a/tensorflow/core/kernels/data/parallel_interleave_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/parallel_interleave_dataset_op_test.cc
@@ -159,7 +159,7 @@ ParallelInterleaveDatasetParams ParallelInterleaveDatasetParams2() {
       /*other_arguments=*/{},
       /*cycle_length=*/2,
       /*block_length=*/1,
-      /*buffer_output_elements=*/0,
+      /*buffer_output_elements=*/1,
       /*prefetch_input_elements=*/0,
       /*num_parallel_calls=*/2,
       /*func=*/
@@ -184,7 +184,7 @@ ParallelInterleaveDatasetParams ParallelInterleaveDatasetParams3() {
       /*other_arguments=*/{},
       /*cycle_length=*/3,
       /*block_length=*/1,
-      /*buffer_output_elements=*/0,
+      /*buffer_output_elements=*/1,
       /*prefetch_input_elements=*/1,
       /*num_parallel_calls=*/2,
       /*func=*/
@@ -488,7 +488,7 @@ ParallelInterleaveDatasetParamsWithInvalidBufferOutputElements() {
       /*other_arguments=*/{},
       /*cycle_length=*/1,
       /*block_length=*/1,
-      /*buffer_output_elements=*/-2,
+      /*buffer_output_elements=*/model::kAutotune,
       /*prefetch_input_elements=*/model::kAutotune,
       /*num_parallel_calls=*/-5,
       /*func=*/
@@ -514,7 +514,7 @@ ParallelInterleaveDatasetParamsWithInvalidPrefetchInputElements() {
       /*other_arguments=*/{},
       /*cycle_length=*/1,
       /*block_length=*/1,
-      /*buffer_output_elements=*/-2,
+      /*buffer_output_elements=*/model::kAutotune,
       /*prefetch_input_elements=*/model::kAutotune,
       /*num_parallel_calls=*/-5,
       /*func=*/

From 5bc13320c9cab1dbca9348412ecc647ef62bebc7 Mon Sep 17 00:00:00 2001
From: Marat Dukhan 
Date: Mon, 15 Jun 2020 16:51:56 -0700
Subject: [PATCH 0236/1390] Re-enable depthwise_conv_2d_test on --config=msan
 in XNNPACK

PiperOrigin-RevId: 316572700
Change-Id: I0fe943907f18aa7d851cfb0e7e66e7b140b3079d
---
 tensorflow/lite/delegates/xnnpack/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index efbaf0cfc42..5736a2995b1 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -318,7 +318,6 @@ cc_test(
         "//tensorflow:emscripten": EMSCRIPTEN_LINKOPTS,
         "//conditions:default": [],
     }),
-    tags = ["nomsan"],  # b/145129478
     deps = [
         ":depthwise_conv_2d_tester",
         ":test_main",

From 6f0425cd06fd66c0686e376e17f3d9c23d773b86 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Mon, 15 Jun 2020 16:58:43 -0700
Subject: [PATCH 0237/1390] Unidirectional sequence LSTM: Use constants for
 tensor indices in lstm_shared.h instead of redefining them.

PiperOrigin-RevId: 316573775
Change-Id: I2de5264e4ab3375344dfe3132b6cde7b6cb0226c
---
 .../kernels/unidirectional_sequence_lstm.cc   | 192 +++++++-----------
 1 file changed, 76 insertions(+), 116 deletions(-)

diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
index a6fe785ce53..0552885f720 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/lstm_eval.h"
+#include "tensorflow/lite/kernels/lstm_shared.h"
 
 namespace tflite {
 namespace ops {
@@ -36,53 +37,6 @@ struct OpData {
   bool compute_row_sums = false;
 };
 
-// Input Tensors of size {max_time, n_batch, n_input}
-constexpr int kInputTensor = 0;
-
-// Input weight tensors of size: {n_cell, n_input}
-constexpr int kInputToInputWeightsTensor = 1;  // Optional
-constexpr int kInputToForgetWeightsTensor = 2;
-constexpr int kInputToCellWeightsTensor = 3;
-constexpr int kInputToOutputWeightsTensor = 4;
-
-// Recurrent weight tensors of size {n_cell, n_output}
-constexpr int kRecurrentToInputWeightsTensor = 5;  // Optional
-constexpr int kRecurrentToForgetWeightsTensor = 6;
-constexpr int kRecurrentToCellWeightsTensor = 7;
-constexpr int kRecurrentToOutputWeightsTensor = 8;
-
-// Peephole weights tensors of size {n_cell}, representing a diagonal matrix.
-constexpr int kCellToInputWeightsTensor = 9;    // Optional
-constexpr int kCellToForgetWeightsTensor = 10;  // Optional
-constexpr int kCellToOutputWeightsTensor = 11;  // Optional
-
-// Gates bias tensors of size {n_cell}
-constexpr int kInputGateBiasTensor = 12;  // Optional
-constexpr int kForgetGateBiasTensor = 13;
-constexpr int kCellGateBiasTensor = 14;
-constexpr int kOutputGateBiasTensor = 15;
-
-// Projection weight tensor of size {n_output, n_cell}
-constexpr int kProjectionWeightsTensor = 16;  // Optional
-// Projection bias tensor of size {n_output}
-constexpr int kProjectionBiasTensor = 17;  // Optional
-
-// Stateful input tensors that are variables and will be modified by the Op.
-// Activation state tensor of size {n_batch, n_output}
-constexpr int kInputActivationStateTensor = 18;
-// Cell state tensor of size {n_batch, n_cell}
-constexpr int kInputCellStateTensor = 19;
-
-// Layer norm coefficient tensors of size {n_cell}, representing a diagonal
-// matrix.
-constexpr int kInputLayerNormCoefficientsTensor = 20;   // Optional
-constexpr int kForgetLayerNormCoefficientsTensor = 21;  // Optional
-constexpr int kCellLayerNormCoefficientsTensor = 22;    // Optional
-constexpr int kOutputLayerNormCoefficientsTensor = 23;  // Optional
-
-// Output tensors.
-constexpr int kOutputTensor = 0;
-
 // Temporary tensors
 enum TemporaryTensor {
   kScratchBuffer = 0,
@@ -122,8 +76,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE(context, params->cell_clip >= 0);
   TF_LITE_ENSURE(context, params->proj_clip >= 0);
 
-  const TfLiteTensor* input_to_input_weights =
-      GetOptionalInputTensor(context, node, kInputToInputWeightsTensor);
+  const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kInputToInputWeightsTensor);
   if (input_to_input_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[0], n_cell);
@@ -131,19 +85,19 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   }
 
   const TfLiteTensor* input_to_forget_weights =
-      GetInput(context, node, kInputToForgetWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToForgetWeightsTensor);
   TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, input_to_forget_weights->dims->data[1], n_input);
 
   const TfLiteTensor* input_to_cell_weights =
-      GetInput(context, node, kInputToCellWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToCellWeightsTensor);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[1], n_input);
 
-  const TfLiteTensor* recurrent_to_input_weights =
-      GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor);
+  const TfLiteTensor* recurrent_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kRecurrentToInputWeightsTensor);
   if (recurrent_to_input_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[0],
@@ -153,7 +107,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   }
 
   const TfLiteTensor* recurrent_to_forget_weights =
-      GetInput(context, node, kRecurrentToForgetWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToForgetWeightsTensor);
   TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->data[0],
                     n_cell);
@@ -161,7 +115,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                     n_output);
 
   const TfLiteTensor* recurrent_to_cell_weights =
-      GetInput(context, node, kRecurrentToCellWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToCellWeightsTensor);
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[1],
@@ -176,22 +130,22 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
        (recurrent_to_input_weights == nullptr));
   TF_LITE_ENSURE(context, cifg_weights_all_or_none == true);
 
-  const TfLiteTensor* cell_to_input_weights =
-      GetOptionalInputTensor(context, node, kCellToInputWeightsTensor);
+  const TfLiteTensor* cell_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToInputWeightsTensor);
   if (cell_to_input_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->data[0], n_cell);
   }
 
-  const TfLiteTensor* cell_to_forget_weights =
-      GetOptionalInputTensor(context, node, kCellToForgetWeightsTensor);
+  const TfLiteTensor* cell_to_forget_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToForgetWeightsTensor);
   if (cell_to_forget_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->data[0], n_cell);
   }
 
-  const TfLiteTensor* cell_to_output_weights =
-      GetOptionalInputTensor(context, node, kCellToOutputWeightsTensor);
+  const TfLiteTensor* cell_to_output_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToOutputWeightsTensor);
   if (cell_to_output_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->data[0], n_cell);
@@ -210,7 +164,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
 
   // Make sure the input gate bias is present only when not a CIFG-LSTM.
   const TfLiteTensor* input_gate_bias =
-      GetOptionalInputTensor(context, node, kInputGateBiasTensor);
+      GetOptionalInputTensor(context, node, lstm::full::kInputGateBiasTensor);
   if (use_cifg) {
     TF_LITE_ENSURE_EQ(context, input_gate_bias, nullptr);
   } else {
@@ -219,21 +173,22 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   }
 
   const TfLiteTensor* forget_gate_bias =
-      GetInput(context, node, kForgetGateBiasTensor);
+      GetInput(context, node, lstm::full::kForgetGateBiasTensor);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell);
 
-  const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
+  const TfLiteTensor* cell_bias =
+      GetInput(context, node, lstm::full::kCellGateBiasTensor);
   TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell);
 
   const TfLiteTensor* output_gate_bias =
-      GetInput(context, node, kOutputGateBiasTensor);
+      GetInput(context, node, lstm::full::kOutputGateBiasTensor);
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->data[0], n_cell);
 
-  const TfLiteTensor* projection_weights =
-      GetOptionalInputTensor(context, node, kProjectionWeightsTensor);
+  const TfLiteTensor* projection_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kProjectionWeightsTensor);
   if (projection_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[0], n_output);
@@ -241,7 +196,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   }
 
   const TfLiteTensor* projection_bias =
-      GetOptionalInputTensor(context, node, kProjectionBiasTensor);
+      GetOptionalInputTensor(context, node, lstm::full::kProjectionBiasTensor);
   if (projection_bias != nullptr) {
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->data[0], n_output);
@@ -258,7 +213,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
 
   if (is_layer_norm_lstm) {
     const TfLiteTensor* input_layer_norm_coefficients = GetOptionalInputTensor(
-        context, node, kInputLayerNormCoefficientsTensor);
+        context, node, lstm::full::kInputLayerNormCoefficientsTensor);
     if (use_cifg) {
       TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients, nullptr);
     } else {
@@ -271,7 +226,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     }
 
     const TfLiteTensor* forget_layer_norm_coefficients =
-        GetInput(context, node, kForgetLayerNormCoefficientsTensor);
+        GetInput(context, node, lstm::full::kForgetLayerNormCoefficientsTensor);
     TF_LITE_ENSURE(context, forget_layer_norm_coefficients != nullptr);
     TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->dims->data[0],
@@ -280,7 +235,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                       kTfLiteFloat32);
 
     const TfLiteTensor* cell_layer_norm_coefficients =
-        GetInput(context, node, kCellLayerNormCoefficientsTensor);
+        GetInput(context, node, lstm::full::kCellLayerNormCoefficientsTensor);
     TF_LITE_ENSURE(context, cell_layer_norm_coefficients != nullptr);
     TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->dims->data[0],
@@ -289,7 +244,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                       kTfLiteFloat32);
 
     const TfLiteTensor* output_layer_norm_coefficients =
-        GetInput(context, node, kOutputLayerNormCoefficientsTensor);
+        GetInput(context, node, lstm::full::kOutputLayerNormCoefficientsTensor);
     TF_LITE_ENSURE(context, output_layer_norm_coefficients != nullptr);
     TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->dims->data[0],
@@ -312,7 +267,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   bool is_layer_norm_lstm = false;
   if (node->inputs->size == 24) {
     const TfLiteTensor* forget_layer_norm_coefficients = GetOptionalInputTensor(
-        context, node, kForgetLayerNormCoefficientsTensor);
+        context, node, lstm::full::kForgetLayerNormCoefficientsTensor);
     if (forget_layer_norm_coefficients == nullptr) {
       is_layer_norm_lstm = false;
     } else {
@@ -332,7 +287,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Inferring batch size, number of outputs and sequence length and
   // number of cells from the input tensors.
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  const TfLiteTensor* input = GetInput(context, node, lstm::full::kInputTensor);
   TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
   TF_LITE_ENSURE(context, input->dims->size > 1);
   const auto* params =
@@ -343,13 +298,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const int n_input = input->dims->data[2];
 
   const TfLiteTensor* input_to_output_weights =
-      GetInput(context, node, kInputToOutputWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToOutputWeightsTensor);
   const int n_cell = input_to_output_weights->dims->data[0];
   TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->data[1], n_input);
 
   const TfLiteTensor* recurrent_to_output_weights =
-      GetInput(context, node, kRecurrentToOutputWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToOutputWeightsTensor);
   TF_LITE_ENSURE_EQ(context, recurrent_to_output_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, recurrent_to_output_weights->dims->data[0],
                     n_cell);
@@ -361,13 +316,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                                                n_cell, is_layer_norm_lstm));
 
   // Get the pointer to output, activation_state and cell_state buffer tensors.
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  TfLiteTensor* output = GetOutput(context, node, lstm::full::kOutputTensor);
 
   TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
+      GetVariableInput(context, node, lstm::full::kInputActivationStateTensor);
   TF_LITE_ENSURE(context, activation_state != nullptr);
   TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
+      GetVariableInput(context, node, lstm::full::kInputCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   // Check the shape of input state tensors.
@@ -395,8 +350,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   scratch_buffer->type = input->type;
   scratch_buffer->allocation_type = kTfLiteArenaRw;
 
-  const TfLiteTensor* input_to_input_weights =
-      GetOptionalInputTensor(context, node, kInputToInputWeightsTensor);
+  const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kInputToInputWeightsTensor);
   const bool use_cifg = (input_to_input_weights == nullptr);
   TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(2);
   scratch_buffer_size->data[0] = n_batch;
@@ -534,8 +489,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     row_sums->type = kTfLiteInt32;
     row_sums->allocation_type = kTfLiteArenaRwPersistent;
     int row_sums_rows = use_cifg ? 6 : 8;
-    const TfLiteTensor* projection_weights =
-        GetOptionalInputTensor(context, node, kProjectionWeightsTensor);
+    const TfLiteTensor* projection_weights = GetOptionalInputTensor(
+        context, node, lstm::full::kProjectionWeightsTensor);
     if (projection_weights != nullptr) {
       row_sums_rows += ceil(static_cast(n_output) / n_cell);
     }
@@ -558,74 +513,79 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const OpData* op_data = reinterpret_cast(node->user_data);
   const bool is_layer_norm_lstm = op_data->is_layer_norm_lstm;
   const bool time_major = params->time_major;
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  const TfLiteTensor* input = GetInput(context, node, lstm::full::kInputTensor);
 
-  const TfLiteTensor* input_to_input_weights =
-      GetOptionalInputTensor(context, node, kInputToInputWeightsTensor);
+  const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kInputToInputWeightsTensor);
   const TfLiteTensor* input_to_forget_weights =
-      GetInput(context, node, kInputToForgetWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToForgetWeightsTensor);
   const TfLiteTensor* input_to_cell_weights =
-      GetInput(context, node, kInputToCellWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToCellWeightsTensor);
   const TfLiteTensor* input_to_output_weights =
-      GetInput(context, node, kInputToOutputWeightsTensor);
+      GetInput(context, node, lstm::full::kInputToOutputWeightsTensor);
 
-  const TfLiteTensor* recurrent_to_input_weights =
-      GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor);
+  const TfLiteTensor* recurrent_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kRecurrentToInputWeightsTensor);
   const TfLiteTensor* recurrent_to_forget_weights =
-      GetInput(context, node, kRecurrentToForgetWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToForgetWeightsTensor);
   const TfLiteTensor* recurrent_to_cell_weights =
-      GetInput(context, node, kRecurrentToCellWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToCellWeightsTensor);
   const TfLiteTensor* recurrent_to_output_weights =
-      GetInput(context, node, kRecurrentToOutputWeightsTensor);
+      GetInput(context, node, lstm::full::kRecurrentToOutputWeightsTensor);
 
-  const TfLiteTensor* cell_to_input_weights =
-      GetOptionalInputTensor(context, node, kCellToInputWeightsTensor);
-  const TfLiteTensor* cell_to_forget_weights =
-      GetOptionalInputTensor(context, node, kCellToForgetWeightsTensor);
-  const TfLiteTensor* cell_to_output_weights =
-      GetOptionalInputTensor(context, node, kCellToOutputWeightsTensor);
+  const TfLiteTensor* cell_to_input_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToInputWeightsTensor);
+  const TfLiteTensor* cell_to_forget_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToForgetWeightsTensor);
+  const TfLiteTensor* cell_to_output_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kCellToOutputWeightsTensor);
 
   const TfLiteTensor* input_gate_bias =
-      GetOptionalInputTensor(context, node, kInputGateBiasTensor);
+      GetOptionalInputTensor(context, node, lstm::full::kInputGateBiasTensor);
   const TfLiteTensor* forget_gate_bias =
-      GetInput(context, node, kForgetGateBiasTensor);
-  const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
+      GetInput(context, node, lstm::full::kForgetGateBiasTensor);
+  const TfLiteTensor* cell_bias =
+      GetInput(context, node, lstm::full::kCellGateBiasTensor);
   const TfLiteTensor* output_gate_bias =
-      GetInput(context, node, kOutputGateBiasTensor);
+      GetInput(context, node, lstm::full::kOutputGateBiasTensor);
 
-  const TfLiteTensor* projection_weights =
-      GetOptionalInputTensor(context, node, kProjectionWeightsTensor);
+  const TfLiteTensor* projection_weights = GetOptionalInputTensor(
+      context, node, lstm::full::kProjectionWeightsTensor);
   const TfLiteTensor* projection_bias =
-      GetOptionalInputTensor(context, node, kProjectionBiasTensor);
+      GetOptionalInputTensor(context, node, lstm::full::kProjectionBiasTensor);
 
   // Index the scratch buffers pointers to the global scratch buffer.
   TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
 
   TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
+      GetVariableInput(context, node, lstm::full::kInputActivationStateTensor);
   TF_LITE_ENSURE(context, activation_state != nullptr);
   TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
+      GetVariableInput(context, node, lstm::full::kInputCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   const TfLiteTensor* input_layer_norm_coefficients =
-      is_layer_norm_lstm ? GetOptionalInputTensor(
-                               context, node, kInputLayerNormCoefficientsTensor)
-                         : nullptr;
+      is_layer_norm_lstm
+          ? GetOptionalInputTensor(
+                context, node, lstm::full::kInputLayerNormCoefficientsTensor)
+          : nullptr;
   const TfLiteTensor* forget_layer_norm_coefficients =
       is_layer_norm_lstm
-          ? GetInput(context, node, kForgetLayerNormCoefficientsTensor)
+          ? GetInput(context, node,
+                     lstm::full::kForgetLayerNormCoefficientsTensor)
           : nullptr;
   const TfLiteTensor* cell_layer_norm_coefficients =
       is_layer_norm_lstm
-          ? GetInput(context, node, kCellLayerNormCoefficientsTensor)
+          ? GetInput(context, node,
+                     lstm::full::kCellLayerNormCoefficientsTensor)
           : nullptr;
   const TfLiteTensor* output_layer_norm_coefficients =
       is_layer_norm_lstm
-          ? GetInput(context, node, kOutputLayerNormCoefficientsTensor)
+          ? GetInput(context, node,
+                     lstm::full::kOutputLayerNormCoefficientsTensor)
           : nullptr;
 
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  TfLiteTensor* output = GetOutput(context, node, lstm::full::kOutputTensor);
 
   // Copy out the LSTM specific params so they can be passed in the function.
   TfLiteLSTMParams lstm_params;

From 304daeb37f56a12e0d0dd26654893e9e6cf840eb Mon Sep 17 00:00:00 2001
From: Jing Pu 
Date: Mon, 15 Jun 2020 17:03:42 -0700
Subject: [PATCH 0238/1390] Fix the tests that would be otherwise broken by
 https://reviews.llvm.org/D80258.

Previously, the tests relied on the misbehavior that col zero pointing to the first column that will be fixed by the LLVM patch.

PiperOrigin-RevId: 316574635
Change-Id: I6a66de318081bc5ebee3f888e1f1601fc3de9af3
---
 .../aot/tests/test_error_message.lit.pbtxt.debug.pbtxt        | 4 ++++
 .../graphdef2mlir/error-message-with-source-info.pbtxt.debug  | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt b/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt
index 7acc8287950..0820ebfd040 100644
--- a/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt
+++ b/tensorflow/compiler/aot/tests/test_error_message.lit.pbtxt.debug.pbtxt
@@ -4,6 +4,7 @@ traces: {
   value: {
     file_line_cols: {
       line: 1
+      col:  1
     }
   }
 }
@@ -12,9 +13,11 @@ traces: {
   value: {
     file_line_cols: {
       line: 3
+      col:  1
     }
     file_line_cols: {
       line: 4
+      col:  1
     }
   }
 }
@@ -23,6 +26,7 @@ traces: {
   value: {
     file_line_cols: {
       line: 2
+      col:  1
     }
   }
 }
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/error-message-with-source-info.pbtxt.debug b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/error-message-with-source-info.pbtxt.debug
index 881a01f3a3a..2e1afeb2329 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/error-message-with-source-info.pbtxt.debug
+++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/error-message-with-source-info.pbtxt.debug
@@ -4,6 +4,7 @@ traces: {
     value: {
       file_line_cols: {
         line      : 1
+        col       : 1
       }
     }
 }
@@ -12,9 +13,11 @@ traces: {
     value: {
       file_line_cols: {
         line      : 3
+        col       : 1
       }
       file_line_cols: {
         line      : 4
+        col       : 1
       }
     }
 }
@@ -23,6 +26,7 @@ traces: {
     value: {
       file_line_cols: {
         line      : 2
+        col       : 1
       }
     }
 }

From 0cc6210daa35247daf7f4cc98c115de611d2d05f Mon Sep 17 00:00:00 2001
From: Yanhui Liang 
Date: Mon, 15 Jun 2020 17:13:45 -0700
Subject: [PATCH 0239/1390] Cache DataHandler in `model.evaluate` to avoid
 function retracing between epochs in `model.fit`.

PiperOrigin-RevId: 316576844
Change-Id: Icf85ce6830b69a003c1f2ebf41f8c70258504afd
---
 tensorflow/python/keras/engine/training.py    | 55 +++++++++++++------
 .../python/keras/engine/training_test.py      | 22 ++++++++
 2 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index b7a4795d768..5567e1733a7 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -1040,6 +1040,10 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
           data_adapter.train_validation_split(
               (x, y, sample_weight), validation_split=validation_split))
 
+    if validation_data:
+      val_x, val_y, val_sample_weight = (
+          data_adapter.unpack_x_y_sample_weight(validation_data))
+
     with self.distribute_strategy.scope(), \
          training_utils.RespectCompiledTrainableState(self):
       # Creates a `tf.data.Dataset` and handles batch and epoch iteration.
@@ -1102,8 +1106,21 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
 
         # Run validation.
         if validation_data and self._should_eval(epoch, validation_freq):
-          val_x, val_y, val_sample_weight = (
-              data_adapter.unpack_x_y_sample_weight(validation_data))
+          # Create data_handler for evaluation and cache it.
+          if getattr(self, '_eval_data_handler', None) is None:
+            self._eval_data_handler = data_adapter.DataHandler(
+                x=val_x,
+                y=val_y,
+                sample_weight=val_sample_weight,
+                batch_size=validation_batch_size or batch_size,
+                steps_per_epoch=validation_steps,
+                initial_epoch=0,
+                epochs=1,
+                max_queue_size=max_queue_size,
+                workers=workers,
+                use_multiprocessing=use_multiprocessing,
+                model=self,
+                steps_per_execution=self._steps_per_execution)
           val_logs = self.evaluate(
               x=val_x,
               y=val_y,
@@ -1123,6 +1140,9 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         if self.stop_training:
           break
 
+      # If eval data_hanlder exists, delete it after all epochs are done.
+      if getattr(self, '_eval_data_handler', None) is not None:
+        del self._eval_data_handler
       callbacks.on_train_end(logs=training_logs)
       return self.history
 
@@ -1318,20 +1338,23 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
     _disallow_inside_tf_function('evaluate')
 
     with self.distribute_strategy.scope():
-      # Creates a `tf.data.Dataset` and handles batch and epoch iteration.
-      data_handler = data_adapter.DataHandler(
-          x=x,
-          y=y,
-          sample_weight=sample_weight,
-          batch_size=batch_size,
-          steps_per_epoch=steps,
-          initial_epoch=0,
-          epochs=1,
-          max_queue_size=max_queue_size,
-          workers=workers,
-          use_multiprocessing=use_multiprocessing,
-          model=self,
-          steps_per_execution=self._steps_per_execution)
+      if getattr(self, '_eval_data_handler', None) is not None:
+        data_handler = self._eval_data_handler
+      else:
+        # Creates a `tf.data.Dataset` and handles batch and epoch iteration.
+        data_handler = data_adapter.DataHandler(
+            x=x,
+            y=y,
+            sample_weight=sample_weight,
+            batch_size=batch_size,
+            steps_per_epoch=steps,
+            initial_epoch=0,
+            epochs=1,
+            max_queue_size=max_queue_size,
+            workers=workers,
+            use_multiprocessing=use_multiprocessing,
+            model=self,
+            steps_per_execution=self._steps_per_execution)
 
       # Container that configures and calls `tf.keras.Callback`s.
       if not isinstance(callbacks, callbacks_module.CallbackList):
diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index aa01463582c..5cf15926bfb 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -59,6 +59,7 @@ from tensorflow.python.ops import template
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training.rmsprop import RMSPropOptimizer
 
 try:
@@ -3538,5 +3539,26 @@ class TestAutoUpdates(keras_parameterized.TestCase):
     self.assertAllEqual(self.evaluate(bn.moving_variance), np.ones((10,)))
 
 
+class TestFunctionTracing(keras_parameterized.TestCase):
+
+  @keras_parameterized.run_all_keras_modes(
+      always_skip_v1=True, always_skip_eager=True)
+  def test_no_tracing_between_epoch(self):
+    if sys.version_info[0] < 3:
+      self.skipTest('self.assertLogs() call is not available in Python 2.')
+
+    model = sequential.Sequential([layers_module.Dense(4, activation='relu')])
+    model.compile(loss='mse', optimizer='rmsprop')
+    x = np.random.random((10, 6))
+    y = np.random.random((10, 4))
+
+    logging.set_verbosity(1)
+    with self.assertLogs(level=1) as logs:
+      model.fit(x, y, epochs=10, batch_size=5, validation_data=(x, y))
+
+    new_func_graph = 'INFO:absl:Creating new FuncGraph for Python function'
+    self.assertEqual(sum(new_func_graph in log for log in logs.output), 9)
+
+
 if __name__ == '__main__':
   test.main()

From ac695a31de1649a3d61e738b87e19ea6c558dada Mon Sep 17 00:00:00 2001
From: Gaurav Jain 
Date: Mon, 15 Jun 2020 17:16:40 -0700
Subject: [PATCH 0240/1390] Add DT_BOOL support to GPU variable ops

This is a follow-on to PR #38848 & PR #39172 and resolves remaining ask
in Issue #35994. The original PR tried to add many variable ops on the
GPU including DT_BOOL. However, this caused testCondModifyBoolPred to
fail and thus the DT_BOOL type was removed. The reason for the test
failure is once DT_BOOL variables are supported on the GPU, we need to
ensure the switch ops are also updated to not have host memory
requirement. Otherwise, a DT_BOOL ref variable is attempted to be
copied to the GPU which fails since we should not be transfering ref
types.

PiperOrigin-RevId: 316577397
Change-Id: Ic0d96ed4cdf8a0ea4674889aaff3a8ecd50991dd
---
 tensorflow/core/kernels/control_flow_ops.cc            | 10 +++++-----
 tensorflow/core/kernels/variable_ops.cc                |  3 +--
 .../debug/lib/debug_graph_reconstruction_test.py       |  6 +++---
 tensorflow/python/ops/control_flow_ops_test.py         |  6 +++---
 4 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/kernels/control_flow_ops.cc b/tensorflow/core/kernels/control_flow_ops.cc
index c8e83b6f672..435de3c5954 100644
--- a/tensorflow/core/kernels/control_flow_ops.cc
+++ b/tensorflow/core/kernels/control_flow_ops.cc
@@ -111,15 +111,17 @@ REGISTER_GPU_SWITCH(uint64);
 TF_CALL_variant(REGISTER_GPU_SWITCH);
 TF_CALL_uint32(REGISTER_GPU_SWITCH);
 TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
+TF_CALL_bool(REGISTER_GPU_SWITCH);
+TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
 
 #undef REGISTER_CPU_SWITCH
 #undef REGISTER_CPU_REF_SWITCH
 #undef REGISTER_GPU_SWITCH
 #undef REGISTER_GPU_REF_SWITCH
 
-// Special GPU kernels for int32 and string.
-// TODO(b/25387198): Also enable int32 in device memory. This kernel
-// registration requires all int32 inputs and outputs to be in host memory.
+// Special GPU kernels for int32, string & resource handles. Requiring all
+// inputs and outputs to be in host memory.
+// TODO(b/25387198): Also enable int32 in device memory.
 #define REGISTER_GPU_HOST_KERNEL(type)                    \
   REGISTER_KERNEL_BUILDER(Name("Switch")                  \
                               .Device(DEVICE_GPU)         \
@@ -149,8 +151,6 @@ TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
 
 REGISTER_GPU_HOST_KERNEL(int32);
 REGISTER_GPU_HOST_REF_KERNEL(int32);
-REGISTER_GPU_HOST_KERNEL(bool);
-REGISTER_GPU_HOST_REF_KERNEL(bool);
 REGISTER_GPU_HOST_KERNEL(tstring);
 REGISTER_GPU_HOST_REF_KERNEL(tstring);
 REGISTER_GPU_HOST_KERNEL(ResourceHandle);
diff --git a/tensorflow/core/kernels/variable_ops.cc b/tensorflow/core/kernels/variable_ops.cc
index 6f5e0b94eca..ccd33e8c75a 100644
--- a/tensorflow/core/kernels/variable_ops.cc
+++ b/tensorflow/core/kernels/variable_ops.cc
@@ -252,8 +252,7 @@ TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SYCL_KERNEL);
 
 TF_CALL_int64(REGISTER_GPU_KERNELS);
 TF_CALL_uint32(REGISTER_GPU_KERNELS);
-TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);
-TF_CALL_COMPLEX_TYPES(REGISTER_GPU_KERNELS);
+TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS);
 #undef REGISTER_GPU_KERNELS
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
diff --git a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
index fb722efab4e..b3baa6e7bc2 100644
--- a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
+++ b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
@@ -73,9 +73,9 @@ class ReconstructNonDebugGraphTest(test_util.TensorFlowTestCase):
           for attr_key in new_node.attr:
             if attr_key == "parallel_iterations":
               new_node.attr[attr_key].i = 1
-        elif new_node.op == "Switch":
-          # We don't check the inputs to Switch ops as their inputs may be
-          # Send/Recv nodes.
+        elif new_node.op == "Switch" or new_node.op == "Identity":
+          # We don't check the inputs to Switch or Identity ops as their inputs
+          # may be Send/Recv nodes.
           del new_node.input[:]
 
     return output_graph_def
diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py
index 9254695d988..3ca9bda82f2 100644
--- a/tensorflow/python/ops/control_flow_ops_test.py
+++ b/tensorflow/python/ops/control_flow_ops_test.py
@@ -396,10 +396,10 @@ class CondTest(test_util.TensorFlowTestCase):
         fn2=lambda: math_ops.add(y, 23))
     self.assertEquals(self.evaluate(z), 24)
 
-  @test_util.run_deprecated_v1
+  @test_util.run_v1_only("Exercises Ref variables")
   def testCondModifyBoolPred(self):
-    # This test in particular used to fail only when running in GPU, hence
-    # use_gpu=True.
+    # We want to use the GPU here because we want to ensure that we can update
+    # a boolean ref variable on the GPU.
     with test_util.use_gpu():
       bool_var = variable_scope.get_variable(
           "bool_var", dtype=dtypes.bool, initializer=True)

From 4ac4683f5ed3dca87afd799447a57ed81b5bfdcf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 17:17:28 -0700
Subject: [PATCH 0241/1390] Go: Update generated wrapper functions for
 TensorFlow ops.

PiperOrigin-RevId: 316577541
Change-Id: I10368c7cc5267b32ac298011dd16abaaab671448
---
 tensorflow/go/op/wrappers.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index f118e2bd494..10acebc7965 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -9495,6 +9495,14 @@ func DebugIdentityV2CircularBufferSize(value int64) DebugIdentityV2Attr {
 	}
 }
 
+// DebugIdentityV2TfdbgRunId sets the optional tfdbg_run_id attribute to value.
+// If not specified, defaults to ""
+func DebugIdentityV2TfdbgRunId(value string) DebugIdentityV2Attr {
+	return func(m optionalAttr) {
+		m["tfdbg_run_id"] = value
+	}
+}
+
 // Debug Identity V2 Op.
 //
 // Provides an identity mapping from input to output, while writing the content of

From c34265b348e2fccff165d33deceb91220d670dad Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 17:35:30 -0700
Subject: [PATCH 0242/1390] Integrate stackdriver support with tensorflow (roll
 forward cl/315012368)

PiperOrigin-RevId: 316580813
Change-Id: I93a66821fd13dc24a45c47f3e67638aed81d59be
---
 .bazelrc                  |  2 ++
 tensorflow/BUILD          |  7 +++++++
 tensorflow/core/BUILD     |  4 ++--
 tensorflow/python/BUILD   |  4 ++++
 tensorflow/tensorflow.bzl | 27 +++++++++++++++++++++++++--
 tensorflow/workspace.bzl  | 10 ++++++++++
 6 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index 15d46120642..5ea8048d5d9 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -200,6 +200,8 @@ build:nogcp --define=no_gcp_support=true
 build:nohdfs --define=no_hdfs_support=true
 build:nonccl --define=no_nccl_support=true
 
+build:stackdriver_support --define=stackdriver_support=true
+
 build --define=use_fast_cpp_protos=true
 build --define=allow_oversize_protos=true
 
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index ce759634232..bd0619b0c05 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -298,6 +298,13 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+# Experimental features
+config_setting(
+    name = "stackdriver_support",
+    define_values = {"stackdriver_support": "true"},
+    visibility = ["//visibility:public"],
+)
+
 # Crosses between platforms and file system libraries not supported on those
 # platforms due to limitations in nested select() statements.
 config_setting(
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 2b16801f6ed..50f1f2527a5 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -100,7 +100,7 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_test_gpu")
 load("//tensorflow:tensorflow.bzl", "tf_cc_tests_gpu")
 
 # buildifier: disable=same-origin-load
-load("//tensorflow:tensorflow.bzl", "tf_monitoring_deps")
+load("//tensorflow:tensorflow.bzl", "tf_monitoring_framework_deps")
 
 # For platform specific build config
 load(
@@ -1950,7 +1950,7 @@ cc_library(
         "@zlib",
         "@double_conversion//:double-conversion",
         "@com_google_protobuf//:protobuf",
-    ] + tf_protos_all_impl() + tf_protos_grappler_impl() + tf_protos_profiler_impl() + tf_monitoring_deps(),
+    ] + tf_protos_all_impl() + tf_protos_grappler_impl() + tf_protos_profiler_impl() + tf_monitoring_framework_deps(),
     # Alwayslink causes a cc_binary to "always link" in the
     # srcs for a given cc_library, even if they are unreferenced, see:
     # https://docs.bazel.build/versions/master/be/c-cpp.html#cc_library.alwayslink
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index a7338772589..87048ba9d40 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -5,6 +5,9 @@
 load("//tensorflow:tensorflow.bzl", "py_strict_library")
 load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cuda_library", "tf_gen_op_wrapper_py")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow:tensorflow.bzl", "tf_monitoring_python_deps")
+
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")
 
@@ -6048,6 +6051,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/core/util/tensor_bundle",
         "//tensorflow/compiler/mlir/python:mlir",
     ] + (tf_additional_lib_deps() +
+         tf_monitoring_python_deps() +
          tf_additional_plugin_deps() +
          tf_additional_profiler_deps()) + if_ngraph([
         "@ngraph_tf//:ngraph_tf",
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index 2609f5a42cf..f97363a919e 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -2864,8 +2864,31 @@ def if_cuda_or_rocm(if_true, if_false = []):
         "//conditions:default": if_false,
     })
 
-def tf_monitoring_deps():
-    return []
+def tf_monitoring_framework_deps(link_to_tensorflow_framework = True):
+    """Get the monitoring libs that will be linked to the tensorflow framework.
+
+      Currently in OSS, the protos must be statically linked to the tensorflow
+      framework, whereas the grpc should not be linked here.
+    """
+    return select({
+        "//tensorflow:stackdriver_support": [
+            "@com_github_googlecloudplatform_tensorflow_gcp_tools//monitoring:stackdriver_exporter_protos",
+        ],
+        "//conditions:default": [],
+    })
+
+def tf_monitoring_python_deps():
+    """Get the monitoring libs that will be linked to the python wrapper.
+
+      Currently in OSS, the grpc must be statically linked to the python wrapper
+      whereas the protos should not be linked here.
+    """
+    return select({
+        "//tensorflow:stackdriver_support": [
+            "@com_github_googlecloudplatform_tensorflow_gcp_tools//monitoring:stackdriver_exporter",
+        ],
+        "//conditions:default": [],
+    })
 
 def tf_jit_compilation_passes_extra_deps():
     return []
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 6b0143e397f..60fef8c0cb9 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -351,6 +351,16 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
         ],
     )
 
+    tf_http_archive(
+        name = "com_github_googlecloudplatform_tensorflow_gcp_tools",
+        sha256 = "5e9ebe17eaa2895eb7f77fefbf52deeda7c4b63f5a616916b823eb74f3a0c542",
+        strip_prefix = "tensorflow-gcp-tools-2643d8caeba6ca2a6a0b46bb123953cb95b7e7d5",
+        urls = [
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/GoogleCloudPlatform/tensorflow-gcp-tools/archive/2643d8caeba6ca2a6a0b46bb123953cb95b7e7d5.tar.gz",
+            "https://github.com/GoogleCloudPlatform/tensorflow-gcp-tools/archive/2643d8caeba6ca2a6a0b46bb123953cb95b7e7d5.tar.gz",
+        ],
+    )
+
     tf_http_archive(
         name = "com_google_googleapis",
         build_file = clean_dep("//third_party/googleapis:googleapis.BUILD"),

From 2f45ee867dffb4dc8accb523d86663b7c9184635 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 17:47:21 -0700
Subject: [PATCH 0243/1390]   Add decomposition for
 ResourceApplyCenteredRMSProp

PiperOrigin-RevId: 316582464
Change-Id: Id61fb84cf6ce26398fcd870f331b51e0af8b9c3d
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    | 43 ++++++++++++
 .../tests/decompose_resource_ops.mlir         | 50 ++++++++++++++
 .../transforms/decompose_resource_ops.td      | 66 +++++++++++++++++++
 3 files changed, 159 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index df8dccb2163..6131a729441 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -7280,6 +7280,49 @@ $$\text{variable} := \text{variable} - \text{lr}_t * m_t / (\sqrt{v_t} + \epsilo
   TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<3>;
 }
 
+def TF_ResourceApplyCenteredRMSPropOp : TF_Op<"ResourceApplyCenteredRMSProp", []> {
+  let summary = "Update '*var' according to the centered RMSProp algorithm.";
+
+  let description = [{
+The centered RMSProp algorithm uses an estimate of the centered second moment
+(i.e., the variance) for normalization, as opposed to regular RMSProp, which
+uses the (uncentered) second moment. This often helps with training, but is
+slightly more expensive in terms of computation and memory.
+
+Note that in dense implementation of this algorithm, mg, ms, and mom will
+update even if the grad is zero, but in this sparse implementation, mg, ms,
+and mom will not update in iterations during which the grad is zero.
+
+mean_square = decay * mean_square + (1-decay) * gradient ** 2
+mean_grad = decay * mean_grad + (1-decay) * gradient
+
+Delta = learning_rate * gradient / sqrt(mean_square + epsilon - mean_grad ** 2)
+
+mg <- rho * mg_{t-1} + (1-rho) * grad
+ms <- rho * ms_{t-1} + (1-rho) * grad * grad
+mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms - mg * mg + epsilon)
+var <- var - mom
+  }];
+
+  let arguments = (ins
+    TF_ResourceTensor:$var,
+    TF_ResourceTensor:$mg,
+    TF_ResourceTensor:$ms,
+    TF_ResourceTensor:$mom,
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64, TF_Qint32, TF_Qint8, TF_Quint8, TF_Uint16, TF_Uint32, TF_Uint64, TF_Uint8]>:$lr,
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64, TF_Qint32, TF_Qint8, TF_Quint8, TF_Uint16, TF_Uint32, TF_Uint64, TF_Uint8]>:$rho,
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64, TF_Qint32, TF_Qint8, TF_Quint8, TF_Uint16, TF_Uint32, TF_Uint64, TF_Uint8]>:$momentum,
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64, TF_Qint32, TF_Qint8, TF_Quint8, TF_Uint16, TF_Uint32, TF_Uint64, TF_Uint8]>:$epsilon,
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64, TF_Qint32, TF_Qint8, TF_Quint8, TF_Uint16, TF_Uint32, TF_Uint64, TF_Uint8]>:$grad,
+
+    DefaultValuedAttr:$use_locking
+  );
+
+  let results = (outs);
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<4>;
+}
+
 def TF_ResourceApplyGradientDescentOp : TF_Op<"ResourceApplyGradientDescent", []> {
   let summary = "Update '*var' by subtracting 'alpha' * 'delta' from it.";
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
index 7a2e5173247..25dfda25358 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
@@ -368,6 +368,56 @@ func @decompose_resource_gather_op(%indices : tensor<5xi32>) -> tensor<2x5x16xi3
 
 // -----
 
+// Tests that composite tf.ResourceApplyCenteredRMSProp operation is decomposed.
+
+// CHECK-LABEL: func @decompose_resource_apply_centered_RMS_prop
+// CHECK-SAME:  [[VAR:%.*]]: tensor, [[MG:%.*]]: tensor, [[MS:%.*]]: tensor, [[MOM:%.*]]: tensor, [[LR:%.*]]: tensor, [[RHO:%.*]]: tensor, [[MOMENTUM:%.*]]: tensor, [[EPSILON:%.*]]: tensor, [[GRAD:%.*]]: tensor
+func @decompose_resource_apply_centered_RMS_prop(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor, %arg6: tensor, %arg7: tensor, %arg8: tensor) -> () {
+  // CHECK: [[ONE:%.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor}
+  // CHECK: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp"
+  // CHECK: [[MG_HANDLE:%.*]] = "tf.VarHandleOp"
+  // CHECK: [[MS_HANDLE:%.*]] = "tf.VarHandleOp"
+  // CHECK: [[MOM_HANDLE:%.*]] = "tf.VarHandleOp"
+  %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf.resource>
+  %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf.resource>
+  %2 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf.resource>
+  %3 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf.resource>
+
+  // CHECK: [[GRADSQ:%.*]] = "tf.Mul"([[GRAD]], [[GRAD]])
+  // CHECK: [[SB:%.*]] = "tf.Sub"([[ONE]], [[RHO]])
+  // CHECK: [[GRAD_SUB:%.*]] = "tf.Mul"([[GRADSQ]], [[SB]])
+  // CHECK: [[MS:%.*]] = "tf.ReadVariableOp"([[MS_HANDLE]])
+  // CHECK: [[MS_RHO:%.*]] = "tf.Mul"([[MS]], [[RHO]])
+  // CHECK: [[MS_NEW:%.*]] = "tf.Add"([[GRAD_SUB]], [[MS_RHO]])
+  // CHECK: "tf.AssignVariableOp"([[MS_HANDLE]], [[MS_NEW]])
+
+  // CHECK: [[SUB_RHO:%.*]] = "tf.Sub"([[ONE]], [[RHO]])
+  // CHECK: [[SUB_GRAD:%.*]] = "tf.Mul"([[GRAD]], [[SUB_RHO]])
+  // CHECK: [[MG:%.*]] = "tf.ReadVariableOp"([[MG_HANDLE]])
+  // CHECK: [[MG_RHO:%.*]] = "tf.Mul"([[MG]], [[RHO]])
+  // CHECK: [[MG_NEW:%.*]] = "tf.Add"([[SUB_GRAD]], [[MG_RHO]])
+  // CHECK: "tf.AssignVariableOp"([[MG_HANDLE]], [[MG_NEW]])
+
+  // CHECK: [[MOM:%.*]] = "tf.ReadVariableOp"([[MOM_HANDLE]])
+  // CHECK: [[MOM_MOM:%.*]] = "tf.Mul"([[MOMENTUM]], [[MOM]])
+  // CHECK: [[LR_GRAD:%.*]] = "tf.Mul"([[LR]], [[GRAD]])
+
+  // CHECK: [[MG_MG:%.*]] = "tf.Mul"([[MG_NEW]], [[MG_NEW]])
+  // CHECK: [[MG_NEW:%.*]] = "tf.Add"([[MG_MG]], [[EPSILON]])
+  // CHECK: [[MG_SUB:%.*]] = "tf.Sub"([[MS_NEW]], [[MG_NEW]])
+  // CHECK: [[MG_SQRT:%.*]] = "tf.Sqrt"([[MG_SUB]])
+  // CHECK: [[MOM_DIV:%.*]] = "tf.Div"([[LR_GRAD]], [[MG_SQRT]])
+  // CHECK: [[MOM_NEW:%.*]] = "tf.Add"([[MOM_MOM]], [[MOM_DIV]])
+
+  // CHECK: [[VAR:%.*]] = "tf.ReadVariableOp"([[VAR_HANDLE]])
+  // CHECK: [[VAR_NEW:%.*]] = "tf.Sub"([[VAR]], [[MOM_NEW]])
+
+  "tf.ResourceApplyCenteredRMSProp"(%0, %1, %2, %3, %arg4, %arg5, %arg6, %arg7, %arg8) {use_locking = false} : (tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor, tensor, tensor, tensor, tensor) -> ()
+  return
+}
+
+// -----
+
 // Tests that composite tf.ResourceScatterUpdate operation is decomposed.
 
 // CHECK-LABEL: @decompose_resource_scatter_update_op
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
index 3869a1a7fa3..0dd7d778e31 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
@@ -327,3 +327,69 @@ def DecomposeVariableShape : Pat<
   (TF_VariableShapeOp:$src_op $resource),
   (TF_ShapeOp (CreateTFReadVariableOpFromResourceHandle $src_op, $resource)),
   [(CheckHasResourceSubtype $resource)]>;
+
+// This decomposition is only correct inside XLA as it ignores use_locking
+// attribute.
+// ms <- rho * ms_{t-1} + (1-rho) * grad * grad
+// mg = grad * (one - rho) + mg * rho;
+// mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms - mg * mg + epsilon)
+//
+def DecomposeResourceApplyCenteredRMSProp :
+  Pattern<
+    (TF_ResourceApplyCenteredRMSPropOp:$src_op
+       $var_resource, $mg_resource, $ms_resource, $mom_resource, $lr, $rho, $momentum, $epsilon,
+       $grad, ConstBoolAttrFalse:$use_locking
+    ),
+    [(TF_ConstOp:$one (GetScalarOfType<1> $grad)),
+     (CreateTFReadVariableOp $src_op, $grad, $ms_resource),
+     (TF_AddOp:$ms_new
+       (TF_MulOp
+         (TF_MulOp $grad, $grad),
+         (TF_SubOp $one, $rho)
+       ),
+       (TF_MulOp
+          (CreateTFReadVariableOp $src_op, $grad, $ms_resource),
+          $rho
+       )
+     ),
+     (TF_AssignVariableOp $ms_resource, $ms_new),
+     // mg = grad * (one - rho) + mg * rho;
+     (TF_AddOp:$mg_new
+       (TF_MulOp
+         $grad,
+         (TF_SubOp $one, $rho)
+       ),
+       (TF_MulOp
+          (CreateTFReadVariableOp $src_op, $grad, $mg_resource),
+          $rho
+       )
+     ),
+     (TF_AssignVariableOp $mg_resource, $mg_new),
+     // mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms - mg * mg + epsilon)
+     (TF_AddOp:$mom_new
+      (TF_MulOp $momentum,
+       (CreateTFReadVariableOp $src_op, $grad, $mom_resource)),
+      (TF_DivOp
+         (TF_MulOp $lr, $grad),
+         (TF_SqrtOp
+           (TF_SubOp
+             $ms_new,
+             (TF_AddOp
+               (TF_MulOp
+                 $mg_new,
+                 $mg_new
+               ),
+               $epsilon
+             )
+           )
+         )
+      )
+     ),
+     (TF_AssignVariableOp $mom_resource, $mom_new),
+     // var <- var - mom
+     (TF_AssignSubVariableOp $var_resource,
+        (TF_SubOp (CreateTFReadVariableOp $src_op, $grad, $var_resource),
+                  $mom_new)
+     )
+   ]
+   >;

From fe6580b4d85b12a2ce7b1a529b70fdbfeedd899e Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Tue, 16 Jun 2020 02:16:09 +0100
Subject: [PATCH 0244/1390] Add comment about ignoring distributed multi
 assignment

---
 .../keras/mixed_precision/experimental/autocast_variable_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 95957f5634e..14f26cdf953 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -311,6 +311,7 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertAllClose(3., self.evaluate(x.assign_sub(3.)))
 
         # Assign multiple times
+        # This currently only works if no strategy is used
         if not ds_context.has_strategy():
           assign = x.assign(1.)
           self.assertAllClose(1., self.evaluate(assign))

From da2046b8a3b1bb79c77bf258aa8a52887bc3703a Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Tue, 16 Jun 2020 02:25:18 +0100
Subject: [PATCH 0245/1390] Use default_strategy instead of dummy scope

---
 .../experimental/autocast_variable_test.py               | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 14f26cdf953..c45015b644e 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -17,7 +17,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import contextlib
 import os
 
 from absl.testing import parameterized
@@ -43,15 +42,9 @@ from tensorflow.python.platform import test
 from tensorflow.python.training import gradient_descent as gradient_descent_v1
 from tensorflow.python.training.tracking import util as trackable_utils
 
-class DummyStrategy(object):
-  @contextlib.contextmanager
-  def scope(self):
-    yield
-
 maybe_distribute = combinations.combine(
     distribution=[
-        combinations.NamedDistribution(
-            "Dummy", lambda: DummyStrategy(), required_gpus=None),
+        strategy_combinations.default_strategy,
         strategy_combinations.mirrored_strategy_with_cpu_1_and_2
     ])
 

From 0292d384f84f55a78312ae0132f700caec245315 Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon 
Date: Mon, 15 Jun 2020 18:23:12 -0700
Subject: [PATCH 0246/1390] Suppress a deprecation warning for Object.finalize

The API is deprecated starting in JDK 9:
https://bugs.openjdk.java.net/browse/JDK-8165641

PiperOrigin-RevId: 316587198
Change-Id: Icafed5879fd50973e955522a176aa11a8773090b
---
 .../java/src/main/java/org/tensorflow/lite/Interpreter.java     | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
index 0f8b7b5c2f2..7c9c5644f47 100644
--- a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
+++ b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
@@ -507,6 +507,8 @@ public final class Interpreter implements AutoCloseable {
     }
   }
 
+  // for Object.finalize, see https://bugs.openjdk.java.net/browse/JDK-8165641
+  @SuppressWarnings("deprecation")
   @Override
   protected void finalize() throws Throwable {
     try {

From 2a96849f478627ad03b8f20f8c71e85d6d7480e3 Mon Sep 17 00:00:00 2001
From: Karim Nosir 
Date: Mon, 15 Jun 2020 18:38:58 -0700
Subject: [PATCH 0247/1390] Update source files with used includes.

PiperOrigin-RevId: 316589177
Change-Id: I0aba0ed1cf9ff478e7890fa53a7749bf844bd26d
---
 tensorflow/lite/kernels/BUILD                 | 406 +++++++++---------
 .../lite/kernels/acceleration_test_util.cc    |  10 +-
 .../lite/kernels/acceleration_test_util.h     |   2 +-
 .../acceleration_test_util_internal.cc        |   8 +
 .../kernels/acceleration_test_util_internal.h |   4 +-
 .../acceleration_test_util_internal_test.cc   |   2 +
 tensorflow/lite/kernels/activations.cc        |  12 +-
 tensorflow/lite/kernels/activations_test.cc   |  18 +-
 tensorflow/lite/kernels/add.cc                |  12 +-
 tensorflow/lite/kernels/add_n.cc              |   2 +
 tensorflow/lite/kernels/add_n_test.cc         |   7 +-
 tensorflow/lite/kernels/add_test.cc           |  10 +-
 tensorflow/lite/kernels/arg_min_max.cc        |   8 +-
 tensorflow/lite/kernels/arg_min_max_test.cc   |  11 +-
 tensorflow/lite/kernels/assign_variable.cc    |   6 +-
 tensorflow/lite/kernels/audio_spectrogram.cc  |   9 +-
 .../lite/kernels/audio_spectrogram_test.cc    |   5 +-
 tensorflow/lite/kernels/basic_rnn_test.cc     |   7 +-
 tensorflow/lite/kernels/batch_matmul.cc       |   6 +
 tensorflow/lite/kernels/batch_matmul_test.cc  |  12 +-
 tensorflow/lite/kernels/batch_to_space_nd.cc  |   8 +-
 .../lite/kernels/batch_to_space_nd_test.cc    |   9 +-
 .../kernels/bidirectional_sequence_lstm.cc    |   3 +
 .../bidirectional_sequence_lstm_test.cc       |   9 +-
 .../kernels/bidirectional_sequence_rnn.cc     |   1 +
 .../bidirectional_sequence_rnn_test.cc        |  11 +-
 tensorflow/lite/kernels/cast.cc               |   6 +-
 tensorflow/lite/kernels/cast_test.cc          |  11 +-
 tensorflow/lite/kernels/ceil.cc               |   1 +
 tensorflow/lite/kernels/ceil_test.cc          |   8 +-
 tensorflow/lite/kernels/comparisons.cc        |   9 +
 tensorflow/lite/kernels/comparisons_test.cc   |  13 +-
 tensorflow/lite/kernels/concatenation.cc      |  13 +-
 tensorflow/lite/kernels/concatenation_test.cc |  13 +-
 tensorflow/lite/kernels/conv.cc               |  13 +-
 tensorflow/lite/kernels/conv_test.cc          |  15 +-
 .../lite/kernels/cpu_backend_context.cc       |   4 +
 tensorflow/lite/kernels/cpu_backend_context.h |   1 +
 tensorflow/lite/kernels/cpu_backend_gemm.h    |   1 +
 .../kernels/cpu_backend_gemm_custom_gemv.h    |   5 +
 .../lite/kernels/cpu_backend_gemm_gemmlowp.h  |   4 +-
 .../lite/kernels/cpu_backend_gemm_ruy.h       |   2 +-
 .../lite/kernels/cpu_backend_gemm_test.cc     |   9 +-
 .../kernels/cpu_backend_threadpool_test.cc    |   2 +
 tensorflow/lite/kernels/custom_ops_register.h |   2 +-
 tensorflow/lite/kernels/densify.cc            |   5 +-
 tensorflow/lite/kernels/densify_test.cc       |   8 +-
 tensorflow/lite/kernels/depth_to_space.cc     |   5 +-
 .../lite/kernels/depth_to_space_test.cc       |  11 +-
 tensorflow/lite/kernels/depthwise_conv.cc     |  15 +-
 .../kernels/depthwise_conv_hybrid_test.cc     |  15 +-
 .../lite/kernels/depthwise_conv_test.cc       |  14 +-
 tensorflow/lite/kernels/dequantize.cc         |   8 +-
 tensorflow/lite/kernels/dequantize.h          |   7 +-
 tensorflow/lite/kernels/dequantize_test.cc    |   9 +-
 .../lite/kernels/detection_postprocess.cc     |  10 +-
 .../kernels/detection_postprocess_test.cc     |   9 +-
 tensorflow/lite/kernels/div.cc                |   9 +-
 tensorflow/lite/kernels/div_test.cc           |  10 +-
 tensorflow/lite/kernels/eigen_support.cc      |   3 +
 tensorflow/lite/kernels/eigen_support_test.cc |   3 +-
 tensorflow/lite/kernels/elementwise.cc        |   4 +
 tensorflow/lite/kernels/elementwise_test.cc   |   8 +-
 tensorflow/lite/kernels/embedding_lookup.cc   |  12 +-
 .../lite/kernels/embedding_lookup_sparse.cc   |   3 +-
 .../kernels/embedding_lookup_sparse_test.cc   |   7 +-
 .../lite/kernels/embedding_lookup_test.cc     |   9 +-
 tensorflow/lite/kernels/exp.cc                |   7 +-
 tensorflow/lite/kernels/exp_test.cc           |  11 +-
 tensorflow/lite/kernels/expand_dims.cc        |   6 +-
 tensorflow/lite/kernels/expand_dims_test.cc   |  14 +-
 tensorflow/lite/kernels/fake_quant.cc         |   7 +-
 tensorflow/lite/kernels/fake_quant_test.cc    |   9 +-
 tensorflow/lite/kernels/fill.cc               |   4 +-
 tensorflow/lite/kernels/fill_test.cc          |  12 +-
 tensorflow/lite/kernels/floor.cc              |   1 +
 tensorflow/lite/kernels/floor_div.cc          |   9 +-
 tensorflow/lite/kernels/floor_div_test.cc     |   9 +-
 tensorflow/lite/kernels/floor_mod.cc          |   7 +-
 tensorflow/lite/kernels/floor_mod_test.cc     |   9 +-
 tensorflow/lite/kernels/floor_test.cc         |   8 +-
 tensorflow/lite/kernels/fully_connected.cc    |   1 -
 .../lite/kernels/fully_connected_test.cc      |  15 +-
 tensorflow/lite/kernels/gather.cc             |   6 +-
 tensorflow/lite/kernels/gather_nd.cc          |   6 +-
 tensorflow/lite/kernels/gather_nd_test.cc     |  13 +-
 tensorflow/lite/kernels/gather_test.cc        |  14 +-
 tensorflow/lite/kernels/hashtable/BUILD       |   2 +-
 tensorflow/lite/kernels/hashtable_lookup.cc   |   9 +-
 .../lite/kernels/hashtable_lookup_test.cc     |  11 +-
 tensorflow/lite/kernels/if.cc                 |   4 +
 tensorflow/lite/kernels/if_test.cc            |   9 +-
 .../kernels/internal/optimized/im2col_utils.h |   2 +
 .../internal/optimized/integer_ops/add.h      |   5 +
 .../internal/optimized/integer_ops/conv.h     |   1 +
 .../optimized/integer_ops/depthwise_conv.h    |   7 +
 .../optimized/integer_ops/fully_connected.h   |   4 +-
 .../internal/optimized/integer_ops/mul.h      |   5 +
 .../internal/optimized/integer_ops/pooling.h  |  13 +-
 .../internal/per_channel_dequantize_test.cc   |   1 +
 .../kernels/internal/reference/batch_matmul.h |   5 +-
 .../lite/kernels/internal/reference/densify.h |   2 +
 .../internal/reference/non_max_suppression.h  |   2 +-
 .../kernels/internal/reference/quantize.h     |   4 +
 .../internal/reference/strided_slice.h        |   2 +
 .../lite/kernels/internal/reference/sub.h     |   8 +-
 .../lite/kernels/internal/reference/svdf.h    |   3 +
 tensorflow/lite/kernels/kernel_util.cc        |   7 +-
 tensorflow/lite/kernels/kernel_util.h         |   4 +-
 tensorflow/lite/kernels/kernel_util_test.cc   |  10 +
 tensorflow/lite/kernels/l2norm.cc             |   5 +-
 tensorflow/lite/kernels/l2norm_test.cc        |  11 +-
 .../lite/kernels/local_response_norm.cc       |   3 +-
 .../lite/kernels/local_response_norm_test.cc  |   9 +-
 tensorflow/lite/kernels/log_softmax_test.cc   |   8 +-
 tensorflow/lite/kernels/logical.cc            |   5 +-
 tensorflow/lite/kernels/logical_test.cc       |   9 +-
 tensorflow/lite/kernels/lsh_projection.cc     |  10 +-
 .../lite/kernels/lsh_projection_test.cc       |   7 +-
 tensorflow/lite/kernels/lstm_eval.cc          |   5 +
 tensorflow/lite/kernels/lstm_eval.h           |   1 -
 tensorflow/lite/kernels/lstm_eval_test.cc     |  12 +-
 tensorflow/lite/kernels/lstm_test.cc          |   8 +-
 tensorflow/lite/kernels/matrix_diag.cc        |   7 +-
 tensorflow/lite/kernels/matrix_diag_test.cc   |  10 +-
 tensorflow/lite/kernels/matrix_set_diag.cc    |   7 +-
 .../lite/kernels/matrix_set_diag_test.cc      |  10 +-
 tensorflow/lite/kernels/maximum_minimum.cc    |  10 +-
 .../lite/kernels/maximum_minimum_test.cc      |  11 +-
 tensorflow/lite/kernels/mfcc.cc               |   9 +-
 tensorflow/lite/kernels/mfcc_test.cc          |   6 +-
 tensorflow/lite/kernels/mirror_pad.cc         |   6 +-
 tensorflow/lite/kernels/mirror_pad_test.cc    |   7 +-
 tensorflow/lite/kernels/mul.cc                |  10 +-
 tensorflow/lite/kernels/mul_test.cc           |  11 +-
 tensorflow/lite/kernels/neg.cc                |   3 +
 tensorflow/lite/kernels/neg_test.cc           |  11 +-
 .../lite/kernels/non_max_suppression.cc       |   8 +-
 .../lite/kernels/non_max_suppression_test.cc  |   9 +-
 tensorflow/lite/kernels/numeric_verify.cc     |   9 +-
 .../lite/kernels/numeric_verify_test.cc       |   6 +-
 tensorflow/lite/kernels/one_hot.cc            |   4 +-
 tensorflow/lite/kernels/one_hot_test.cc       |   7 +-
 tensorflow/lite/kernels/op_macros.h           |   2 -
 .../lite/kernels/optional_tensor_test.cc      |  10 +-
 tensorflow/lite/kernels/pack.cc               |   4 +
 tensorflow/lite/kernels/pack_test.cc          |  11 +-
 tensorflow/lite/kernels/pad.cc                |  11 +-
 tensorflow/lite/kernels/pad_test.cc           |   8 +-
 tensorflow/lite/kernels/pooling.cc            |  13 +-
 tensorflow/lite/kernels/pooling_test.cc       |  11 +-
 tensorflow/lite/kernels/pow.cc                |   5 +-
 tensorflow/lite/kernels/pow_test.cc           |  10 +-
 .../lite/kernels/quant_basic_lstm_test.cc     |   7 +-
 tensorflow/lite/kernels/quantize.cc           |   1 +
 tensorflow/lite/kernels/quantize_test.cc      |   7 +-
 tensorflow/lite/kernels/range.cc              |   8 +-
 tensorflow/lite/kernels/range_test.cc         |   8 +-
 tensorflow/lite/kernels/rank.cc               |   5 +-
 tensorflow/lite/kernels/rank_test.cc          |   8 +-
 tensorflow/lite/kernels/read_variable.cc      |   5 +-
 tensorflow/lite/kernels/reduce.cc             |   9 +-
 tensorflow/lite/kernels/reduce_test.cc        |  10 +-
 tensorflow/lite/kernels/register.cc           |   2 +
 tensorflow/lite/kernels/register.h            |   1 -
 tensorflow/lite/kernels/register_ref.cc       |   4 +
 tensorflow/lite/kernels/register_ref.h        |   2 +-
 tensorflow/lite/kernels/reshape.cc            |   2 +-
 tensorflow/lite/kernels/reshape_test.cc       |   7 +-
 tensorflow/lite/kernels/reshape_test_common.h |   8 +
 tensorflow/lite/kernels/resize_bilinear.cc    |   7 +-
 .../lite/kernels/resize_bilinear_test.cc      |  10 +-
 .../lite/kernels/resize_nearest_neighbor.cc   |   9 +-
 .../kernels/resize_nearest_neighbor_test.cc   |  10 +-
 tensorflow/lite/kernels/reverse.cc            |   4 +-
 tensorflow/lite/kernels/reverse_sequence.cc   |   3 +
 .../lite/kernels/reverse_sequence_test.cc     |   9 +-
 tensorflow/lite/kernels/reverse_test.cc       |   9 +-
 tensorflow/lite/kernels/rfft2d.cc             |  12 +-
 tensorflow/lite/kernels/rfft2d_test.cc        |   8 +-
 tensorflow/lite/kernels/round.cc              |   3 +
 tensorflow/lite/kernels/round_test.cc         |   7 +-
 tensorflow/lite/kernels/scatter_nd.cc         |   7 +-
 tensorflow/lite/kernels/scatter_nd_test.cc    |  12 +-
 tensorflow/lite/kernels/segment_sum.cc        |   4 +-
 tensorflow/lite/kernels/segment_sum_test.cc   |   8 +-
 tensorflow/lite/kernels/select.cc             |   6 +-
 tensorflow/lite/kernels/select_test.cc        |  10 +-
 tensorflow/lite/kernels/shape.cc              |   5 +-
 tensorflow/lite/kernels/shape_test.cc         |   7 +-
 tensorflow/lite/kernels/skip_gram.cc          |   2 -
 tensorflow/lite/kernels/skip_gram_test.cc     |   7 +-
 tensorflow/lite/kernels/slice.cc              |  12 +-
 tensorflow/lite/kernels/slice_test.cc         |  11 +-
 tensorflow/lite/kernels/softmax_test.cc       |  11 +-
 tensorflow/lite/kernels/space_to_batch_nd.cc  |   9 +-
 .../lite/kernels/space_to_batch_nd_test.cc    |  10 +-
 tensorflow/lite/kernels/space_to_depth.cc     |   5 +-
 .../lite/kernels/space_to_depth_test.cc       |  10 +-
 tensorflow/lite/kernels/sparse_to_dense.cc    |  12 +-
 .../lite/kernels/sparse_to_dense_test.cc      |  11 +-
 tensorflow/lite/kernels/split.cc              |   7 +-
 tensorflow/lite/kernels/split_test.cc         |  12 +-
 tensorflow/lite/kernels/split_v.cc            |   5 +-
 tensorflow/lite/kernels/split_v_test.cc       |  11 +-
 tensorflow/lite/kernels/squared_difference.cc |   8 +-
 .../lite/kernels/squared_difference_test.cc   |  10 +-
 tensorflow/lite/kernels/squeeze.cc            |   3 -
 tensorflow/lite/kernels/squeeze_test.cc       |  11 +-
 tensorflow/lite/kernels/strided_slice.cc      |   8 +-
 tensorflow/lite/kernels/strided_slice_test.cc |  10 +-
 tensorflow/lite/kernels/sub.cc                |  11 +-
 tensorflow/lite/kernels/sub_test.cc           |  10 +-
 tensorflow/lite/kernels/subgraph_test_util.cc |  13 +-
 tensorflow/lite/kernels/subgraph_test_util.h  |   5 +
 .../lite/kernels/subgraph_test_util_test.cc   |   8 +-
 tensorflow/lite/kernels/svdf.cc               |   1 +
 tensorflow/lite/kernels/svdf_test.cc          |  10 +-
 tensorflow/lite/kernels/test_util.cc          |  23 +-
 tensorflow/lite/kernels/test_util.h           |  24 +-
 tensorflow/lite/kernels/test_util_test.cc     |   6 +
 tensorflow/lite/kernels/tile.cc               |  11 +-
 tensorflow/lite/kernels/tile_test.cc          |  13 +-
 tensorflow/lite/kernels/topk_v2.cc            |  11 +-
 tensorflow/lite/kernels/topk_v2_test.cc       |  12 +-
 tensorflow/lite/kernels/transpose.cc          |   9 +-
 tensorflow/lite/kernels/transpose_conv.cc     |  13 +-
 .../lite/kernels/transpose_conv_test.cc       |  16 +-
 tensorflow/lite/kernels/transpose_test.cc     |  13 +-
 .../kernels/unidirectional_sequence_lstm.cc   |   2 +
 .../unidirectional_sequence_lstm_test.cc      |   6 +-
 .../unidirectional_sequence_rnn_test.cc       |   7 +-
 tensorflow/lite/kernels/unique.cc             |   6 +
 tensorflow/lite/kernels/unique_test.cc        |   9 +-
 tensorflow/lite/kernels/unpack.cc             |   4 +
 tensorflow/lite/kernels/unpack_test.cc        |  12 +-
 tensorflow/lite/kernels/variable_ops_test.cc  |   4 +-
 tensorflow/lite/kernels/where.cc              |   2 +
 tensorflow/lite/kernels/where_test.cc         |   8 +-
 tensorflow/lite/kernels/while.cc              |   3 +
 tensorflow/lite/kernels/while_test.cc         |  10 +-
 tensorflow/lite/kernels/zeros_like.cc         |   4 +
 tensorflow/lite/kernels/zeros_like_test.cc    |   9 +-
 .../lite/micro/kernels/cmsis-nn/pooling.cc    |   1 +
 244 files changed, 1457 insertions(+), 809 deletions(-)

diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
index aad79ffbc89..b16a85c65d8 100644
--- a/tensorflow/lite/kernels/BUILD
+++ b/tensorflow/lite/kernels/BUILD
@@ -115,10 +115,11 @@ cc_test(
     size = "small",
     srcs = ["optional_tensor_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/testing:util",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -158,7 +159,6 @@ cc_test(
     ],
     deps = [
         ":acceleration_test_util_internal",
-        "@com_google_absl//absl/types:optional",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -175,8 +175,11 @@ cc_library(
         "//tensorflow/lite:framework",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite:schema_fbs_version",
+        "//tensorflow/lite:string",
         "//tensorflow/lite:string_util",
+        "//tensorflow/lite:type_to_tflitetype",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/api",
         "//tensorflow/lite/delegates/nnapi:acceleration_test_util",
         "//tensorflow/lite/delegates/nnapi:nnapi_delegate",
         "//tensorflow/lite/kernels/internal:tensor_utils",
@@ -187,6 +190,7 @@ cc_library(
         "//tensorflow/lite/tools/optimize/sparsity:format_converter",
         "//tensorflow/lite/tools/versioning",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -217,6 +221,7 @@ cc_library(
         "//tensorflow/lite:arena_planner",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels/internal:optimized",
+        "//third_party/eigen3",
     ],
 )
 
@@ -226,7 +231,9 @@ cc_test(
     srcs = ["eigen_support_test.cc"],
     deps = [
         ":eigen_support",
+        "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels/internal:optimized",
+        "//third_party/eigen3",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -291,6 +298,7 @@ cc_library(
         # gemmlowp_context_ and ruy_context_ members.
         "@ruy//ruy:context",
         "@gemmlowp",
+        "//tensorflow/lite/c:common",
         "//tensorflow/lite:external_cpu_backend_context",
     ],
 )
@@ -333,17 +341,18 @@ cc_library(
         "cpu_backend_gemm_eigen.cc",
         "cpu_backend_gemm_eigen.h",
         "cpu_backend_gemm_gemmlowp.h",
-        "cpu_backend_gemm_ruy.h",
     ],
     hdrs = [
         "cpu_backend_gemm.h",
         "cpu_backend_gemm_params.h",
+        "cpu_backend_gemm_ruy.h",
     ],
     copts = tflite_copts(),
     deps = [
         ":tflite_with_ruy_only",
         "//tensorflow/lite/kernels/internal:common",
         "//tensorflow/lite/kernels/internal:compatibility",
+        "//tensorflow/lite/kernels/internal:cpu_check",
         "//tensorflow/lite/kernels/internal:types",
         ":cpu_backend_context",
         ":cpu_backend_threadpool",
@@ -369,6 +378,7 @@ cc_test(
         ":cpu_backend_context",
         ":cpu_backend_gemm",
         "@com_google_googletest//:gtest",
+        "@ruy//ruy:matrix",
         # ruy:reference_mul provides the reference implementation
         # that this test compares against.
         "@ruy//ruy:reference_mul",
@@ -382,6 +392,7 @@ cc_library(
     ],
     build_for_embedded = True,
     copts = tflite_copts(),
+    deps = ["//tensorflow/lite/micro:debug_log"],
 )
 
 cc_library(
@@ -408,6 +419,7 @@ cc_test(
     srcs = ["kernel_util_test.cc"],
     deps = [
         ":kernel_util",
+        "//tensorflow/lite/c:common",
         "//tensorflow/lite/testing:util",
         "@com_google_googletest//:gtest",
     ],
@@ -569,7 +581,12 @@ cc_library(
     ],
     copts = tflite_copts() + tf_opts_nortti_if_android() + EXTRA_EIGEN_COPTS,
     visibility = ["//visibility:private"],
-    deps = BUILTIN_KERNEL_DEPS + ["@farmhash_archive//:farmhash"],
+    deps = BUILTIN_KERNEL_DEPS + [
+        "@ruy//ruy/profiler:instrumentation",
+        "//tensorflow/lite/kernels/internal:cppmath",
+        "//tensorflow/lite:string",
+        "@farmhash_archive//:farmhash",
+    ],
 )
 
 # Creates a target where Ruy is unconditionally enabled along with caching
@@ -583,7 +600,12 @@ cc_library(
     ],
     copts = tflite_copts() + tf_opts_nortti_if_android() + EXTRA_EIGEN_COPTS,
     visibility = ["//visibility:private"],
-    deps = BUILTIN_KERNEL_DEPS + ["@farmhash_archive//:farmhash"] + [":tflite_with_ruy_only_and_caching_enabled"],
+    deps = BUILTIN_KERNEL_DEPS + [
+        "@ruy//ruy/profiler:instrumentation",
+        "//tensorflow/lite/kernels/internal:cppmath",
+        "//tensorflow/lite:string",
+        "@farmhash_archive//:farmhash",
+    ] + [":tflite_with_ruy_only_and_caching_enabled"],
 )
 
 cc_library(
@@ -611,7 +633,6 @@ cc_test(
     ],
     deps = [
         ":test_main",
-        ":test_util",
         ":variable_op_kernels",  # buildcleaner: keep
         "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:tensor",
@@ -632,6 +653,7 @@ cc_library(
         "//tensorflow/lite/kernels/hashtable:hashtable_op_kernels",
         "//tensorflow/lite/kernels/internal:kernel_utils",
         "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
         "//third_party/fft2d:fft2d_headers",
         "@fft2d",
         "@ruy//ruy/profiler:instrumentation",
@@ -681,6 +703,7 @@ cc_library(
         ":builtin_op_kernels",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
     ],
     alwayslink = 1,
 )
@@ -697,6 +720,7 @@ cc_library(
         ":builtin_op_kernels",
         "//tensorflow/lite:framework_lib",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
     ],
 )
 
@@ -713,6 +737,7 @@ cc_library(
         ":builtin_op_kernels_ruy_and_caching",
         "//tensorflow/lite:framework_lib",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
     ],
 )
 
@@ -728,6 +753,7 @@ cc_library(
         "//tensorflow/lite:framework",
         "//tensorflow/lite:util",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
     ],
 )
 
@@ -736,10 +762,10 @@ cc_test(
     size = "small",
     srcs = ["audio_spectrogram_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
     ],
@@ -750,10 +776,10 @@ cc_test(
     size = "small",
     srcs = ["mfcc_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
     ],
@@ -764,10 +790,10 @@ cc_test(
     size = "small",
     srcs = ["detection_postprocess_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
     ],
@@ -782,12 +808,15 @@ cc_test(
         "tflite_xnnpack",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -800,11 +829,11 @@ cc_test(
         "tflite_xnnpack",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -813,11 +842,11 @@ cc_test(
     size = "small",
     srcs = ["add_n_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -827,11 +856,11 @@ cc_test(
     srcs = ["arg_min_max_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -841,11 +870,11 @@ cc_test(
     srcs = ["div_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -855,11 +884,11 @@ cc_test(
     srcs = ["sub_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -869,13 +898,15 @@ cc_test(
     srcs = ["transpose_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/kernels/internal:compatibility",
         "//tensorflow/lite/kernels/internal:reference",
         "//tensorflow/lite/kernels/internal:reference_base",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -884,11 +915,11 @@ cc_test(
     size = "small",
     srcs = ["space_to_batch_nd_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -898,11 +929,12 @@ cc_test(
     srcs = ["batch_to_space_nd_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -911,11 +943,11 @@ cc_test(
     size = "small",
     srcs = ["batch_matmul_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -925,11 +957,11 @@ cc_test(
     srcs = ["cast_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -939,11 +971,11 @@ cc_test(
     srcs = ["concatenation_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -953,10 +985,11 @@ cc_test(
     srcs = ["conv_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
     ],
@@ -968,15 +1001,11 @@ cc_test(
     srcs = ["densify_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels/internal:types",
         "//tensorflow/lite/schema:schema_fbs",
-        "//tensorflow/lite/tools/optimize/sparsity:format_converter",
-        "//third_party/eigen3",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
     ],
@@ -987,13 +1016,16 @@ cc_test(
     size = "small",
     srcs = ["depthwise_conv_hybrid_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core/api",
         "//tensorflow/lite/kernels/internal:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1003,13 +1035,16 @@ cc_test(
     srcs = ["depthwise_conv_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core/api",
         "//tensorflow/lite/kernels/internal:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1019,14 +1054,16 @@ cc_test(
     srcs = ["dequantize_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/api",
         "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "//third_party/eigen3",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1036,11 +1073,11 @@ cc_test(
     srcs = ["numeric_verify_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "//third_party/eigen3",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
@@ -1053,11 +1090,11 @@ cc_test(
     srcs = ["basic_rnn_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1067,12 +1104,11 @@ cc_test(
     srcs = ["bidirectional_sequence_lstm_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1082,10 +1118,9 @@ cc_test(
     srcs = ["floor_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1098,10 +1133,9 @@ cc_test(
         "tflite_not_portable_ios",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1114,10 +1148,9 @@ cc_test(
         "tflite_not_portable_ios",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1128,10 +1161,9 @@ cc_test(
     srcs = ["elementwise_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1142,11 +1174,11 @@ cc_test(
     srcs = ["unidirectional_sequence_lstm_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1155,11 +1187,11 @@ cc_test(
     size = "small",
     srcs = ["bidirectional_sequence_rnn_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1169,11 +1201,11 @@ cc_test(
     srcs = ["unidirectional_sequence_rnn_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1184,11 +1216,11 @@ cc_test(
     # TODO(b/143912164): Enable NNAPI test when fix nnapi.
     # tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1198,11 +1230,11 @@ cc_test(
     srcs = ["exp_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1211,11 +1243,11 @@ cc_test(
     size = "small",
     srcs = ["fake_quant_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1225,10 +1257,9 @@ cc_test(
     srcs = ["maximum_minimum_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1239,11 +1270,11 @@ cc_test(
     srcs = ["reduce_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1256,11 +1287,11 @@ cc_test(
         "tflite_xnnpack",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1270,11 +1301,12 @@ cc_test(
     srcs = ["pad_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1286,6 +1318,8 @@ cc_library(
     ],
     deps = [
         ":test_util",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
     ],
 )
 
@@ -1295,11 +1329,10 @@ cc_test(
     srcs = ["reshape_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":reshape_test_common",
         ":test_main",
-        ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1310,12 +1343,13 @@ cc_test(
     srcs = ["gather_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1324,12 +1358,13 @@ cc_test(
     size = "small",
     srcs = ["gather_nd_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1338,12 +1373,12 @@ cc_test(
     size = "small",
     srcs = ["scatter_nd_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1353,11 +1388,10 @@ cc_test(
     srcs = ["topk_v2_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1368,11 +1402,11 @@ cc_test(
     srcs = ["resize_bilinear_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1382,11 +1416,11 @@ cc_test(
     srcs = ["resize_nearest_neighbor_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1396,11 +1430,11 @@ cc_test(
     srcs = ["svdf_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1410,11 +1444,11 @@ cc_test(
     srcs = ["embedding_lookup_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1425,12 +1459,13 @@ cc_test(
     srcs = ["embedding_lookup_sparse_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1444,9 +1479,13 @@ cc_test(
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core/api",
         "//tensorflow/lite/kernels/internal:tensor_utils",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1456,11 +1495,11 @@ cc_test(
     srcs = ["local_response_norm_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1473,11 +1512,11 @@ cc_test(
         "tflite_xnnpack",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1490,12 +1529,13 @@ cc_test(
         "tflite_xnnpack",
     ],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:reference_base",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1505,12 +1545,13 @@ cc_test(
     srcs = ["log_softmax_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:reference_base",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1520,11 +1561,11 @@ cc_test(
     srcs = ["lsh_projection_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1534,12 +1575,13 @@ cc_test(
     srcs = ["hashtable_lookup_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite:string_util",
         "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1550,11 +1592,12 @@ cc_test(
     srcs = ["lstm_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1563,11 +1606,9 @@ cc_test(
     size = "small",
     srcs = ["lstm_eval_test.cc"],
     deps = [
-        ":builtin_ops",
+        ":cpu_backend_context",
         ":lstm_eval",
         ":test_main",
-        ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
         "@com_google_googletest//:gtest",
     ],
@@ -1578,12 +1619,14 @@ cc_test(
     size = "small",
     srcs = ["skip_gram_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite:string_util",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1593,11 +1636,11 @@ cc_test(
     srcs = ["space_to_depth_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1607,11 +1650,11 @@ cc_test(
     srcs = ["depth_to_space_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1621,11 +1664,11 @@ cc_test(
     srcs = ["split_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1634,11 +1677,11 @@ cc_test(
     size = "small",
     srcs = ["split_v_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1648,11 +1691,11 @@ cc_test(
     srcs = ["squeeze_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1662,10 +1705,9 @@ cc_test(
     srcs = ["strided_slice_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1676,11 +1718,11 @@ cc_test(
     srcs = ["tile_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1693,11 +1735,12 @@ cc_test(
     ],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1707,11 +1750,11 @@ cc_test(
     srcs = ["neg_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1723,11 +1766,11 @@ cc_test(
     ],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1739,10 +1782,10 @@ cc_test(
     ],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1753,10 +1796,11 @@ cc_test(
     srcs = ["transpose_conv_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
     ],
@@ -1768,11 +1812,11 @@ cc_test(
     srcs = ["expand_dims_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1782,11 +1826,10 @@ cc_test(
     size = "small",
     srcs = ["sparse_to_dense_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1796,11 +1839,11 @@ cc_test(
     size = "small",
     srcs = ["shape_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1810,12 +1853,13 @@ cc_test(
     size = "small",
     srcs = ["rank_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1825,12 +1869,11 @@ cc_test(
     srcs = ["pow_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels/internal:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1840,11 +1883,10 @@ cc_test(
     size = "small",
     srcs = ["pack_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1854,10 +1896,10 @@ cc_test(
     size = "small",
     srcs = ["one_hot_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1868,12 +1910,12 @@ cc_test(
     srcs = ["logical_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1882,11 +1924,9 @@ cc_test(
     size = "small",
     srcs = ["unpack_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1896,11 +1936,9 @@ cc_test(
     size = "small",
     srcs = ["floor_div_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1910,12 +1948,11 @@ cc_test(
     size = "small",
     srcs = ["where_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1924,12 +1961,12 @@ cc_test(
     size = "small",
     srcs = ["zeros_like_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1938,11 +1975,9 @@ cc_test(
     size = "small",
     srcs = ["floor_mod_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1952,11 +1987,9 @@ cc_test(
     size = "small",
     srcs = ["range_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1966,12 +1999,11 @@ cc_test(
     size = "small",
     srcs = ["squared_difference_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:builtin_op_data",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -1981,15 +2013,11 @@ cc_test(
     srcs = ["if_test.cc"],
     tags = ["tflite_not_portable_ios"],
     deps = [
-        ":builtin_ops",
         ":kernel_util",
         ":subgraph_test_util",
         ":test_main",
-        ":test_util",
-        "//tensorflow/lite:builtin_op_data",
         "//tensorflow/lite:framework",
         "@com_google_googletest//:gtest",
-        "@flatbuffers",
     ],
 )
 
@@ -1999,15 +2027,10 @@ cc_test(
     srcs = ["while_test.cc"],
     tags = ["tflite_not_portable_ios"],
     deps = [
-        ":builtin_ops",
-        ":kernel_util",
         ":subgraph_test_util",
         ":test_main",
-        ":test_util",
-        "//tensorflow/lite:builtin_op_data",
         "//tensorflow/lite:framework",
         "@com_google_googletest//:gtest",
-        "@flatbuffers",
     ],
 )
 
@@ -2017,10 +2040,10 @@ cc_test(
     srcs = ["fill_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2029,10 +2052,9 @@ cc_test(
     name = "unique_test",
     srcs = ["unique_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2042,10 +2064,9 @@ cc_test(
     size = "small",
     srcs = ["reverse_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2059,6 +2080,8 @@ cc_test(
         ":test_util",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/testing:util",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2068,11 +2091,11 @@ cc_test(
     size = "small",
     srcs = ["non_max_suppression_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -2092,10 +2115,9 @@ cc_test(
     name = "mirror_pad_test",
     srcs = ["mirror_pad_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2111,6 +2133,7 @@ cc_library(
         ":test_util",
         "//tensorflow/lite:builtin_op_data",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:common",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
     ],
@@ -2123,8 +2146,8 @@ cc_test(
     deps = [
         ":kernel_util",
         ":subgraph_test_util",
-        ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/testing:util",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2134,10 +2157,9 @@ cc_test(
     size = "small",
     srcs = ["reverse_sequence_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2147,10 +2169,10 @@ cc_test(
     size = "small",
     srcs = ["matrix_diag_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2161,12 +2183,12 @@ cc_test(
     srcs = ["quantize_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -2175,10 +2197,10 @@ cc_test(
     size = "small",
     srcs = ["matrix_set_diag_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2189,12 +2211,11 @@ cc_test(
     srcs = ["quant_basic_lstm_test.cc"],
     tags = ["tflite_nnapi"],
     deps = [
-        ":builtin_ops",
-        ":kernel_util",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
+        "@flatbuffers",
     ],
 )
 
@@ -2202,10 +2223,9 @@ cc_test(
     name = "segment_sum_test",
     srcs = ["segment_sum_test.cc"],
     deps = [
-        ":builtin_ops",
         ":test_main",
         ":test_util",
-        "//tensorflow/lite:framework",
+        "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
     ],
 )
diff --git a/tensorflow/lite/kernels/acceleration_test_util.cc b/tensorflow/lite/kernels/acceleration_test_util.cc
index 0dffd22fa26..741c34d9672 100644
--- a/tensorflow/lite/kernels/acceleration_test_util.cc
+++ b/tensorflow/lite/kernels/acceleration_test_util.cc
@@ -14,19 +14,11 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/acceleration_test_util.h"
 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
 #include 
 #include 
-#include 
 
+#include 
 #include "absl/types/optional.h"
-#include "tensorflow/lite/kernels/acceleration_test_util_internal.h"
-#include "tensorflow/lite/minimal_logging.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/acceleration_test_util.h b/tensorflow/lite/kernels/acceleration_test_util.h
index 75b5d79c8c4..a6a88d5f131 100644
--- a/tensorflow/lite/kernels/acceleration_test_util.h
+++ b/tensorflow/lite/kernels/acceleration_test_util.h
@@ -15,7 +15,7 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_H_
 #define TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_H_
 
-#include 
+#include 
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/acceleration_test_util_internal.cc b/tensorflow/lite/kernels/acceleration_test_util_internal.cc
index f4a1f5cdc87..a6ad8234f59 100644
--- a/tensorflow/lite/kernels/acceleration_test_util_internal.cc
+++ b/tensorflow/lite/kernels/acceleration_test_util_internal.cc
@@ -14,6 +14,14 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/acceleration_test_util_internal.h"
 
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+
 namespace tflite {
 
 void ReadAccelerationConfig(
diff --git a/tensorflow/lite/kernels/acceleration_test_util_internal.h b/tensorflow/lite/kernels/acceleration_test_util_internal.h
index 8999af7e7ad..24fc2383f9e 100644
--- a/tensorflow/lite/kernels/acceleration_test_util_internal.h
+++ b/tensorflow/lite/kernels/acceleration_test_util_internal.h
@@ -18,14 +18,12 @@ limitations under the License.
 #include 
 #include 
 #include 
-#include 
+#include 
 #include 
-#include 
 #include 
 
 #include "absl/types/optional.h"
 #include "re2/re2.h"
-#include "tensorflow/lite/minimal_logging.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/acceleration_test_util_internal_test.cc b/tensorflow/lite/kernels/acceleration_test_util_internal_test.cc
index 71e0c9e9912..82d21fd9332 100644
--- a/tensorflow/lite/kernels/acceleration_test_util_internal_test.cc
+++ b/tensorflow/lite/kernels/acceleration_test_util_internal_test.cc
@@ -14,7 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/acceleration_test_util_internal.h"
 
+#include 
 #include 
+#include 
 #include 
 
 #include 
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 749a0d69ef9..2b2428f3f92 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -12,25 +12,35 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
 #include 
 #include 
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/log_softmax.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/logistic.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/tanh.h"
 #include "tensorflow/lite/kernels/internal/reference/logistic.h"
+#include "tensorflow/lite/kernels/internal/reference/prelu.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/softmax.h"
+#include "tensorflow/lite/kernels/internal/reference/tanh.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 #if __aarch64__ && __clang__
 #include 
diff --git a/tensorflow/lite/kernels/activations_test.cc b/tensorflow/lite/kernels/activations_test.cc
index 5a679147469..50b1c041e34 100644
--- a/tensorflow/lite/kernels/activations_test.cc
+++ b/tensorflow/lite/kernels/activations_test.cc
@@ -12,16 +12,28 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+#include 
+#include 
+
+#include 
+#include 
 #include 
+#include 
+#include 
 #include 
+#include 
+#include 
+#include 
 
 #include 
 #include "absl/memory/memory.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/add.cc b/tensorflow/lite/kernels/add.cc
index d9b8c87eeb7..279f6aa12ce 100644
--- a/tensorflow/lite/kernels/add.cc
+++ b/tensorflow/lite/kernels/add.cc
@@ -14,16 +14,26 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/add.h"
 
+#include 
+#include 
+
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/add.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/add_n.cc b/tensorflow/lite/kernels/add_n.cc
index 5f6437fe331..7b4d52c5272 100644
--- a/tensorflow/lite/kernels/add_n.cc
+++ b/tensorflow/lite/kernels/add_n.cc
@@ -12,6 +12,8 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
diff --git a/tensorflow/lite/kernels/add_n_test.cc b/tensorflow/lite/kernels/add_n_test.cc
index ac6ccec2b66..4db646ace9c 100644
--- a/tensorflow/lite/kernels/add_n_test.cc
+++ b/tensorflow/lite/kernels/add_n_test.cc
@@ -12,13 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
 #include 
 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/add_test.cc b/tensorflow/lite/kernels/add_test.cc
index 267b80564c9..bb883dd9b05 100644
--- a/tensorflow/lite/kernels/add_test.cc
+++ b/tensorflow/lite/kernels/add_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/arg_min_max.cc b/tensorflow/lite/kernels/arg_min_max.cc
index e99f59ba703..4a3902ac57c 100644
--- a/tensorflow/lite/kernels/arg_min_max.cc
+++ b/tensorflow/lite/kernels/arg_min_max.cc
@@ -12,13 +12,19 @@ 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.
 ==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/arg_min_max.h"
+
+#include 
+
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/arg_min_max_test.cc b/tensorflow/lite/kernels/arg_min_max_test.cc
index d028bac3fb2..957d3473b8d 100644
--- a/tensorflow/lite/kernels/arg_min_max_test.cc
+++ b/tensorflow/lite/kernels/arg_min_max_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/assign_variable.cc b/tensorflow/lite/kernels/assign_variable.cc
index 41ddcdda6f7..4cb4e08e43a 100644
--- a/tensorflow/lite/kernels/assign_variable.cc
+++ b/tensorflow/lite/kernels/assign_variable.cc
@@ -13,17 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/experimental/resource/resource_variable.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/audio_spectrogram.cc b/tensorflow/lite/kernels/audio_spectrogram.cc
index 99457ea11b1..29c9eeef3d0 100644
--- a/tensorflow/lite/kernels/audio_spectrogram.cc
+++ b/tensorflow/lite/kernels/audio_spectrogram.cc
@@ -13,15 +13,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+#include 
+
+#include 
+
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/spectrogram.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/audio_spectrogram_test.cc b/tensorflow/lite/kernels/audio_spectrogram_test.cc
index 0f4182ea728..cdb77303748 100644
--- a/tensorflow/lite/kernels/audio_spectrogram_test.cc
+++ b/tensorflow/lite/kernels/audio_spectrogram_test.cc
@@ -13,16 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
-#include 
 #include 
 
 #include 
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/basic_rnn_test.cc b/tensorflow/lite/kernels/basic_rnn_test.cc
index f7cbaa5a814..2146d086c9a 100644
--- a/tensorflow/lite/kernels/basic_rnn_test.cc
+++ b/tensorflow/lite/kernels/basic_rnn_test.cc
@@ -14,17 +14,14 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite RNN op.
 
-#include 
 #include 
-#include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/batch_matmul.cc b/tensorflow/lite/kernels/batch_matmul.cc
index 9cbad101bab..d2115f96e1c 100644
--- a/tensorflow/lite/kernels/batch_matmul.cc
+++ b/tensorflow/lite/kernels/batch_matmul.cc
@@ -15,15 +15,21 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/internal/reference/batch_matmul.h"
 
+#include 
+
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/batch_matmul.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
diff --git a/tensorflow/lite/kernels/batch_matmul_test.cc b/tensorflow/lite/kernels/batch_matmul_test.cc
index aec031015c0..5e52479f49b 100644
--- a/tensorflow/lite/kernels/batch_matmul_test.cc
+++ b/tensorflow/lite/kernels/batch_matmul_test.cc
@@ -12,12 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/batch_to_space_nd.cc b/tensorflow/lite/kernels/batch_to_space_nd.cc
index d7d796ebec1..9d6492e0fcb 100644
--- a/tensorflow/lite/kernels/batch_to_space_nd.cc
+++ b/tensorflow/lite/kernels/batch_to_space_nd.cc
@@ -12,17 +12,15 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/batch_to_space_nd_test.cc b/tensorflow/lite/kernels/batch_to_space_nd_test.cc
index cffa1036c84..e675faafd74 100644
--- a/tensorflow/lite/kernels/batch_to_space_nd_test.cc
+++ b/tensorflow/lite/kernels/batch_to_space_nd_test.cc
@@ -13,11 +13,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
+#include 
+#include 
+
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
index 8ccc7a68eb7..a984ff5124f 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
@@ -13,6 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc b/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
index c468c4c09fb..3a52de130e3 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
@@ -14,17 +14,12 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite Bidirectional LSTM op.
 
-#include 
-#include 
-#include 
+#include 
 #include 
 
-#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc b/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
index 58a2ef9c1ea..abaf6df9fa8 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include 
 #include 
 #include 
 
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_rnn_test.cc b/tensorflow/lite/kernels/bidirectional_sequence_rnn_test.cc
index 4a7cc9a016d..870b99d7437 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_rnn_test.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_rnn_test.cc
@@ -14,15 +14,18 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite Bidirectional RNN op.
 
-#include 
+#include 
+#include 
+#include 
+#include 
+#include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/cast.cc b/tensorflow/lite/kernels/cast.cc
index a24dadb5279..415f1270328 100644
--- a/tensorflow/lite/kernels/cast.cc
+++ b/tensorflow/lite/kernels/cast.cc
@@ -12,18 +12,14 @@ 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.
 ==============================================================================*/
-#include 
-
 #include 
 #include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
-#include "tensorflow/lite/string_util.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/cast_test.cc b/tensorflow/lite/kernels/cast_test.cc
index 8f1cb44f1c9..a615edbd085 100644
--- a/tensorflow/lite/kernels/cast_test.cc
+++ b/tensorflow/lite/kernels/cast_test.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/ceil.cc b/tensorflow/lite/kernels/ceil.cc
index 3b1df4f6c2f..9914dbe09ce 100644
--- a/tensorflow/lite/kernels/ceil.cc
+++ b/tensorflow/lite/kernels/ceil.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/ceil_test.cc b/tensorflow/lite/kernels/ceil_test.cc
index 36486087fcf..bb23a5b6197 100644
--- a/tensorflow/lite/kernels/ceil_test.cc
+++ b/tensorflow/lite/kernels/ceil_test.cc
@@ -13,11 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/comparisons.cc b/tensorflow/lite/kernels/comparisons.cc
index 91dbc447c35..7d1c6b7804e 100644
--- a/tensorflow/lite/kernels/comparisons.cc
+++ b/tensorflow/lite/kernels/comparisons.cc
@@ -12,10 +12,19 @@ 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.
 ==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/comparisons.h"
+
+#include 
+
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/string_util.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/comparisons_test.cc b/tensorflow/lite/kernels/comparisons_test.cc
index 986600ccd1a..f8cf6dee74c 100644
--- a/tensorflow/lite/kernels/comparisons_test.cc
+++ b/tensorflow/lite/kernels/comparisons_test.cc
@@ -12,11 +12,18 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/concatenation.cc b/tensorflow/lite/kernels/concatenation.cc
index 8beb962b1b1..61748e5ce58 100644
--- a/tensorflow/lite/kernels/concatenation.cc
+++ b/tensorflow/lite/kernels/concatenation.cc
@@ -12,20 +12,19 @@ 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.
 ==============================================================================*/
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
+#include "tensorflow/lite/kernels/internal/reference/concatenation.h"
+
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/concatenation_test.cc b/tensorflow/lite/kernels/concatenation_test.cc
index 8f4abe0bcda..4e362598aae 100644
--- a/tensorflow/lite/kernels/concatenation_test.cc
+++ b/tensorflow/lite/kernels/concatenation_test.cc
@@ -12,13 +12,18 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
+#include 
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc
index 1d610b2e068..fa6caff5baa 100644
--- a/tensorflow/lite/kernels/conv.cc
+++ b/tensorflow/lite/kernels/conv.cc
@@ -14,19 +14,17 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h"
 
-#include 
-#include 
-#include 
+#include 
+
 #include 
-#include 
-#include 
-#include 
-#include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/eigen_support.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 // b/131835803 forces us to include multithreaded_conv.h before optimized_ops.h
 #ifndef TFLITE_WITH_RUY_ONLY
 #include "tensorflow/lite/kernels/internal/optimized/multithreaded_conv.h"
@@ -39,7 +37,6 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/kernels/padding.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/conv_test.cc b/tensorflow/lite/kernels/conv_test.cc
index ef1d5366255..a1fd34eb1cb 100644
--- a/tensorflow/lite/kernels/conv_test.cc
+++ b/tensorflow/lite/kernels/conv_test.cc
@@ -12,15 +12,22 @@ 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.
 ==============================================================================*/
-#include 
-#include 
+#include 
+#include 
 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include 
 #include 
 #include "absl/memory/memory.h"
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/cpu_backend_context.cc b/tensorflow/lite/kernels/cpu_backend_context.cc
index 7a16bed0ead..a99d08769ea 100644
--- a/tensorflow/lite/kernels/cpu_backend_context.cc
+++ b/tensorflow/lite/kernels/cpu_backend_context.cc
@@ -15,8 +15,12 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 
+#include 
+
 #include "public/gemmlowp.h"
 #include "ruy/context.h"  // from @ruy
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/external_cpu_backend_context.h"
 #include "tensorflow/lite/kernels/op_macros.h"
 
 namespace {
diff --git a/tensorflow/lite/kernels/cpu_backend_context.h b/tensorflow/lite/kernels/cpu_backend_context.h
index b4973feb56f..19ef88bf8e3 100644
--- a/tensorflow/lite/kernels/cpu_backend_context.h
+++ b/tensorflow/lite/kernels/cpu_backend_context.h
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "public/gemmlowp.h"
 #include "ruy/context.h"  // from @ruy
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/external_cpu_backend_context.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/cpu_backend_gemm.h b/tensorflow/lite/kernels/cpu_backend_gemm.h
index 8e324c8b515..f4d20d8970a 100644
--- a/tensorflow/lite/kernels/cpu_backend_gemm.h
+++ b/tensorflow/lite/kernels/cpu_backend_gemm.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include 
 
+#include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
diff --git a/tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h b/tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h
index 1c3c0ca39c4..224f8ecea41 100644
--- a/tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h
+++ b/tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h
@@ -32,6 +32,9 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_CUSTOM_GEMV_H_
 #define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_CUSTOM_GEMV_H_
 
+#include 
+
+#include 
 #include 
 #include 
 
@@ -40,6 +43,8 @@ limitations under the License.
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
 #include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 
 namespace tflite {
 namespace cpu_backend_gemm {
diff --git a/tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h b/tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h
index 1e1074523ab..77d37aac291 100644
--- a/tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h
+++ b/tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h
@@ -16,13 +16,15 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_GEMMLOWP_H_
 #define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_GEMMLOWP_H_
 
+#include 
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #ifndef TFLITE_WITH_RUY_ONLY
 
 #include 
 #include 
 
 #include "public/gemmlowp.h"
-#include "ruy/ruy.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
diff --git a/tensorflow/lite/kernels/cpu_backend_gemm_ruy.h b/tensorflow/lite/kernels/cpu_backend_gemm_ruy.h
index 8aaedb6a13a..07ae2ff08b7 100644
--- a/tensorflow/lite/kernels/cpu_backend_gemm_ruy.h
+++ b/tensorflow/lite/kernels/cpu_backend_gemm_ruy.h
@@ -17,10 +17,10 @@ limitations under the License.
 #define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_RUY_H_
 
 #include "ruy/matrix.h"  // from @ruy
-#include "ruy/path.h"  // from @ruy
 #include "ruy/ruy.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 
 namespace tflite {
 namespace cpu_backend_gemm {
diff --git a/tensorflow/lite/kernels/cpu_backend_gemm_test.cc b/tensorflow/lite/kernels/cpu_backend_gemm_test.cc
index 20334947dde..d79d1357696 100644
--- a/tensorflow/lite/kernels/cpu_backend_gemm_test.cc
+++ b/tensorflow/lite/kernels/cpu_backend_gemm_test.cc
@@ -15,19 +15,26 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/cpu_backend_gemm.h"
 
+#include 
+#include 
+#include 
+
 #include 
-#include 
+#include 
 #include 
 #include 
 #include 
 #include 
 #include 
 #include 
+#include 
 
 #include 
+#include "ruy/matrix.h"  // from @ruy
 #include "ruy/reference_mul.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/cpu_backend_threadpool_test.cc b/tensorflow/lite/kernels/cpu_backend_threadpool_test.cc
index 5089323070a..fafe4c40067 100644
--- a/tensorflow/lite/kernels/cpu_backend_threadpool_test.cc
+++ b/tensorflow/lite/kernels/cpu_backend_threadpool_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
 
+#include 
+
 #include 
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 
diff --git a/tensorflow/lite/kernels/custom_ops_register.h b/tensorflow/lite/kernels/custom_ops_register.h
index ca9fac81889..3abc893243b 100644
--- a/tensorflow/lite/kernels/custom_ops_register.h
+++ b/tensorflow/lite/kernels/custom_ops_register.h
@@ -15,7 +15,7 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_CUSTOM_OPS_REGISTER_H_
 #define TENSORFLOW_LITE_KERNELS_CUSTOM_OPS_REGISTER_H_
 
-#include "tensorflow/lite/context.h"
+#include "tensorflow/lite/c/common.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/densify.cc b/tensorflow/lite/kernels/densify.cc
index 0c2742d8696..cc3ac67464d 100644
--- a/tensorflow/lite/kernels/densify.cc
+++ b/tensorflow/lite/kernels/densify.cc
@@ -14,16 +14,13 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/reference/densify.h"
 
-#include 
+#include 
 
 #include 
-#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/densify_test.cc b/tensorflow/lite/kernels/densify_test.cc
index 5cb90932069..d453606cf2e 100644
--- a/tensorflow/lite/kernels/densify_test.cc
+++ b/tensorflow/lite/kernels/densify_test.cc
@@ -14,18 +14,16 @@ limitations under the License.
 ==============================================================================*/
 #include 
 #include 
+#include 
+#include 
 
+#include 
 #include 
 #include "absl/memory/memory.h"
-#include "third_party/eigen3/Eigen/Core"
 #include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/types.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 #include "tensorflow/lite/schema/schema_generated.h"
-#include "tensorflow/lite/tools/optimize/sparsity/format_converter.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc
index d6fe8c7ab1c..8a81ea932bf 100644
--- a/tensorflow/lite/kernels/depth_to_space.cc
+++ b/tensorflow/lite/kernels/depth_to_space.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/depth_to_space_test.cc b/tensorflow/lite/kernels/depth_to_space_test.cc
index 8d59a1ad82f..4429faf9909 100644
--- a/tensorflow/lite/kernels/depth_to_space_test.cc
+++ b/tensorflow/lite/kernels/depth_to_space_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/depthwise_conv.cc b/tensorflow/lite/kernels/depthwise_conv.cc
index 8500b5cd39b..1897d14a065 100644
--- a/tensorflow/lite/kernels/depthwise_conv.cc
+++ b/tensorflow/lite/kernels/depthwise_conv.cc
@@ -15,27 +15,28 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h"
 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
+#include 
+#include 
+
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_multithread.h"
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/depthwiseconv_float.h"
 #include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/depthwise_conv.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/kernels/padding.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/depthwise_conv_hybrid_test.cc b/tensorflow/lite/kernels/depthwise_conv_hybrid_test.cc
index 09cd7cec4d7..c5158eac3d0 100644
--- a/tensorflow/lite/kernels/depthwise_conv_hybrid_test.cc
+++ b/tensorflow/lite/kernels/depthwise_conv_hybrid_test.cc
@@ -12,18 +12,25 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+
 #include 
 #include 
-#include 
+#include 
+#include 
+#include 
+#include 
 
+#include 
 #include 
 #include "absl/memory/memory.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/test_util.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/depthwise_conv_test.cc b/tensorflow/lite/kernels/depthwise_conv_test.cc
index 5d85eac4aa9..f410476d983 100644
--- a/tensorflow/lite/kernels/depthwise_conv_test.cc
+++ b/tensorflow/lite/kernels/depthwise_conv_test.cc
@@ -12,17 +12,25 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+
 #include 
 #include 
+#include 
+#include 
+#include 
+#include 
 
+#include 
 #include 
 #include "absl/memory/memory.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/test_util.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/dequantize.cc b/tensorflow/lite/kernels/dequantize.cc
index 272662d9c48..a2a1bd495cf 100644
--- a/tensorflow/lite/kernels/dequantize.cc
+++ b/tensorflow/lite/kernels/dequantize.cc
@@ -14,15 +14,11 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/dequantize.h"
 
-#include 
+#include 
 
-#include 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/dequantize.h b/tensorflow/lite/kernels/dequantize.h
index 3d9e7ccb135..30739eb2c57 100644
--- a/tensorflow/lite/kernels/dequantize.h
+++ b/tensorflow/lite/kernels/dequantize.h
@@ -15,16 +15,17 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_DEQUANTIZE_H_
 #define TENSORFLOW_LITE_KERNELS_DEQUANTIZE_H_
 
+#include 
+
 #include "third_party/eigen3/Eigen/Core"
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/dequantize.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
-#include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/dequantize_test.cc b/tensorflow/lite/kernels/dequantize_test.cc
index f55a23e138d..da795474400 100644
--- a/tensorflow/lite/kernels/dequantize_test.cc
+++ b/tensorflow/lite/kernels/dequantize_test.cc
@@ -13,15 +13,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include 
+#include 
+#include 
+#include 
 
+#include 
 #include 
 #include "absl/memory/memory.h"
 #include "third_party/eigen3/Eigen/Core"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/types.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/detection_postprocess.cc b/tensorflow/lite/kernels/detection_postprocess.cc
index 5d848bc9eab..c0b5b2ddf7c 100644
--- a/tensorflow/lite/kernels/detection_postprocess.cc
+++ b/tensorflow/lite/kernels/detection_postprocess.cc
@@ -12,19 +12,23 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+#include 
+#include 
 
+#include 
+#include 
 #include 
 #include 
 
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/detection_postprocess_test.cc b/tensorflow/lite/kernels/detection_postprocess_test.cc
index 348ea45a515..cf0d3ba2f3d 100644
--- a/tensorflow/lite/kernels/detection_postprocess_test.cc
+++ b/tensorflow/lite/kernels/detection_postprocess_test.cc
@@ -12,16 +12,17 @@ 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.
 ==============================================================================*/
-#include 
-#include 
+#include 
+
+#include 
 #include 
 
+#include 
 #include 
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/div.cc b/tensorflow/lite/kernels/div.cc
index 731fb3c2fe2..cdd02277ec9 100644
--- a/tensorflow/lite/kernels/div.cc
+++ b/tensorflow/lite/kernels/div.cc
@@ -12,15 +12,22 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/div_test.cc b/tensorflow/lite/kernels/div_test.cc
index e72565f84a0..b1a691d2452 100644
--- a/tensorflow/lite/kernels/div_test.cc
+++ b/tensorflow/lite/kernels/div_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/eigen_support.cc b/tensorflow/lite/kernels/eigen_support.cc
index 3f71b0a20fc..c911222f26e 100644
--- a/tensorflow/lite/kernels/eigen_support.cc
+++ b/tensorflow/lite/kernels/eigen_support.cc
@@ -14,9 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/eigen_support.h"
 
+#include 
+#include 
 #include 
 
 #include "tensorflow/lite/arena_planner.h"
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h"
 #include "tensorflow/lite/kernels/op_macros.h"
 
diff --git a/tensorflow/lite/kernels/eigen_support_test.cc b/tensorflow/lite/kernels/eigen_support_test.cc
index a8c8dc0a5e4..08b58446c16 100644
--- a/tensorflow/lite/kernels/eigen_support_test.cc
+++ b/tensorflow/lite/kernels/eigen_support_test.cc
@@ -15,9 +15,10 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/eigen_support.h"
 
-#include 
+#include 
 
 #include 
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 78ded8f932c..95b791be3f2 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -13,11 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
 #include 
 
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/elementwise_test.cc b/tensorflow/lite/kernels/elementwise_test.cc
index ec00d3e071d..9495be0e590 100644
--- a/tensorflow/lite/kernels/elementwise_test.cc
+++ b/tensorflow/lite/kernels/elementwise_test.cc
@@ -13,11 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/embedding_lookup.cc b/tensorflow/lite/kernels/embedding_lookup.cc
index 0334c0daf84..36e0737c7e2 100644
--- a/tensorflow/lite/kernels/embedding_lookup.cc
+++ b/tensorflow/lite/kernels/embedding_lookup.cc
@@ -29,19 +29,13 @@ limitations under the License.
 //   When indices are out of bound, the ops will not succeed.
 //
 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
+#include 
+
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/embedding_lookup_sparse.cc b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
index 8ecf427d4c1..92574817e3b 100644
--- a/tensorflow/lite/kernels/embedding_lookup_sparse.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
@@ -62,6 +62,8 @@ limitations under the License.
 //
 //   When indices are out of bound, the op will not succeed.
 
+#include 
+
 #include 
 #include 
 
@@ -70,7 +72,6 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/embedding_lookup_sparse_test.cc b/tensorflow/lite/kernels/embedding_lookup_sparse_test.cc
index 1714ff52e70..d9b2d523297 100644
--- a/tensorflow/lite/kernels/embedding_lookup_sparse_test.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_sparse_test.cc
@@ -15,15 +15,18 @@ limitations under the License.
 // Unit test for TFLite sparse lookup op.
 
 #include 
+#include 
+#include 
+#include 
 #include 
 
 #include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/embedding_lookup_test.cc b/tensorflow/lite/kernels/embedding_lookup_test.cc
index 58e77afeeef..cbca3af0e8e 100644
--- a/tensorflow/lite/kernels/embedding_lookup_test.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_test.cc
@@ -14,17 +14,20 @@ License.
 ==============================================================================*/
 // Unit test for TFLite Lookup op.
 
+#include 
+
+#include 
 #include 
-#include 
+#include 
 #include 
 
 #include 
 #include 
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/exp.cc b/tensorflow/lite/kernels/exp.cc
index b53d2f9f56a..764ae6ce7a3 100644
--- a/tensorflow/lite/kernels/exp.cc
+++ b/tensorflow/lite/kernels/exp.cc
@@ -12,16 +12,11 @@ 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.
 ==============================================================================*/
-#include 
-
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/exp_test.cc b/tensorflow/lite/kernels/exp_test.cc
index b6f73169c4b..97cb591d353 100644
--- a/tensorflow/lite/kernels/exp_test.cc
+++ b/tensorflow/lite/kernels/exp_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/expand_dims.cc b/tensorflow/lite/kernels/expand_dims.cc
index 5c7bd167425..721ab3d510a 100644
--- a/tensorflow/lite/kernels/expand_dims.cc
+++ b/tensorflow/lite/kernels/expand_dims.cc
@@ -1,4 +1,5 @@
 
+#include 
 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,14 +16,11 @@ limitations under the License.
 ==============================================================================*/
 #include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/expand_dims_test.cc b/tensorflow/lite/kernels/expand_dims_test.cc
index 5bb1d76f00f..6d231d4cb63 100644
--- a/tensorflow/lite/kernels/expand_dims_test.cc
+++ b/tensorflow/lite/kernels/expand_dims_test.cc
@@ -1,4 +1,3 @@
-
 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,12 +12,17 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/fake_quant.cc b/tensorflow/lite/kernels/fake_quant.cc
index 79ce8416a0a..e8ea090c7b3 100644
--- a/tensorflow/lite/kernels/fake_quant.cc
+++ b/tensorflow/lite/kernels/fake_quant.cc
@@ -12,16 +12,13 @@ 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.
 ==============================================================================*/
-#include 
-
-#include 
-
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/fake_quant_test.cc b/tensorflow/lite/kernels/fake_quant_test.cc
index 3a2dc258fe2..94c03ff1fbd 100644
--- a/tensorflow/lite/kernels/fake_quant_test.cc
+++ b/tensorflow/lite/kernels/fake_quant_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/fill.cc b/tensorflow/lite/kernels/fill.cc
index 19ff1de4939..68ec3e9eca3 100644
--- a/tensorflow/lite/kernels/fill.cc
+++ b/tensorflow/lite/kernels/fill.cc
@@ -13,12 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/string_util.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/fill_test.cc b/tensorflow/lite/kernels/fill_test.cc
index 0717a31b9d7..4fc753ba36a 100644
--- a/tensorflow/lite/kernels/fill_test.cc
+++ b/tensorflow/lite/kernels/fill_test.cc
@@ -12,11 +12,17 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/floor.cc b/tensorflow/lite/kernels/floor.cc
index a01a3199ffc..2e341218700 100644
--- a/tensorflow/lite/kernels/floor.cc
+++ b/tensorflow/lite/kernels/floor.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/floor_div.cc b/tensorflow/lite/kernels/floor_div.cc
index 05b6d9cfeae..5677dc4d9b7 100644
--- a/tensorflow/lite/kernels/floor_div.cc
+++ b/tensorflow/lite/kernels/floor_div.cc
@@ -12,11 +12,18 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+#include 
+
+#include 
+
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/floor_div_test.cc b/tensorflow/lite/kernels/floor_div_test.cc
index b7dcb3babe2..d219f8913ce 100644
--- a/tensorflow/lite/kernels/floor_div_test.cc
+++ b/tensorflow/lite/kernels/floor_div_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/floor_mod.cc b/tensorflow/lite/kernels/floor_mod.cc
index a91ab7f07ab..a4bc9fa9841 100644
--- a/tensorflow/lite/kernels/floor_mod.cc
+++ b/tensorflow/lite/kernels/floor_mod.cc
@@ -12,14 +12,15 @@ 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.
 ==============================================================================*/
-#include 
-#include 
+#include 
+#include 
 
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 // TODO(b/117523611): We should factor out a binary_op and put binary ops there.
 namespace tflite {
diff --git a/tensorflow/lite/kernels/floor_mod_test.cc b/tensorflow/lite/kernels/floor_mod_test.cc
index f96988a855a..33b3834c972 100644
--- a/tensorflow/lite/kernels/floor_mod_test.cc
+++ b/tensorflow/lite/kernels/floor_mod_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/floor_test.cc b/tensorflow/lite/kernels/floor_test.cc
index e66158ba7ba..8602786f439 100644
--- a/tensorflow/lite/kernels/floor_test.cc
+++ b/tensorflow/lite/kernels/floor_test.cc
@@ -13,11 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/fully_connected.cc b/tensorflow/lite/kernels/fully_connected.cc
index 37ccf8459bb..a1893878232 100644
--- a/tensorflow/lite/kernels/fully_connected.cc
+++ b/tensorflow/lite/kernels/fully_connected.cc
@@ -22,7 +22,6 @@ limitations under the License.
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
-#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/optimized/sparse_ops/fully_connected.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
diff --git a/tensorflow/lite/kernels/fully_connected_test.cc b/tensorflow/lite/kernels/fully_connected_test.cc
index 364d54b0c8e..7f02ed079bd 100644
--- a/tensorflow/lite/kernels/fully_connected_test.cc
+++ b/tensorflow/lite/kernels/fully_connected_test.cc
@@ -16,19 +16,28 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/fully_connected.h"
 
+#include 
+#include 
+
+#include 
 #include 
-#include 
+#include 
+#include 
+#include 
 #include 
+#include 
 #include 
 
 #include 
 #include 
 #include "absl/memory/memory.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/gather.cc b/tensorflow/lite/kernels/gather.cc
index b1485397291..1de49f7c486 100644
--- a/tensorflow/lite/kernels/gather.cc
+++ b/tensorflow/lite/kernels/gather.cc
@@ -12,14 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/gather_nd.cc b/tensorflow/lite/kernels/gather_nd.cc
index 4ca0864b94f..fd31b8c4ddd 100644
--- a/tensorflow/lite/kernels/gather_nd.cc
+++ b/tensorflow/lite/kernels/gather_nd.cc
@@ -12,12 +12,14 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/gather_nd_test.cc b/tensorflow/lite/kernels/gather_nd_test.cc
index 7e2714dac5e..33dce89917d 100644
--- a/tensorflow/lite/kernels/gather_nd_test.cc
+++ b/tensorflow/lite/kernels/gather_nd_test.cc
@@ -12,13 +12,18 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
 #include 
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/gather_test.cc b/tensorflow/lite/kernels/gather_test.cc
index 483b59fb533..01be7f01935 100644
--- a/tensorflow/lite/kernels/gather_test.cc
+++ b/tensorflow/lite/kernels/gather_test.cc
@@ -12,12 +12,18 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/hashtable/BUILD b/tensorflow/lite/kernels/hashtable/BUILD
index 4ec3abe77ee..d141abf4f95 100644
--- a/tensorflow/lite/kernels/hashtable/BUILD
+++ b/tensorflow/lite/kernels/hashtable/BUILD
@@ -25,7 +25,6 @@ cc_library(
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/experimental/resource",
         "//tensorflow/lite/kernels:kernel_util",
-        "//tensorflow/lite/kernels:op_macros",
         "//tensorflow/lite/kernels/internal:tensor",
         "//tensorflow/lite/schema:schema_fbs",
         "@flatbuffers",
@@ -49,6 +48,7 @@ cc_test(
         "//tensorflow/lite/testing:util",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
         "@flatbuffers",
     ],
 )
diff --git a/tensorflow/lite/kernels/hashtable_lookup.cc b/tensorflow/lite/kernels/hashtable_lookup.cc
index a432dcb8e22..65e50fe41c2 100644
--- a/tensorflow/lite/kernels/hashtable_lookup.cc
+++ b/tensorflow/lite/kernels/hashtable_lookup.cc
@@ -31,18 +31,13 @@ limitations under the License.
 //   Each item indicates whether the corresponding lookup has a returned value.
 //   0 for missing key, 1 for found key.
 
-#include 
-#include 
-#include 
+#include 
+
 #include 
 #include 
-#include 
-#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/hashtable_lookup_test.cc b/tensorflow/lite/kernels/hashtable_lookup_test.cc
index 638d82ea167..8f90de3a71a 100644
--- a/tensorflow/lite/kernels/hashtable_lookup_test.cc
+++ b/tensorflow/lite/kernels/hashtable_lookup_test.cc
@@ -14,16 +14,21 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite Lookup op.
 
-#include 
+#include 
+
+#include 
+#include 
+#include 
+#include 
 #include 
 
 #include 
 #include 
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/if.cc b/tensorflow/lite/kernels/if.cc
index 731b9d4c82f..d3f92a92b08 100644
--- a/tensorflow/lite/kernels/if.cc
+++ b/tensorflow/lite/kernels/if.cc
@@ -13,7 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
+#include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
diff --git a/tensorflow/lite/kernels/if_test.cc b/tensorflow/lite/kernels/if_test.cc
index c81300e5d1d..0bef77ef7a7 100644
--- a/tensorflow/lite/kernels/if_test.cc
+++ b/tensorflow/lite/kernels/if_test.cc
@@ -13,14 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
+#include 
+#include 
+
 #include 
-#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/subgraph_test_util.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/internal/optimized/im2col_utils.h b/tensorflow/lite/kernels/internal/optimized/im2col_utils.h
index 42aa4825771..ca4ce6bbea3 100644
--- a/tensorflow/lite/kernels/internal/optimized/im2col_utils.h
+++ b/tensorflow/lite/kernels/internal/optimized/im2col_utils.h
@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
 
+#include 
+
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/internal/types.h"
 
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/add.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/add.h
index 44479d93a31..a63763b755a 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/add.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/add.h
@@ -15,9 +15,14 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
 
+#include 
+
+#include "fixedpoint/fixedpoint.h"
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h
index 61f848c888e..c426ceb3a67 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/cpu_backend_gemm.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h
index 0ff153da977..c84e7dc04d9 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h
@@ -15,13 +15,20 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_H_
 
+#include 
+
+#include 
+#include 
+
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h"
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h"
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h
index 8de99c1a564..d234c5bb4a1 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h
@@ -18,9 +18,11 @@ limitations under the License.
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/cpu_backend_gemm.h"
-#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/fully_connected.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 
 namespace tflite {
 namespace optimized_integer_ops {
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h
index 15c3d291ec3..45c27ab026a 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h
@@ -15,9 +15,14 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
 
+#include 
+
+#include "fixedpoint/fixedpoint.h"
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/mul.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h b/tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h
index 060845f4a10..f2696500ab9 100644
--- a/tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h
+++ b/tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h
@@ -15,23 +15,16 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
 
-#include 
-#include 
-#include 
+#include 
 
 #include 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
 
-#include "fixedpoint/fixedpoint.h"
 #include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/cppmath.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
diff --git a/tensorflow/lite/kernels/internal/per_channel_dequantize_test.cc b/tensorflow/lite/kernels/internal/per_channel_dequantize_test.cc
index 3ad125b86e4..89710b99b95 100644
--- a/tensorflow/lite/kernels/internal/per_channel_dequantize_test.cc
+++ b/tensorflow/lite/kernels/internal/per_channel_dequantize_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include 
 #include 
 
+#include 
 #include 
 #include "tensorflow/lite/kernels/internal/reference/dequantize.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/internal/reference/batch_matmul.h b/tensorflow/lite/kernels/internal/reference/batch_matmul.h
index c8d6d6a0e29..1394bd9da64 100644
--- a/tensorflow/lite/kernels/internal/reference/batch_matmul.h
+++ b/tensorflow/lite/kernels/internal/reference/batch_matmul.h
@@ -15,8 +15,11 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_BATCH_MATMUL_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_BATCH_MATMUL_H_
 
-#include "tensorflow/lite/c/common.h"
+#include 
+#include 
+
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/reference/portable_tensor_utils_impl.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 
diff --git a/tensorflow/lite/kernels/internal/reference/densify.h b/tensorflow/lite/kernels/internal/reference/densify.h
index d1fd488700a..71a9a26cc09 100644
--- a/tensorflow/lite/kernels/internal/reference/densify.h
+++ b/tensorflow/lite/kernels/internal/reference/densify.h
@@ -15,6 +15,8 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_DENSIFY_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_DENSIFY_H_
 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/internal/reference/non_max_suppression.h b/tensorflow/lite/kernels/internal/reference/non_max_suppression.h
index 5d3823788ef..64c27c174fa 100644
--- a/tensorflow/lite/kernels/internal/reference/non_max_suppression.h
+++ b/tensorflow/lite/kernels/internal/reference/non_max_suppression.h
@@ -17,7 +17,7 @@ limitations under the License.
 
 #include 
 #include 
-#include 
+#include 
 #include 
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/internal/reference/quantize.h b/tensorflow/lite/kernels/internal/reference/quantize.h
index 58d19c0a14c..d36db06f2e0 100644
--- a/tensorflow/lite/kernels/internal/reference/quantize.h
+++ b/tensorflow/lite/kernels/internal/reference/quantize.h
@@ -15,7 +15,11 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_QUANTIZE_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_QUANTIZE_H_
 
+#include 
+#include 
+
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/cppmath.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 
diff --git a/tensorflow/lite/kernels/internal/reference/strided_slice.h b/tensorflow/lite/kernels/internal/reference/strided_slice.h
index ba6d4c22554..8b6f0c13da1 100644
--- a/tensorflow/lite/kernels/internal/reference/strided_slice.h
+++ b/tensorflow/lite/kernels/internal/reference/strided_slice.h
@@ -16,8 +16,10 @@ limitations under the License.
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_STRIDED_SLICE_H_
 
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/strided_slice_logic.h"
 #include "tensorflow/lite/kernels/internal/types.h"
+
 namespace tflite {
 
 namespace reference_ops {
diff --git a/tensorflow/lite/kernels/internal/reference/sub.h b/tensorflow/lite/kernels/internal/reference/sub.h
index 48d03de02ee..6191eaac558 100644
--- a/tensorflow/lite/kernels/internal/reference/sub.h
+++ b/tensorflow/lite/kernels/internal/reference/sub.h
@@ -15,9 +15,15 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SUB_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SUB_H_
 
-#include "fixedpoint/fixedpoint.h"
+#include 
+
+#include 
+#include 
+
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/internal/reference/svdf.h b/tensorflow/lite/kernels/internal/reference/svdf.h
index f03c0392ff4..ffa46b8f422 100644
--- a/tensorflow/lite/kernels/internal/reference/svdf.h
+++ b/tensorflow/lite/kernels/internal/reference/svdf.h
@@ -15,7 +15,10 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SVDF_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_SVDF_H_
 
+#include 
+
 #include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index ded536ab3a7..032726a7860 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -14,10 +14,15 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/kernel_util.h"
 
+#include 
+#include 
+
 #include 
-#include 
+#include 
 #include 
 
+#include "tensorflow/lite/c/builtin_op_data.h"
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/cppmath.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 
diff --git a/tensorflow/lite/kernels/kernel_util.h b/tensorflow/lite/kernels/kernel_util.h
index d57234afa77..6fc69fa1629 100644
--- a/tensorflow/lite/kernels/kernel_util.h
+++ b/tensorflow/lite/kernels/kernel_util.h
@@ -15,10 +15,10 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_KERNEL_UTIL_H_
 #define TENSORFLOW_LITE_KERNELS_KERNEL_UTIL_H_
 
-#include 
+#include 
+
 #include 
 
-#include "flatbuffers/flatbuffers.h"
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 
diff --git a/tensorflow/lite/kernels/kernel_util_test.cc b/tensorflow/lite/kernels/kernel_util_test.cc
index 7a7467ee0d4..db0cc3cb39c 100644
--- a/tensorflow/lite/kernels/kernel_util_test.cc
+++ b/tensorflow/lite/kernels/kernel_util_test.cc
@@ -14,8 +14,18 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/kernel_util.h"
 
+#include 
+#include 
+#include 
+#include 
+
+#include 
+#include 
+
 #include 
 #include 
+#include "tensorflow/lite/c/builtin_op_data.h"
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/testing/util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/l2norm.cc b/tensorflow/lite/kernels/l2norm.cc
index ab009f337de..a7fb35ed594 100644
--- a/tensorflow/lite/kernels/l2norm.cc
+++ b/tensorflow/lite/kernels/l2norm.cc
@@ -14,12 +14,15 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/l2normalization.h"
+#include "tensorflow/lite/kernels/internal/reference/l2normalization.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/l2norm_test.cc b/tensorflow/lite/kernels/l2norm_test.cc
index e4793dc5c74..968bcc556b3 100644
--- a/tensorflow/lite/kernels/l2norm_test.cc
+++ b/tensorflow/lite/kernels/l2norm_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/local_response_norm.cc b/tensorflow/lite/kernels/local_response_norm.cc
index 85d0796cc1c..f4b996c45a1 100644
--- a/tensorflow/lite/kernels/local_response_norm.cc
+++ b/tensorflow/lite/kernels/local_response_norm.cc
@@ -17,8 +17,9 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/local_response_norm_test.cc b/tensorflow/lite/kernels/local_response_norm_test.cc
index 701da5ceb3d..353cce3db8a 100644
--- a/tensorflow/lite/kernels/local_response_norm_test.cc
+++ b/tensorflow/lite/kernels/local_response_norm_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/log_softmax_test.cc b/tensorflow/lite/kernels/log_softmax_test.cc
index bc265915279..a65ee528a30 100644
--- a/tensorflow/lite/kernels/log_softmax_test.cc
+++ b/tensorflow/lite/kernels/log_softmax_test.cc
@@ -14,17 +14,17 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite LOG_SOFTMAX op.
 
-#include 
+#include 
 #include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/logical.cc b/tensorflow/lite/kernels/logical.cc
index 397964cfd19..ec650dd4210 100644
--- a/tensorflow/lite/kernels/logical.cc
+++ b/tensorflow/lite/kernels/logical.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/logical_test.cc b/tensorflow/lite/kernels/logical_test.cc
index 276d5d91cd1..cd85e320069 100644
--- a/tensorflow/lite/kernels/logical_test.cc
+++ b/tensorflow/lite/kernels/logical_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/lsh_projection.cc b/tensorflow/lite/kernels/lsh_projection.cc
index 68d0719c4e2..b809748c59c 100644
--- a/tensorflow/lite/kernels/lsh_projection.cc
+++ b/tensorflow/lite/kernels/lsh_projection.cc
@@ -50,20 +50,16 @@ limitations under the License.
 //     Output.Dim == { Tensor[0].Dim[0] * Tensor[0].Dim[1] }
 //     A flattened tensor represents projected bit vectors.
 
-#include 
-#include 
-#include 
-#include 
+#include 
+#include 
+
 #include 
-#include 
-#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include 
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/lsh_projection_test.cc b/tensorflow/lite/kernels/lsh_projection_test.cc
index 1b75992de6d..008a5c45aaa 100644
--- a/tensorflow/lite/kernels/lsh_projection_test.cc
+++ b/tensorflow/lite/kernels/lsh_projection_test.cc
@@ -13,13 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
 #include 
 
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 7fa3d85687c..b285ed1030f 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -14,8 +14,13 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/lstm_eval.h"
 
+#include 
+#include 
+
 #include 
 #include 
+#include 
+#include 
 
 #include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/c/builtin_op_data.h"
diff --git a/tensorflow/lite/kernels/lstm_eval.h b/tensorflow/lite/kernels/lstm_eval.h
index 877cfd70a89..91f47b18df6 100644
--- a/tensorflow/lite/kernels/lstm_eval.h
+++ b/tensorflow/lite/kernels/lstm_eval.h
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include 
 #include 
-#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
diff --git a/tensorflow/lite/kernels/lstm_eval_test.cc b/tensorflow/lite/kernels/lstm_eval_test.cc
index 885ae250ae7..baf2e5e83df 100644
--- a/tensorflow/lite/kernels/lstm_eval_test.cc
+++ b/tensorflow/lite/kernels/lstm_eval_test.cc
@@ -14,17 +14,17 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/lstm_eval.h"
 
+#include 
+#include 
+
 #include 
-#include 
+#include 
 #include 
 
-#include 
 #include 
+#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/lstm_test.cc b/tensorflow/lite/kernels/lstm_test.cc
index 74ec8d324c6..ba5ee6508cc 100644
--- a/tensorflow/lite/kernels/lstm_test.cc
+++ b/tensorflow/lite/kernels/lstm_test.cc
@@ -17,15 +17,17 @@ limitations under the License.
 // TODO(alanchiao): add unit test with invalid input dimensions for this and its
 // variants.
 
-#include 
+#include 
+
+#include 
 #include 
 
 #include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/matrix_diag.cc b/tensorflow/lite/kernels/matrix_diag.cc
index a4137c1e0b7..c921650926f 100644
--- a/tensorflow/lite/kernels/matrix_diag.cc
+++ b/tensorflow/lite/kernels/matrix_diag.cc
@@ -12,17 +12,14 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/matrix_diag_test.cc b/tensorflow/lite/kernels/matrix_diag_test.cc
index 09a72e9b726..d0c2a45b3b3 100644
--- a/tensorflow/lite/kernels/matrix_diag_test.cc
+++ b/tensorflow/lite/kernels/matrix_diag_test.cc
@@ -12,10 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
+#include 
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/matrix_set_diag.cc b/tensorflow/lite/kernels/matrix_set_diag.cc
index 4602ca0228c..e9c17f985d3 100644
--- a/tensorflow/lite/kernels/matrix_set_diag.cc
+++ b/tensorflow/lite/kernels/matrix_set_diag.cc
@@ -12,17 +12,14 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/matrix_set_diag_test.cc b/tensorflow/lite/kernels/matrix_set_diag_test.cc
index 46b314735b3..8fdb381f2b9 100644
--- a/tensorflow/lite/kernels/matrix_set_diag_test.cc
+++ b/tensorflow/lite/kernels/matrix_set_diag_test.cc
@@ -12,10 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
+#include 
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/maximum_minimum.cc b/tensorflow/lite/kernels/maximum_minimum.cc
index cad86acd8dd..ae1920e53db 100644
--- a/tensorflow/lite/kernels/maximum_minimum.cc
+++ b/tensorflow/lite/kernels/maximum_minimum.cc
@@ -12,17 +12,19 @@ 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.
 ==============================================================================*/
-#include 
+#include "tensorflow/lite/kernels/internal/reference/maximum_minimum.h"
 
-#include 
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/maximum_minimum_test.cc b/tensorflow/lite/kernels/maximum_minimum_test.cc
index b22435d3e97..2c036e369bd 100644
--- a/tensorflow/lite/kernels/maximum_minimum_test.cc
+++ b/tensorflow/lite/kernels/maximum_minimum_test.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/mfcc.cc b/tensorflow/lite/kernels/mfcc.cc
index cba7cb132eb..5fe5b948a87 100644
--- a/tensorflow/lite/kernels/mfcc.cc
+++ b/tensorflow/lite/kernels/mfcc.cc
@@ -14,16 +14,21 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/mfcc.h"
 
+#include 
+#include 
+
+#include 
+
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/mfcc_dct.h"
 #include "tensorflow/lite/kernels/internal/mfcc_mel_filterbank.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/mfcc_test.cc b/tensorflow/lite/kernels/mfcc_test.cc
index a6b769ccd37..abe9b7b9dad 100644
--- a/tensorflow/lite/kernels/mfcc_test.cc
+++ b/tensorflow/lite/kernels/mfcc_test.cc
@@ -13,16 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
-#include 
 #include 
 
+#include 
 #include 
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/mirror_pad.cc b/tensorflow/lite/kernels/mirror_pad.cc
index 17756113069..8f4f02f7848 100644
--- a/tensorflow/lite/kernels/mirror_pad.cc
+++ b/tensorflow/lite/kernels/mirror_pad.cc
@@ -13,9 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
+#include 
 #include 
 #include 
 
+#include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
@@ -26,7 +31,6 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/mirror_pad_test.cc b/tensorflow/lite/kernels/mirror_pad_test.cc
index 91e48fa68aa..fc8a7e68c49 100644
--- a/tensorflow/lite/kernels/mirror_pad_test.cc
+++ b/tensorflow/lite/kernels/mirror_pad_test.cc
@@ -12,11 +12,12 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/mul.cc b/tensorflow/lite/kernels/mul.cc
index 4140a1ac5b2..0ab378e278d 100644
--- a/tensorflow/lite/kernels/mul.cc
+++ b/tensorflow/lite/kernels/mul.cc
@@ -14,16 +14,24 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h"
 
+#include 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/mul.h"
+#include "tensorflow/lite/kernels/internal/reference/mul.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/mul_test.cc b/tensorflow/lite/kernels/mul_test.cc
index db8dca9f4b3..9499fd40bea 100644
--- a/tensorflow/lite/kernels/mul_test.cc
+++ b/tensorflow/lite/kernels/mul_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/neg.cc b/tensorflow/lite/kernels/neg.cc
index 135f888a9e5..4a4ce8fcbd5 100644
--- a/tensorflow/lite/kernels/neg.cc
+++ b/tensorflow/lite/kernels/neg.cc
@@ -15,9 +15,12 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/internal/reference/neg.h"
 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/neg_test.cc b/tensorflow/lite/kernels/neg_test.cc
index 0cdf5161628..87326846ddd 100644
--- a/tensorflow/lite/kernels/neg_test.cc
+++ b/tensorflow/lite/kernels/neg_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/non_max_suppression.cc b/tensorflow/lite/kernels/non_max_suppression.cc
index f57ee1bc5d2..d6e13cdbd33 100644
--- a/tensorflow/lite/kernels/non_max_suppression.cc
+++ b/tensorflow/lite/kernels/non_max_suppression.cc
@@ -14,16 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/reference/non_max_suppression.h"
 
-#include 
+#include 
 
-#include 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/non_max_suppression_test.cc b/tensorflow/lite/kernels/non_max_suppression_test.cc
index 454bb5a0959..9b7baa147e5 100644
--- a/tensorflow/lite/kernels/non_max_suppression_test.cc
+++ b/tensorflow/lite/kernels/non_max_suppression_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/numeric_verify.cc b/tensorflow/lite/kernels/numeric_verify.cc
index fa6324086e1..bbd2448ece0 100644
--- a/tensorflow/lite/kernels/numeric_verify.cc
+++ b/tensorflow/lite/kernels/numeric_verify.cc
@@ -12,25 +12,24 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+#include 
+#include 
 
 #include 
-#include 
 #include 
 #include 
 #include 
 
-#include "third_party/eigen3/Eigen/Core"
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/dequantize.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/numeric_verify_test.cc b/tensorflow/lite/kernels/numeric_verify_test.cc
index 7dcf2436b32..9fb2e559c37 100644
--- a/tensorflow/lite/kernels/numeric_verify_test.cc
+++ b/tensorflow/lite/kernels/numeric_verify_test.cc
@@ -12,17 +12,19 @@ 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.
 ==============================================================================*/
+#include 
+
 #include 
 #include 
+#include 
 
 #include 
 #include "absl/memory/memory.h"
 #include "third_party/eigen3/Eigen/Core"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/types.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/one_hot.cc b/tensorflow/lite/kernels/one_hot.cc
index 750d6dac2ef..76d53c6396f 100644
--- a/tensorflow/lite/kernels/one_hot.cc
+++ b/tensorflow/lite/kernels/one_hot.cc
@@ -12,11 +12,13 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/one_hot_test.cc b/tensorflow/lite/kernels/one_hot_test.cc
index 96b549cb6eb..e94854612ab 100644
--- a/tensorflow/lite/kernels/one_hot_test.cc
+++ b/tensorflow/lite/kernels/one_hot_test.cc
@@ -13,13 +13,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
+#include 
+#include 
 
 #include 
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/op_macros.h b/tensorflow/lite/kernels/op_macros.h
index 8c1a6b1be16..5c190f1c595 100644
--- a/tensorflow/lite/kernels/op_macros.h
+++ b/tensorflow/lite/kernels/op_macros.h
@@ -20,7 +20,6 @@ limitations under the License.
 #ifdef TF_LITE_MCU_DEBUG_LOG
 
 #include "tensorflow/lite/micro/debug_log.h"
-#include "tensorflow/lite/micro/micro_error_reporter.h"
 
 #define DEBUG_LOG(x) \
   do {               \
@@ -37,7 +36,6 @@ inline void InfiniteLoop() {
 
 #else  // TF_LITE_MCU_DEBUG_LOG
 
-#include 
 #include 
 #include 
 
diff --git a/tensorflow/lite/kernels/optional_tensor_test.cc b/tensorflow/lite/kernels/optional_tensor_test.cc
index a09f8601589..26d619276aa 100644
--- a/tensorflow/lite/kernels/optional_tensor_test.cc
+++ b/tensorflow/lite/kernels/optional_tensor_test.cc
@@ -14,16 +14,14 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite LSTM op.
 
-#include 
-#include 
+#include 
 #include 
 
-#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/pack.cc b/tensorflow/lite/kernels/pack.cc
index ebc3381dae8..fc7a87692c4 100644
--- a/tensorflow/lite/kernels/pack.cc
+++ b/tensorflow/lite/kernels/pack.cc
@@ -13,10 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/pack_test.cc b/tensorflow/lite/kernels/pack_test.cc
index bc6758c7249..c15a0b7fe1f 100644
--- a/tensorflow/lite/kernels/pack_test.cc
+++ b/tensorflow/lite/kernels/pack_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/pad.cc b/tensorflow/lite/kernels/pad.cc
index cc735e4eede..2239511b60a 100644
--- a/tensorflow/lite/kernels/pad.cc
+++ b/tensorflow/lite/kernels/pad.cc
@@ -12,17 +12,20 @@ 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.
 ==============================================================================*/
-#include 
+#include "tensorflow/lite/kernels/internal/reference/pad.h"
 
-#include 
+#include 
+
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/pad_test.cc b/tensorflow/lite/kernels/pad_test.cc
index 6e1e00cc3b8..983642298be 100644
--- a/tensorflow/lite/kernels/pad_test.cc
+++ b/tensorflow/lite/kernels/pad_test.cc
@@ -13,11 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include 
+#include 
+#include 
+
+#include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc
index 0dcb667e901..1dc5cbb6199 100644
--- a/tensorflow/lite/kernels/pooling.cc
+++ b/tensorflow/lite/kernels/pooling.cc
@@ -14,21 +14,22 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h"
 
-#include 
-#include 
-#include 
+#include 
+#include 
+
 #include 
-#include 
-#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
+#include "tensorflow/lite/kernels/internal/reference/pooling.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/kernels/padding.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/pooling_test.cc b/tensorflow/lite/kernels/pooling_test.cc
index e609f04e21d..e614fedccfd 100644
--- a/tensorflow/lite/kernels/pooling_test.cc
+++ b/tensorflow/lite/kernels/pooling_test.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/pow.cc b/tensorflow/lite/kernels/pow.cc
index 8b72e7c0418..a76c77a3f9f 100644
--- a/tensorflow/lite/kernels/pow.cc
+++ b/tensorflow/lite/kernels/pow.cc
@@ -12,12 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/pow_test.cc b/tensorflow/lite/kernels/pow_test.cc
index e106cbb1b91..fa7b6d2ef9a 100644
--- a/tensorflow/lite/kernels/pow_test.cc
+++ b/tensorflow/lite/kernels/pow_test.cc
@@ -12,12 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/test_util.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/quant_basic_lstm_test.cc b/tensorflow/lite/kernels/quant_basic_lstm_test.cc
index 0baae569d24..3e081c221c5 100644
--- a/tensorflow/lite/kernels/quant_basic_lstm_test.cc
+++ b/tensorflow/lite/kernels/quant_basic_lstm_test.cc
@@ -14,16 +14,13 @@ limitations under the License.
 ==============================================================================*/
 #include 
 #include 
-#include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/quantize.cc b/tensorflow/lite/kernels/quantize.cc
index 21cc73278a0..1779500e6a2 100644
--- a/tensorflow/lite/kernels/quantize.cc
+++ b/tensorflow/lite/kernels/quantize.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/requantize.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
diff --git a/tensorflow/lite/kernels/quantize_test.cc b/tensorflow/lite/kernels/quantize_test.cc
index cfd9c1f434e..d7392b3e3ea 100644
--- a/tensorflow/lite/kernels/quantize_test.cc
+++ b/tensorflow/lite/kernels/quantize_test.cc
@@ -13,13 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include 
+#include 
+#include 
 
 #include 
-#include "tensorflow/lite/interpreter.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/internal/types.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/range.cc b/tensorflow/lite/kernels/range.cc
index ae6db1b601f..55cc543d745 100644
--- a/tensorflow/lite/kernels/range.cc
+++ b/tensorflow/lite/kernels/range.cc
@@ -13,7 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+#include 
+#include 
+
+#include 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
diff --git a/tensorflow/lite/kernels/range_test.cc b/tensorflow/lite/kernels/range_test.cc
index 7de4fe3cb76..52f7231def9 100644
--- a/tensorflow/lite/kernels/range_test.cc
+++ b/tensorflow/lite/kernels/range_test.cc
@@ -12,11 +12,13 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/rank.cc b/tensorflow/lite/kernels/rank.cc
index 53fd92f1682..2202f6dd953 100644
--- a/tensorflow/lite/kernels/rank.cc
+++ b/tensorflow/lite/kernels/rank.cc
@@ -12,11 +12,12 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/rank_test.cc b/tensorflow/lite/kernels/rank_test.cc
index 5373a0a66fe..760560c5a92 100644
--- a/tensorflow/lite/kernels/rank_test.cc
+++ b/tensorflow/lite/kernels/rank_test.cc
@@ -13,13 +13,17 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
+#include 
+#include 
 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/read_variable.cc b/tensorflow/lite/kernels/read_variable.cc
index f4762303b8b..ad6e8d43858 100644
--- a/tensorflow/lite/kernels/read_variable.cc
+++ b/tensorflow/lite/kernels/read_variable.cc
@@ -13,17 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
 #include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/experimental/resource/resource_variable.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/reduce.cc b/tensorflow/lite/kernels/reduce.cc
index f0222a08fe3..af42b2a369c 100644
--- a/tensorflow/lite/kernels/reduce.cc
+++ b/tensorflow/lite/kernels/reduce.cc
@@ -12,16 +12,20 @@ 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.
 ==============================================================================*/
-#include 
+#include "tensorflow/lite/kernels/internal/reference/reduce.h"
+
+#include 
 
 #include 
 #include 
-#include 
 
+#include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/mean.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/mean.h"
@@ -30,7 +34,6 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/reduce_test.cc b/tensorflow/lite/kernels/reduce_test.cc
index ddbd5106063..2c83369ea37 100644
--- a/tensorflow/lite/kernels/reduce_test.cc
+++ b/tensorflow/lite/kernels/reduce_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/register.cc b/tensorflow/lite/kernels/register.cc
index 8ca58e6a309..90688a2aa1f 100644
--- a/tensorflow/lite/kernels/register.cc
+++ b/tensorflow/lite/kernels/register.cc
@@ -15,7 +15,9 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/register.h"
 
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/builtin_op_kernels.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/register.h b/tensorflow/lite/kernels/register.h
index 3e5bd298baf..a2a41ea9428 100644
--- a/tensorflow/lite/kernels/register.h
+++ b/tensorflow/lite/kernels/register.h
@@ -15,7 +15,6 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_REGISTER_H_
 #define TENSORFLOW_LITE_KERNELS_REGISTER_H_
 
-#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/model.h"
 #include "tensorflow/lite/mutable_op_resolver.h"
 
diff --git a/tensorflow/lite/kernels/register_ref.cc b/tensorflow/lite/kernels/register_ref.cc
index 426f8a8e896..233520e2165 100644
--- a/tensorflow/lite/kernels/register_ref.cc
+++ b/tensorflow/lite/kernels/register_ref.cc
@@ -14,6 +14,10 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/lite/kernels/register_ref.h"
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 #include "tensorflow/lite/util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/register_ref.h b/tensorflow/lite/kernels/register_ref.h
index 5d9cb2c0b95..406fad89673 100644
--- a/tensorflow/lite/kernels/register_ref.h
+++ b/tensorflow/lite/kernels/register_ref.h
@@ -16,8 +16,8 @@ limitations under the License.
 #define TENSORFLOW_LITE_KERNELS_REGISTER_REF_H_
 
 #include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/model.h"
 #include "tensorflow/lite/mutable_op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/reshape.cc b/tensorflow/lite/kernels/reshape.cc
index 6afc3c8a670..ab6f0d8577d 100644
--- a/tensorflow/lite/kernels/reshape.cc
+++ b/tensorflow/lite/kernels/reshape.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include 
 #include 
 
 #include 
@@ -20,7 +21,6 @@ limitations under the License.
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/reshape_test.cc b/tensorflow/lite/kernels/reshape_test.cc
index 310d594698c..09f8ce6a3d5 100644
--- a/tensorflow/lite/kernels/reshape_test.cc
+++ b/tensorflow/lite/kernels/reshape_test.cc
@@ -12,14 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
 #include 
 
 #include 
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/reshape_test_common.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/reshape_test_common.h b/tensorflow/lite/kernels/reshape_test_common.h
index 9dbf028e7be..662c163b7c0 100644
--- a/tensorflow/lite/kernels/reshape_test_common.h
+++ b/tensorflow/lite/kernels/reshape_test_common.h
@@ -15,7 +15,15 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_RESHAPE_TEST_COMMON_H_
 #define TENSORFLOW_LITE_KERNELS_RESHAPE_TEST_COMMON_H_
 
+#include 
+
+#include 
+#include 
+#include 
+
 #include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 // There are three ways to specify the output shape of a Reshape
diff --git a/tensorflow/lite/kernels/resize_bilinear.cc b/tensorflow/lite/kernels/resize_bilinear.cc
index dfd58255491..b0488a0b464 100644
--- a/tensorflow/lite/kernels/resize_bilinear.cc
+++ b/tensorflow/lite/kernels/resize_bilinear.cc
@@ -12,13 +12,18 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/resize_bilinear_test.cc b/tensorflow/lite/kernels/resize_bilinear_test.cc
index d4d414ae29c..6dedc0d169d 100644
--- a/tensorflow/lite/kernels/resize_bilinear_test.cc
+++ b/tensorflow/lite/kernels/resize_bilinear_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/resize_nearest_neighbor.cc b/tensorflow/lite/kernels/resize_nearest_neighbor.cc
index 1b58e5245ee..fff45ac13cc 100644
--- a/tensorflow/lite/kernels/resize_nearest_neighbor.cc
+++ b/tensorflow/lite/kernels/resize_nearest_neighbor.cc
@@ -12,13 +12,20 @@ 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.
 ==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/resize_nearest_neighbor.h"
+
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/resize_nearest_neighbor_test.cc b/tensorflow/lite/kernels/resize_nearest_neighbor_test.cc
index 656bd6ee750..b22ad48afb9 100644
--- a/tensorflow/lite/kernels/resize_nearest_neighbor_test.cc
+++ b/tensorflow/lite/kernels/resize_nearest_neighbor_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/reverse.cc b/tensorflow/lite/kernels/reverse.cc
index 760236ad6a7..9ce845b4b7b 100644
--- a/tensorflow/lite/kernels/reverse.cc
+++ b/tensorflow/lite/kernels/reverse.cc
@@ -13,10 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/reverse_sequence.cc b/tensorflow/lite/kernels/reverse_sequence.cc
index 8e976cccd90..7390876d39b 100644
--- a/tensorflow/lite/kernels/reverse_sequence.cc
+++ b/tensorflow/lite/kernels/reverse_sequence.cc
@@ -13,10 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/reverse_sequence_test.cc b/tensorflow/lite/kernels/reverse_sequence_test.cc
index f1fefdd4856..3d9dcdc0b69 100644
--- a/tensorflow/lite/kernels/reverse_sequence_test.cc
+++ b/tensorflow/lite/kernels/reverse_sequence_test.cc
@@ -12,10 +12,13 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include 
+
+#include 
+
+#include 
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/reverse_test.cc b/tensorflow/lite/kernels/reverse_test.cc
index 02101ab172c..f1fcf67fd42 100644
--- a/tensorflow/lite/kernels/reverse_test.cc
+++ b/tensorflow/lite/kernels/reverse_test.cc
@@ -12,10 +12,13 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include 
+
+#include 
+
+#include 
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/rfft2d.cc b/tensorflow/lite/kernels/rfft2d.cc
index fa201153daf..9aeee53f637 100644
--- a/tensorflow/lite/kernels/rfft2d.cc
+++ b/tensorflow/lite/kernels/rfft2d.cc
@@ -13,13 +13,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+#include 
+#include 
+
+#include 
+#include 
+
 #include "third_party/fft2d/fft2d.h"
 #include "ruy/profiler/instrumentation.h"  // from @ruy
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/rfft2d_test.cc b/tensorflow/lite/kernels/rfft2d_test.cc
index d4b6a0a9d83..e9b23bacf0c 100644
--- a/tensorflow/lite/kernels/rfft2d_test.cc
+++ b/tensorflow/lite/kernels/rfft2d_test.cc
@@ -13,13 +13,17 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
+#include 
+
+#include 
+#include 
 
 #include 
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/custom_ops_register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/round.cc b/tensorflow/lite/kernels/round.cc
index fd16605abeb..341d2880705 100644
--- a/tensorflow/lite/kernels/round.cc
+++ b/tensorflow/lite/kernels/round.cc
@@ -13,9 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "tensorflow/lite/kernels/internal/reference/round.h"
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/round_test.cc b/tensorflow/lite/kernels/round_test.cc
index baa614347d2..3402014f25f 100644
--- a/tensorflow/lite/kernels/round_test.cc
+++ b/tensorflow/lite/kernels/round_test.cc
@@ -13,11 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/scatter_nd.cc b/tensorflow/lite/kernels/scatter_nd.cc
index 32f83357f8c..4e904f66692 100644
--- a/tensorflow/lite/kernels/scatter_nd.cc
+++ b/tensorflow/lite/kernels/scatter_nd.cc
@@ -13,14 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/context.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/scatter_nd_test.cc b/tensorflow/lite/kernels/scatter_nd_test.cc
index e25ba9b93f3..9fdf176fe1f 100644
--- a/tensorflow/lite/kernels/scatter_nd_test.cc
+++ b/tensorflow/lite/kernels/scatter_nd_test.cc
@@ -12,13 +12,15 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/segment_sum.cc b/tensorflow/lite/kernels/segment_sum.cc
index db8aa688ebe..8185359321e 100644
--- a/tensorflow/lite/kernels/segment_sum.cc
+++ b/tensorflow/lite/kernels/segment_sum.cc
@@ -13,10 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/segment_sum_test.cc b/tensorflow/lite/kernels/segment_sum_test.cc
index d083feb44aa..ec531ffd92d 100644
--- a/tensorflow/lite/kernels/segment_sum_test.cc
+++ b/tensorflow/lite/kernels/segment_sum_test.cc
@@ -12,11 +12,13 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/select.cc b/tensorflow/lite/kernels/select.cc
index 89fac10c869..281425253c5 100644
--- a/tensorflow/lite/kernels/select.cc
+++ b/tensorflow/lite/kernels/select.cc
@@ -12,12 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
-#include "tensorflow/lite/string_util.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/select_test.cc b/tensorflow/lite/kernels/select_test.cc
index 36935b0b6dc..56ed994d805 100644
--- a/tensorflow/lite/kernels/select_test.cc
+++ b/tensorflow/lite/kernels/select_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/shape.cc b/tensorflow/lite/kernels/shape.cc
index d979f083f70..afeadc38c20 100644
--- a/tensorflow/lite/kernels/shape.cc
+++ b/tensorflow/lite/kernels/shape.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/shape_test.cc b/tensorflow/lite/kernels/shape_test.cc
index 3eeb83f5000..292b40ab7cc 100644
--- a/tensorflow/lite/kernels/shape_test.cc
+++ b/tensorflow/lite/kernels/shape_test.cc
@@ -13,13 +13,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
+#include 
+#include 
 
 #include 
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/skip_gram.cc b/tensorflow/lite/kernels/skip_gram.cc
index 3eb415a55a8..8348a25bba7 100644
--- a/tensorflow/lite/kernels/skip_gram.cc
+++ b/tensorflow/lite/kernels/skip_gram.cc
@@ -31,13 +31,11 @@ limitations under the License.
 
 #include 
 
-#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/skip_gram_test.cc b/tensorflow/lite/kernels/skip_gram_test.cc
index 12d631660ee..9a5b541c077 100644
--- a/tensorflow/lite/kernels/skip_gram_test.cc
+++ b/tensorflow/lite/kernels/skip_gram_test.cc
@@ -13,13 +13,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
 #include 
 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 #include "tensorflow/lite/string_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc
index 1d2cc588abc..c99e1b573b9 100644
--- a/tensorflow/lite/kernels/slice.cc
+++ b/tensorflow/lite/kernels/slice.cc
@@ -13,17 +13,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
+#include 
 
-#include 
+#include 
+#include 
 #include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc
index 1a31ae44a5d..f950f3346ab 100644
--- a/tensorflow/lite/kernels/slice_test.cc
+++ b/tensorflow/lite/kernels/slice_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/softmax_test.cc b/tensorflow/lite/kernels/softmax_test.cc
index 79f3608fe9a..b8e3b3076b3 100644
--- a/tensorflow/lite/kernels/softmax_test.cc
+++ b/tensorflow/lite/kernels/softmax_test.cc
@@ -14,17 +14,18 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite SOFTMAX op.
 
-#include 
+#include "tensorflow/lite/kernels/internal/reference/softmax.h"
+
+#include 
 #include 
 #include 
 
-#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/space_to_batch_nd.cc b/tensorflow/lite/kernels/space_to_batch_nd.cc
index e8756ef5f2e..7fc58e7ee6b 100644
--- a/tensorflow/lite/kernels/space_to_batch_nd.cc
+++ b/tensorflow/lite/kernels/space_to_batch_nd.cc
@@ -12,17 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/space_to_batch_nd_test.cc b/tensorflow/lite/kernels/space_to_batch_nd_test.cc
index d34989f0fb7..0591265f73b 100644
--- a/tensorflow/lite/kernels/space_to_batch_nd_test.cc
+++ b/tensorflow/lite/kernels/space_to_batch_nd_test.cc
@@ -13,11 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/space_to_depth.cc b/tensorflow/lite/kernels/space_to_depth.cc
index 527b7c83adb..e4c7efaaf99 100644
--- a/tensorflow/lite/kernels/space_to_depth.cc
+++ b/tensorflow/lite/kernels/space_to_depth.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/space_to_depth_test.cc b/tensorflow/lite/kernels/space_to_depth_test.cc
index ad2f95d82ba..523dc60a37f 100644
--- a/tensorflow/lite/kernels/space_to_depth_test.cc
+++ b/tensorflow/lite/kernels/space_to_depth_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/sparse_to_dense.cc b/tensorflow/lite/kernels/sparse_to_dense.cc
index 29c87734748..bdf0f4e703a 100644
--- a/tensorflow/lite/kernels/sparse_to_dense.cc
+++ b/tensorflow/lite/kernels/sparse_to_dense.cc
@@ -12,20 +12,16 @@ 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.
 ==============================================================================*/
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
+#include 
+
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/op_macros.h"
-#include "tensorflow/lite/kernels/padding.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/sparse_to_dense_test.cc b/tensorflow/lite/kernels/sparse_to_dense_test.cc
index ad040b2ce04..add4b53ca48 100644
--- a/tensorflow/lite/kernels/sparse_to_dense_test.cc
+++ b/tensorflow/lite/kernels/sparse_to_dense_test.cc
@@ -1,4 +1,3 @@
-
 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,12 +12,14 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/split.cc b/tensorflow/lite/kernels/split.cc
index da239e6ecab..3b7781f409e 100644
--- a/tensorflow/lite/kernels/split.cc
+++ b/tensorflow/lite/kernels/split.cc
@@ -12,17 +12,16 @@ 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.
 ==============================================================================*/
-#include 
-
-#include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/split_test.cc b/tensorflow/lite/kernels/split_test.cc
index 7952396880c..ae7c5cf76e8 100644
--- a/tensorflow/lite/kernels/split_test.cc
+++ b/tensorflow/lite/kernels/split_test.cc
@@ -12,11 +12,17 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/split_v.cc b/tensorflow/lite/kernels/split_v.cc
index b5529b98ecb..7d60086a91d 100644
--- a/tensorflow/lite/kernels/split_v.cc
+++ b/tensorflow/lite/kernels/split_v.cc
@@ -12,6 +12,8 @@ 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.
 ==============================================================================*/
+#include 
+
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
@@ -19,8 +21,9 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/split_v_test.cc b/tensorflow/lite/kernels/split_v_test.cc
index 4e143cabe58..a10e277d653 100644
--- a/tensorflow/lite/kernels/split_v_test.cc
+++ b/tensorflow/lite/kernels/split_v_test.cc
@@ -12,13 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/squared_difference.cc b/tensorflow/lite/kernels/squared_difference.cc
index fbea2403a53..e17ff8e3191 100644
--- a/tensorflow/lite/kernels/squared_difference.cc
+++ b/tensorflow/lite/kernels/squared_difference.cc
@@ -12,14 +12,18 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/c/builtin_op_data.h"
+#include 
+#include 
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/squared_difference_test.cc b/tensorflow/lite/kernels/squared_difference_test.cc
index 249590f37e9..efac1969144 100644
--- a/tensorflow/lite/kernels/squared_difference_test.cc
+++ b/tensorflow/lite/kernels/squared_difference_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/squeeze.cc b/tensorflow/lite/kernels/squeeze.cc
index fa5656e1d59..c4dc51026a6 100644
--- a/tensorflow/lite/kernels/squeeze.cc
+++ b/tensorflow/lite/kernels/squeeze.cc
@@ -14,13 +14,10 @@ limitations under the License.
 ==============================================================================*/
 #include 
 
-#include 
-
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/squeeze_test.cc b/tensorflow/lite/kernels/squeeze_test.cc
index 575a02a70f8..4239ae43e1c 100644
--- a/tensorflow/lite/kernels/squeeze_test.cc
+++ b/tensorflow/lite/kernels/squeeze_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/strided_slice.cc b/tensorflow/lite/kernels/strided_slice.cc
index e2ca812d193..50c2255e526 100644
--- a/tensorflow/lite/kernels/strided_slice.cc
+++ b/tensorflow/lite/kernels/strided_slice.cc
@@ -15,14 +15,20 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/internal/reference/strided_slice.h"
 
+#include 
+#include 
+
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/strided_slice_logic.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/strided_slice_test.cc b/tensorflow/lite/kernels/strided_slice_test.cc
index e97eab5b7c4..5f625d3f201 100644
--- a/tensorflow/lite/kernels/strided_slice_test.cc
+++ b/tensorflow/lite/kernels/strided_slice_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/sub.cc b/tensorflow/lite/kernels/sub.cc
index a2282a0545b..aa628fa5408 100644
--- a/tensorflow/lite/kernels/sub.cc
+++ b/tensorflow/lite/kernels/sub.cc
@@ -14,18 +14,27 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/internal/reference/sub.h"
 
+#include 
+#include 
+
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/add.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/sub_test.cc b/tensorflow/lite/kernels/sub_test.cc
index adda1b810ce..21f2dc7cabd 100644
--- a/tensorflow/lite/kernels/sub_test.cc
+++ b/tensorflow/lite/kernels/sub_test.cc
@@ -12,11 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/subgraph_test_util.cc b/tensorflow/lite/kernels/subgraph_test_util.cc
index 00f947a9e38..8f1964ad10f 100644
--- a/tensorflow/lite/kernels/subgraph_test_util.cc
+++ b/tensorflow/lite/kernels/subgraph_test_util.cc
@@ -15,13 +15,18 @@ limitations under the License.
 
 #include "tensorflow/lite/kernels/subgraph_test_util.h"
 
-#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include 
+#include 
+#include 
+
+#include 
+
+#include 
+#include "tensorflow/lite/c/builtin_op_data.h"
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/kernels/builtin_op_kernels.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/register.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 
 namespace tflite {
 namespace subgraph_test_util {
diff --git a/tensorflow/lite/kernels/subgraph_test_util.h b/tensorflow/lite/kernels/subgraph_test_util.h
index 95b7206fc29..7306f82344d 100644
--- a/tensorflow/lite/kernels/subgraph_test_util.h
+++ b/tensorflow/lite/kernels/subgraph_test_util.h
@@ -20,6 +20,11 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_SUBGRAPH_TEST_UTIL_H_
 #define TENSORFLOW_LITE_KERNELS_SUBGRAPH_TEST_UTIL_H_
 
+#include 
+
+#include 
+#include 
+
 #include 
 #include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/interpreter.h"
diff --git a/tensorflow/lite/kernels/subgraph_test_util_test.cc b/tensorflow/lite/kernels/subgraph_test_util_test.cc
index 4bd0482da17..39e013294f7 100644
--- a/tensorflow/lite/kernels/subgraph_test_util_test.cc
+++ b/tensorflow/lite/kernels/subgraph_test_util_test.cc
@@ -14,10 +14,16 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/lite/kernels/subgraph_test_util.h"
+
+#include 
+
+#include 
+#include 
+
 #include 
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/testing/util.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/svdf.cc b/tensorflow/lite/kernels/svdf.cc
index 57eedb6b204..1b8bf904b8a 100644
--- a/tensorflow/lite/kernels/svdf.cc
+++ b/tensorflow/lite/kernels/svdf.cc
@@ -24,6 +24,7 @@ limitations under the License.
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
diff --git a/tensorflow/lite/kernels/svdf_test.cc b/tensorflow/lite/kernels/svdf_test.cc
index 68963b784f4..b0ac2011948 100644
--- a/tensorflow/lite/kernels/svdf_test.cc
+++ b/tensorflow/lite/kernels/svdf_test.cc
@@ -14,15 +14,15 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite SVDF op.
 
-#include 
+#include 
+
+#include 
 #include 
 
-#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/test_util.cc b/tensorflow/lite/kernels/test_util.cc
index 010aed31dc6..24f6e4f11ca 100644
--- a/tensorflow/lite/kernels/test_util.cc
+++ b/tensorflow/lite/kernels/test_util.cc
@@ -14,18 +14,37 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/test_util.h"
 
-#include 
+#include 
+#include 
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
 
 #include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/core/subgraph.h"
 #include "tensorflow/lite/delegates/nnapi/acceleration_test_util.h"
 #include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/acceleration_test_util.h"
+#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/model.h"
 #include "tensorflow/lite/nnapi/nnapi_implementation.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/string_util.h"
 #include "tensorflow/lite/tools/versioning/op_version.h"
 #include "tensorflow/lite/version.h"
 
diff --git a/tensorflow/lite/kernels/test_util.h b/tensorflow/lite/kernels/test_util.h
index f0f02d25add..bc93bdae58a 100644
--- a/tensorflow/lite/kernels/test_util.h
+++ b/tensorflow/lite/kernels/test_util.h
@@ -15,25 +15,39 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_TEST_UTIL_H_
 #define TENSORFLOW_LITE_KERNELS_TEST_UTIL_H_
 
+#include 
+#include 
+#include 
+#include 
+
+#include 
 #include 
 #include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
 #include 
+#include 
 #include 
 
 #include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/delegates/nnapi/nnapi_delegate_kernel.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_utils.h"
-#include "tensorflow/lite/kernels/register.h"
-#include "tensorflow/lite/model.h"
 #include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 #include "tensorflow/lite/string_util.h"
-#include "tensorflow/lite/testing/util.h"
+#include "tensorflow/lite/testing/util.h"  // IWYU pragma: keep
 #include "tensorflow/lite/tools/optimize/quantization_utils.h"
 #include "tensorflow/lite/tools/optimize/sparsity/format_converter.h"
+#include "tensorflow/lite/type_to_tflitetype.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/test_util_test.cc b/tensorflow/lite/kernels/test_util_test.cc
index 7abb7011f9d..e6f865f6cd6 100644
--- a/tensorflow/lite/kernels/test_util_test.cc
+++ b/tensorflow/lite/kernels/test_util_test.cc
@@ -13,7 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/kernels/test_util.h"
+
+#include 
+
+#include 
+
 #include 
+#include "tensorflow/lite/testing/util.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/tile.cc b/tensorflow/lite/kernels/tile.cc
index 64f6bd05485..884456fcbf2 100644
--- a/tensorflow/lite/kernels/tile.cc
+++ b/tensorflow/lite/kernels/tile.cc
@@ -12,16 +12,19 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
+#include 
+#include 
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/string_util.h"
+
 namespace tflite {
 namespace ops {
 namespace builtin {
diff --git a/tensorflow/lite/kernels/tile_test.cc b/tensorflow/lite/kernels/tile_test.cc
index 0df1f33c2bc..7a5203937a5 100644
--- a/tensorflow/lite/kernels/tile_test.cc
+++ b/tensorflow/lite/kernels/tile_test.cc
@@ -12,12 +12,17 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/topk_v2.cc b/tensorflow/lite/kernels/topk_v2.cc
index 3e313481725..6a5bd392086 100644
--- a/tensorflow/lite/kernels/topk_v2.cc
+++ b/tensorflow/lite/kernels/topk_v2.cc
@@ -12,13 +12,18 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+
+#include 
+#include 
+#include 
 
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
+
 namespace tflite {
 namespace ops {
 namespace builtin {
diff --git a/tensorflow/lite/kernels/topk_v2_test.cc b/tensorflow/lite/kernels/topk_v2_test.cc
index 72ed82c1449..5d96fe06bf4 100644
--- a/tensorflow/lite/kernels/topk_v2_test.cc
+++ b/tensorflow/lite/kernels/topk_v2_test.cc
@@ -1,4 +1,3 @@
-
 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,12 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc
index 511b45e9f4a..27f5cf6f065 100644
--- a/tensorflow/lite/kernels/transpose.cc
+++ b/tensorflow/lite/kernels/transpose.cc
@@ -12,17 +12,16 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
-#include 
-
-#include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc
index 494433159d4..33e122ba037 100644
--- a/tensorflow/lite/kernels/transpose_conv.cc
+++ b/tensorflow/lite/kernels/transpose_conv.cc
@@ -13,27 +13,26 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include 
-#include 
-#include 
-#include 
-#include 
-#include 
+#include 
+#include 
+
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/cpu_backend_context.h"
 #include "tensorflow/lite/kernels/eigen_support.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 // NOLINTNEXTLINE - This header file should't go to the top.
 #include "tensorflow/lite/kernels/internal/optimized/integer_ops/transpose_conv.h"
 #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
 // NOLINTNEXTLINE - This header file should't go to the top.
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/transpose_conv.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/kernels/padding.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/transpose_conv_test.cc b/tensorflow/lite/kernels/transpose_conv_test.cc
index 25c55c95412..b57bc047f62 100644
--- a/tensorflow/lite/kernels/transpose_conv_test.cc
+++ b/tensorflow/lite/kernels/transpose_conv_test.cc
@@ -12,14 +12,24 @@ 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.
 ==============================================================================*/
-#include 
+#include 
+#include 
 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include 
 #include 
 #include "absl/memory/memory.h"
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/transpose_test.cc b/tensorflow/lite/kernels/transpose_test.cc
index 30594da5d51..a88abec7161 100644
--- a/tensorflow/lite/kernels/transpose_test.cc
+++ b/tensorflow/lite/kernels/transpose_test.cc
@@ -12,13 +12,20 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
index 0552885f720..b8b9396f436 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
index 4ea018c0cab..43cc75f894b 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
@@ -14,15 +14,13 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite Sequential LSTM op.
 
-#include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_rnn_test.cc b/tensorflow/lite/kernels/unidirectional_sequence_rnn_test.cc
index 8b6f102acdb..f1486267c17 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_rnn_test.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_rnn_test.cc
@@ -14,15 +14,14 @@ limitations under the License.
 ==============================================================================*/
 // Unit test for TFLite Sequential RNN op.
 
-#include 
+#include 
 #include 
 
 #include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/unique.cc b/tensorflow/lite/kernels/unique.cc
index d0d277ecaa4..dd5c801b468 100644
--- a/tensorflow/lite/kernels/unique.cc
+++ b/tensorflow/lite/kernels/unique.cc
@@ -13,12 +13,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
 #include 
+#include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/unique_test.cc b/tensorflow/lite/kernels/unique_test.cc
index b18fcbed654..d01616025bf 100644
--- a/tensorflow/lite/kernels/unique_test.cc
+++ b/tensorflow/lite/kernels/unique_test.cc
@@ -12,11 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/unpack.cc b/tensorflow/lite/kernels/unpack.cc
index 285e01b3558..8d307acb268 100644
--- a/tensorflow/lite/kernels/unpack.cc
+++ b/tensorflow/lite/kernels/unpack.cc
@@ -13,10 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/unpack_test.cc b/tensorflow/lite/kernels/unpack_test.cc
index 0c6b8fa157c..14d9b9e66b7 100644
--- a/tensorflow/lite/kernels/unpack_test.cc
+++ b/tensorflow/lite/kernels/unpack_test.cc
@@ -12,11 +12,17 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+#include 
+#include 
+
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/variable_ops_test.cc b/tensorflow/lite/kernels/variable_ops_test.cc
index 2efac9d7d8f..077a03df21d 100644
--- a/tensorflow/lite/kernels/variable_ops_test.cc
+++ b/tensorflow/lite/kernels/variable_ops_test.cc
@@ -12,13 +12,11 @@ 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.
 ==============================================================================*/
-#include 
+#include 
 
 #include 
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/where.cc b/tensorflow/lite/kernels/where.cc
index 867d3069f2c..a20efa8baaa 100644
--- a/tensorflow/lite/kernels/where.cc
+++ b/tensorflow/lite/kernels/where.cc
@@ -12,6 +12,8 @@ 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.
 ==============================================================================*/
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
diff --git a/tensorflow/lite/kernels/where_test.cc b/tensorflow/lite/kernels/where_test.cc
index f52c9aad487..ba93bed6e74 100644
--- a/tensorflow/lite/kernels/where_test.cc
+++ b/tensorflow/lite/kernels/where_test.cc
@@ -12,13 +12,15 @@ 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.
 ==============================================================================*/
+#include 
+
 #include 
 
+#include 
 #include 
-#include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/kernels/while.cc b/tensorflow/lite/kernels/while.cc
index 81b6c0c6634..99d6d2cc1c8 100644
--- a/tensorflow/lite/kernels/while.cc
+++ b/tensorflow/lite/kernels/while.cc
@@ -13,7 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
diff --git a/tensorflow/lite/kernels/while_test.cc b/tensorflow/lite/kernels/while_test.cc
index 324519e32a0..b0b63f8c643 100644
--- a/tensorflow/lite/kernels/while_test.cc
+++ b/tensorflow/lite/kernels/while_test.cc
@@ -12,14 +12,14 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+#include 
+
 #include 
-#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/subgraph_test_util.h"
-#include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
 
 namespace tflite {
 
diff --git a/tensorflow/lite/kernels/zeros_like.cc b/tensorflow/lite/kernels/zeros_like.cc
index ad6d03649f6..8586c945e7c 100644
--- a/tensorflow/lite/kernels/zeros_like.cc
+++ b/tensorflow/lite/kernels/zeros_like.cc
@@ -13,8 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+#include 
+
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/kernels/zeros_like_test.cc b/tensorflow/lite/kernels/zeros_like_test.cc
index 09a233c2c30..0be6336ce3e 100644
--- a/tensorflow/lite/kernels/zeros_like_test.cc
+++ b/tensorflow/lite/kernels/zeros_like_test.cc
@@ -12,11 +12,16 @@ 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.
 ==============================================================================*/
+#include 
+
+#include 
+
+#include 
 #include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/kernels/test_util.h"
-#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
index 381e3632b3a..001b0feaef2 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/reference/pooling.h"
 
 #include "cmsis/CMSIS/NN/Include/arm_nnfunctions.h"
+#include "flatbuffers/base.h"  // from @flatbuffers
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"

From cd7da16dd6c17df428dc9ec105c0c8f11e5fd4f5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 18:39:25 -0700
Subject: [PATCH 0248/1390] Stops calling std::copy_n over zero-sided vec,
 which indeed dereferences null. Also added vector::reserve() call before the
 for loop.

PiperOrigin-RevId: 316589226
Change-Id: Iae6edd6f3bf5c8b737cbad782d45c978e622df43
---
 tensorflow/core/kernels/ctc_decoder_ops.cc | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/ctc_decoder_ops.cc b/tensorflow/core/kernels/ctc_decoder_ops.cc
index 517612eecb6..d62aef2d03b 100644
--- a/tensorflow/core/kernels/ctc_decoder_ops.cc
+++ b/tensorflow/core/kernels/ctc_decoder_ops.cc
@@ -154,7 +154,14 @@ class CTCDecodeHelper {
         auto& p_batch = sequences[b][p];
         int64 num_decoded = p_batch.size();
         max_decoded = std::max(max_decoded, num_decoded);
-        std::copy_n(p_batch.begin(), num_decoded, &values_t(offset));
+        if (num_decoded > 0) {
+          DCHECK_NE(values_t.data(), nullptr)
+              << "values_t should not be nullptr: p_num=" << p_num
+              << " num_decoded=" << num_decoded;
+          DCHECK_LT(offset, values_t.size())
+              << "offset should be smaller than values_t.size()";
+          std::copy_n(p_batch.begin(), num_decoded, &values_t(offset));
+        }
         for (int64 t = 0; t < num_decoded; ++t, ++offset) {
           indices_t(offset, 0) = b;
           indices_t(offset, 1) = t;
@@ -203,6 +210,7 @@ class CTCGreedyDecoderOp : public OpKernel {
 
     auto inputs_t = inputs->tensor();
 
+    input_list_t.reserve(max_time);
     for (std::size_t t = 0; t < max_time; ++t) {
       input_list_t.emplace_back(inputs_t.data() + t * batch_size * num_classes,
                                 batch_size, num_classes);
@@ -305,6 +313,7 @@ class CTCBeamSearchDecoderOp : public OpKernel {
 
     std::vector::UnalignedConstMatrix> input_list_t;
 
+    input_list_t.reserve(max_time);
     for (std::size_t t = 0; t < max_time; ++t) {
       input_list_t.emplace_back(inputs_t.data() + t * batch_size * num_classes,
                                 batch_size, num_classes);

From 0a0161836d350f62e26fbb30bd234a9017ebfd59 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 18:41:44 -0700
Subject: [PATCH 0249/1390] Require Keras Preprocessing Layers be restored from
 config.

PiperOrigin-RevId: 316589499
Change-Id: Id8f5ffcc84d60735c5d174581077007845717b78
---
 .../keras/engine/base_preprocessing_layer.py    |  1 +
 .../engine/base_preprocessing_layer_test.py     | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/tensorflow/python/keras/engine/base_preprocessing_layer.py b/tensorflow/python/keras/engine/base_preprocessing_layer.py
index e6b75033d60..08df07e33e3 100644
--- a/tensorflow/python/keras/engine/base_preprocessing_layer.py
+++ b/tensorflow/python/keras/engine/base_preprocessing_layer.py
@@ -41,6 +41,7 @@ from tensorflow.python.util.tf_export import keras_export
 class PreprocessingLayer(Layer):
   """Base class for PreprocessingLayers."""
   __metaclass__ = abc.ABCMeta
+  _must_restore_from_config = True
 
   @abc.abstractmethod
   def adapt(self, data, reset_state=True):
diff --git a/tensorflow/python/keras/engine/base_preprocessing_layer_test.py b/tensorflow/python/keras/engine/base_preprocessing_layer_test.py
index f35871a2f00..70d088cf3d3 100644
--- a/tensorflow/python/keras/engine/base_preprocessing_layer_test.py
+++ b/tensorflow/python/keras/engine/base_preprocessing_layer_test.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 import json
+import os
 
 from absl.testing import parameterized
 import numpy as np
@@ -35,6 +36,7 @@ from tensorflow.python.keras.engine import base_preprocessing_layer
 from tensorflow.python.keras.engine import base_preprocessing_layer_v1
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import variables
 from tensorflow.python.ops.ragged import ragged_factory_ops
 from tensorflow.python.platform import test
 from tensorflow.python.util import compat
@@ -349,6 +351,21 @@ class PreprocessingLayerTest(keras_parameterized.TestCase):
     layer_2.adapt(np.array([1, 2]), reset_state=False)
     self.assertAllEqual([[19], [20], [21]], model_2.predict([1., 2., 3.]))
 
+  def test_loading_without_providing_class_fails(self):
+    input_data = keras.Input(shape=(1,))
+    layer = get_layer()
+    output = layer(input_data)
+    model = keras.Model(input_data, output)
+
+    if not context.executing_eagerly():
+      self.evaluate(variables.variables_initializer(model.variables))
+
+    output_path = os.path.join(self.get_temp_dir(), "tf_keras_saved_model")
+    model.save(output_path, save_format="tf")
+
+    with self.assertRaisesRegex(RuntimeError, "Unable to restore a layer of"):
+      _ = keras.models.load_model(output_path)
+
 
 @keras_parameterized.run_all_keras_modes
 class ConvertToListTest(keras_parameterized.TestCase):

From 20bfd347119d5ff8dc3d6d58526e8bc683231660 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 18:48:16 -0700
Subject: [PATCH 0250/1390] Update ops-related pbtxt files.

PiperOrigin-RevId: 316590217
Change-Id: Id3e8cb3e4dc0e9c930d1dc457312c55f64a02470
---
 .../ops_history_v2/DebugIdentityV2.pbtxt      | 66 +++++++++++++++++++
 tensorflow/core/ops/ops.pbtxt                 |  7 ++
 2 files changed, 73 insertions(+)

diff --git a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
index 9502555a20e..ea92aaa4943 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/DebugIdentityV2.pbtxt
@@ -109,3 +109,69 @@ op {
   }
   is_stateful: true
 }
+op {
+  name: "DebugIdentityV2"
+  input_arg {
+    name: "input"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+  }
+  attr {
+    name: "tfdbg_context_id"
+    type: "string"
+    default_value {
+      s: ""
+    }
+  }
+  attr {
+    name: "op_name"
+    type: "string"
+    default_value {
+      s: ""
+    }
+  }
+  attr {
+    name: "output_slot"
+    type: "int"
+    default_value {
+      i: -1
+    }
+  }
+  attr {
+    name: "tensor_debug_mode"
+    type: "int"
+    default_value {
+      i: -1
+    }
+  }
+  attr {
+    name: "debug_urls"
+    type: "list(string)"
+    default_value {
+      list {
+      }
+    }
+  }
+  attr {
+    name: "circular_buffer_size"
+    type: "int"
+    default_value {
+      i: 1000
+    }
+  }
+  attr {
+    name: "tfdbg_run_id"
+    type: "string"
+    default_value {
+      s: ""
+    }
+  }
+  is_stateful: true
+}
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index c67b347af04..1d92a0671b8 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -10860,6 +10860,13 @@ op {
       i: 1000
     }
   }
+  attr {
+    name: "tfdbg_run_id"
+    type: "string"
+    default_value {
+      s: ""
+    }
+  }
   is_stateful: true
 }
 op {

From 72e98dfe9675e4b2a78203743904d97586d5775a Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Mon, 15 Jun 2020 19:06:11 -0700
Subject: [PATCH 0251/1390] Simple update for passing dynamic indexing with
 integers test

Also copy the _slice_helper logic into tf numpy so we can iterate more quickly
there.

PiperOrigin-RevId: 316592568
Change-Id: I1adb870385872fd407e49381cadcb117755f95a0
---
 tensorflow/python/ops/numpy_ops/np_arrays.py | 134 ++++++++++++++++++-
 1 file changed, 133 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index 0e320d415b6..a7696ad31c2 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -17,14 +17,135 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import numbers
 import numpy as np
 import six
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.util import nest
+
+
+_SLICE_TYPE_ERROR = (
+    'Only integers, slices (`:`), ellipsis (`...`), '
+    'tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid '
+    'indices')
+
+_SUPPORTED_SLICE_DTYPES = (dtypes.int32, dtypes.int32_ref, dtypes.int64,
+                           dtypes.int64_ref)
+
+
+def _check_index(idx):
+  """Check if a given value is a valid index into a tensor."""
+  if isinstance(idx, (numbers.Integral, tensor_shape.Dimension)):
+    return
+
+  # Optimistic check. Assumptions:
+  # * any object with a dtype is supported
+  # * any object with a dtype has a sizeable shape attribute.
+  dtype = getattr(idx, 'dtype', None)
+  if (dtype is None or dtypes.as_dtype(dtype) not in _SUPPORTED_SLICE_DTYPES or
+      idx.shape and len(idx.shape) == 1):
+    # TODO(slebedev): IndexError seems more appropriate here, but it
+    # will break `_slice_helper` contract.
+    raise TypeError(_SLICE_TYPE_ERROR + ', got {!r}'.format(idx))
+
+
+def _is_undefined_dimension(d):
+  return isinstance(d, tensor_shape.Dimension) and d.value is None
+
+
+def _slice_helper(tensor, slice_spec, var=None):
+  """Copied from array_ops._slice_helper, will be merged back later."""
+  if isinstance(slice_spec, bool) or \
+  (isinstance(slice_spec, ops.Tensor) and slice_spec.dtype == dtypes.bool) or \
+  (isinstance(slice_spec, np.ndarray) and slice_spec.dtype == bool):
+    return array_ops.boolean_mask(tensor=tensor, mask=slice_spec)
+
+  if not isinstance(slice_spec, (list, tuple)):
+    slice_spec = [slice_spec]
+
+  begin, end, strides = [], [], []
+  index = 0
+
+  new_axis_mask, shrink_axis_mask = 0, 0
+  begin_mask, end_mask = 0, 0
+  ellipsis_mask = 0
+  for s in slice_spec:
+    if isinstance(s, slice):
+      if s.start is not None and not _is_undefined_dimension(s.start):
+        _check_index(s.start)
+        begin.append(s.start)
+      else:
+        begin.append(0)
+        begin_mask |= (1 << index)
+      if s.stop is not None and not _is_undefined_dimension(s.stop):
+        _check_index(s.stop)
+        end.append(s.stop)
+      else:
+        end.append(0)
+        end_mask |= (1 << index)
+      if s.step is not None and not _is_undefined_dimension(s.step):
+        _check_index(s.step)
+        strides.append(s.step)
+      else:
+        strides.append(1)
+    elif s is Ellipsis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      ellipsis_mask |= (1 << index)
+    elif s is array_ops.newaxis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      new_axis_mask |= (1 << index)
+    else:
+      _check_index(s)
+      begin.append(s)
+      end.append(s + 1)
+      strides.append(1)
+      shrink_axis_mask |= (1 << index)
+    index += 1
+
+  # stack possibly involves no tensors, so we must use op_scope correct graph.
+  with ops.name_scope(
+      None,
+      'strided_slice', [tensor] + begin + end + strides,
+      skip_on_eager=False) as name:
+    if begin:
+      packed_begin, packed_end, packed_strides = (array_ops.stack(begin),
+                                                  array_ops.stack(end),
+                                                  array_ops.stack(strides))
+      if (packed_begin.dtype == dtypes.int64 or
+          packed_end.dtype == dtypes.int64 or
+          packed_strides.dtype == dtypes.int64):
+        if packed_begin.dtype != dtypes.int64:
+          packed_begin = math_ops.cast(packed_begin, dtypes.int64)
+        if packed_end.dtype != dtypes.int64:
+          packed_end = math_ops.cast(packed_end, dtypes.int64)
+        if packed_strides.dtype != dtypes.int64:
+          packed_strides = math_ops.cast(packed_strides, dtypes.int64)
+    else:
+      var_empty = constant_op.constant([], dtype=dtypes.int32)
+      packed_begin = packed_end = packed_strides = var_empty
+    return array_ops.strided_slice(
+        tensor,
+        packed_begin,
+        packed_end,
+        packed_strides,
+        begin_mask=begin_mask,
+        end_mask=end_mask,
+        shrink_axis_mask=shrink_axis_mask,
+        new_axis_mask=new_axis_mask,
+        ellipsis_mask=ellipsis_mask,
+        var=var,
+        name=name)
 
 
 def convert_to_tensor(value, dtype=None, dtype_hint=None):
@@ -184,7 +305,18 @@ class ndarray(object):  # pylint: disable=invalid-name
 
   def __getitem__(self, slice_spec):
     # TODO(srbs): Need to support better indexing.
-    result_t = self.data.__getitem__(slice_spec)
+    def _gettensor(x):
+      if isinstance(x, ndarray):
+        x = x.data
+      if isinstance(x, ops.Tensor) and x.dtype not in (
+          dtypes.int32, dtypes.int64):
+        # Currently _slice_helper will only work with int32/int64 tensors, but
+        # type inference by numpy can create {u,}int{8,16}, so just cast.
+        x = math_ops.cast(x, dtypes.int32)
+      return x
+    slice_spec = nest.map_structure(_gettensor, slice_spec)
+
+    result_t = _slice_helper(self.data, slice_spec)
     return tensor_to_ndarray(result_t)
 
   def __iter__(self):

From 43de71078e7a1e1b8267937448da107d214cd81c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 19:27:05 -0700
Subject: [PATCH 0252/1390] More compatibility fixes for typing.Generic:  *
 types.new_class is required in some distributions  * avoid calling isinstance
 on some function objects in python 3.6 Required for #40132.

PiperOrigin-RevId: 316594891
Change-Id: I7c67bc04d6cfe4706a85be5b9c2271b4a4f0b97b
---
 tensorflow/python/framework/test_util.py |  7 -------
 tensorflow/python/util/tf_should_use.py  | 18 +++++-------------
 2 files changed, 5 insertions(+), 20 deletions(-)

diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index 950e17d0d8c..2967bb3de84 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -33,7 +33,6 @@ import tempfile
 import threading
 import time
 import unittest
-import weakref
 
 from absl.testing import parameterized
 import numpy as np
@@ -733,12 +732,6 @@ def assert_no_new_tensors(f):
     """Finds existing Tensors, runs the test, checks for new Tensors."""
 
     def _is_tensorflow_object(obj):
-      if isinstance(obj, weakref.ReferenceType):
-        obj = obj()
-        if obj is None:
-          return False
-      if not hasattr(obj, "__class__"):
-        return False
       try:
         return isinstance(obj,
                           (ops.Tensor, variables.Variable,
diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py
index 41c3220f5ca..1671b078fa3 100644
--- a/tensorflow/python/util/tf_should_use.py
+++ b/tensorflow/python/util/tf_should_use.py
@@ -21,12 +21,15 @@ import copy
 import sys
 import textwrap
 import traceback
-import types
+
+import six  # pylint: disable=unused-import
+
 
 from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util import tf_decorator
+# pylint: enable=g-bad-import-order,g-import-not-at-top
 
 
 class _TFShouldUseHelper(object):
@@ -151,18 +154,7 @@ def _get_wrapper(x, tf_should_use_helper):
   tx = copy.deepcopy(type_x)
   # Prefer using __orig_bases__, which preserve generic type arguments.
   bases = getattr(tx, '__orig_bases__', tx.__bases__)
-
-  # Use types.new_class when available, which is preferred over plain type in
-  # some distributions.
-  if sys.version_info >= (3, 5):
-    def set_body(ns):
-      ns.update(tx.__dict__)
-      return ns
-
-    copy_tx = types.new_class(tx.__name__, bases, exec_body=set_body)
-  else:
-    copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
-
+  copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
   copy_tx.__init__ = _new__init__
   copy_tx.__getattribute__ = _new__getattribute__
   copy_tx.mark_used = _new_mark_used

From 71d1fe5f14f1fec658fde5bcee83fa486ca05355 Mon Sep 17 00:00:00 2001
From: Yunlu Li 
Date: Mon, 15 Jun 2020 19:31:43 -0700
Subject: [PATCH 0253/1390] Make model sparsification work for Conv,
 DepthwiseConv and TransposeConv.

PiperOrigin-RevId: 316595363
Change-Id: Iad873c74599eda46b6699b0a3cd479209b8a4ca7
---
 tensorflow/compiler/mlir/lite/ir/tfl_ops.td   | 20 ++++++-
 .../mlir/lite/transforms/dense_to_sparse.cc   | 55 ++++++++++---------
 2 files changed, 48 insertions(+), 27 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
index ab48a0f8f92..509c13ae161 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
@@ -436,7 +436,7 @@ class TFL_Op traits = []> :
 class TFL_ConvOp :
     TFL_Op,
     TFL_ChannelDimIndexInterface, AffineOpCoefficient,
-    TFL_GpuTargetOp]> {
+    TFL_GpuTargetOp, TFL_SparseOp]> {
   let summary = opSummary # " operator";
 
   let description = [{
@@ -571,7 +571,8 @@ def TFL_TransposeConvOp: TFL_Op<"transpose_conv", [
     TFL_OperandHasRank<2, 4>,
     PredOpTrait<"input and output must have same element type",
       TFL_TCresVTEtIsSameAsOp<0, 2>>,
-    TFL_GpuTargetOp]> {
+    TFL_GpuTargetOp,
+    TFL_SparseOp]> {
   let summary = "Transpose convolution operator";
 
   let description = [{
@@ -593,6 +594,13 @@ def TFL_TransposeConvOp: TFL_Op<"transpose_conv", [
   let hasOptions = 1;
 
   let verifier = [{ return Verify(*this); }];
+
+  let extraClassDeclaration = [{
+    // SparseOpInterface:
+    std::vector GetSparseOperands() { return {1}; }
+    std::vector> GetFloatBlockSize() { return {}; }
+    std::vector> GetQuantizedBlockSize() { return {}; }
+  }];
 }
 
 def TFL_AveragePool2DOp:
@@ -826,6 +834,10 @@ def TFL_Conv2DOp : TFL_ConvOp<"conv_2d", "Convolution", 0> {
   let extraClassDeclaration = [{
     // ChannelDimIndexInterface:
     int GetChannelDimIndex() { return 0; }
+    // SparseOpInterface:
+    std::vector GetSparseOperands() { return {1}; }
+    std::vector> GetFloatBlockSize() { return {}; }
+    std::vector> GetQuantizedBlockSize() { return {}; }
   }];
 }
 
@@ -866,6 +878,10 @@ def TFL_DepthwiseConv2DOp :
   let extraClassDeclaration = [{
     // ChannelDimIndexInterface:
     int GetChannelDimIndex() { return 3; }
+    // SparseOpInterface:
+    std::vector GetSparseOperands() { return {1}; }
+    std::vector> GetFloatBlockSize() { return {}; }
+    std::vector> GetQuantizedBlockSize() { return {}; }
   }];
 }
 
diff --git a/tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc b/tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc
index 9b526f40277..f5ef2585be5 100644
--- a/tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc
@@ -40,14 +40,22 @@ void PopulateEncodingParams(const std::vector& block_size,
                             std::vector* traversal_order,
                             std::vector* format,
                             std::vector* b_map, std::vector* b_size) {
-  *traversal_order = {0, 1};
-  *format = {kTfLiteDimDense, kTfLiteDimSparseCSR};
+  const int dims_count = block_size.size();
+  traversal_order->resize(dims_count);
+  format->resize(dims_count);
+  for (int i = 0; i < dims_count; i++) {
+    (*traversal_order)[i] = i;
+  }
+  for (int i = 0; i < dims_count - 1; i++) {
+    (*format)[i] = kTfLiteDimDense;
+  }
+  (*format)[dims_count - 1] = kTfLiteDimSparseCSR;
   *b_map = {};
   *b_size = {};
   int block_rank = 0;
-  for (int i = 0; i < 2; i++) {
+  for (int i = 0; i < dims_count; i++) {
     if (block_size[i] != 1) {
-      traversal_order->push_back(block_rank + 2);
+      traversal_order->push_back(block_rank + dims_count);
       format->push_back(kTfLiteDimDense);
       block_rank++;
       b_map->push_back(i);
@@ -58,27 +66,18 @@ void PopulateEncodingParams(const std::vector& block_size,
 
 float CalculateRandomSparsity(const ElementsAttr& attr,
                               const ShapedType& type) {
-  int num_elements = 1;
-  for (int i = 0; i < 2; i++) {
-    num_elements *= type.getDimSize(i);
-  }
+  int num_elements = type.getNumElements();
   int num_zeros = 0;
 
   if (type.getElementType().isF32()) {
-    std::vector data;
-    data.reserve(type.getNumElements());
-    for (const auto val : attr.getValues()) data.push_back(val);
-    for (int i = 0; i < data.size(); i++) {
-      if (data[i] == 0) {
+    for (const auto val : attr.getValues()) {
+      if (val == 0.f) {
         num_zeros++;
       }
     }
   } else if (type.getElementType().isa()) {
-    std::vector data;
-    data.reserve(type.getNumElements());
-    for (const auto val : attr.getValues()) data.push_back(val);
-    for (int i = 0; i < data.size(); i++) {
-      if (data[i] == 0) {
+    for (const auto val : attr.getValues()) {
+      if (val == 0) {
         num_zeros++;
       }
     }
@@ -150,9 +149,10 @@ InspectResult InspectWeight(
     type = cst.getType().cast();
   }
 
-  // TODO(b/147449640): Add ability to encode weights more than 2-D, e.g. Conv
-  // weights.
-  if (type.getRank() != 2) {
+  // Currently we only support compressing weights of ops:
+  //   Conv, DepthwiseConv, TransposeConv, whose filter has rank 4, and
+  //   FullyConnected, whose filter has rank 2.
+  if (type.getRank() != 2 && type.getRank() != 4) {
     result.can_compress = false;
     return result;
   }
@@ -195,9 +195,11 @@ std::vector BuildSparsityParameterAttribute(
     attr = cst.value();
     type = cst.getType().cast();
   }
-  std::vector shape(2);
-  shape[0] = type.getDimSize(0);
-  shape[1] = type.getDimSize(1);
+  const int dims_count = type.getRank();
+  std::vector shape(dims_count);
+  for (int i = 0; i < dims_count; i++) {
+    shape[i] = type.getDimSize(i);
+  }
 
   std::vector traversal_order = {};
   std::vector format = {};
@@ -271,10 +273,13 @@ void DenseToSparse::runOnFunction() {
         continue;
       }
 
+      ShapedType type;
       if (isa(inst)) {
         supported_block_size = sparse_op.GetFloatBlockSize();
+        type = dyn_cast(inst).getType().cast();
       } else if (isa(inst)) {
         supported_block_size = sparse_op.GetQuantizedBlockSize();
+        type = dyn_cast(inst).getType().cast();
       } else {
         continue;
       }
@@ -286,7 +291,7 @@ void DenseToSparse::runOnFunction() {
 
       // The weight is not block sparse. Encode with random sparsity.
       if (result.selected_block_size.empty()) {
-        result.selected_block_size = {1, 1};
+        result.selected_block_size = std::vector(type.getRank(), 1);
       }
 
       builder.setInsertionPoint(op);

From 552604a2efb050dc5f79f75d2e9c286ba7725c5c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 19:46:08 -0700
Subject: [PATCH 0254/1390] Go: Update generated wrapper functions for
 TensorFlow ops.

PiperOrigin-RevId: 316596872
Change-Id: I3eed9544a8c70c2c7d1ffe5307ba9d37890a8d8c
---
 tensorflow/go/op/wrappers.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 10acebc7965..485baa16f39 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -32348,11 +32348,11 @@ func StridedSliceShrinkAxisMask(value int64) StridedSliceAttr {
 // begin = [1, 2, x, x, 0, x] # x denotes don't care (usually 0)
 // end = [2, 4, x, x, -3, x]
 // strides = [1, 1, x, x, -1, 1]
-// begin_mask = 1<<4 | 1 << 5 = 48
+// begin_mask = 1<<4 | 1<<5 = 48
 // end_mask = 1<<5 = 32
 // ellipsis_mask = 1<<3 = 8
-// new_axis_mask = 1<<2 4
-// shrink_axis_mask = 1<<0
+// new_axis_mask = 1<<2 = 4
+// shrink_axis_mask = 1<<0 = 1
 // ```
 //
 // In this case if `foo.shape` is (5, 5, 5, 5, 5, 5) the final shape of

From c118bc4b5e1bb3a2760a026621d7a19053b0626c Mon Sep 17 00:00:00 2001
From: Prakalp Srivastava 
Date: Mon, 15 Jun 2020 19:50:34 -0700
Subject: [PATCH 0255/1390] Fix decomposition of ResourceApplyCenteredRMSProp
 op pattern.

AssignSubVariableOp only needs the resource handle and the tensor to be subtracted from it.

PiperOrigin-RevId: 316597384
Change-Id: I26a99cb6fe210cc66429ebca13765af3109748ae
---
 .../mlir/tensorflow/tests/decompose_resource_ops.mlir        | 1 +
 .../mlir/tensorflow/transforms/decompose_resource_ops.td     | 5 +----
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
index 25dfda25358..ff4dbf41221 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir
@@ -411,6 +411,7 @@ func @decompose_resource_apply_centered_RMS_prop(%arg0: tensor, %arg1: tens
 
   // CHECK: [[VAR:%.*]] = "tf.ReadVariableOp"([[VAR_HANDLE]])
   // CHECK: [[VAR_NEW:%.*]] = "tf.Sub"([[VAR]], [[MOM_NEW]])
+  // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[VAR_NEW]])
 
   "tf.ResourceApplyCenteredRMSProp"(%0, %1, %2, %3, %arg4, %arg5, %arg6, %arg7, %arg8) {use_locking = false} : (tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor<*x!tf.resource>, tensor, tensor, tensor, tensor, tensor) -> ()
   return
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
index 0dd7d778e31..40339cebd31 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/decompose_resource_ops.td
@@ -387,9 +387,6 @@ def DecomposeResourceApplyCenteredRMSProp :
      ),
      (TF_AssignVariableOp $mom_resource, $mom_new),
      // var <- var - mom
-     (TF_AssignSubVariableOp $var_resource,
-        (TF_SubOp (CreateTFReadVariableOp $src_op, $grad, $var_resource),
-                  $mom_new)
-     )
+     (TF_AssignSubVariableOp $var_resource, $mom_new)
    ]
    >;

From 5de12f4d96c48898f58495f81156651c64eed275 Mon Sep 17 00:00:00 2001
From: Francois Chollet 
Date: Mon, 15 Jun 2020 20:32:58 -0700
Subject: [PATCH 0256/1390] Update tf.keras version number.

PiperOrigin-RevId: 316602175
Change-Id: I94fa651c3005f578abafcafe86382a2dc75451e6
---
 tensorflow/python/keras/__init__.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/tensorflow/python/keras/__init__.py b/tensorflow/python/keras/__init__.py
index 6f79c219867..47f207329d7 100644
--- a/tensorflow/python/keras/__init__.py
+++ b/tensorflow/python/keras/__init__.py
@@ -32,13 +32,6 @@ from tensorflow.python.keras.engine.training import Model
 
 from tensorflow.python.util.tf_export import keras_export
 
-if tf2.enabled():
-  __version__ = '2.3.0-tf'
-else:
-  __version__ = '2.2.4-tf'
+__version__ = '2.4.0'
 
 keras_export('keras.__version__').export_constant(__name__, '__version__')
-
-del absolute_import
-del division
-del print_function

From 13d59c2d1fccecc6343965cef89464229d00db21 Mon Sep 17 00:00:00 2001
From: Chen Lei 
Date: Tue, 16 Jun 2020 11:50:49 +0800
Subject: [PATCH 0257/1390] Update stale.yml

---
 .github/stale.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/stale.yml b/.github/stale.yml
index 5f8dd12f477..7eef5309ecd 100644
--- a/.github/stale.yml
+++ b/.github/stale.yml
@@ -23,7 +23,7 @@
 daysUntilStale: 7
 # Number of days of inactivity before a stale Issue or Pull Request is closed
 daysUntilClose: 7
-# Only issues or pull requests with all of these labels are check if stale. Defaults to `[]` (disabled)
+# Only issues or pull requests with all of these labels are checked if stale. Defaults to `[]` (disabled)
 onlyLabels:
  - stat:awaiting response
 # Comment to post when marking as stale. Set to `false` to disable

From e0b43845d711e9dc520b9b6716ff89c3b4cd631f Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Tue, 16 Jun 2020 03:51:57 +0000
Subject: [PATCH 0258/1390] Sign compare warning fixes batch 2

---
 .../core/grappler/costs/graph_memory.cc       |  5 +-
 .../core/grappler/costs/graph_properties.cc   | 59 +++++++++++--------
 .../grappler/costs/op_level_cost_estimator.cc |  8 +--
 .../core/grappler/costs/virtual_scheduler.cc  | 14 +++--
 .../optimizers/common_subgraph_elimination.cc |  2 +-
 .../grappler/optimizers/debug_stripper.cc     |  2 +-
 .../grappler/optimizers/function_optimizer.cc | 10 ++--
 .../core/grappler/optimizers/model_pruner.cc  |  3 +-
 .../optimizers/pin_to_host_optimizer.cc       |  3 +-
 .../grappler/optimizers/shape_optimizer.cc    |  3 +-
 tensorflow/core/grappler/utils.cc             |  2 +-
 tensorflow/core/grappler/utils/graph_view.cc  | 37 +++++++-----
 .../core/grappler/utils/graph_view_internal.h | 34 +++++++----
 .../core/grappler/utils/topological_sort.cc   |  4 +-
 .../sql/sqlite_query_connection.cc            |  2 +-
 tensorflow/python/grappler/model_analyzer.cc  |  4 +-
 16 files changed, 113 insertions(+), 79 deletions(-)

diff --git a/tensorflow/core/grappler/costs/graph_memory.cc b/tensorflow/core/grappler/costs/graph_memory.cc
index 020e8cf1d1f..768a025b0e6 100644
--- a/tensorflow/core/grappler/costs/graph_memory.cc
+++ b/tensorflow/core/grappler/costs/graph_memory.cc
@@ -255,7 +255,8 @@ void GraphMemory::InferFromTrace(const StepStats& timeline) {
     std::unordered_set live_at_peak;
     size_t current = 0;
     std::unordered_set currently_live;
-    for (int i = 0; i < events.size(); ++i) {
+    int events_size = events.size();
+    for (int i = 0; i < events_size; ++i) {
       const auto& event = events[i];
 
       if (event.allocated) {
@@ -271,7 +272,7 @@ void GraphMemory::InferFromTrace(const StepStats& timeline) {
         current -= event.tensor->memory_used;
         currently_live.erase(event.tensor);
       }
-      if (i + 1 == events.size() ||
+      if (i + 1 == events_size ||
           event.timestamp != events[i + 1].timestamp) {
         if (current > peak) {
           peak = current;
diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index ee691e7a081..0c14607e9e2 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -363,7 +363,7 @@ void VerboseLogUnknownDimensionSources(
 std::vector ReplaceUnknownDimFromConstWithUnknownDim(
     InferenceContext* ic, const std::vector& shapes) {
   std::vector converted_shapes(shapes.size());
-  for (int i = 0; i < shapes.size(); i++) {
+  for (int i = 0, shapes_size = shapes.size(); i < shapes_size; i++) {
     const auto& shape = shapes[i];
     if (!ic->RankKnown(shape)) {
       converted_shapes[i] = shape;
@@ -502,7 +502,7 @@ class TopoQueue {
       const std::vector& topo_order) const {
     absl::flat_hash_map map;
     map.reserve(topo_order.size());
-    for (int i = 0; i < topo_order.size(); ++i) {
+    for (int i = 0, topo_order_size = topo_order.size(); i < topo_order_size; ++i) {
       map.emplace(topo_order[i], i);
     }
     return map;
@@ -680,14 +680,16 @@ class SymbolicShapeRefiner {
                         ", shape: ", ic->DebugString(ic->input(i)),
                         ", tensor: ");
         Tensor t1;
-        if (input_tensor_protos.size() > i &&
+        int input_tensor_protos_size = input_tensor_protos.size();
+        if (input_tensor_protos_size > i &&
             input_tensor_protos.at(i) != nullptr &&
             t1.FromProto(*input_tensor_protos.at(i))) {
           absl::StrAppend(&output, t1.DebugString(), ", tensor_as_shape: ");
         } else {
           absl::StrAppend(&output, " null, tensor_as_shape: ");
         }
-        if (input_tensors_as_shapes_to_propagate.size() > i) {
+        int input_tensors_as_shapes_to_propagate_size = input_tensors_as_shapes_to_propagate.size();
+        if (input_tensors_as_shapes_to_propagate_size > i) {
           absl::StrAppend(
               &output,
               StringifyShapeHandle(input_tensors_as_shapes_to_propagate.at(i)),
@@ -702,14 +704,16 @@ class SymbolicShapeRefiner {
                         ", shape: ", ic->DebugString(ic->output(i)),
                         ", tensor: ");
         Tensor t2;
-        if (output_tensor_protos.size() > i &&
+        int output_tensor_protos_size = output_tensor_protos.size();
+        if (output_tensor_protos_size > i &&
             output_tensor_protos.at(i) != nullptr &&
             t2.FromProto(*output_tensor_protos.at(i))) {
           absl::StrAppend(&output, t2.DebugString(), ", tensor_as_shape: ");
         } else {
           absl::StrAppend(&output, " null, tensor_as_shape: ");
         }
-        if (output_tensors_as_shapes.size() > i) {
+        int output_tensors_as_shapes_size = output_tensors_as_shapes.size();
+        if (output_tensors_as_shapes_size > i) {
           absl::StrAppend(&output,
                           StringifyShapeHandle(output_tensors_as_shapes.at(i)),
                           "\n");
@@ -779,7 +783,7 @@ class SymbolicShapeRefiner {
     MutableGraphView gv(&grappler_function_item.graph);
 
     // Forward shapes from function input nodes to argument nodes.
-    for (int i = 0; i < grappler_function_item.inputs().size(); ++i) {
+    for (int i = 0, iter_limit = grappler_function_item.inputs().size(); i < iter_limit; ++i) {
       auto& fun_input = grappler_function_item.input(i);
       NodeDef* fun_node = gv.GetNode(fun_input.node_name);
       const TensorId input_tensor = ParseTensorName(function_node->input(i));
@@ -858,13 +862,13 @@ class SymbolicShapeRefiner {
       if (IsConstant(*input_node)) {
         TF_CHECK_OK(
             ReplaceInputWithConst(*input_node, i, &grappler_function_item));
-      } else if (ctx->input_tensor_protos.size() > i &&
+      } else if (static_cast(ctx->input_tensor_protos.size()) > i &&
                  ctx->input_tensor_protos[i] != nullptr) {
         NodeDef const_input_node = MakeConstNodeDefFromTensorProto(
             ic, *ctx->input_tensor_protos[i], ctx->input_types[i]);
         TF_CHECK_OK(ReplaceInputWithConst(const_input_node, i,
                                           &grappler_function_item));
-      } else if (ic->input_tensors_as_shapes().size() > i &&
+      } else if (static_cast(ic->input_tensors_as_shapes().size()) > i &&
                  IsShapeFullyDefinedIntegerVectorOrScalar(
                      ic, ic->input(i), ic->input_tensors_as_shapes()[i],
                      ctx->input_types[i])) {
@@ -912,7 +916,8 @@ class SymbolicShapeRefiner {
       }
 
       auto output_properties = gp.GetOutputProperties(retnode->name());
-      if (out_tensor.index() >= output_properties.size()) {
+      int output_properties_size = output_properties.size();
+      if (out_tensor.index() >= output_properties_size) {
         return errors::InvalidArgument(
             out_tensor.ToString(), " has invalid position ", out_tensor.index(),
             " (output_properties.size() = ", output_properties.size(), ").");
@@ -975,12 +980,12 @@ class SymbolicShapeRefiner {
       // NodeContext:
       // output_tensor_protos to input_tensor_protos and input_tensors, and
       // output_tensors_as_shapes to input_tensors_as_shapes.
-      if (src_ctx->output_tensors_as_shapes.size() > src_output) {
+      if (static_cast(src_ctx->output_tensors_as_shapes.size()) > src_output) {
         ctx->input_tensors_as_shapes_to_propagate[dst_input] =
             src_ctx->output_tensors_as_shapes[src_output];
       }
 
-      if (src_ctx->output_tensor_protos.size() > src_output) {
+      if (static_cast(src_ctx->output_tensor_protos.size()) > src_output) {
         const auto* tensor_proto = src_ctx->output_tensor_protos[src_output];
         if (tensor_proto != nullptr) {
           ctx->input_tensor_protos[dst_input] = tensor_proto;
@@ -1233,7 +1238,7 @@ class SymbolicShapeRefiner {
     if (st1.size() != st2.size()) {
       return false;
     }
-    for (int i = 0; i < st1.size(); ++i) {
+    for (int i = 0, st1_size = st1.size(); i < st1_size; ++i) {
       const ShapeAndType& s1 = st1[i];
       const ShapeAndType& s2 = st2[i];
       if (s1.dtype != s2.dtype) {
@@ -1268,13 +1273,13 @@ class SymbolicShapeRefiner {
       return Status::OK();
     }
 
-    if (grappler_function_item.inputs().size() > function_node->input_size()) {
+    if (static_cast(grappler_function_item.inputs().size()) > function_node->input_size()) {
       return errors::FailedPrecondition(
           "Function input size should be smaller than node input size.");
     }
 
-    for (int i = grappler_function_item.inputs().size();
-         i < function_node->input_size(); ++i) {
+    for (int i = grappler_function_item.inputs().size(), iter_limit = function_node->input_size();
+         i < iter_limit; ++i) {
       const string& input = function_node->input(i);
       if (!IsControlInput(input)) {
         return errors::FailedPrecondition(
@@ -1357,18 +1362,20 @@ class SymbolicShapeRefiner {
   // Returns true if all the output tensors have known values.
   bool AllOutputValuesKnown(NodeContext* c) {
     InferenceContext* ic = c->inference_context.get();
-    if (c->output_tensors_as_shapes.size() < ic->num_outputs() &&
-        c->output_tensor_protos.size() < ic->num_outputs()) {
+    int c_output_tensors_as_shapes_size = c->output_tensors_as_shapes.size();
+    int c_output_tensor_protos_size = c->output_tensor_protos.size();
+    if (c_output_tensors_as_shapes_size < ic->num_outputs() &&
+        c_output_tensor_protos_size < ic->num_outputs()) {
       return false;
     } else {
       // Checks if we can get output value via either output_tensor_proto or
       // output_tensors_as_shapes.
       for (int i = 0; i < ic->num_outputs(); i++) {
-        if (c->output_tensor_protos.size() > i &&
+        if (c_output_tensor_protos_size > i &&
             c->output_tensor_protos[i] != nullptr) {
           continue;
         }
-        if (c->output_tensors_as_shapes.size() > i &&
+        if (c_output_tensors_as_shapes_size > i &&
             ic->FullyDefined(c->output_tensors_as_shapes[i])) {
           bool no_unknown_dim_from_const = true;
           for (int32 j = 0; j < ic->Rank(c->output_tensors_as_shapes[i]); ++j) {
@@ -1539,7 +1546,7 @@ class SymbolicShapeRefiner {
                                     &resource_mgr_, &outputs));
     c->output_tensors_as_shapes.resize(outputs.size());
     c->output_tensor_protos.resize(outputs.size(), nullptr);
-    for (int k = 0; k < outputs.size(); k++) {
+    for (int k = 0, outputs_size = outputs.size(); k < outputs_size; k++) {
       const auto& t = outputs[k];
       // Override output shape.
       ShapeHandle output_shape;
@@ -2297,7 +2304,7 @@ Status GraphProperties::UpdateEnqueue(
 
   // TODO(bsteiner): handle EnqueueMany as well.
   std::vector shapes_and_types;
-  for (int i = 1; i < ctx->input_types.size(); ++i) {
+  for (int i = 1, iter_limit = ctx->input_types.size(); i < iter_limit; ++i) {
     GraphView::InputPort inp(enqueue_node, i);
     GraphView::OutputPort fanin = shape_refiner->graph().GetRegularFanin(inp);
     InferenceContext* in = shape_refiner->GetContext(fanin.node);
@@ -2490,10 +2497,10 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds,
             const TensorProto& raw_val =
                 fanin.node->attr().at("value").tensor();
             *input_properties[i].mutable_value() = raw_val;
-          } else if (ctx->input_tensor_protos.size() > i &&
+          } else if (static_cast(ctx->input_tensor_protos.size()) > i &&
                      ctx->input_tensor_protos[i] != nullptr) {
             *input_properties[i].mutable_value() = *ctx->input_tensor_protos[i];
-          } else if (ic->input_tensors_as_shapes().size() > i &&
+          } else if (static_cast(ic->input_tensors_as_shapes().size()) > i &&
                      IsShapeFullyDefinedIntegerVectorOrScalar(
                          ic, ic->input(i), ic->input_tensors_as_shapes()[i],
                          ctx->input_types[i])) {
@@ -2525,11 +2532,11 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds,
             // TODO(rmlarsen): Eliminate this copy.
             const TensorProto& raw_val = node.attr().at("value").tensor();
             *output_properties[i].mutable_value() = raw_val;
-          } else if (ctx->output_tensor_protos.size() > i &&
+          } else if (static_cast(ctx->output_tensor_protos.size()) > i &&
                      ctx->output_tensor_protos[i] != nullptr) {
             *output_properties[i].mutable_value() =
                 *ctx->output_tensor_protos[i];
-          } else if (converted_output_tensors_as_shapes.size() > i &&
+          } else if (static_cast(converted_output_tensors_as_shapes.size()) > i &&
                      IsShapeFullyDefinedIntegerVectorOrScalar(
                          ic, ic->output(i),
                          converted_output_tensors_as_shapes[i],
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index b8b62cbd6e5..a62359025be 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -1470,8 +1470,8 @@ Costs OpLevelCostEstimator::PredictEinsum(const OpContext& op_context) const {
                          (a_input.shape().dim_size() < matrix_rank) ||
                          (b_input.shape().dim_size() < matrix_rank);
 
-  if (a_input_str.size() != a_input_shape.dim_size() ||
-      b_input_str.size() != b_input_shape.dim_size()) {
+  if (a_input_str.size() != static_cast(a_input_shape.dim_size()) ||
+      b_input_str.size() != static_cast(b_input_shape.dim_size())) {
     VLOG(1) << "Missing accurate estimator for op: " << op_info.op()
             << ", equation subscripts don't match tensor rank.";
     return PredictCostOfAnUnknownOp(op_context);
@@ -1513,7 +1513,7 @@ Costs OpLevelCostEstimator::PredictEinsum(const OpContext& op_context) const {
   n_dim.set_size(1);
   k_dim.set_size(1);
 
-  for (int i_idx = 0; i_idx < a_input_str.size(); ++i_idx) {
+  for (int i_idx = 0, a_input_str_size = a_input_str.size(); i_idx < a_input_str_size; ++i_idx) {
     if (b_input_str.find(a_input_str[i_idx]) == std::string::npos) {
       if (rhs_str.find(a_input_str[i_idx]) == std::string::npos) {
         VLOG(1) << "Missing accurate estimator for op: " << op_info.op();
@@ -1533,7 +1533,7 @@ Costs OpLevelCostEstimator::PredictEinsum(const OpContext& op_context) const {
     *(a_matrix_shape->add_dim()) = a_input_shape.dim(i_idx);
     *(b_matrix_shape->add_dim()) = a_input_shape.dim(i_idx);
   }
-  for (int i_idx = 0; i_idx < b_input_str.size(); ++i_idx) {
+  for (int i_idx = 0, b_input_str_size = b_input_str.size(); i_idx < b_input_str_size; ++i_idx) {
     if (a_input_str.find(b_input_str[i_idx]) == std::string::npos) {
       if (rhs_str.find(b_input_str[i_idx]) == std::string::npos) {
         VLOG(1) << "Missing accurate estimator for op: " << op_info.op();
diff --git a/tensorflow/core/grappler/costs/virtual_scheduler.cc b/tensorflow/core/grappler/costs/virtual_scheduler.cc
index 5339b00627e..b20fad8b41c 100644
--- a/tensorflow/core/grappler/costs/virtual_scheduler.cc
+++ b/tensorflow/core/grappler/costs/virtual_scheduler.cc
@@ -522,8 +522,8 @@ Status SchedulerState::Init(const GrapplerItem* item,
 
     if (IsPersistent(*curr_node)) {
       auto& device_state = device_[curr_node_device];
-      for (int port_num = 0;
-           port_num < curr_node_state.output_properties.size(); ++port_num) {
+      for (int port_num = 0, port_num_iter_limit = curr_node_state.output_properties.size();
+           port_num < port_num_iter_limit; ++port_num) {
         device_state.persistent_nodes.insert(
             std::make_pair(curr_node, port_num));
       }
@@ -795,7 +795,8 @@ void SchedulerState::GetOutputNodes(const NodeDef* node,
       // Execute a node as soon as all its inputs are ready. Merge nodes are
       // special since they run as soon as one of their inputs becomes
       // available.
-      if (output_state.num_inputs_ready == output_state.inputs.size() ||
+      int output_state_inputs_size = output_state.inputs.size();
+      if (output_state.num_inputs_ready == output_state_inputs_size ||
           IsMerge(*output_node)) {
         // This output node is now ready.
         output_state.time_ready = curr_time;
@@ -900,8 +901,9 @@ std::vector SchedulerState::MarkNodeExecuted(
     auto port = input_port.second;
     auto& input_state = node_map_[input];
     input_state.num_outputs_executed[port]++;
-    if (input_state.num_outputs_executed[port] ==
-            input_state.outputs[port].size() &&
+    int input_state_outputs_size_ = input_state.outputs[port].size();
+    if (input_state.num_outputs_executed[port] == input_state_outputs_size_
+             &&
         !IsPersistent(*input)) {
       // All the outputs are executed; no reference to this output port of
       // input node.
@@ -1119,7 +1121,7 @@ void SchedulerState::GenerateRunMetadata(RunMetadata* metadata) {
       const NodeState& nodestate = node_map_.at(node_def);
       NodeExecStats* node_stats = device_stepstats->add_node_stats();
       uint64 total_output_size = 0;
-      for (int slot = 0; slot < nodestate.output_properties.size(); slot++) {
+      for (int slot = 0, slot_iter_limit = nodestate.output_properties.size(); slot < slot_iter_limit; slot++) {
         const auto& properties = nodestate.output_properties[slot];
         NodeOutput* no = node_stats->add_output();
         no->set_slot(slot);
diff --git a/tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc b/tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc
index af323e913a7..4f385797f20 100644
--- a/tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc
+++ b/tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc
@@ -73,7 +73,7 @@ class UniqueNodes {
     if (it == memoized_signatures_.end()) return;
 
     std::vector& candidates = rep_[it->second];
-    for (int i = 0; i < candidates.size(); ++i) {
+    for (int i = 0, candidates_size = candidates.size(); i < candidates_size; ++i) {
       if (candidates[i] == node) {
         std::swap(candidates[i], candidates[candidates.size() - 1]);
         candidates.resize(candidates.size() - 1);
diff --git a/tensorflow/core/grappler/optimizers/debug_stripper.cc b/tensorflow/core/grappler/optimizers/debug_stripper.cc
index d4b3bf395c3..de62e8fe6b9 100644
--- a/tensorflow/core/grappler/optimizers/debug_stripper.cc
+++ b/tensorflow/core/grappler/optimizers/debug_stripper.cc
@@ -63,7 +63,7 @@ Status DebugStripper::Optimize(Cluster* cluster, const GrapplerItem& item,
       node.mutable_attr()->swap(new_attr);
       // As Identity op only takes one input, mark redundant inputs as control
       // input.
-      for (size_t i = 1; i < node.input_size(); ++i) {
+      for (int i = 1, node_input_size = node.input_size(); i < node_input_size; ++i) {
         if (!IsControlInput(node.input(i))) {
           *node.mutable_input(i) = AsControlDependency(NodeName(node.input(i)));
         }
diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index ed3af955c13..5c703b18a6d 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -438,8 +438,8 @@ bool HasUnusedOutputs(const NodeDef& func_node, const FunctionDef& func,
   int num_outputs = func.signature().output_arg_size();
   const absl::flat_hash_set active_outputs =
       GetActiveOutputs(func_node, ctx, /*size_hind*/ num_outputs);
-
-  return active_outputs.size() != num_outputs;
+  int active_outputs_size = active_outputs.size(); 
+  return active_outputs_size != num_outputs;
 }
 
 // Return pruned FunctionDefLibrary with functions that are reachable from
@@ -563,7 +563,8 @@ void RemoveUnusedOutputsTypes(const FunctionSpecialization& specialization,
   if (tout == nullptr || !tout->has_list()) return;
 
   // Nothing to do if all outputs are active.
-  if (specialization.active_outputs.size() == tout->list().type_size()) return;
+  int specialization_active_outputs_size = specialization.active_outputs.size();
+  if (specialization_active_outputs_size == tout->list().type_size()) return;
 
   // Clear input types for the specialized node.
   auto* attr = specialized_func_node->mutable_attr();
@@ -1142,7 +1143,8 @@ void AddFrameForwardingControlEdge(const std::vector& info,
                                    Node* caller, Graph* g) {
   // All nodes added to the graph by v2 control flow lowering and function
   // inlining are guaranteed to have control edges to nested function calls.
-  if (caller->id() >= info.size()) return;
+  int info_size = info.size();
+  if (caller->id() >= info_size ) return;
 
   // Check if a lowered node is executing inside a while loop.
   const Node* frame = info[caller->id()].frame;
diff --git a/tensorflow/core/grappler/optimizers/model_pruner.cc b/tensorflow/core/grappler/optimizers/model_pruner.cc
index 20db4360f73..634ef35ab21 100644
--- a/tensorflow/core/grappler/optimizers/model_pruner.cc
+++ b/tensorflow/core/grappler/optimizers/model_pruner.cc
@@ -401,9 +401,10 @@ Status SplitIdentityNInputs(GraphDef* graph,
     }
 
     const int num_non_control_inputs = NumNonControlInputs(*node);
+    int terminal_second_size = terminal.second.size(); 
     if (node->attr().count("T") == 0 ||
         node->attr().at("T").list().type_size() != num_non_control_inputs ||
-        terminal.second.size() >= num_non_control_inputs) {
+        terminal_second_size >= num_non_control_inputs) {
       continue;
     }
 
diff --git a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
index ec16de1294b..35d0c5b0e40 100644
--- a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
@@ -107,7 +107,8 @@ Status IsNodeOutputPortHostFriendly(const GraphView& graph,
         /*include_tensor_values=*/false));
   }
   const auto& output_properties = properties->GetOutputProperties(node.name());
-  if (port_id >= output_properties.size()) {
+  int output_properties_size = output_properties.size();
+  if (port_id >= output_properties_size) {
     LOG(WARNING) << "port_id=" << port_id
                  << " but output_properties.size()=" << output_properties.size()
                  << "\n"
diff --git a/tensorflow/core/grappler/optimizers/shape_optimizer.cc b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
index 69de1cde4ca..656c1a1db1c 100644
--- a/tensorflow/core/grappler/optimizers/shape_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
@@ -99,7 +99,8 @@ Status ShapeOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
         }
         const auto& prop =
             properties.GetOutputProperties(reduce_indices.node->name());
-        if (prop.size() <= reduce_indices.port_id) {
+        int prop_size = prop.size();
+        if (prop_size <= reduce_indices.port_id) {
           continue;
         }
         const TensorShapeProto& reduction_indices_shape =
diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index cd6b4855583..240a52d1c6b 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -357,7 +357,7 @@ void PermuteNodesInPlace(GraphDef* graph, std::vector* permutation,
     }
     permutation->swap(inv_perm);
   }
-  for (std::size_t n = 0; n + 1 < permutation->size(); ++n) {
+  for (int n = 0, permutation_size = permutation->size(); n + 1 < permutation_size; ++n) {
     while (n != (*permutation)[n]) {
       std::size_t r = (*permutation)[n];
       graph->mutable_node()->SwapElements(n, r);
diff --git a/tensorflow/core/grappler/utils/graph_view.cc b/tensorflow/core/grappler/utils/graph_view.cc
index 5a9a1cd2abb..c19e600e869 100644
--- a/tensorflow/core/grappler/utils/graph_view.cc
+++ b/tensorflow/core/grappler/utils/graph_view.cc
@@ -63,7 +63,7 @@ bool NodeView::HasFanout(const FaninView& fanout) const {
     return false;
   } else if (fanout.index() == Graph::kControlSlot) {
     return view->fanins_set_.contains({this->node(), Graph::kControlSlot});
-  } else if (fanout.index() >= view->regular_fanins_.size()) {
+  } else if (fanout.index() >= static_cast(view->regular_fanins_.size())) {
     return false;
   }
   return view->regular_fanins_[fanout.index()].node_index_ == node_index_;
@@ -152,7 +152,8 @@ Status GraphView::CheckAndAddFaninsInternal(NodeView* node_view) {
                                      Graph::kControlSlot);
       has_observed_control = true;
     } else {
-      if (fanin_node_view.regular_fanouts_by_port_.size() <
+      int fanin_node_view_regular_fanouts_by_port_size = fanin_node_view.regular_fanouts_by_port_.size();
+      if (fanin_node_view_regular_fanouts_by_port_size <
           fanin_id.index() + 1) {
         fanin_node_view.regular_fanouts_by_port_.resize(fanin_id.index() + 1);
       }
@@ -197,7 +198,7 @@ bool MutableNodeView::HasFanout(const MutableFaninView& fanout) const {
     return false;
   } else if (fanout.index() == Graph::kControlSlot) {
     return view->fanins_count_.contains({this->node(), Graph::kControlSlot});
-  } else if (fanout.index() >= view->regular_fanins_.size()) {
+  } else if (fanout.index() >= static_cast(view->regular_fanins_.size())) {
     return false;
   }
   return view->regular_fanins_[fanout.index()].node_index_ == node_index_;
@@ -279,7 +280,8 @@ void Mutation::AddMutation(
 void Mutation::RemoveNode(MutableNodeView* node) {
   auto& update_index = node->update_index_;
   if (update_index != internal::kMissingIndex) {
-    if (update_index < updated_nodes_.size() - 1) {
+    int updated_nodes_size = updated_nodes_.size();
+    if (update_index < updated_nodes_size - 1) {
       graph_view_->nodes_[updated_nodes_.back().node_index].update_index_ =
           update_index;
       std::swap(updated_nodes_[update_index], updated_nodes_.back());
@@ -574,7 +576,8 @@ void MutableGraphView::AddFaninsInternal(
           --last_pos;
         }
       } else {
-        if (fanin_node_view.regular_fanouts_by_port_.size() <
+      	int fanin_node_view_regular_fanouts_by_port_size = fanin_node_view.regular_fanouts_by_port_.size();
+        if (fanin_node_view_regular_fanouts_by_port_size <
             fanin_id.index() + 1) {
           fanin_node_view.regular_fanouts_by_port_.resize(fanin_id.index() + 1);
         }
@@ -852,8 +855,8 @@ template 
 void MutableGraphView::ReplaceNodeFanouts(MutableNodeView* node, T* fanouts) {
   node->num_regular_fanouts_ = fanouts->num_regular_fanouts_;
   node->regular_fanouts_by_port_ = std::move(fanouts->regular_fanouts_by_port_);
-  for (int i = 0; i < node->regular_fanouts_by_port_.size(); ++i) {
-    for (int j = 0; j < node->regular_fanouts_by_port_[i].size(); ++j) {
+  for (int i = 0, i_max = node->regular_fanouts_by_port_.size(); i < i_max; ++i) {
+    for (int j = 0, j_max = node->regular_fanouts_by_port_[i].size(); j < j_max; ++j) {
       auto& fanout = node->regular_fanouts_by_port_[i][j];
       auto* fanout_node_view = fanout.node_view();
       auto& fanout_fanin = fanout_node_view->regular_fanins_[fanout.index()];
@@ -868,7 +871,7 @@ void MutableGraphView::ReplaceNodeFanouts(MutableNodeView* node, T* fanouts) {
     }
   }
   node->controlled_fanouts_ = std::move(fanouts->controlled_fanouts_);
-  for (int i = 0; i < node->controlled_fanouts_.size(); ++i) {
+  for (int i = 0, i_max = node->controlled_fanouts_.size(); i < i_max; ++i) {
     auto& fanout = node->controlled_fanouts_[i];
     auto* fanout_node_view = fanout.node_view();
     auto& fanout_fanin =
@@ -1017,7 +1020,8 @@ inline void MutableGraphView::RemoveRegularFaninFanoutInternal(
                       {&graph_->node(fanin.node_index_), fanin.index()});
   auto* fanin_node_view = fanin.node_view();
   auto& fanouts = fanin_node_view->regular_fanouts_by_port_[fanin.index()];
-  if (fanin.fanout_index_ < fanouts.size() - 1) {
+  int fanouts_size = fanouts.size();
+  if (fanin.fanout_index_ < fanouts_size - 1) {
     // Swap fanout with last fanout in vector, and update it's associated fanin
     // index.
     MutableFaninView& last_fanout = fanouts.back();
@@ -1043,7 +1047,8 @@ inline void MutableGraphView::RemoveRegularFaninFanoutInternal(
       break;
     }
   }
-  if (last_fanout_index < fanin_node_view->regular_fanouts_by_port_.size()) {
+  int fanin_node_view_regular_fanouts_by_port_size = fanin_node_view->regular_fanouts_by_port_.size();
+  if (last_fanout_index < fanin_node_view_regular_fanouts_by_port_size) {
     fanin_node_view->regular_fanouts_by_port_.resize(last_fanout_index);
   }
 }
@@ -1052,7 +1057,8 @@ inline void MutableGraphView::AddRegularFaninInternal(
     MutableNodeView* node_view, const SafeTensorId& fanin_id) {
   MutableNodeView* fanin_node_view = GetNode(fanin_id.node());
   // Resize fanouts to include new output port index.
-  if (fanin_node_view->regular_fanouts_by_port_.size() < fanin_id.index() + 1) {
+  int fanin_node_view_regular_fanouts_by_port_size = fanin_node_view->regular_fanouts_by_port_.size();
+  if (fanin_node_view_regular_fanouts_by_port_size < fanin_id.index() + 1) {
     fanin_node_view->regular_fanouts_by_port_.resize(fanin_id.index() + 1);
   }
 
@@ -1078,7 +1084,8 @@ inline void MutableGraphView::UpdateRegularFaninInternal(
 
   MutableNodeView* fanin_node_view = GetNode(fanin_id.node());
   // Resize fanouts to include new output port index.
-  if (fanin_node_view->regular_fanouts_by_port_.size() < fanin_id.index() + 1) {
+  int fanin_node_view_regular_fanouts_by_port_size = fanin_node_view->regular_fanouts_by_port_.size();
+  if (fanin_node_view_regular_fanouts_by_port_size < fanin_id.index() + 1) {
     fanin_node_view->regular_fanouts_by_port_.resize(fanin_id.index() + 1);
   }
 
@@ -1110,8 +1117,9 @@ inline void MutableGraphView::RemoveControllingFaninFanoutInternal(
     // controlled fanout in controlling fanin with controlled fanout to be
     // removed.
     auto* control_to_remove_view = control_to_remove.node_view();
+    int control_to_remove_view_controlled_fanouts_size = control_to_remove_view->controlled_fanouts_.size();
     if (control_to_remove.fanout_index_ <
-        control_to_remove_view->controlled_fanouts_.size() - 1) {
+        control_to_remove_view_controlled_fanouts_size - 1) {
       auto& control_to_remove_view_last_control =
           control_to_remove_view->controlled_fanouts_.back();
       control_to_remove_view_last_control.node_view()
@@ -1137,7 +1145,8 @@ inline void MutableGraphView::RemoveControllingFaninInternal(
     RemoveControllingFaninFanoutInternal(node_view, control_index);
 
     // Swap last controlling fanin in node with controlling fanin to be removed.
-    if (control_index < node_view->controlling_fanins_.size() - 1) {
+    int node_view_controlling_fanins_size = node_view->controlling_fanins_.size();
+    if (control_index < node_view_controlling_fanins_size - 1) {
       auto& last_control = node_view->controlling_fanins_.back();
       auto* last_control_view = last_control.node_view();
       last_control_view->controlled_fanouts_[last_control.fanout_index_]
diff --git a/tensorflow/core/grappler/utils/graph_view_internal.h b/tensorflow/core/grappler/utils/graph_view_internal.h
index d07f9f71640..9b142444d8a 100644
--- a/tensorflow/core/grappler/utils/graph_view_internal.h
+++ b/tensorflow/core/grappler/utils/graph_view_internal.h
@@ -172,7 +172,8 @@ class NodeViewInternal {
   // Returns a regular fanin based on input index. If no such fanin exist, a
   // missing fanin is returned, with no NodeView set and an index of -2.
   const FanoutViewT& GetRegularFanin(int i) const {
-    if (i < 0 || i >= regular_fanins_.size()) {
+    int regular_fanins_size = regular_fanins_.size(); 
+    if (i < 0 || i >= regular_fanins_size) {
       return GetMissingFanin();
     }
     return regular_fanins_[i];
@@ -191,7 +192,8 @@ class NodeViewInternal {
   // Returns a regular fanout(s) based on output index. If no such output index
   // exists, no fanouts will be returned.
   const std::vector& GetRegularFanout(int i) const {
-    if (i < 0 || i >= regular_fanouts_by_port_.size()) {
+    int regular_fanouts_by_port_size = regular_fanouts_by_port_.size();
+    if (i < 0 || i >= regular_fanouts_by_port_size) {
       return GetMissingFanout();
     }
     return regular_fanouts_by_port_[i];
@@ -289,14 +291,16 @@ class GraphViewInternal {
   // Finds node by index in the graph. If no such node exists in the graph, a
   // `nullptr` is returned.
   const NodeViewT* GetNode(int node_index) const {
-    if (node_index < 0 || node_index >= nodes_.size()) {
+    int nodes_size = nodes_.size(); 
+    if (node_index < 0 || node_index >= nodes_size) {
       return nullptr;
     }
     return &nodes_[node_index];
   }
 
   NodeViewT* GetNode(int node_index) {
-    if (node_index < 0 || node_index >= nodes_.size()) {
+    int nodes_size = nodes_.size(); 
+    if (node_index < 0 || node_index >= nodes_size) {
       return nullptr;
     }
     return &nodes_[node_index];
@@ -444,13 +448,14 @@ inline bool UpdateDevice(NodeViewDiff* diff,
 template 
 inline bool AddOrUpdateAtIndex(std::vector* v, int i, const U& value,
                                const T& default_value) {
-  if (i > v->size()) {
+  int v_size = v->size(); 
+  if (i > v_size) {
     // Resize to include `value`, filling the newly introduced gap with
     // `default_value` for later checks of validity (gaps in vector).
     v->reserve(i + 1);
     v->resize(i, default_value);
     v->push_back({value});
-  } else if (i == v->size()) {
+  } else if (i == v_size) {
     // Vector is large enough, simply append `value` to the end.
     v->push_back({value});
   } else {
@@ -494,7 +499,8 @@ inline bool AddOrUpdateRegularFanin(NodeViewDiff* diff, int index,
     // index from beginning of regular fanins.
     const int relative_removal_index = num_regular_fanins - index - 1;
     // Check if at relative index fanin was already marked for removal.
-    if (relative_removal_index < diff->regular_inputs_to_remove.size() &&
+    int diff_regular_inputs_to_remove_size = diff->regular_inputs_to_remove.size();
+    if (relative_removal_index < diff_regular_inputs_to_remove_size &&
         diff->regular_inputs_to_remove[relative_removal_index]) {
       // Unmark fanin for removal.
       diff->regular_inputs_to_remove[relative_removal_index] = false;
@@ -543,7 +549,8 @@ inline bool RemoveRegularFanin(NodeViewDiff* diff, int index) {
   } else {
     // Relative index from end of regular fanins.
     const int relative_add_index = index - num_regular_fanins;
-    if (relative_add_index >= diff->regular_inputs_to_add.size() ||
+    int diff_regular_inputs_to_add_size = diff->regular_inputs_to_add.size();
+    if (relative_add_index >= diff_regular_inputs_to_add_size ||
         IsEmptyTensorId(diff->regular_inputs_to_add[relative_add_index])) {
       // At relative index, appended regular fanin was already marked for
       // removal.
@@ -671,7 +678,8 @@ inline bool IsWellFormed(
     const absl::flat_hash_map& updated_node_names) {
   ResizeByTrimmingEndForValue(&diff->regular_inputs_to_remove, false);
   ResizeByTrimmingEndForValue(&diff->regular_inputs_to_add, EmptyTensorId());
-  if (diff->regular_inputs_to_add.size() != diff->num_regular_inputs_to_add) {
+  int diff_regular_inputs_to_add_size = diff->regular_inputs_to_add.size();
+  if (diff_regular_inputs_to_add_size != diff->num_regular_inputs_to_add) {
     // Missing regular fanins in between appended fanins.
     return false;
   } else if (diff->num_regular_inputs_to_add > 0 &&
@@ -679,7 +687,7 @@ inline bool IsWellFormed(
     // Appending new fanins while removing existing fanins, resulting in missing
     // regular fanins in between.
     return false;
-  } else if (diff->regular_inputs_to_remove.size() !=
+  } else if ( static_cast(diff->regular_inputs_to_remove.size()) !=
              diff->num_regular_inputs_to_remove) {
     // Regular fanins exist in between removed fanins.
     return false;
@@ -830,7 +838,8 @@ inline void AddOrUpdateRegularFanin(NewNode* new_node, int index,
 // remove existing fanins and updated/added fanins via AddOrUpdateRegularFanins.
 template 
 inline void RemoveRegularFanin(NewNode* new_node, int index) {
-  if (index < 0 || index >= new_node->regular_fanins.size() ||
+  int new_node_regular_fanins_size = new_node->regular_fanins.size();
+  if (index < 0 || index >= new_node_regular_fanins_size ||
       IsEmptyTensorId(new_node->regular_fanins[index])) {
     return;
   }
@@ -874,7 +883,8 @@ inline bool IsWellFormed(
     NewNode* new_node,
     const absl::flat_hash_map& updated_node_names) {
   ResizeByTrimmingEndForValue(&new_node->regular_fanins, EmptyTensorId());
-  if (new_node->regular_fanins.size() != new_node->num_regular_fanins) {
+  int new_node_regular_fanins_size = new_node->regular_fanins.size();
+  if (new_node_regular_fanins_size != new_node->num_regular_fanins) {
     return false;
   }
 
diff --git a/tensorflow/core/grappler/utils/topological_sort.cc b/tensorflow/core/grappler/utils/topological_sort.cc
index e24a457593a..5ed292d1983 100644
--- a/tensorflow/core/grappler/utils/topological_sort.cc
+++ b/tensorflow/core/grappler/utils/topological_sort.cc
@@ -81,7 +81,7 @@ Status ComputeTopologicalOrder(
     int ready_node = (*ready_nodes)[front];
     for (int fanout : graph_view.GetFanout(ready_node)) {
       ++num_ready_inputs[fanout];
-      if (num_ready_inputs[fanout] == graph_view.GetFanin(fanout).size()) {
+      if (num_ready_inputs[fanout] == static_cast(graph_view.GetFanin(fanout).size())) {
         ready_nodes->push_back(fanout);
         ++back;
       }
@@ -95,7 +95,7 @@ Status ComputeTopologicalOrder(
                  "at node = "
               << graph.node(back).DebugString();
       for (int i = 0; i < graph_view.num_nodes(); ++i) {
-        if (num_ready_inputs[i] != graph_view.GetFanin(i).size()) {
+        if (num_ready_inputs[i] != static_cast(graph_view.GetFanin(i).size())) {
           VLOG(1) << "Node not ready: " << graph.node(i).DebugString();
         }
       }
diff --git a/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc b/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
index e86cbc7684c..ada94be15bf 100644
--- a/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
+++ b/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
@@ -68,7 +68,7 @@ Status SqliteQueryConnection::GetNext(IteratorContext* ctx,
 
 Status SqliteQueryConnection::PrepareQuery() {
   TF_RETURN_IF_ERROR(db_->Prepare(query_, &stmt_));
-  int column_count = stmt_.ColumnCount();
+  size_t column_count = stmt_.ColumnCount();
   if (column_count != output_types_.size()) {
     stmt_ = SqliteStatement();
     return errors::InvalidArgument(tensorflow::strings::Printf(
diff --git a/tensorflow/python/grappler/model_analyzer.cc b/tensorflow/python/grappler/model_analyzer.cc
index 5a76cdd8fb2..250010c0fed 100644
--- a/tensorflow/python/grappler/model_analyzer.cc
+++ b/tensorflow/python/grappler/model_analyzer.cc
@@ -48,7 +48,7 @@ void ModelAnalyzer::PrintNodeInfo(const NodeDef* node,
   if (properties.HasOutputProperties(node->name())) {
     const std::vector& props =
         properties.GetOutputProperties(node->name());
-    for (int i = 0; i < props.size(); ++i) {
+    for (int i = 0, props_size = props.size(); i < props_size; ++i) {
       const OpInfo::TensorProperties& prop = props[i];
       os << "\t"
          << "output " << i << " (" << DataTypeString(prop.dtype())
@@ -88,7 +88,7 @@ void ModelAnalyzer::PrintNodeInfo(const NodeDef* node,
     } else if (properties.HasInputProperties(node->name())) {
       const std::vector& props =
           properties.GetInputProperties(node->name());
-      for (int i = 0; i < props.size(); ++i) {
+      for (int i = 0, props_size = props.size(); i < props_size; ++i) {
         const OpInfo::TensorProperties& prop = props[i];
         if (prop.has_value()) {
           os << "\t"

From 6166444602c0ddab5e5e7ff129113341a99bd98c Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Mon, 15 Jun 2020 20:56:58 -0700
Subject: [PATCH 0259/1390] [tf.data] Add prefetch benchmark.

This CL adds a benchmark to test the performance of the prefetch dataset transformation.

PiperOrigin-RevId: 316605096
Change-Id: Iaf5d8f8c3afba6e51a53805afe5bc978916ff01e
---
 tensorflow/python/data/benchmarks/BUILD       |  9 +++++
 .../data/benchmarks/prefetch_benchmark.py     | 40 +++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 tensorflow/python/data/benchmarks/prefetch_benchmark.py

diff --git a/tensorflow/python/data/benchmarks/BUILD b/tensorflow/python/data/benchmarks/BUILD
index 256f10dcefd..3f0faf5364a 100644
--- a/tensorflow/python/data/benchmarks/BUILD
+++ b/tensorflow/python/data/benchmarks/BUILD
@@ -85,6 +85,15 @@ tf_py_test(
     ],
 )
 
+tf_py_test(
+    name = "prefetch_benchmark",
+    srcs = ["prefetch_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)
+
 tf_py_test(
     name = "range_benchmark",
     srcs = ["range_benchmark.py"],
diff --git a/tensorflow/python/data/benchmarks/prefetch_benchmark.py b/tensorflow/python/data/benchmarks/prefetch_benchmark.py
new file mode 100644
index 00000000000..31177508397
--- /dev/null
+++ b/tensorflow/python/data/benchmarks/prefetch_benchmark.py
@@ -0,0 +1,40 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.prefetch()`."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+
+
+class PrefetchBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.prefetch()`."""
+
+  def benchmark_prefetch(self):
+    num_elements = 1000000
+    for prefetch_buffer in [1, 5, 10, 20, 100]:
+      dataset = dataset_ops.Dataset.range(num_elements)
+      dataset = dataset.prefetch(prefetch_buffer)
+
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=num_elements,
+          name="prefetch_{}".format(prefetch_buffer))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

From 52e1dba6b14da82ddd30344526e13557cf33cc32 Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Tue, 16 Jun 2020 04:02:50 +0000
Subject: [PATCH 0260/1390] getting rid of pesky lingering commits

---
 .../mlir/lite/quantization/import_quant_stats_pass.cc  |  4 ++--
 .../mlir/lite/quantization/quantization_config.cc      |  8 ++++----
 .../mlir/lite/quantization/quantization_driver.cc      |  4 ++--
 .../mlir/lite/quantization/quantization_utils.cc       | 10 +++++-----
 .../compiler/mlir/tensorflow/utils/dump_mlir_util.cc   |  2 +-
 tensorflow/compiler/mlir/xla/ir/chlo_ops.cc            |  2 +-
 tensorflow/compiler/mlir/xla/ir/hlo_ops.cc             |  6 +++---
 7 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc
index e00a088c38c..d924a3e82ac 100644
--- a/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc
@@ -76,7 +76,7 @@ class ImportQuantStatsPass
   // If the index is out of range, this method returns false. Otherwise it
   // returns true if the value is a float tensor.
   bool IsQuantizableResult(Operation *op, int index) {
-    if (index < 0 || index >= static_cast(op->getNumResults())) return false;
+    if (index < 0 || index >= op->getNumResults()) return false;
     Value res = op->getResult(index);
     return res.getType().isa() &&
            res.getType().cast().getElementType().isa();
@@ -158,7 +158,7 @@ void ImportQuantStatsPass::ImportAsStatsOps(OpBuilder b, Operation *op,
     InsertStatsOpAtResult(b, op->getResult(index), layer_stats, axis_stats,
                           axis);
   } else {
-    for (int i = 0, e = op->getNumResults(); i < e; ++i) {
+    for (int i = 0; i < op->getNumResults(); ++i) {
       if (IsQuantizableResult(op, i)) {
         InsertStatsOpAtResult(b, op->getResult(i), layer_stats, axis_stats,
                               axis);
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
index cdff93502f2..6b897bd5608 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
@@ -48,9 +48,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names,
   std::vector node_mins;
   if (!min_values.empty()) {
     std::vector node_mins_str = absl::StrSplit(min_values, ',');
-    for (const std::string&node_min : node_mins_str.size()) {
+    for (int i = 0; i < node_mins_str.size(); i++) {
       double value;
-      if (!absl::SimpleAtod(node_min, &value)) {
+      if (!absl::SimpleAtod(node_mins_str[i], &value)) {
         return true;
       }
       node_mins.push_back(value);
@@ -60,9 +60,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names,
   std::vector node_maxs;
   if (!max_values.empty()) {
     std::vector node_maxs_str = absl::StrSplit(max_values, ',');
-    for (const std::string&node_max : node_maxs_str.size()) {
+    for (int i = 0; i < node_maxs_str.size(); i++) {
       double value;
-      if (!absl::SimpleAtod(node_max, &value)) {
+      if (!absl::SimpleAtod(node_maxs_str[i], &value)) {
         llvm::errs() << "Unexpected mins: " << node_maxs_str[i] << "\n";
         return true;
       }
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
index a9f4eb78431..2964a3e79f8 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
@@ -294,7 +294,7 @@ class QuantizationDriver {
         return;
       if (current_op == op) llvm::errs() << "===>>>";
       llvm::errs() << op->getName() << " : (";
-      for (int i = 0, e = op->getNumOperands(); i < e; ++i) {
+      for (auto i = 0; i < op->getNumOperands(); ++i) {
         if (auto params = GetOperandQuantState(op, i).params)
           params.print(llvm::errs());
         else
@@ -303,7 +303,7 @@ class QuantizationDriver {
         llvm::errs() << ",";
       }
       llvm::errs() << ") -> (";
-      for (int i = 0, e = op->getNumResults(); i < e; ++i) {
+      for (auto i = 0; i < op->getNumResults(); ++i) {
         if (auto params = GetResultQuantState(op, i).params)
           params.print(llvm::errs());
         else
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
index 57b24eb8772..3d50f280d0f 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
@@ -54,7 +54,7 @@ static Type GetQuantizedType(Builder builder, Type input_type,
   } else if (min.size() == max.size()) {
     auto shape = input_type.dyn_cast();
     if (!shape || shape.getRank() <= quant_dim ||
-        static_cast(min.size()) != shape.getDimSize(quant_dim)) {
+        min.size() != shape.getDimSize(quant_dim)) {
       return {};
     }
     // TODO(b/141508873): the quantization dim is set to the last dimension.
@@ -75,7 +75,7 @@ TypeAttr RescaleQuantizedType(Type input, Attribute factor) {
   if (auto qtype = ele_type.dyn_cast()) {
     ArrayRef scales = qtype.getScales();
     // Broadcasting hasn't been implemented yet.
-    if (static_cast(scales.size()) != factor_values.getNumElements()) return {};
+    if (scales.size() != factor_values.getNumElements()) return {};
     SmallVector new_scales;
     new_scales.reserve(scales.size());
     auto scales_iter = scales.begin();
@@ -269,7 +269,7 @@ Type GetUniformQuantizedPerAxisTypeForWeight(ElementsAttr attr, int quant_dim,
                                              bool narrow_range) {
   Builder builder(attr.getContext());
   auto shape = attr.getType().cast().getShape();
-  if (static_cast(shape.size()) <= quant_dim) return {};
+  if (shape.size() <= quant_dim) return {};
   // `symmetric` can only be used when it is `signed` and `narrow_range`.
   if (symmetric && (!is_signed || !narrow_range)) return {};
 
@@ -334,7 +334,7 @@ quant::QuantizedType GetUniformQuantizedTypeForBias(
     const std::vector& op_types) {
   if (op_types.empty()) return {};
 
-  size_t axis_size = 1;
+  int axis_size = 1;
   int32_t quant_dim = -1;
   Type expressed_type;
   // Requires all the op types are valid UniformQuantizedTypes or
@@ -368,7 +368,7 @@ quant::QuantizedType GetUniformQuantizedTypeForBias(
         scales[index_scale.index()] *= index_scale.value();
       }
     } else if (auto type = op_type.dyn_cast()) {
-      for (int index = 0, e = axis_size; index != e; ++index) {
+      for (int index = 0; index != axis_size; ++index) {
         scales[index] *= type.getScale();
       }
     }
diff --git a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
index febf2bc096d..797687ea658 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
@@ -41,7 +41,7 @@ std::string MakeUniqueFilename(string name) {
   static NameCounts& instance = *new NameCounts;
 
   // Remove illegal characters from `name`.
-  for (int i = 0, e = name.size(); i < e; ++i) {
+  for (int i = 0; i < name.size(); ++i) {
     char ch = name[i];
     if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' ||
         ch == '\\') {
diff --git a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc
index 3408f3ed0cc..26db4549a2a 100644
--- a/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc
+++ b/tensorflow/compiler/mlir/xla/ir/chlo_ops.cc
@@ -49,7 +49,7 @@ static Type GetBroadcastType(Type x, Type y, Type element_type,
 
   if (shape_x.size() == shape_y.size()) {
     llvm::SmallVector out_shape(shape_x.size());
-    for (int i = 0, e = shape_x.size(); i < e; i++) {
+    for (int i = 0; i < shape_x.size(); i++) {
       auto x_val = shape_x[i];
       auto y_val = shape_y[i];
       if (x_val == -1 || y_val == -1) {
diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
index 7f313b56925..d20f1713eba 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
@@ -143,7 +143,7 @@ DenseIntElementsAttr BuildConvPaddingAttrs(
 
   int rank = padding_low.size();
   SmallVector padding;
-  for (unsigned i = 0, e = rank; i < e; ++i) {
+  for (unsigned i = 0; i < rank; ++i) {
     padding.push_back(GetPaddingValue(padding_attr, {i, 0}) + padding_low[i]);
     padding.push_back(GetPaddingValue(padding_attr, {i, 1}) + padding_high[i]);
   }
@@ -853,7 +853,7 @@ static Attribute foldConcatenateHelper(ConcatenateOp* op,
   auto shape = type.getShape();
 
   size_t top_size = 1;
-  for (int i = 0, e = axis; i < e; i++) {
+  for (int i = 0; i < axis; i++) {
     top_size = top_size * shape[i];
   }
 
@@ -1118,7 +1118,7 @@ static LogicalResult Verify(MapOp op) {
   // increasing.
   auto values = op.dimensions().getValues();
   auto dimensions = std::vector{values.begin(), values.end()};
-  for (int i = 0, e = dimensions.size(); i < e; ++i) {
+  for (int i = 0; i < dimensions.size(); ++i) {
     if (dimensions[i] != i)
       return op.emitOpError() << "requires monotonically increasing dimension "
                                  "numbers, but got: "

From 7ee9571a8e127b39a4a8a01016e40105b7613bbf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 21:21:15 -0700
Subject: [PATCH 0261/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/1a7f115dce22

PiperOrigin-RevId: 316608084
Change-Id: I6df721c782e850371267e6a2fba7eda96d3b1610
---
 third_party/mlir/BUILD      | 22 ++++++++++++++++++++++
 third_party/mlir/test.BUILD |  1 +
 2 files changed, 23 insertions(+)

diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index 5f3f0a4b99b..06e0ed8d4b4 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -756,6 +756,25 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ShapeToSCF",
+    srcs = glob([
+        "lib/Conversion/ShapeToSCF/*.cpp",
+        "lib/Conversion/ShapeToSCF/*.h",
+    ]) + ["lib/Conversion/PassDetail.h"],
+    hdrs = ["include/mlir/Conversion/ShapeToSCF/ShapeToSCF.h"],
+    includes = ["include"],
+    deps = [
+        ":ConversionPassIncGen",
+        ":IR",
+        ":Pass",
+        ":SCFDialect",
+        ":Shape",
+        ":StandardOps",
+        ":Transforms",
+    ],
+)
+
 gentbl(
     name = "ShapeTransformsPassIncGen",
     strip_include_prefix = "include",
@@ -2613,6 +2632,7 @@ cc_library(
         ":Parser",
         ":Pass",
         ":SCFTransforms",
+        ":ShapeToSCF",
         ":ShapeToStandard",
         ":ShapeTransforms",
         ":StandardOpsTransforms",
@@ -2713,6 +2733,7 @@ cc_library(
         ":SPIRVPassIncGen",
         ":SPIRVToLLVM",
         ":Shape",
+        ":ShapeToSCF",
         ":ShapeToStandard",
         ":ShapeTransforms",
         ":ShapeTransformsPassIncGen",
@@ -3282,6 +3303,7 @@ cc_library(
         ":LinalgTransforms",
         ":Pass",
         ":StandardOps",
+        ":StandardOpsTransforms",
         ":Support",
         ":Transforms",
         ":VectorToLLVM",
diff --git a/third_party/mlir/test.BUILD b/third_party/mlir/test.BUILD
index 23287ce28d6..14c2ba7778e 100644
--- a/third_party/mlir/test.BUILD
+++ b/third_party/mlir/test.BUILD
@@ -166,6 +166,7 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:StandardOps",
+        "@llvm-project//mlir:StandardOpsTransforms",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TargetNVVMIR",
         "@llvm-project//mlir:TargetROCDLIR",

From e74010b4e803f5f47f6c013a2932172e4beae72a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 22:24:35 -0700
Subject: [PATCH 0262/1390] Clarify the documentation for PermuteDimensions

PiperOrigin-RevId: 316616603
Change-Id: Iccfbd986276688bdc25b6757cdee6806f0b587d6
---
 tensorflow/compiler/xla/shape_util.cc | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc
index 52cbb8f95ac..ab46e49b181 100644
--- a/tensorflow/compiler/xla/shape_util.cc
+++ b/tensorflow/compiler/xla/shape_util.cc
@@ -968,17 +968,18 @@ Status ForEachMutableSubshapeHelper(
   // `shape`'s list of dimensions is isomorphic to the identity I.
   //
   // Let `shape`'s layout be L.  A layout is a permutation which maps a
-  // minor-to-major physical layout to the order of a shape's logical dims.
-  // Therefore inverse of a layout maps from logical to physical dims, and so
-  // the physical layout of I is simply L'.I = L', where L' is the inverse of L.
+  // minor-to-major physical dimension ordering to a shape's logical dimension
+  // ordering.  Therefore the inverse of a layout maps from logical to physical
+  // dims, and so the physical ordering of I is simply L'.I = L', where L' is
+  // the inverse of L.
   //
   // Let the argument `permutation` be P.  This is a permutation over `shape`'s
   // dimensions, so our return value will be a shape with dims P.I = P.  Our
-  // goal is to construct a layout permutation L* that we can apply to P such
-  // that the physical dimension ordering of the returned shape is the same
-  // as that of the original shape, namely L'.
+  // goal is to construct a layout permutation L* for this shape. The physical
+  // dimension ordering of this returned shape must be the same as that of the
+  // original shape, namely L'.
   //
-  // Our returned shape has dims P and layout L*, so its in-memory layout is
+  // Our returned shape has dims P and layout L*, so its in-memory ordering is
   // L*'.P.  Setting this equal to L' and solving for L*, we get:
   //
   //   L*'.P = L'    =>

From 42a734170dae2942fcf553ccf5480fd48840795a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 15 Jun 2020 23:00:10 -0700
Subject: [PATCH 0263/1390] tf.numpy: Change a bunch of ops to handle unknown
 shapes. Fix logic in sort_ops to handle 1D values.

PiperOrigin-RevId: 316620637
Change-Id: Iedc2ba8aad7673bbe210661bb741bf0660f047aa
---
 .../python/ops/numpy_ops/np_array_ops.py      | 36 ++++++--
 tensorflow/python/ops/numpy_ops/np_arrays.py  | 51 ++++++++---
 .../python/ops/numpy_ops/np_math_ops.py       | 90 +++++++++++++------
 tensorflow/python/ops/numpy_ops/np_utils.py   | 32 ++++---
 tensorflow/python/ops/sort_ops.py             | 16 +---
 5 files changed, 153 insertions(+), 72 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index fbf67a46e31..e97bb61613b 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -223,9 +223,10 @@ def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
   Raises:
     ValueError: if `fill_value` can not be broadcast to shape `shape`.
   """
+  if not isinstance(shape, np_arrays.ndarray):
+    shape = asarray(np_arrays.convert_to_tensor(shape, dtype_hint=np.int32))
+  shape = atleast_1d(shape).data
   fill_value = asarray(fill_value, dtype=dtype)
-  if np_utils.isscalar(shape):
-    shape = array_ops.reshape(shape, [1])
   return np_arrays.tensor_to_ndarray(
       array_ops.broadcast_to(fill_value.data, shape))
 
@@ -808,16 +809,21 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: d
 @np_utils.np_doc(np.std)
 def std(a, axis=None, keepdims=None):  # pylint: disable=missing-function-docstring
   return _reduce(
-      math_ops.reduce_std, a, axis=axis, dtype=None, keepdims=keepdims,
+      math_ops.reduce_std,
+      a,
+      axis=axis,
+      dtype=None,
+      keepdims=keepdims,
       promote_int=_TO_FLOAT)
 
 
 @np_utils.np_doc(np.ravel)
 def ravel(a):  # pylint: disable=missing-docstring
   a = asarray(a)
-  if a.ndim == 1:
-    return a
-  return np_utils.tensor_to_ndarray(array_ops.reshape(a.data, [-1]))
+  out = np_utils.cond(
+      math_ops.equal(a.ndim, 1), lambda: a.data,
+      lambda: array_ops.reshape(a.data, [-1]))
+  return np_utils.tensor_to_ndarray(out)
 
 
 setattr(np_arrays.ndarray, 'ravel', ravel)
@@ -846,7 +852,8 @@ def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
   a = asarray(a).data
   original_shape = a._shape_as_list()  # pylint: disable=protected-access
   # Best effort recovery of the shape.
-  if original_shape is not None and None not in original_shape:
+  known_shape = original_shape is not None and None not in original_shape
+  if known_shape:
     if not original_shape:
       original_shape = (repeats,)
     else:
@@ -865,7 +872,8 @@ def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
 
   repeats = asarray(repeats).data
   result = array_ops.repeat(a, repeats, axis)
-  result.set_shape(original_shape)
+  if known_shape:
+    result.set_shape(original_shape)
 
   return np_utils.tensor_to_ndarray(result)
 
@@ -1287,7 +1295,13 @@ def broadcast_to(array, shape):  # pylint: disable=redefined-outer-name
 
 
 @np_utils.np_doc(np.stack)
-def stack(arrays, axis=0):
+def stack(arrays, axis=0):  # pylint: disable=missing-function-docstring
+  if isinstance(arrays, (np_arrays.ndarray, ops.Tensor)):
+    arrays = asarray(arrays)
+    if axis == 0:
+      return arrays
+    else:
+      return swapaxes(arrays, 0, axis)
   arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
   unwrapped_arrays = [
       a.data if isinstance(a, np_arrays.ndarray) else a for a in arrays
@@ -1450,6 +1464,8 @@ def tri(N, M=None, k=0, dtype=None):  # pylint: disable=invalid-name,missing-doc
 @np_utils.np_doc(np.tril)
 def tril(m, k=0):  # pylint: disable=missing-docstring
   m = asarray(m).data
+  if m.shape.ndims is None:
+    raise ValueError('Argument to tril should have known rank')
   m_shape = m.shape.as_list()
 
   if len(m_shape) < 2:
@@ -1470,6 +1486,8 @@ def tril(m, k=0):  # pylint: disable=missing-docstring
 @np_utils.np_doc(np.triu)
 def triu(m, k=0):  # pylint: disable=missing-docstring
   m = asarray(m).data
+  if m.shape.ndims is None:
+    raise ValueError('Argument to triu should have known rank')
   m_shape = m.shape.as_list()
 
   if len(m_shape) < 2:
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index a7696ad31c2..e2f73100909 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 """ndarray class."""
+
+# pylint: disable=g-direct-tensorflow-import
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -151,13 +154,16 @@ def _slice_helper(tensor, slice_spec, var=None):
 def convert_to_tensor(value, dtype=None, dtype_hint=None):
   """Wrapper over `tf.convert_to_tensor`.
 
-     Args:
-       value: value to convert
-       dtype: (optional) the type we would like it to be converted to.
-       dtype_hint: (optional) soft preference for the type we would like it to
-         be converted to. `tf.convert_to_tensor` will attempt to convert value
-         to this type first, but will not fail if conversion is not possible
-         falling back to inferring the type instead.
+  Args:
+    value: value to convert
+    dtype: (optional) the type we would like it to be converted to.
+    dtype_hint: (optional) soft preference for the type we would like it to be
+      converted to. `tf.convert_to_tensor` will attempt to convert value to this
+      type first, but will not fail if conversion is not possible falling back
+      to inferring the type instead.
+
+  Returns:
+    Value converted to tf.Tensor.
   """
   # A safer version of `tf.convert_to_tensor` to work around b/149876037.
   # TODO(wangpeng): Remove this function once the bug is fixed.
@@ -250,8 +256,12 @@ class ndarray(object):  # pylint: disable=invalid-name
 
   @property
   def shape(self):
-    """Returns a tuple of array dimensions."""
-    return self.data._shape_tuple()  # pylint: disable=protected-access
+    """Returns a tuple or tf.Tensor of array dimensions."""
+    shape = self.data.shape
+    if shape.is_fully_defined():
+      return tuple(shape.as_list())
+    else:
+      return array_ops.shape(self.data)
 
   @property
   def dtype(self):
@@ -259,19 +269,30 @@ class ndarray(object):  # pylint: disable=invalid-name
 
   @property
   def ndim(self):
-    return self.data.shape.ndims
+    ndims = self.data.shape.ndims
+    if ndims is None:
+      return array_ops.rank(self.data)
+    else:
+      return ndims
 
   @property
   def size(self):
     """Returns the number of elements in the array."""
-    return np.prod(self.shape)
+    shape = self.shape
+    if isinstance(shape, ops.Tensor):
+      return array_ops.size(self.data)
+    else:
+      return np.prod(self.shape)
 
   @property
   def T(self):  # pylint: disable=invalid-name
     return self.transpose()
 
   def __len__(self):
-    if self.shape:
+    shape = self.shape
+    if isinstance(shape, ops.Tensor):
+      raise TypeError('len() of symbolic tensor undefined')
+    elif shape:
       return self.shape[0]
     else:
       raise TypeError('len() of unsized object.')
@@ -320,6 +341,8 @@ class ndarray(object):  # pylint: disable=invalid-name
     return tensor_to_ndarray(result_t)
 
   def __iter__(self):
+    if not isinstance(self.data, ops.EagerTensor):
+      raise TypeError('Iteration over symbolic tensor is not allowed')
     for i in range(self.shape[0]):
       result_t = self.data[i]
       yield tensor_to_ndarray(result_t)
@@ -356,6 +379,8 @@ class ndarray(object):  # pylint: disable=invalid-name
       ValueError: If the array does not have size 1.
     """
     # TODO(wangpeng): Handle graph mode
+    if not isinstance(self.data, ops.EagerTensor):
+      raise TypeError('Indexing using symbolic tensor is not allowed')
     return np.asscalar(self.data.numpy())
 
   def tolist(self):
@@ -384,5 +409,3 @@ def ndarray_to_tensor(arr, dtype=None, name=None, as_ref=False):
 
 
 ops.register_tensor_conversion_function(ndarray, ndarray_to_tensor)
-
-
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 02d37b3a3a4..b32f78bee5a 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -50,9 +50,9 @@ def dot(a, b):  # pylint: disable=missing-docstring
             math_ops.equal(array_ops.rank(b), 0)),
         lambda: a * b,
         lambda: np_utils.cond(  # pylint: disable=g-long-lambda
-            math_ops.equal(array_ops.rank(b), 1), lambda: math_ops.tensordot(
-                a, b, axes=[[-1], [-1]]), lambda: math_ops.tensordot(
-                    a, b, axes=[[-1], [-2]])))
+            math_ops.equal(array_ops.rank(b), 1),
+            lambda: math_ops.tensordot(a, b, axes=[[-1], [-1]]),
+            lambda: math_ops.tensordot(a, b, axes=[[-1], [-2]])))
 
   return _bin_op(f, a, b)
 
@@ -204,8 +204,8 @@ def matmul(x1, x2):  # pylint: disable=missing-docstring
       return np_utils.cond(
           math_ops.equal(array_ops.rank(x2), 1),
           lambda: math_ops.tensordot(x1, x2, axes=1),
-          lambda: np_utils.cond(
-              math_ops.equal(array_ops.rank(x1), 1),  # pylint: disable=g-long-lambda
+          lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+              math_ops.equal(array_ops.rank(x1), 1),
               lambda: math_ops.tensordot(  # pylint: disable=g-long-lambda
                   x1, x2, axes=[[0], [-2]]),
               lambda: math_ops.matmul(x1, x2)))
@@ -352,14 +352,30 @@ def hypot(x1, x2):
 def kron(a, b):  # pylint: disable=missing-function-docstring
   # pylint: disable=protected-access,g-complex-comprehension
   a, b = np_array_ops._promote_dtype(a, b)
-  ndim = max(a.ndim, b.ndim)
-  if a.ndim < ndim:
-    a = np_array_ops.reshape(a, np_array_ops._pad_left_to(ndim, a.shape))
-  if b.ndim < ndim:
-    b = np_array_ops.reshape(b, np_array_ops._pad_left_to(ndim, b.shape))
-  a_reshaped = np_array_ops.reshape(a, [i for d in a.shape for i in (d, 1)])
-  b_reshaped = np_array_ops.reshape(b, [i for d in b.shape for i in (1, d)])
-  out_shape = tuple(np.multiply(a.shape, b.shape))
+  t_a = np_utils.cond(
+      a.ndim < b.ndim,
+      lambda: np_array_ops.reshape(  # pylint: disable=g-long-lambda
+          a.data, np_array_ops._pad_left_to(b.ndim, a.shape)),
+      lambda: a.data)
+  t_b = np_utils.cond(
+      b.ndim < a.ndim,
+      lambda: np_array_ops.reshape(  # pylint: disable=g-long-lambda
+          b.data, np_array_ops._pad_left_to(a.ndim, b.shape)),
+      lambda: b.data)
+
+  def _make_shape(shape, prepend):
+    ones = array_ops.ones_like(shape)
+    if prepend:
+      shapes = [ones, shape]
+    else:
+      shapes = [shape, ones]
+    return array_ops.reshape(array_ops.stack(shapes, axis=1), [-1])
+
+  a_shape = array_ops.shape(t_a)
+  b_shape = array_ops.shape(t_b)
+  a_reshaped = np_array_ops.reshape(t_a, _make_shape(a_shape, False))
+  b_reshaped = np_array_ops.reshape(t_b, _make_shape(b_shape, True))
+  out_shape = a_shape * b_shape
   return np_array_ops.reshape(a_reshaped * b_reshaped, out_shape)
 
 
@@ -454,7 +470,8 @@ def _tf_gcd(x1, x2):  # pylint: disable=missing-function-docstring
   if (not np.issubdtype(x1.dtype.as_numpy_dtype, np.integer) or
       not np.issubdtype(x2.dtype.as_numpy_dtype, np.integer)):
     raise ValueError('Arguments to gcd must be integers.')
-  shape = array_ops.broadcast_static_shape(x1.shape, x2.shape)
+  shape = array_ops.broadcast_dynamic_shape(
+      array_ops.shape(x1), array_ops.shape(x2))
   x1 = array_ops.broadcast_to(x1, shape)
   x2 = array_ops.broadcast_to(x2, shape)
   value, _ = control_flow_ops.while_loop(_gcd_cond_fn, _gcd_body_fn,
@@ -607,7 +624,7 @@ def signbit(x):
 
   def f(x):
     if x.dtype == dtypes.bool:
-      return array_ops.fill(x.shape, False)
+      return array_ops.fill(array_ops.shape(x), False)
     return x < 0
 
   return _scalar(f, x)
@@ -866,7 +883,11 @@ def square(x):
 def diff(a, n=1, axis=-1):  # pylint: disable=missing-function-docstring
 
   def f(a):
+    # TODO(agarwal): transpose and reshape to N, H, 1 and do a 1D convolution
+    # TODO(agarwal): avoid depending on static rank.
     nd = a.shape.rank
+    if nd is None:
+      raise ValueError('diff currently requires known rank for input `a`')
     if (axis + nd if axis < 0 else axis) >= nd:
       raise ValueError('axis %s is out of bounds for array of dimension %s' %
                        (axis, nd))
@@ -887,8 +908,10 @@ def diff(a, n=1, axis=-1):  # pylint: disable=missing-function-docstring
 
 
 def _flip_args(f):
+
   def _f(a, b):
     return f(b, a)
+
   return _f
 
 
@@ -910,6 +933,7 @@ setattr(np_arrays.ndarray, '__rtruediv__', _flip_args(true_divide))
 
 
 def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):
+  """Helper function for comparision."""
   dtype = np_utils.result_type(x1, x2)
   # Cast x1 and x2 to the result_type if needed.
   x1 = np_array_ops.array(x1, dtype=dtype)
@@ -953,12 +977,18 @@ def less_equal(x1, x2):
 
 
 @np_utils.np_doc(np.array_equal)
-def array_equal(a1, a2):
+def array_equal(a1, a2):  # pylint: disable=missing-function-docstring
 
-  def f(a1, a2):
-    if a1.shape != a2.shape:
-      return constant_op.constant(False)
-    return math_ops.reduce_all(math_ops.equal(a1, a2))
+  def f(x1, x2):
+    return np_utils.cond(
+        math_ops.equal(array_ops.rank(x1), array_ops.rank(x2)),
+        lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+            np_utils.reduce_all(
+                math_ops.equal(array_ops.shape(x1), array_ops.shape(x2))
+            ),
+            lambda: math_ops.reduce_all(math_ops.equal(x1, x2)),
+            lambda: constant_op.constant(False)),
+        lambda: constant_op.constant(False))
 
   return _comparison(f, a1, a2)
 
@@ -1001,7 +1031,13 @@ setattr(np_arrays.ndarray, '__ne__', not_equal)
 
 @np_utils.np_doc(np.linspace)
 def linspace(  # pylint: disable=missing-docstring
-    start, stop, num=50, endpoint=True, retstep=False, dtype=float, axis=0):
+    start,
+    stop,
+    num=50,
+    endpoint=True,
+    retstep=False,
+    dtype=float,
+    axis=0):
   if dtype:
     dtype = np_utils.result_type(dtype)
   start = np_array_ops.array(start, dtype=dtype).data
@@ -1054,10 +1090,14 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):  # pylint
   start_sign = 1 - np_array_ops.sign(np_array_ops.real(start))
   stop_sign = 1 - np_array_ops.sign(np_array_ops.real(stop))
   signflip = 1 - start_sign * stop_sign // 2
-  res = signflip * logspace(log10(signflip * start),
-                            log10(signflip * stop), num,
-                            endpoint=endpoint, base=10.0,
-                            dtype=computation_dtype, axis=0)
+  res = signflip * logspace(
+      log10(signflip * start),
+      log10(signflip * stop),
+      num,
+      endpoint=endpoint,
+      base=10.0,
+      dtype=computation_dtype,
+      axis=0)
   if axis != 0:
     res = np_array_ops.moveaxis(res, 0, axis)
   return np_utils.tensor_to_ndarray(math_ops.cast(res, dtype))
diff --git a/tensorflow/python/ops/numpy_ops/np_utils.py b/tensorflow/python/ops/numpy_ops/np_utils.py
index 47b45b171fb..186e56816fe 100644
--- a/tensorflow/python/ops/numpy_ops/np_utils.py
+++ b/tensorflow/python/ops/numpy_ops/np_utils.py
@@ -47,7 +47,7 @@ def _canonicalize_axes(axes, rank):
     canonicalizer = (
         lambda axis: cond(axis < 0, lambda: axis + rank, lambda: axis))
   else:
-    canonicalizer = lambda axis: axis+rank if axis < 0 else axis
+    canonicalizer = lambda axis: axis + rank if axis < 0 else axis
 
   return [canonicalizer(axis) for axis in axes]
 
@@ -100,9 +100,16 @@ def finfo(dtype):
 
 def isscalar(val):
   """Returns whether `val` is a scalar value or scalar Tensor."""
-  if isinstance(val, (np.ndarray, np_arrays.ndarray, ops.Tensor)):
-    return len(val.shape) == 0  # pylint: disable=g-explicit-length-test
-  return np.isscalar(val)
+  if isinstance(val, np_arrays.ndarray):
+    val = val.data
+  if isinstance(val, ops.Tensor):
+    ndims = val.shape.ndims
+    if ndims is not None:
+      return ndims == 0
+    else:
+      return math_ops.equal(array_ops.rank(val), 0)
+  else:
+    return np.isscalar(val)
 
 
 # Can't use np_doc because np.result_type is a builtin function.
@@ -119,8 +126,8 @@ def result_type(*arrays_and_dtypes):
   def maybe_get_dtype(x):
     # Don't put np.ndarray in this list, because np.result_type looks at the
     # value (not just dtype) of np.ndarray to decide the result type.
-    if isinstance(x, (np_arrays.ndarray, ops.Tensor,
-                      indexed_slices.IndexedSlices)):
+    if isinstance(
+        x, (np_arrays.ndarray, ops.Tensor, indexed_slices.IndexedSlices)):
       return _to_numpy_type(x.dtype)
     elif isinstance(x, dtypes.DType):
       return _to_numpy_type(x)
@@ -277,8 +284,11 @@ def np_doc(np_fun, np_fun_name=None):
       #   for name in np_sig.parameters:
       #     if name not in sig.parameters:
       #       unsupported_params.append(name)
-    f.__doc__ = _np_doc_helper(f, np_fun, np_fun_name=np_fun_name,
-                               unsupported_params=unsupported_params)
+    f.__doc__ = _np_doc_helper(
+        f,
+        np_fun,
+        np_fun_name=np_fun_name,
+        unsupported_params=unsupported_params)
     return f
 
   return decorator
@@ -287,9 +297,9 @@ def np_doc(np_fun, np_fun_name=None):
 def _np_doc_helper(f, np_f, np_fun_name=None, unsupported_params=None):
   """Helper to get docs."""
   if not unsupported_params and not _has_docstring(f) and _has_docstring(np_f):
-      # TODO(wangpeng): It looks like code snippets in numpy doc don't work
-      # correctly with doctest. Fix that and remove the reformatting of the np_f
-      # comment, here and below.
+    # TODO(wangpeng): It looks like code snippets in numpy doc don't work
+    # correctly with doctest. Fix that and remove the reformatting of the np_f
+    # comment, here and below.
     return np_f.__doc__.replace('>>>', '>')
   assert np_f or np_fun_name
   if not np_fun_name:
diff --git a/tensorflow/python/ops/sort_ops.py b/tensorflow/python/ops/sort_ops.py
index 4e66a80bc01..d711516cb86 100644
--- a/tensorflow/python/ops/sort_ops.py
+++ b/tensorflow/python/ops/sort_ops.py
@@ -134,7 +134,7 @@ def _sort_or_argsort(values, axis, direction, return_argsort):
   # Axis must be an integer, not a Tensor.
   axis = framework_ops.convert_to_tensor(axis, name='axis')
   axis_static = tensor_util.constant_value(axis)
-  if axis.shape.ndims != 0 or axis_static is None:
+  if axis.shape.ndims not in (None, 0) or axis_static is None:
     raise ValueError('axis must be a constant scalar')
   axis_static = int(axis_static)  # Avoids NumPy casting error
 
@@ -184,18 +184,8 @@ def _descending_sort(values, axis, return_argsort=False):
           name='transposition')
     else:
       # Generate the transposition array from the tensors.
-      transposition = array_ops.concat(
-          [
-              # Axes up to axis are unchanged.
-              math_ops.range(axis),
-              # Swap axis and rank - 1.
-              [rank - 1],
-              # Axes in [axis + 1, rank - 1) are unchanged.
-              math_ops.range(axis + 1, rank - 1),
-              # Swap axis and rank - 1.
-              [axis]
-          ],
-          axis=0)
+      transposition = array_ops.tensor_scatter_update(
+          math_ops.range(rank), [[axis], [rank-1]], [rank-1, axis])
     top_k_input = array_ops.transpose(values, transposition)
 
   values, indices = nn_ops.top_k(top_k_input, k)

From 3fa71a90282e984050d25f4ea61e4bbe8b75bd12 Mon Sep 17 00:00:00 2001
From: Mehdi Amini 
Date: Mon, 15 Jun 2020 23:20:01 -0700
Subject: [PATCH 0264/1390] Bump open source llvm revision to
 1a7f115dce22b2c09fdd4f7f79d24da5de6eaef8

PiperOrigin-RevId: 316622447
Change-Id: I4f0e46be7b6a3e2624034cab02e06a1f8e86a1a0
---
 tensorflow/workspace.bzl                  | 4 ++--
 third_party/llvm/llvm.autogenerated.BUILD | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 60fef8c0cb9..31daf2249a0 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "ec02635d104c5f42840c63ed41d0cea774d649fe"
-    LLVM_SHA256 = "2e1155b29bd84b7382bcd7c00a88899f4156e9b09bc443debfdd7d3a5931e929"
+    LLVM_COMMIT = "1a7f115dce22b2c09fdd4f7f79d24da5de6eaef8"
+    LLVM_SHA256 = "3ace55744a86211c9c837915b88c18e1ed3e5cd839aaeade6aa88b02bc86e47e"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
diff --git a/third_party/llvm/llvm.autogenerated.BUILD b/third_party/llvm/llvm.autogenerated.BUILD
index 88f007dff1d..d05fa841420 100644
--- a/third_party/llvm/llvm.autogenerated.BUILD
+++ b/third_party/llvm/llvm.autogenerated.BUILD
@@ -3602,6 +3602,7 @@ cc_library(
     ]),
     copts = llvm_copts,
     deps = [
+        ":Core",
         ":MC",
         ":MCDisassembler",
         ":Object",

From b03a4dbbd6a679d02437708e210f115525ed20f1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 00:03:48 -0700
Subject: [PATCH 0265/1390] tf.numpy: Add module comments.

PiperOrigin-RevId: 316626561
Change-Id: I2a750832ccb5461f01c5bf0a948bc15190f4fd39
---
 tensorflow/python/ops/numpy_ops/__init__.py | 108 ++++++++++++++++++--
 1 file changed, 102 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/__init__.py b/tensorflow/python/ops/numpy_ops/__init__.py
index d20b171205e..10ace06df9a 100644
--- a/tensorflow/python/ops/numpy_ops/__init__.py
+++ b/tensorflow/python/ops/numpy_ops/__init__.py
@@ -12,30 +12,126 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tensorflow numpy API."""
+"""numpy_ops.
+
+This module provides a subset of numpy API, built on top of TensorFlow
+operations. APIs are based on numpy 1.16 version.
+
+The set of supported APIs may be expanded over time. Also future releases may
+change the baseline version of numpy API being supported. A list of some
+systematic differences with numpy are listed later in the "Differences with
+Numpy" section.
+
+Types
+-----
+
+The module provide an `ndarray` class which wraps an immutable `tf.Tensor`.
+Additional functions are provided which accept array-like objects. Here
+array-like objects includes `ndarrays` as defined by this module, as well as
+`tf.Tensor`, in addition to types accepted by `numpy`.
+
+A subset of `numpy` dtypes are supported, along with `tf.bfloat16`.
+Additionally, support is provided for selecting the default float type
+(`np.float32` vs `np.float64`) given that some applications may prefer lower
+precision.
+
+Device Support
+-------------
+
+Given that `ndarray` and functions wrap TensorFlow constructs, the code will
+have GPU and TPU support on par with TensorFlow. Also the code can be wrapped
+with `tf.function` and XLA compiled. Device placement can be controlled by using
+`with tf.device` scopes.
+
+Graph and Eager Modes
+--------------------
+
+Eager mode execution should typically match numpy semantics of executing
+op-by-op. However the same code can be executed in graph mode, by putting it
+inside a `tf.function`. This can change behavior of certain operations since
+symbolic execution may not have information that is computed during runtime.
+
+Some differences are:
+  * Shapes can be incomplete or unknown. This means that `ndarray.shape`,
+    `ndarray.size` and `ndarray.ndim` can return `ndarray` objects instead of
+    returning integer (or tuple of integer) values.
+  * Python control flow based on `ndarray` values may not work and may have to
+    be rewritten to use `tf.cond` or `tf.while_loop`. Note that autograph
+    conversion as part of `tf.function` should still work.
+  * `__len__`, `__iter__` and `__index__` properties of `ndarray` may similarly
+    not work in graph mode.
+
+Mutation and Variables
+---------------------
+
+`ndarrays` currently wrap immutable `tf.Tensor`. Also currently mutation
+operations like slice assigns are not supported. This may change in the future.
+
+There is currently no explict construct on par with tf.Variable. However one can
+directly construct a `tf.Variable` and use that with the numpy APIs in this
+module. See section on Interop.
+
+Interop
+------
+
+The numpy API calls can be interleaved with TensorFlow calls without incurring
+Tensor data copies. This is true even if the `ndarray` or `tf.Tensor` is placed
+on a non-CPU device.
+
+Additionally, one could put these calls in a `with tf.GradientTape()` context to
+compute gradients through the numpy API calls. Similarly, code vectorization can
+be done using `tf.vectorized_map()`.
+
+In general, the expected behavior should be on par with that of code involving
+`tf.Tensor` and running stateless TensorFlow functions on them.
+
+Array Interface
+--------------
+
+The `ndarray` class implements the `__array__ interface. This should allow these
+objects to be passed into contexts that expect a `numpy` or array-like object
+(e.g. matplotlib).
+
+
+Differences with Numpy
+---------------------
+
+Here is a non-exhaustive list of differences:
+  * Not all dtypes are currently supported. e.g. `np.float96`, `np.float128`.
+    `np.object`, `np.str`, `np.recarray` types are not supported.
+  * `ndarray` storage is in C order only. Fortran order, views, stride_tricks
+    are not supported.
+  * Only a subset of functions and modules are supported. This set would be
+    expanded over time. For supported functions, some arguments or argument
+    values may not be supported. This differences are listed in the function
+    comments.
+  * Buffer mutation is currently not supported. `ndarrays` wrap immutable
+    tensors. This means that output buffer arguments (e..g `out` in ufuncs) are
+    not supported
+  * full `ufunc` support is not provided.
+  * Numpy C API is not supported. Numpy's Cython and Swig integration are not
+    supported.
+"""
 # pylint: disable=g-direct-tensorflow-import
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.ops.array_ops import newaxis
 from tensorflow.python.ops.numpy_ops import np_random as random
-
 # pylint: disable=wildcard-import
-
 from tensorflow.python.ops.numpy_ops.np_array_ops import *
 from tensorflow.python.ops.numpy_ops.np_arrays import ndarray
 from tensorflow.python.ops.numpy_ops.np_dtypes import *
 from tensorflow.python.ops.numpy_ops.np_math_ops import *
+# pylint: enable=wildcard-import
 from tensorflow.python.ops.numpy_ops.np_utils import finfo
 from tensorflow.python.ops.numpy_ops.np_utils import promote_types
 from tensorflow.python.ops.numpy_ops.np_utils import result_type
-# pylint: enable=wildcard-import
 
 # pylint: disable=redefined-builtin,undefined-variable
 max = amax
 min = amin
 round = around
 # pylint: enable=redefined-builtin,undefined-variable
-
-from tensorflow.python.ops.array_ops import newaxis

From 5107743c47cff6980ebd68d61931bc8b5c3c6a87 Mon Sep 17 00:00:00 2001
From: Anjali Sridhar 
Date: Tue, 16 Jun 2020 00:34:30 -0700
Subject: [PATCH 0266/1390] Add VariablePolicy field to the DistributedVariable
 class as part of an internal refactor that will allow us to attach different
 policies to a DistributedVariable.

PiperOrigin-RevId: 316629999
Change-Id: I20160480b0678657198112adaa61ad7a47823cbd
---
 tensorflow/python/distribute/values.py      | 603 +++++++++++++++-----
 tensorflow/python/distribute/values_util.py | 152 +++++
 2 files changed, 615 insertions(+), 140 deletions(-)

diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index 5c038c01999..c6e0eb34a7b 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -41,6 +41,36 @@ from tensorflow.python.types import core
 from tensorflow.python.util.tf_export import tf_export
 
 
+def _on_write_update_replica(var, update_fn, value, **kwargs):
+  """Updates variables with ON_WRITE synchronization in replica context."""
+  if var.aggregation == vs.VariableAggregation.NONE:
+    return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access
+
+  def merge_fn(strategy, value, **kwargs):
+    """Aggregate values and update all variables in cross replica context."""
+    # Don't allow MEAN with non float dtype, since it may cause unexpected
+    # precision loss. Python3 and NumPy automatically upcast integers to
+    # float in division, but we should always preserve the type.
+    #
+    # Note that to be backward compatible we allow the case when the value
+    # is *always* the same on each replica. I.E. value is not a
+    # PerReplica. Refer to regroup() to see how values are grouped.
+    if var.aggregation == vs.VariableAggregation.MEAN and (
+        not var.dtype.is_floating) and isinstance(value, PerReplica):
+      raise ValueError(
+          "Cannot update non-float variables with "
+          "tf.VariableAggregation.MEAN aggregation in replica context. "
+          "Either change the variable dtype to float or update it in "
+          "cross-replica context.")
+
+    assert strategy == var.distribute_strategy
+    v = values_util.apply_aggregation(strategy, value, var.aggregation, var)
+    return var._update_cross_replica(update_fn, v, **kwargs)  # pylint: disable=protected-access
+
+  return ds_context.get_replica_context().merge_call(
+      merge_fn, args=(value,), kwargs=kwargs)
+
+
 @tf_export("distribute.DistributedValues", v1=[])
 class DistributedValues(object):
   """Base class for representing distributed values.
@@ -409,10 +439,7 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
                           core.Tensor):
   """Holds a map from replica to variables."""
 
-  # TODO(josh11b): Support changing the set of variables if e.g. if new
-  # devices are joining or a device is to leave.
-
-  def __init__(self, strategy, values, aggregation):
+  def __init__(self, strategy, values, aggregation, var_policy=None):
     self._distribute_strategy = strategy
     self._aggregation = aggregation
     super(DistributedVariable, self).__init__(values)
@@ -439,6 +466,9 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
     # when restoring from a checkpoint, we may set the _initializer_op
     # property on the entire `DistributedVariable`.
     self._initializer_op = None
+    # Set a VariablePolicy which decides how we replicate/aggregate the given
+    # variable.
+    self._var_policy = var_policy
 
   def is_initialized(self, name=None):
     """Identifies if all the component variables are initialized.
@@ -580,6 +610,8 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       return array_ops.identity(self._get())
 
   def value(self):
+    if self._var_policy:
+      return self._var_policy.value(self)
     return self._get_on_device_or_primary().value()
 
   def numpy(self):
@@ -590,87 +622,104 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
           "numpy() is only available when eager execution is enabled.")
 
   def assign_sub(self, value, use_locking=False, name=None, read_value=True):
-    assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)
-    return self._update(
-        update_fn=assign_sub_fn,
-        value=value,
-        use_locking=use_locking,
-        name=name,
-        read_value=read_value)
+    if self._var_policy:
+      return self._var_policy.assign_sub(self, value, use_locking=use_locking,
+                                         name=name, read_value=read_value)
+    return values_util.on_write_assign_sub(self, value, use_locking=use_locking,
+                                           name=name, read_value=read_value)
 
   def assign_add(self, value, use_locking=False, name=None, read_value=True):
-    assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw)
-    return self._update(
-        update_fn=assign_add_fn,
-        value=value,
-        use_locking=use_locking,
-        name=name,
-        read_value=read_value)
+    if self._var_policy:
+      return self._var_policy.assign_add(self, value, use_locking=use_locking,
+                                         name=name, read_value=read_value)
+    return values_util.on_write_assign_add(self, value, use_locking=use_locking,
+                                           name=name, read_value=read_value)
 
   def assign(self, value, use_locking=False, name=None, read_value=True):
-    assign_fn = lambda var, *a, **kw: var.assign(*a, **kw)
-    return self._update(
-        update_fn=assign_fn,
-        value=value,
-        use_locking=use_locking,
-        name=name,
-        read_value=read_value)
+    if self._var_policy:
+      return self._var_policy.assign(self, value, use_locking=use_locking,
+                                     name=name, read_value=read_value)
+    return values_util.on_write_assign(self, value, use_locking=use_locking,
+                                       name=name, read_value=read_value)
 
   def scatter_sub(self, sparse_delta, use_locking=False, name=None):
-    scatter_sub_fn = lambda var, *a, **kw: var.scatter_sub(*a, **kw)
-    return self._update(
-        update_fn=scatter_sub_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_sub(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_sub(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_add(self, sparse_delta, use_locking=False, name=None):
-    scatter_add_fn = lambda var, *a, **kw: var.scatter_add(*a, **kw)
-    return self._update(
-        update_fn=scatter_add_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_add(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_add(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_mul(self, sparse_delta, use_locking=False, name=None):
-    scatter_mul_fn = lambda var, *a, **kw: var.scatter_mul(*a, **kw)
-    return self._update(
-        update_fn=scatter_mul_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_mul(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_mul(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_div(self, sparse_delta, use_locking=False, name=None):
-    scatter_div_fn = lambda var, *a, **kw: var.scatter_div(*a, **kw)
-    return self._update(
-        update_fn=scatter_div_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_div(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_div(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_min(self, sparse_delta, use_locking=False, name=None):
-    scatter_min_fn = lambda var, *a, **kw: var.scatter_min(*a, **kw)
-    return self._update(
-        update_fn=scatter_min_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_min(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_min(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_max(self, sparse_delta, use_locking=False, name=None):
-    scatter_max_fn = lambda var, *a, **kw: var.scatter_max(*a, **kw)
-    return self._update(
-        update_fn=scatter_max_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_max(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
+    return values_util.scatter_max(self, sparse_delta, use_locking=use_locking,
+                                   name=name)
 
   def scatter_update(self, sparse_delta, use_locking=False, name=None):
-    scatter_update_fn = lambda var, *a, **kw: var.scatter_update(*a, **kw)
-    return self._update(
-        update_fn=scatter_update_fn,
-        value=sparse_delta,
-        use_locking=use_locking,
-        name=name)
+    if self._var_policy:
+      self._var_policy.scatter_update(self, sparse_delta,
+                                      use_locking=use_locking, name=name)
+    return values_util.scatter_update(self, sparse_delta,
+                                      use_locking=use_locking,
+                                      name=name)
+
+  def _gather_saveables_for_checkpoint(self):
+    """Overrides Trackable method.
+
+    This allows both name-based and object-based save and restore of
+    DistributedVariables.
+
+    Returns:
+      A dictionary mapping attribute names to `SaveableObject` factories.
+    """
+
+    def _saveable_factory(name=self._common_name):
+      return _DistributedVariableSaveable(self, self._primary, name)
+
+    return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}
+
+  def _as_graph_element(self):
+    if self._var_policy:
+      return self._var_policy._as_graph_element(self)  # pylint: disable=protected-access
+
+    raise NotImplementedError("No policy set for calling _as_graph_element.")
+
+  def _get_cross_replica(self):
+    if self._var_policy:
+      return self._var_policy._get_cross_replica(self)  # pylint: disable=protected-access
+
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes which support cross-"
+        "replica accesses.")
 
   def _update_cross_replica(self, update_fn, value, **kwargs):
     """Applies updates across replicas.
@@ -699,6 +748,8 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
     Returns:
       Updated variable or `tf.Operation`.
     """
+    if self._var_policy:
+      return self._var_policy._update_replica(self, update_fn, value, **kwargs)  # pylint: disable=protected-access
     raise NotImplementedError("should be implemented by subclass.")
 
   def _update(self, update_fn, value, **kwargs):
@@ -735,6 +786,31 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
     """Pass resource_variable_ops.is_resource_variable check."""
     pass
 
+  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
+    """Converts a variable to a tensor."""
+    with ds_context.enter_or_assert_strategy(self._distribute_strategy):
+      return ops.convert_to_tensor(
+          self._get(), dtype=dtype, name=name, as_ref=as_ref)
+
+
+class _DistributedVariableSaveable(saveable_object.SaveableObject):
+  """Class for defining how to restore a DistributedVariable."""
+
+  def __init__(self, distributed_variable, primary_variable, name):
+    self._distributed_variable = distributed_variable
+    if not self._distributed_variable._var_policy:
+      raise ValueError("VariablePolicy has not been set for the distributed "
+                       "variable.")
+    tensor, spec = distributed_variable._var_policy.get_saveable(
+        distributed_variable, primary_variable, name)
+    super(_DistributedVariableSaveable, self).__init__(tensor, spec, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restore the same value into all variables."""
+    tensor, = restored_tensors
+    return self._distributed_variable._var_policy.get_restore_ops(  # pylint: disable=protected-access
+        self._distributed_variable, tensor)
+
 
 class _MirroredSaveable(saveable_object_util.ResourceVariableSaveable):
   """Class for defining how to restore a MirroredVariable."""
@@ -756,61 +832,27 @@ class MirroredVariable(DistributedVariable, Mirrored):
   """Holds a map from replica to variables whose values are kept in sync."""
 
   def _update_replica(self, update_fn, value, **kwargs):
-    if self.aggregation == vs.VariableAggregation.NONE:
-      return update_fn(self._get_on_device_or_primary(), value, **kwargs)
-
-    def merge_fn(strategy, value, **kwargs):
-      """Aggregate values and update all variables in cross replica context."""
-      # Don't allow MEAN with non float dtype, since it may cause unexpected
-      # precision loss. Python3 and NumPy automatically upcast integers to
-      # float in division, but we should always preserve the type.
-      #
-      # Note that to be backward compatible we allow the case when the value
-      # is *always* the same on each replica. I.E. value is not a
-      # PerReplica. Refer to regroup() to see how values are grouped.
-      if self._aggregation == vs.VariableAggregation.MEAN and (
-          not self.dtype.is_floating) and isinstance(value, PerReplica):
-        raise ValueError(
-            "Cannot update non-float variables with "
-            "tf.VariableAggregation.MEAN aggregation in replica context. "
-            "Either change the variable dtype to float or update it in "
-            "cross-replica context.")
-
-      assert strategy == self.distribute_strategy
-      v = values_util.apply_aggregation(strategy, value, self.aggregation, self)
-      return self._update_cross_replica(update_fn, v, **kwargs)
-
-    return ds_context.get_replica_context().merge_call(
-        merge_fn, args=(value,), kwargs=kwargs)
+    return _on_write_update_replica(self, update_fn, value, **kwargs)
 
   def scatter_min(self, *args, **kwargs):
     if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
         self._aggregation != vs.VariableAggregation.NONE):
-      raise NotImplementedError("scatter_min is only supported for mirrored "
-                                "variable (variable created within certain "
-                                "`tf.distribute.Strategy` scope) with NONE or "
-                                "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
-                                self._aggregation)
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_min", aggregation=self._aggregation))
     return super(MirroredVariable, self).scatter_min(*args, **kwargs)
 
   def scatter_max(self, *args, **kwargs):
     if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
         self._aggregation != vs.VariableAggregation.NONE):
-      raise NotImplementedError("scatter_max is only supported for mirrored "
-                                "variable (variable created within certain "
-                                "`tf.distribute.Strategy` scope) with NONE or "
-                                "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
-                                self._aggregation)
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_min", aggregation=self._aggregation))
     return super(MirroredVariable, self).scatter_max(*args, **kwargs)
 
   def scatter_update(self, *args, **kwargs):
     if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
         self._aggregation != vs.VariableAggregation.NONE):
-      raise NotImplementedError("scatter_update is only supported for mirrored "
-                                "variable (variable created within certain "
-                                "`tf.distribute.Strategy` scope) with NONE or "
-                                "`ONLY_FIRST_REPLICA` aggregation, got: %s" %
-                                self._aggregation)
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_min", aggregation=self._aggregation))
     return super(MirroredVariable, self).scatter_update(*args, **kwargs)
 
   def _get_cross_replica(self):
@@ -893,28 +935,13 @@ class SyncOnReadVariable(DistributedVariable):
   def _update_replica(self, update_fn, value, **kwargs):
     return update_fn(self._get_on_device_or_primary(), value, **kwargs)
 
-  def _assign_on_each_device(self, assign_func, value, read_value):
-    update = control_flow_ops.group(
-        tuple(
-            assign_func(v.device, v, value)
-            for v in self._values))
-    if not read_value:
-      return update
-    with ops.control_dependencies([update] if update else []):
-      return self.read_value()
-
   # TODO(b/154017756): Make assign behaivor in cross replica context consistent
   # with MirroredVariable.
   def assign_sub(self, value, use_locking=False, name=None, read_value=True):
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       if ds_context.in_cross_replica_context():
-        if self._aggregation == vs.VariableAggregation.SUM:
-          raise ValueError(
-              "SyncOnReadVariable does not support `assign_sub` in "
-              "cross-replica context when aggregation is set to "
-              "`tf.VariableAggregation.SUM`.")
-        return self._assign_on_each_device(values_util.assign_sub_on_device,
-                                           value, read_value)
+        return values_util.on_read_assign_sub_cross_replica(
+            self, value, read_value=read_value)
       else:
         return super(SyncOnReadVariable,
                      self).assign_sub(value, use_locking, name, read_value)
@@ -922,13 +949,8 @@ class SyncOnReadVariable(DistributedVariable):
   def assign_add(self, value, use_locking=False, name=None, read_value=True):
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       if ds_context.in_cross_replica_context():
-        if self._aggregation == vs.VariableAggregation.SUM:
-          raise ValueError(
-              "SyncOnReadVariable does not support `assign_add` in "
-              "cross-replica context when aggregation is set to "
-              "`tf.VariableAggregation.SUM`.")
-        return self._assign_on_each_device(values_util.assign_add_on_device,
-                                           value, read_value)
+        return values_util.on_read_assign_add_cross_replica(
+            self, value, read_value=read_value)
       else:
         return super(SyncOnReadVariable,
                      self).assign_add(value, use_locking, name, read_value)
@@ -936,13 +958,8 @@ class SyncOnReadVariable(DistributedVariable):
   def assign(self, value, use_locking=False, name=None, read_value=True):
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       if ds_context.in_cross_replica_context():
-        # To preserve the sum across save and restore, we have to divide the
-        # total across all devices when restoring a variable that was summed
-        # when saving.
-        if self._aggregation == vs.VariableAggregation.SUM:
-          value = math_ops.cast(value / len(self._values), self.dtype)
-        return self._assign_on_each_device(values_util.assign_on_device, value,
-                                           read_value)
+        return values_util.on_read_assign_cross_replica(
+            self, value, read_value=read_value)
       else:
         return super(SyncOnReadVariable,
                      self).assign(value, use_locking, name, read_value)
@@ -987,7 +1004,7 @@ class SyncOnReadVariable(DistributedVariable):
 
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       return self._distribute_strategy.reduce(
-          reduce_util.ReduceOp.from_variable_aggregation(self.aggregation),
+          reduce_util.ReduceOp.from_variable_aggregation(self._aggregation),
           self,
           axis=None)
 
@@ -1022,6 +1039,16 @@ class SyncOnReadVariable(DistributedVariable):
 
 # Register a conversion functions which reads the value of the variable,
 # allowing instances of the class to be used as tensors.
+# DistributedVariable
+def _tensor_conversion_distributed_var(var, dtype=None, name=None,
+                                       as_ref=False):
+  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access
+
+
+ops.register_tensor_conversion_function(DistributedVariable,
+                                        _tensor_conversion_distributed_var)
+
+
 # MirroredVariables
 def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False):
   return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access
@@ -1048,3 +1075,299 @@ def _tensor_conversion_sync_on_read(var, dtype=None, name=None, as_ref=False):
 
 ops.register_tensor_conversion_function(SyncOnReadVariable,
                                         _tensor_conversion_sync_on_read)
+
+
+class VariablePolicy(object):
+  """Policy defining synchronization and aggregation of a distributed variable.
+
+  Given `synchronization` and `aggregation` parameters set on a `tf.Variable`
+  during variable creation within `tf.distribute` scope, `tf.distribute` creates
+  an appropriate policy object and assigns it to the distributed variable. All
+  variable operations are delegated to the respective policy object.
+  """
+
+  def __init__(self, aggregation):
+    self._aggregation = aggregation
+
+  def value(self):
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes.")
+
+  def _is_mirrored(self):
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes.")
+
+  def _as_graph_element(self, _):
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes.")
+
+  def _get_cross_replica(self, var):
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes.")
+
+  def _update_replica(self, var, update_fn, value, **kwargs):
+    raise NotImplementedError(
+        "This method should be overridden by sub-classes.")
+
+
+class OnReadPolicy(VariablePolicy):
+  """Policy defined for `tf.VariableSynchronization.ON_READ` synchronization.
+
+  This policy is created when `synchronization` is set to
+  `tf.VariableSynchronization.ON_READ` and `aggregation` is set to any of the
+  values allowed by the `tf.VariableAggregation` enum such as `NONE`, `SUM`,
+  `MEAN` or `ONLY_FIRST_REPLICA`when creating a `tf.Variable` in `tf.distribute`
+  scope.
+  """
+
+  def _is_mirrored(self):
+    return False
+
+  def value(self, var):
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      if ds_context.in_cross_replica_context():
+        return var._get_cross_replica()  # pylint: disable=protected-access
+      else:
+        return var._get_on_device_or_primary().value()  # pylint: disable=protected-access
+
+  def _as_graph_element(self, var):
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      if ds_context.in_cross_replica_context():
+        return ops.convert_to_tensor(var._get_cross_replica())   # pylint: disable=protected-access
+    return var._get()._as_graph_element()   # pylint: disable=protected-access
+
+  def _get_cross_replica(self, var):
+    if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
+      return var._primary  # pylint: disable=protected-access
+
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      return  var.distribute_strategy.reduce(
+          reduce_util.ReduceOp.from_variable_aggregation(self._aggregation),
+          var,
+          axis=None)
+
+  def _update_replica(self, var, update_fn, value, **kwargs):
+    return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access
+
+  def _scatter_not_implemented(self, method):
+    raise NotImplementedError(
+        "ON_READ variables doesn't support `%s` in cross replica context" %
+        method)
+
+  def assign_sub(self, var, value, use_locking=False, name=None,
+                 read_value=True):
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      if ds_context.in_cross_replica_context():
+        return values_util.on_read_assign_sub_cross_replica(
+            var, value, read_value=read_value)
+      else:
+        return values_util.on_write_assign_sub(
+            var, value, use_locking=use_locking, name=name,
+            read_value=read_value)
+
+  def assign_add(self, var, value, use_locking=False, name=None,
+                 read_value=True):
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      if ds_context.in_cross_replica_context():
+        return values_util.on_read_assign_add_cross_replica(
+            var, value, read_value=read_value)
+      else:
+        return values_util.on_write_assign_add(
+            var, value, use_locking=use_locking, name=name,
+            read_value=read_value)
+
+  def assign(self, var, value, use_locking=False, name=None, read_value=True):
+    with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+      if ds_context.in_cross_replica_context():
+        return values_util.on_read_assign_cross_replica(var, value,
+                                                        read_value=read_value)
+      else:
+        return values_util.on_write_assign(var, value,
+                                           use_locking=use_locking,
+                                           name=name,
+                                           read_value=read_value)
+
+  def scatter_sub(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_sub")
+
+  def scatter_add(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_add")
+
+  def scatter_mul(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_mul")
+
+  def scatter_div(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_div")
+
+  def scatter_min(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_min")
+
+  def scatter_max(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_max")
+
+  def scatter_update(self, *args, **kwargs):
+    del args, kwargs
+    self._scatter_not_implemented("scatter_update")
+
+  def get_saveable(self, var, primary_var, name):
+    """Create a saveable object for the given variable."""
+    # We use a callable so that we don't have to evaluate this expression
+    # in the case where we are trying to restore instead of save.
+    def tensor():
+      strategy = var.distribute_strategy
+      return strategy.extended.read_var(var)
+
+    spec = saveable_object.SaveSpec(
+        tensor=tensor,
+        slice_spec="",
+        name=name,
+        dtype=var.dtype,
+        device=primary_var.device)
+
+    return tensor, [spec]
+
+  def get_restore_ops(self, var, tensor):
+    """Restore the same value into all variables."""
+    # To preserve the sum across save and restore, we have to divide the
+    # total across all devices when restoring a variable that was summed
+    # when saving.
+    if self._aggregation == vs.VariableAggregation.SUM:
+      tensor = math_ops.cast(tensor / len(var._devices),  # pylint: disable=protected-access
+                             var.dtype)
+    return control_flow_ops.group(
+        tuple(
+            values_util.assign_on_device(v.device, v, tensor)
+            for v in var.values))
+
+
+class AutoPolicy(VariablePolicy):
+  """Policy defined for `tf.VariableSynchronization.AUTO` synchronization.
+
+  This policy is created when `synchronization` is set to
+  `tf.VariableSynchronization.AUTO` and `aggregation` is set to
+  `tf.VariableAggregation.NONE` when creating a `tf.Variable` in `tf.distribute`
+  scope.
+  """
+
+  def _is_mirrored(self):
+    return True
+
+  def value(self, var):
+    return var._get_on_device_or_primary().value()  # pylint: disable=protected-access
+
+  def _as_graph_element(self, var):
+    return var._get_on_device_or_primary()._as_graph_element()  # pylint: disable=protected-access
+
+  def _get_cross_replica(self, var):
+    # Return identity, to avoid directly exposing the variable to the user and
+    # allowing it to be modified by mistake.
+    return array_ops.identity(Mirrored._get_cross_replica(var))  # pylint: disable=protected-access
+
+  def _update_replica(self, var, update_fn, value, **kwargs):
+    return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access
+
+  def assign(self, var, value, use_locking=False, name=None, read_value=True):
+    return values_util.on_write_assign(var, value, use_locking=use_locking,
+                                       name=name, read_value=read_value)
+
+  def assign_add(self, var, value, use_locking=False, name=None,
+                 read_value=True):
+    return values_util.on_write_assign_add(var, value, use_locking=use_locking,
+                                           name=name, read_value=read_value)
+
+  def assign_sub(self, var, value, use_locking=False, name=None,
+                 read_value=True):
+    return values_util.on_write_assign_sub(var, value, use_locking=use_locking,
+                                           name=name, read_value=read_value)
+
+  def scatter_sub(self, var, sparse_delta, use_locking=False, name=None):
+    return values_util.scatter_sub(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_add(self, var, sparse_delta, use_locking=False, name=None):
+    return values_util.scatter_add(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_mul(self, var, sparse_delta, use_locking=False, name=None):
+    return values_util.scatter_mul(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_div(self, var, sparse_delta, use_locking=False, name=None):
+    return values_util.scatter_div(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_min(self, var, sparse_delta, use_locking=False, name=None):
+    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
+        self._aggregation != vs.VariableAggregation.NONE):
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_min", aggregation=self._aggregation))
+    return values_util.scatter_min(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_max(self, var, sparse_delta, use_locking=False, name=None):
+    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
+        self._aggregation != vs.VariableAggregation.NONE):
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_max", aggregation=self._aggregation))
+    return values_util.scatter_max(var, sparse_delta, use_locking=use_locking,
+                                   name=name)
+
+  def scatter_update(self, var, sparse_delta, use_locking=False, name=None):
+    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
+        self._aggregation != vs.VariableAggregation.NONE):
+      raise NotImplementedError(values_util.scatter_error_msg.format(
+          op_name="scatter_update", aggregation=self._aggregation))
+    return values_util.scatter_update(var, sparse_delta,
+                                      use_locking=use_locking,
+                                      name=name)
+
+  def get_saveable(self, var, primary_var, name):
+    del var, name
+    return primary_var, ""
+
+  def get_restore_ops(self, var, tensor):
+    return control_flow_ops.group(
+        tuple(
+            values_util.assign_on_device(v.device, v, tensor)
+            for v in var.values))
+
+
+class OnWritePolicy(AutoPolicy):
+  """Policy defined for `tf.VariableSynchronization.ON_WRITE` synchronization.
+
+  This policy is created when the following `synchronization` and
+  `aggregation` parameters are specified when creating a `tf.Variable` in
+  `tf.distribute` scope:
+  * `synchronization` is equal to `tf.VariableSynchronization.AUTO` and
+  aggregation can be any of the following `tf.VariableAggregation` enum
+  values such as `SUM`, `MEAN` or `ONLY_FIRST_REPLICA`.
+  * `synchronization` is equal to `tf.VariableSynchronization.ON_WRITE` and
+  aggregation can be any of the following `tf.VariableAggregation` enum
+  values such as `NONE`, `SUM`, `MEAN` or `ONLY_FIRST_REPLICA`.
+  """
+
+  def _update_replica(self, var, update_fn, value, **kwargs):
+    return _on_write_update_replica(var, update_fn, value, **kwargs)
+
+
+# Utility functions
+# Return True if the Value is Mirrored or the Variable is replicated and kept in
+# sync.
+def _is_mirrored(val):
+  if isinstance(val, DistributedVariable):
+    if val._var_policy:  # pylint: disable=protected-access
+      return val._var_policy._is_mirrored()  # pylint: disable=protected-access
+  return isinstance(val, Mirrored)
+
+
+def _is_sync_on_read(val):
+  if isinstance(val, DistributedVariable):
+    if val._var_policy:  # pylint: disable=protected-access
+      return not val._var_policy._is_mirrored()  # pylint: disable=protected-access
+  return not isinstance(val, Mirrored)
diff --git a/tensorflow/python/distribute/values_util.py b/tensorflow/python/distribute/values_util.py
index c42ac9e4de1..ddb0d2d0401 100644
--- a/tensorflow/python/distribute/values_util.py
+++ b/tensorflow/python/distribute/values_util.py
@@ -23,9 +23,155 @@ from tensorflow.python.distribute import distribution_strategy_context as ds_con
 from tensorflow.python.distribute import reduce_util
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variable_scope as vs
 
 
+def on_write_assign(var, value, use_locking=False, name=None, read_value=True):
+  assign_fn = lambda var, *a, **kw: var.assign(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=assign_fn,
+      value=value,
+      use_locking=use_locking,
+      name=name,
+      read_value=read_value)
+
+
+def on_write_assign_add(var, value, use_locking=False, name=None,
+                        read_value=True):
+  assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=assign_add_fn,
+      value=value,
+      use_locking=use_locking,
+      name=name,
+      read_value=read_value)
+
+
+def on_write_assign_sub(var, value, use_locking=False, name=None,
+                        read_value=True):
+  assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=assign_sub_fn,
+      value=value,
+      use_locking=use_locking,
+      name=name,
+      read_value=read_value)
+
+
+def assign_on_each_device(var, assign_func, value, read_value):
+  update = control_flow_ops.group(
+      tuple(assign_func(v.device, v, value) for v in var._values))  # pylint: disable=protected-access
+  if not read_value:
+    return update
+  with ops.control_dependencies([update] if update else []):
+    return var.read_value()
+
+
+def on_read_assign_sub_cross_replica(var, value, read_value=True):
+  with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+    if ds_context.in_cross_replica_context():
+      if var.aggregation == vs.VariableAggregation.SUM:
+        raise ValueError(
+            "SyncOnReadVariable does not support `assign_sub` in "
+            "cross-replica context when aggregation is set to "
+            "`tf.VariableAggregation.SUM`.")
+      return assign_on_each_device(var, assign_sub_on_device,
+                                   value, read_value)
+
+
+def on_read_assign_add_cross_replica(var, value, read_value=True):
+  with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+    if ds_context.in_cross_replica_context():
+      if var.aggregation == vs.VariableAggregation.SUM:
+        raise ValueError(
+            "SyncOnReadVariable does not support `assign_add` in "
+            "cross-replica context when aggregation is set to "
+            "`tf.VariableAggregation.SUM`.")
+      return assign_on_each_device(var, assign_add_on_device,
+                                   value, read_value)
+
+
+def on_read_assign_cross_replica(var, value, read_value=True):
+  """Return the value of the variable in cross replica context."""
+  with ds_context.enter_or_assert_strategy(var.distribute_strategy):
+    if ds_context.in_cross_replica_context():
+      # To preserve the sum across save and restore, we have to divide the
+      # total across all devices when restoring a variable that was summed
+      # when saving.
+      tensor = value
+      # TODO(anjs): Should this be over all the replicas in sync since we
+      # call `reduce` on the variable during read?
+      if var.aggregation == vs.VariableAggregation.SUM:
+        tensor = math_ops.cast(tensor / len(var._values), var.dtype)  # pylint: disable=protected-access
+      return assign_on_each_device(var, assign_on_device, tensor,
+                                   read_value)
+
+
+def scatter_sub(var, sparse_delta, use_locking=False, name=None):
+  scatter_sub_fn = lambda var, *a, **kw: var.scatter_sub(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_sub_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_add(var, sparse_delta, use_locking=False, name=None):
+  scatter_add_fn = lambda var, *a, **kw: var.scatter_add(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_add_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_mul(var, sparse_delta, use_locking=False, name=None):
+  scatter_mul_fn = lambda var, *a, **kw: var.scatter_mul(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_mul_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_div(var, sparse_delta, use_locking=False, name=None):
+  scatter_div_fn = lambda var, *a, **kw: var.scatter_div(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_div_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_min(var, sparse_delta, use_locking=False, name=None):
+  scatter_min_fn = lambda var, *a, **kw: var.scatter_min(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_min_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_max(var, sparse_delta, use_locking=False, name=None):
+  scatter_max_fn = lambda var, *a, **kw: var.scatter_max(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_max_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
+def scatter_update(var, sparse_delta, use_locking=False, name=None):
+  scatter_update_fn = lambda var, *a, **kw: var.scatter_update(*a, **kw)
+  return var._update(  # pylint: disable=protected-access
+      update_fn=scatter_update_fn,
+      value=sparse_delta,
+      use_locking=use_locking,
+      name=name)
+
+
 def get_current_replica_id_as_int():
   """Returns the current replica ID as an integer, or `None`."""
   replica_context = ds_context.get_replica_context()
@@ -89,3 +235,9 @@ aggregation_error_msg = (
     "`tf.distribute.get_replica_context().merge_call(merge_fn, ..)`."
     "Inside `merge_fn`, you can then update the {variable_type} "
     "using `tf.distribute.StrategyExtended.update()`.")
+
+
+scatter_error_msg = ("{op_name} is only supported for mirrored "
+                     "variable (variable created within certain "
+                     "`tf.distribute.Strategy` scope) with NONE or "
+                     "`ONLY_FIRST_REPLICA` aggregation, got: {aggregation}.")

From e887f933e4fd4965742944668f94255e9a603533 Mon Sep 17 00:00:00 2001
From: Thai Nguyen 
Date: Tue, 16 Jun 2020 00:39:41 -0700
Subject: [PATCH 0267/1390] Add the missing versions to runtime version

This cl also add a new unit test to ensure new version
got reflected in runtime version.

PiperOrigin-RevId: 316630739
Change-Id: I139ecf3077b2eec9bcc13ea5bb9030199d29203a
---
 .../lite/tests/mlir2flatbuffer/if_op.mlir     |  2 +-
 .../tests/mlir2flatbuffer/quantization.mlir   |  2 +-
 .../tests/mlir2flatbuffer/tfl_while_op.mlir   |  2 +-
 .../lite/tests/mlir2flatbuffer/while_op.mlir  |  2 +-
 tensorflow/lite/tools/versioning/BUILD        |  2 +
 .../lite/tools/versioning/runtime_version.cc  | 98 ++++++++++++++++---
 .../lite/tools/versioning/runtime_version.h   |  6 ++
 .../tools/versioning/runtime_version_test.cc  | 23 ++++-
 8 files changed, 118 insertions(+), 19 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/if_op.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/if_op.mlir
index b89ba0fe400..7290209cc4a 100644
--- a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/if_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/if_op.mlir
@@ -157,7 +157,7 @@
 // CHECK-NEXT:   }, {
 // CHECK-EMPTY:
 // CHECK-NEXT:   }, {
-// CHECK-NEXT:   data: [ 49, 46, 49, 52, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+// CHECK-NEXT:   data: [ 49, 46, 49, 53, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
 // CHECK-NEXT:   } ],
 // CHECK-NEXT:   metadata: [ {
 // CHECK-NEXT:   name: "min_runtime_version",
diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/quantization.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/quantization.mlir
index 9bbbdca1c97..dbe10a3f90c 100644
--- a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/quantization.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/quantization.mlir
@@ -154,7 +154,7 @@ func @main(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x1001xf32> {
 // CHECK-NEXT:  }, {
 // CHECK-EMPTY:
 // CHECK-NEXT:  }, {
-// CHECK-NEXT:    data: [ 49, 46, 49, 51, 46, 49, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+// CHECK-NEXT:    data: [ 49, 46, 49, 52, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
 // CHECK-NEXT:  } ],
 // CHECK-NEXT:  metadata: [ {
 // CHECK-NEXT:  name: "min_runtime_version",
diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/tfl_while_op.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/tfl_while_op.mlir
index 1d3a70f0996..996543cc9c7 100644
--- a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/tfl_while_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/tfl_while_op.mlir
@@ -190,7 +190,7 @@
 // CHECK-NEXT:   }, {
 // CHECK-EMPTY:
 // CHECK-NEXT:   }, {
-// CHECK-NEXT:     data: [ 49, 46, 49, 52, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+// CHECK-NEXT:     data: [ 49, 46, 49, 53, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
 // CHECK-NEXT:   } ],
 // CHECK-NEXT:  metadata: [ {
 // CHECK-NEXT:  name: "min_runtime_version",
diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/while_op.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/while_op.mlir
index a76fbbeb871..d69e8f40311 100644
--- a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/while_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/while_op.mlir
@@ -190,7 +190,7 @@
 // CHECK-NEXT:   }, {
 // CHECK-EMPTY:
 // CHECK-NEXT:   }, {
-// CHECK-NEXT:     data: [ 49, 46, 49, 52, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+// CHECK-NEXT:     data: [ 49, 46, 49, 53, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
 // CHECK-NEXT:   } ],
 // CHECK-NEXT:   metadata: [ {
 // CHECK-NEXT:   name: "min_runtime_version",
diff --git a/tensorflow/lite/tools/versioning/BUILD b/tensorflow/lite/tools/versioning/BUILD
index 1ba221d3fa9..34d63bd9645 100644
--- a/tensorflow/lite/tools/versioning/BUILD
+++ b/tensorflow/lite/tools/versioning/BUILD
@@ -38,6 +38,8 @@ tf_cc_test(
     ],
     deps = [
         ":versioning",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/schema:schema_fbs",
         "//tensorflow/lite/schema:schema_fbs_with_mutable",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/tensorflow/lite/tools/versioning/runtime_version.cc b/tensorflow/lite/tools/versioning/runtime_version.cc
index 702f9fc7f85..92a7001606f 100644
--- a/tensorflow/lite/tools/versioning/runtime_version.cc
+++ b/tensorflow/lite/tools/versioning/runtime_version.cc
@@ -22,6 +22,12 @@ limitations under the License.
 #include "tensorflow/lite/schema/mutable/schema_generated.h"
 
 namespace tflite {
+namespace {
+// Use this as the placeholder string if a particular op is not yet included
+// in any Tensorflow's RC/Final release source package. Once that op is
+// included in the release, please update this with the real version string.
+static constexpr char kPendingReleaseVersion[] = "UNKNOWN";
+}  // namespace
 
 bool CompareRuntimeVersion(const std::string& v1, const std::string& v2) {
   const std::vector vec1 = absl::StrSplit(v1, '.');
@@ -40,11 +46,8 @@ bool CompareRuntimeVersion(const std::string& v1, const std::string& v2) {
   return i < vec2.size();
 }
 
-void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
-  // Use this as the placeholder string if a particular op is not yet included
-  // in any Tensorflow's RC/Final release source package. Once that op is
-  // included in the release, please update this with the real version string.
-  static constexpr char kPendingReleaseOpVersion[] = "UNKNOWN";
+std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
+                                           int op_version) {
   // A map from the version key of an op to its minimum runtime version.
   // For example, {{kAveragePool, 1}, "1.5.0"},  means the 1st version of
   // AveragePool requires a minimum TF Lite runtime version '1.5.0`.
@@ -53,27 +56,41 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
           new std::map, std::string>({
               {{BuiltinOperator_AVERAGE_POOL_2D, 1}, "1.5.0"},
               {{BuiltinOperator_AVERAGE_POOL_2D, 2}, "1.14.0"},
+              {{BuiltinOperator_AVERAGE_POOL_2D, 3}, kPendingReleaseVersion},
+              {{BuiltinOperator_BATCH_MATMUL, 1}, kPendingReleaseVersion},
               {{BuiltinOperator_CONV_2D, 1}, "1.5.0"},
               {{BuiltinOperator_CONV_2D, 2}, "1.14.0"},
               {{BuiltinOperator_CONV_2D, 3}, "1.14.0"},
+              {{BuiltinOperator_CONV_2D, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 1}, "1.5.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 2}, "1.12.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 3}, "1.14.0"},
+              {{BuiltinOperator_DEPTHWISE_CONV_2D, 4}, "2.2.0"},
+              {{BuiltinOperator_DEPTHWISE_CONV_2D, 5}, kPendingReleaseVersion},
               {{BuiltinOperator_ADD, 1}, "1.5.0"},
               {{BuiltinOperator_ADD, 2}, "1.14.0"},
               {{BuiltinOperator_ADD_N, 1}, "1.14.0"},
               {{BuiltinOperator_SPACE_TO_BATCH_ND, 1}, "1.6.0"},
               {{BuiltinOperator_SPACE_TO_BATCH_ND, 2}, "1.14.0"},
+              {{BuiltinOperator_SPACE_TO_BATCH_ND, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_SUB, 1}, "1.6.0"},
               {{BuiltinOperator_SUB, 2}, "1.14.0"},
+              {{BuiltinOperator_SUB, 3}, kPendingReleaseVersion},
+              {{BuiltinOperator_DENSIFY, 1}, "2.2.0"},
               {{BuiltinOperator_DIV, 1}, "1.6.0"},
-              {{BuiltinOperator_DIV, 2}, kPendingReleaseOpVersion},
+              {{BuiltinOperator_DIV, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_BATCH_TO_SPACE_ND, 1}, "1.6.0"},
               {{BuiltinOperator_BATCH_TO_SPACE_ND, 2}, "1.14.0"},
+              {{BuiltinOperator_BATCH_TO_SPACE_ND, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_CAST, 1}, "1.5.0"},
               {{BuiltinOperator_CONCATENATION, 1}, "1.5.0"},
               {{BuiltinOperator_CONCATENATION, 2}, "1.14.0"},
+              {{BuiltinOperator_CONCATENATION, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_DEPTH_TO_SPACE, 1}, "2.1.0"},
+              {{BuiltinOperator_EMBEDDING_LOOKUP, 1}, "1.13.0"},
+              {{BuiltinOperator_EMBEDDING_LOOKUP, 2}, "1.14.0"},
+              {{BuiltinOperator_EMBEDDING_LOOKUP, 3}, "1.14.0"},
+              {{BuiltinOperator_EMBEDDING_LOOKUP_SPARSE, 1}, "1.5.0"},
               {{BuiltinOperator_FAKE_QUANT, 1}, "1.5.0"},
               {{BuiltinOperator_FAKE_QUANT, 2}, "1.10.0"},
               {{BuiltinOperator_FULLY_CONNECTED, 1}, "1.5.0"},
@@ -82,10 +99,14 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_FULLY_CONNECTED, 4}, "1.14.0"},
               {{BuiltinOperator_FULLY_CONNECTED, 5}, "2.0.0"},
               {{BuiltinOperator_FULLY_CONNECTED, 6}, "2.1.0"},
+              {{BuiltinOperator_FULLY_CONNECTED, 7}, kPendingReleaseVersion},
+              {{BuiltinOperator_FULLY_CONNECTED, 8}, kPendingReleaseVersion},
               {{BuiltinOperator_GATHER, 1}, "1.6.0"},
               {{BuiltinOperator_GATHER, 2}, "1.14.0"},
               {{BuiltinOperator_GATHER, 3}, "1.15.0"},
               {{BuiltinOperator_GATHER_ND, 1}, "1.14.0"},
+              {{BuiltinOperator_GATHER_ND, 2}, kPendingReleaseVersion},
+              {{BuiltinOperator_HASHTABLE_LOOKUP, 1}, "1.5.0"},
               {{BuiltinOperator_SVDF, 1}, "1.5.0"},
               {{BuiltinOperator_SVDF, 2}, "1.14.0"},
               {{BuiltinOperator_SVDF, 3}, "2.2.0"},
@@ -95,13 +116,21 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, 1}, "1.5.0"},
               {{BuiltinOperator_MAX_POOL_2D, 1}, "1.5.0"},
               {{BuiltinOperator_MAX_POOL_2D, 2}, "1.14.0"},
+              {{BuiltinOperator_MAX_POOL_2D, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_MAXIMUM, 1}, "1.14.0"},
               {{BuiltinOperator_MAXIMUM, 2}, "1.14.0"},
+              {{BuiltinOperator_MAXIMUM, 3}, kPendingReleaseVersion},
+              {{BuiltinOperator_MAXIMUM, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_MINIMUM, 1}, "1.14.0"},
               {{BuiltinOperator_MINIMUM, 2}, "1.14.0"},
+              {{BuiltinOperator_MINIMUM, 3}, kPendingReleaseVersion},
+              {{BuiltinOperator_MINIMUM, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_MUL, 1}, "1.5.0"},
               {{BuiltinOperator_MUL, 2}, "1.14.0"},
               {{BuiltinOperator_MUL, 3}, "1.15.0"},
+              {{BuiltinOperator_MUL, 4}, kPendingReleaseVersion},
+              {{BuiltinOperator_NON_MAX_SUPPRESSION_V4, 1}, "2.1.0"},
+              {{BuiltinOperator_NON_MAX_SUPPRESSION_V5, 1}, "2.1.0"},
               {{BuiltinOperator_PAD, 1}, "1.5.0"},
               {{BuiltinOperator_PAD, 2}, "1.14.0"},
               {{BuiltinOperator_TILE, 1}, "1.10.1"},
@@ -111,18 +140,23 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_RESHAPE, 1}, "1.5.0"},
               {{BuiltinOperator_SOFTMAX, 1}, "1.5.0"},
               {{BuiltinOperator_SOFTMAX, 2}, "1.14.0"},
+              {{BuiltinOperator_SOFTMAX, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_SPACE_TO_DEPTH, 1}, "1.5.0"},
               {{BuiltinOperator_SPACE_TO_DEPTH, 2}, "1.14.0"},
               {{BuiltinOperator_TRANSPOSE, 1}, "1.6.0"},
               {{BuiltinOperator_TRANSPOSE, 2}, "1.14.0"},
               {{BuiltinOperator_TRANSPOSE, 3}, "1.15.0"},
+              {{BuiltinOperator_TRANSPOSE, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_LSTM, 1}, "1.7.0"},
               {{BuiltinOperator_LSTM, 2}, "1.10.0"},
               {{BuiltinOperator_LSTM, 3}, "1.14.0"},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 1}, "1.13.1"},
-              {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 1}, "1.14.0"},
+              {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 2}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 1}, "1.14.0"},
+              {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 2}, "1.14.0"},
+              {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 3}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, 1}, "1.14.0"},
+              {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, 2}, "1.14.0"},
               {{BuiltinOperator_MEAN, 1}, "1.6.0"},
               {{BuiltinOperator_MEAN, 2}, "1.14.0"},
               {{BuiltinOperator_SUM, 1}, "1.10.0"},
@@ -140,14 +174,22 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_RESIZE_BILINEAR, 3}, "2.2.0"},
               {{BuiltinOperator_RESIZE_NEAREST_NEIGHBOR, 1}, "1.13.1"},
               {{BuiltinOperator_RESIZE_NEAREST_NEIGHBOR, 2}, "1.14.0"},
+              {{BuiltinOperator_RESIZE_NEAREST_NEIGHBOR, 3},
+               kPendingReleaseVersion},
+              {{BuiltinOperator_RNN, 1}, "1.5.0"},
+              {{BuiltinOperator_RNN, 2}, "1.14.0"},
+              {{BuiltinOperator_SKIP_GRAM, 1}, "1.5.0"},
               {{BuiltinOperator_SQUEEZE, 1}, "1.6.0"},
               {{BuiltinOperator_SPLIT, 1}, "1.5.0"},
               {{BuiltinOperator_SPLIT, 2}, "1.14.0"},
               {{BuiltinOperator_SPLIT, 3}, "1.14.0"},
+              {{BuiltinOperator_SPLIT, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_SPLIT_V, 1}, "1.13.1"},
+              {{BuiltinOperator_SPLIT_V, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_STRIDED_SLICE, 1}, "1.6.0"},
               {{BuiltinOperator_STRIDED_SLICE, 2}, "1.14.0"},
               {{BuiltinOperator_STRIDED_SLICE, 3}, "2.1.0"},
+              {{BuiltinOperator_STRIDED_SLICE, 4}, "2.2.0"},
               {{BuiltinOperator_TOPK_V2, 1}, "1.7.0"},
               {{BuiltinOperator_TOPK_V2, 2}, "1.14.0"},
               {{BuiltinOperator_ARG_MAX, 1}, "1.9.0"},
@@ -155,40 +197,53 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_ARG_MIN, 1}, "1.9.0"},
               {{BuiltinOperator_ARG_MIN, 2}, "1.14.0"},
               {{BuiltinOperator_TRANSPOSE_CONV, 1}, "1.9.0"},
+              {{BuiltinOperator_TRANSPOSE_CONV, 2}, "2.2.0"},
+              {{BuiltinOperator_TRANSPOSE_CONV, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_SPARSE_TO_DENSE, 1}, "1.9.0"},
               {{BuiltinOperator_SPARSE_TO_DENSE, 2}, "1.14.0"},
               {{BuiltinOperator_SPARSE_TO_DENSE, 3}, "1.15.0"},
               {{BuiltinOperator_EXPAND_DIMS, 1}, "1.10.0"},
               {{BuiltinOperator_PACK, 1}, "1.11.0"},
               {{BuiltinOperator_PACK, 2}, "1.14.0"},
+              {{BuiltinOperator_PACK, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_SHAPE, 1}, "1.10.0"},
               {{BuiltinOperator_SLICE, 1}, "1.14.0"},
               {{BuiltinOperator_SLICE, 2}, "1.14.0"},
               {{BuiltinOperator_SLICE, 3}, "1.14.0"},
               {{BuiltinOperator_TANH, 1}, "1.14.0"},
               {{BuiltinOperator_TANH, 2}, "1.14.0"},
+              {{BuiltinOperator_TANH, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_ONE_HOT, 1}, "1.11.0"},
               {{BuiltinOperator_UNPACK, 1}, "1.11.0"},
               {{BuiltinOperator_UNPACK, 2}, "1.14.0"},
               {{BuiltinOperator_UNPACK, 3}, "2.2.0"},
+              {{BuiltinOperator_UNPACK, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_LEAKY_RELU, 1}, "1.13.1"},
+              {{BuiltinOperator_LEAKY_RELU, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_LOGISTIC, 1}, "1.14.0"},
               {{BuiltinOperator_LOGISTIC, 2}, "1.14.0"},
+              {{BuiltinOperator_LOGISTIC, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_LOG_SOFTMAX, 1}, "1.14.0"},
               {{BuiltinOperator_LOG_SOFTMAX, 2}, "1.14.0"},
+              {{BuiltinOperator_LSH_PROJECTION, 1}, "1.5.0"},
               {{BuiltinOperator_SQUARED_DIFFERENCE, 1}, "1.13.1"},
               {{BuiltinOperator_MIRROR_PAD, 1}, "1.13.1"},
+              {{BuiltinOperator_MIRROR_PAD, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_UNIQUE, 1}, "1.14.0"},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, 1}, "1.14.0"},
+              {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, 2}, "1.14.0"},
               {{BuiltinOperator_WHERE, 1}, "1.14.0"},
               {{BuiltinOperator_DEQUANTIZE, 1}, "1.13.1"},
               {{BuiltinOperator_DEQUANTIZE, 2}, "1.14.0"},
               {{BuiltinOperator_DEQUANTIZE, 3}, "1.15.0"},
+              {{BuiltinOperator_DEQUANTIZE, 4}, "2.2.0"},
               {{BuiltinOperator_REVERSE_SEQUENCE, 1}, "1.14.0"},
               {{BuiltinOperator_EQUAL, 1}, "1.14.0"},
               {{BuiltinOperator_EQUAL, 2}, "1.14.0"},
+              {{BuiltinOperator_EQUAL, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_NOT_EQUAL, 1}, "1.14.0"},
               {{BuiltinOperator_NOT_EQUAL, 2}, "1.14.0"},
+              {{BuiltinOperator_NOT_EQUAL, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_GREATER, 1}, "1.14.0"},
               {{BuiltinOperator_GREATER, 2}, "1.14.0"},
               {{BuiltinOperator_GREATER_EQUAL, 1}, "1.14.0"},
@@ -197,10 +252,12 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_LESS, 2}, "1.14.0"},
               {{BuiltinOperator_LESS_EQUAL, 1}, "1.14.0"},
               {{BuiltinOperator_LESS_EQUAL, 2}, "1.14.0"},
+              {{BuiltinOperator_SCATTER_ND, 1}, "2.1.0"},
               {{BuiltinOperator_SEGMENT_SUM, 1}, "2.2.0"},
               {{BuiltinOperator_SELECT, 1}, "1.14.0"},
               {{BuiltinOperator_SELECT, 2}, "1.14.0"},
               {{BuiltinOperator_SELECT_V2, 1}, "2.2.0"},
+              {{BuiltinOperator_IF, 1}, "1.15.0"},
               {{BuiltinOperator_FLOOR_DIV, 1}, "1.14.0"},
               {{BuiltinOperator_FLOOR_DIV, 2}, "1.14.0"},
               {{BuiltinOperator_FLOOR, 1}, "1.9.0"},
@@ -208,6 +265,8 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_MATRIX_DIAG, 1}, "1.14.0"},
               {{BuiltinOperator_MATRIX_SET_DIAG, 1}, "1.14.0"},
               {{BuiltinOperator_ELU, 1}, "1.14.0"},
+              {{BuiltinOperator_QUANTIZE, 1}, "1.14.0"},
+              {{BuiltinOperator_QUANTIZE, 2}, "1.15.0"},
               {{BuiltinOperator_ROUND, 1}, "1.14.0"},
               {{BuiltinOperator_RELU, 1}, "1.5.0"},
               {{BuiltinOperator_RELU, 2}, "2.1.0"},
@@ -224,17 +283,29 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
               {{BuiltinOperator_RANGE, 1}, "1.13.0"},
               {{BuiltinOperator_SIN, 1}, "1.9.0"},
               {{BuiltinOperator_LOG, 1}, "1.14.0"},
+              {{BuiltinOperator_SQRT, 1}, "1.10.0"},
               {{BuiltinOperator_RSQRT, 1}, "1.10.0"},
               {{BuiltinOperator_SQUARE, 1}, "1.12.0"},
               {{BuiltinOperator_ZEROS_LIKE, 1}, "1.12.0"},
               {{BuiltinOperator_ABS, 1}, "1.13.0"},
               {{BuiltinOperator_HARD_SWISH, 1}, "1.15.0"},
               {{BuiltinOperator_FILL, 1}, "1.13.0"},
+              {{BuiltinOperator_FILL, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_REVERSE_V2, 1}, "1.14.0"},
               {{BuiltinOperator_REVERSE_V2, 2}, "2.2.0"},
               {{BuiltinOperator_RANK, 1}, "1.14.0"},
+              {{BuiltinOperator_WHILE, 1}, "1.15.0"},
           });
 
+  std::pair version_key = {op_code, op_version};
+  auto it = op_version_map->find(version_key);
+  if (it == op_version_map->end()) {
+    return std::string();
+  }
+  return it->second;
+}
+
+void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
   auto model = GetMutableModel(model_buffer_pointer);
   std::string model_min_version;
   auto subgraphs = model->subgraphs();
@@ -244,19 +315,18 @@ void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer) {
       const Operator* op = subgraph->operators()->Get(j);
       const OperatorCode* op_code =
           model->operator_codes()->Get(op->opcode_index());
-      std::pair version_key = {op_code->builtin_code(),
-                                                     op_code->version()};
-      auto it = op_version_map->find(version_key);
-      if (it == op_version_map->end() ||
-          it->second == kPendingReleaseOpVersion) {
+      std::string runtime_version = FindMinimumRuntimeVersionForOp(
+          op_code->builtin_code(), op_code->version());
+      if (runtime_version.empty() ||
+          runtime_version == kPendingReleaseVersion) {
         // In case we didn't find the current op in the map, or the operator
         // doesn't have a minimum runtime version associated, continue.
         continue;
       }
-      if (CompareRuntimeVersion(model_min_version, it->second)) {
+      if (CompareRuntimeVersion(model_min_version, runtime_version)) {
         // Current min model runtime version should be bumped if we see a higher
         // op version.
-        model_min_version = it->second;
+        model_min_version = runtime_version;
       }
     }
   }
diff --git a/tensorflow/lite/tools/versioning/runtime_version.h b/tensorflow/lite/tools/versioning/runtime_version.h
index ad88bd2ab89..64329eb1118 100644
--- a/tensorflow/lite/tools/versioning/runtime_version.h
+++ b/tensorflow/lite/tools/versioning/runtime_version.h
@@ -18,11 +18,17 @@ limitations under the License.
 #include 
 
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/schema/mutable/schema_generated.h"
 
 namespace tflite {
 // Update minimum runtime version of the given TFL flatbuffer model.
 void UpdateMinimumRuntimeVersionForModel(uint8_t* model_buffer_pointer);
 
+// Find the minimum runtime version of a given op version. Return an empty
+// string the version is not registered.
+std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
+                                           int op_version);
+
 // Returns true if the first version string precedes the second.
 // For example, '1.9' should precede '1.14', also '1.14' should precede
 // '1.14.1'. If two version string is equal, then false will be returned.
diff --git a/tensorflow/lite/tools/versioning/runtime_version_test.cc b/tensorflow/lite/tools/versioning/runtime_version_test.cc
index c7b70552340..c32de228cc3 100644
--- a/tensorflow/lite/tools/versioning/runtime_version_test.cc
+++ b/tensorflow/lite/tools/versioning/runtime_version_test.cc
@@ -16,7 +16,8 @@ limitations under the License.
 
 #include 
 #include 
-
+#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 namespace tflite {
 
 TEST(OpVersionTest, CompareRuntimeVersion) {
@@ -31,4 +32,24 @@ TEST(OpVersionTest, CompareRuntimeVersion) {
   EXPECT_FALSE(CompareRuntimeVersion("", ""));
 }
 
+// This test will fail if an op version is added to a builtin op, but not
+// registered to runtime version.
+TEST(OpVersionTest, OpversionMissing) {
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+
+  for (int id = BuiltinOperator_MIN; id <= BuiltinOperator_MAX; ++id) {
+    for (int version = 1;; ++version) {
+      auto op_code = static_cast(id);
+      if (resolver.FindOp(op_code, version) == nullptr) break;
+      // Throw error if the version is not registered in runtime version.
+      std::string runtime_version =
+          FindMinimumRuntimeVersionForOp(op_code, version);
+      EXPECT_NE(runtime_version, "")
+          << "Please add the version " << version << " of "
+          << tflite::EnumNamesBuiltinOperator()[op_code]
+          << " runtime_version.cc";
+    }
+  }
+}
+
 }  // namespace tflite

From ec0e105c6fe537969a736ddb546c277ae18b9282 Mon Sep 17 00:00:00 2001
From: Thai Nguyen 
Date: Tue, 16 Jun 2020 00:47:28 -0700
Subject: [PATCH 0268/1390] Fix build failure of list_flex_ops_main in OSS

The cc_binary required --config=monolithic which can't be passed
into a native.genrule. Using tf_cc_binary solves the build failure.

PiperOrigin-RevId: 316631689
Change-Id: Ia706d532578ccbf5bc8f172f6344f166d05531fb
---
 tensorflow/lite/tools/BUILD                 |  7 ++++---
 tensorflow/lite/tools/list_flex_ops_test.cc | 23 +++++++++++----------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/tensorflow/lite/tools/BUILD b/tensorflow/lite/tools/BUILD
index c34453e0809..89d3da1ec6a 100644
--- a/tensorflow/lite/tools/BUILD
+++ b/tensorflow/lite/tools/BUILD
@@ -1,5 +1,6 @@
 load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
 load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
 
 package(
     default_visibility = [
@@ -271,7 +272,7 @@ cc_library(
 
 # This tool list flex ops and kernels inside a TFLite file.
 # It is used to generate header file for selective registration.
-cc_binary(
+tf_cc_binary(
     name = "list_flex_ops_main",
     srcs = ["list_flex_ops_main.cc"],
     visibility = ["//visibility:public"],
@@ -282,7 +283,7 @@ cc_binary(
     ],
 )
 
-cc_test(
+tf_cc_test(
     name = "list_flex_ops_test",
     srcs = ["list_flex_ops_test.cc"],
     data = [
@@ -293,7 +294,6 @@ cc_test(
         "//tensorflow/lite:testdata/test_model_broken.bin",
     ],
     tags = [
-        "no_oss",  # Currently requires --config=monolithic, b/118895218.
         "tflite_not_portable_android",
         "tflite_not_portable_ios",
     ],
@@ -301,6 +301,7 @@ cc_test(
         ":list_flex_ops",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:resource_loader",
         "//tensorflow/lite/kernels:test_util",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
diff --git a/tensorflow/lite/tools/list_flex_ops_test.cc b/tensorflow/lite/tools/list_flex_ops_test.cc
index 67ddc06325a..872d7509d0c 100644
--- a/tensorflow/lite/tools/list_flex_ops_test.cc
+++ b/tensorflow/lite/tools/list_flex_ops_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/resource_loader.h"
 #include "tensorflow/lite/kernels/test_util.h"
 
 namespace tflite {
@@ -31,8 +32,9 @@ class FlexOpsListTest : public ::testing::Test {
  protected:
   FlexOpsListTest() {}
 
-  void ReadOps(const string& model_path) {
-    auto model = FlatBufferModel::BuildFromFile(model_path.data());
+  void ReadOps(const string& path) {
+    std::string full_path = tensorflow::GetDataDependencyFilepath(path);
+    auto model = FlatBufferModel::BuildFromFile(full_path.data());
     AddFlexOpsFromModel(model->GetModel(), &flex_ops_);
     output_text_ = OpListToJSONString(flex_ops_);
   }
@@ -84,30 +86,29 @@ class FlexOpModel : public SingleOpModel {
 };
 
 TEST_F(FlexOpsListTest, TestModelsNoFlex) {
-  ReadOps("third_party/tensorflow/lite/testdata/test_model.bin");
+  ReadOps("tensorflow/lite/testdata/test_model.bin");
   EXPECT_EQ(output_text_, "[]");
 }
 
 TEST_F(FlexOpsListTest, TestBrokenModel) {
   EXPECT_DEATH_IF_SUPPORTED(
-      ReadOps("third_party/tensorflow/lite/testdata/test_model_broken.bin"),
-      "");
+      ReadOps("tensorflow/lite/testdata/test_model_broken.bin"), "");
 }
 
 TEST_F(FlexOpsListTest, TestZeroSubgraphs) {
-  ReadOps("third_party/tensorflow/lite/testdata/0_subgraphs.bin");
+  ReadOps("tensorflow/lite/testdata/0_subgraphs.bin");
   EXPECT_EQ(output_text_, "[]");
 }
 
 TEST_F(FlexOpsListTest, TestFlexAdd) {
-  ReadOps("third_party/tensorflow/lite/testdata/multi_add_flex.bin");
+  ReadOps("tensorflow/lite/testdata/multi_add_flex.bin");
   EXPECT_EQ(output_text_,
             "[[\"Add\", \"BinaryOp>\"]]");
 }
 
 TEST_F(FlexOpsListTest, TestTwoModel) {
-  ReadOps("third_party/tensorflow/lite/testdata/multi_add_flex.bin");
-  ReadOps("third_party/tensorflow/lite/testdata/softplus_flex.bin");
+  ReadOps("tensorflow/lite/testdata/multi_add_flex.bin");
+  ReadOps("tensorflow/lite/testdata/softplus_flex.bin");
   EXPECT_EQ(output_text_,
             "[[\"Add\", \"BinaryOp>\"],\n[\"Softplus\", \"SoftplusOp>\"]]");
 }

From d43a0150f891b938dfa4247744e4d18e2e696e06 Mon Sep 17 00:00:00 2001
From: Jing Pu 
Date: Tue, 16 Jun 2020 00:57:18 -0700
Subject: [PATCH 0269/1390] Add a pattern to legalize hlo.reduce to tf.Min.

PiperOrigin-RevId: 316632939
Change-Id: I7fbc90c1a75e5cc8bafb6b87475284de6ebe91a7
---
 .../mlir/tensorflow/tests/legalize_hlo.mlir   | 21 +++++++
 .../tensorflow/transforms/legalize_hlo.cc     | 57 +++++++++++++++----
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/legalize_hlo.mlir b/tensorflow/compiler/mlir/tensorflow/tests/legalize_hlo.mlir
index 6a32bcb5254..c8542ab3bae 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/legalize_hlo.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/legalize_hlo.mlir
@@ -773,6 +773,20 @@ func @convert_reduce_to_max(%arg0: tensor<1x256xf32>) -> tensor<1xf32> {
   return %1 : tensor<1xf32>
 }
 
+
+func @convert_reduce_to_min(%arg0: tensor<1x256xf32>) -> tensor<1xf32> {
+  // "0x7F800000" represents INF for f32.
+  %0 = xla_hlo.constant dense<0x7F800000> : tensor
+  %1 = "xla_hlo.reduce"(%arg0, %0) ( {
+  ^bb0(%arg1: tensor, %arg2: tensor):
+    %2 = xla_hlo.minimum %arg1, %arg2 : tensor
+    "xla_hlo.return"(%2) : (tensor) -> ()
+  }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<1x256xf32>, tensor) -> tensor<1xf32>
+  return %1 : tensor<1xf32>
+}
+
+
+
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 
 // CHECK-LABEL:   func @biasAdd_NHWC(
@@ -1689,3 +1703,10 @@ func @convert_reduce_to_max(%arg0: tensor<1x256xf32>) -> tensor<1xf32> {
 // CHECK:           [[VAL_418:%.*]] = "tf.Max"([[VAL_416:%.*]], [[VAL_417:%.*]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32>
 // CHECK:           return [[VAL_418]] : tensor<1xf32>
 // CHECK:         }
+
+// CHECK-LABEL:   func @convert_reduce_to_min(
+// CHECK-SAME:                  [[VAL_419:%.*]]: tensor<1x256xf32>) -> tensor<1xf32> {
+// CHECK:           [[VAL_420:%.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64>
+// CHECK:           [[VAL_421:%.*]] = "tf.Min"([[VAL_419:%.*]], [[VAL_420:%.*]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32>
+// CHECK:           return [[VAL_421]] : tensor<1xf32>
+// CHECK:         }
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc b/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
index ad7abc08d94..9bb23213919 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
@@ -489,9 +489,9 @@ LogicalResult MatchReduceOpInput(xla_hlo::ReduceOp reduce_op) {
   return success();
 }
 
-// TODO(b/157192370): This "xla_hlo::ReduceOp" can corresponds to many TF ops
-// with different ops in reduce_op.body. Now we only match to "tf.Max" and
-// "tf.Sum".
+// TODO(jingpu): This "xla_hlo::ReduceOp" can corresponds to many TF ops
+// with different ops in reduce_op.body. Now we only match to "tf.Max", "tf.Min"
+// and "tf.Sum".
 class ConvertReduceOpToTfSum : public OpConversionPattern {
  public:
   using OpConversionPattern::OpConversionPattern;
@@ -504,15 +504,13 @@ class ConvertReduceOpToTfSum : public OpConversionPattern {
     Operation *first_op = &reduce_op.body().front().front();
     if (!llvm::isa(first_op)) return failure();
 
-    // In `MatchReduceOpInput` function, we only match that the
+    // In `MatchReduceOpInput` function, we already match that the
     // "xla_hlo::ReduceOp" only has one input, one init_value and one result.
     auto input = reduce_op.operands()[0];
     // Get reduction dimension.
     DenseIntElementsAttr dimension = reduce_op.dimensions();
     SmallVector reduce_dims;
-    const int64_t input_rank = input.getType().cast().getRank();
     for (const int64_t &dim : dimension.getValues()) {
-      if (dim < 0 || dim >= input_rank) return failure();
       reduce_dims.emplace_back(dim);
     }
 
@@ -545,15 +543,13 @@ class ConvertReduceOpToTfMax : public OpConversionPattern {
     Operation *first_op = &reduce_op.body().front().front();
     if (!llvm::isa(first_op)) return failure();
 
-    // In `MatchReduceOpInput` function, we only match that the
+    // In `MatchReduceOpInput` function, we already match that the
     // "xla_hlo::ReduceOp" only has one input, one init_value and one result.
     auto input = reduce_op.operands()[0];
     // Get reduction dimension.
     DenseIntElementsAttr dimension = reduce_op.dimensions();
     SmallVector reduce_dims;
-    const int64_t input_rank = input.getType().cast().getRank();
     for (const int64_t &dim : dimension.getValues()) {
-      if (dim < 0 || dim >= input_rank) return failure();
       reduce_dims.emplace_back(dim);
     }
 
@@ -576,6 +572,47 @@ class ConvertReduceOpToTfMax : public OpConversionPattern {
   };
 };
 
+class ConvertReduceOpToTfMin : public OpConversionPattern {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      xla_hlo::ReduceOp reduce_op, ArrayRef args,
+      ConversionPatternRewriter &rewriter) const final {
+    if (failed(MatchReduceOpInput(reduce_op))) return failure();
+
+    Operation *first_op = &reduce_op.body().front().front();
+    if (!llvm::isa(first_op)) return failure();
+
+    // In `MatchReduceOpInput` function, we already match that the
+    // "xla_hlo::ReduceOp" only has one input, one init_value and one result.
+    Value input = reduce_op.operands()[0];
+    // Get reduction dimension.
+    DenseIntElementsAttr dimension = reduce_op.dimensions();
+    SmallVector reduce_dims;
+    for (const int64_t &dim : dimension.getValues()) {
+      reduce_dims.emplace_back(dim);
+    }
+
+    // Check initial value is +INF.
+    DenseFPElementsAttr init_value;
+    if (!matchPattern(reduce_op.init_values()[0], m_Constant(&init_value)) ||
+        !init_value.isSplat() ||
+        !init_value.getSplatValue().isInfinity() ||
+        init_value.getSplatValue().isNegative())
+      return failure();
+
+    auto dim_type = RankedTensorType::get(
+        {static_cast(reduce_dims.size())}, rewriter.getI64Type());
+    auto reduction_indices = rewriter.create(
+        reduce_op.getLoc(), dim_type, rewriter.getI64TensorAttr(reduce_dims));
+    rewriter.replaceOpWithNewOp(
+        reduce_op, reduce_op.getType(0), input, reduction_indices,
+        /*keep_dim=*/rewriter.getBoolAttr(false));
+    return success();
+  };
+};
+
 class LegalizeHloToTf : public PassWrapper {
  public:
   LegalizeHloToTf() = default;
@@ -709,7 +746,7 @@ void LegalizeHloToTf::runOnFunction() {
   OwningRewritePatternList patterns;
   populateWithGenerated(&context, &patterns);
   patterns.insert(&context);
+                  ConvertReduceOpToTfMin, ConvertReduceOpToTfSum>(&context);
 
   ConversionTarget target(context);
   target.addLegalDialect();

From 44db81e3241f98e61d386aeb8b1ee0dee33e04b6 Mon Sep 17 00:00:00 2001
From: YoungSeok Yoon 
Date: Tue, 16 Jun 2020 01:16:13 -0700
Subject: [PATCH 0270/1390] Fix broken hyperlinks in guide docs

PiperOrigin-RevId: 316635046
Change-Id: I604b94075e2e10520bbfb2885089e534f3a649cd
---
 tensorflow/lite/g3doc/guide/ops_select.md           | 2 +-
 tensorflow/lite/g3doc/performance/best_practices.md | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/g3doc/guide/ops_select.md b/tensorflow/lite/g3doc/guide/ops_select.md
index 2226d86b1c9..0fa608cfa96 100644
--- a/tensorflow/lite/g3doc/guide/ops_select.md
+++ b/tensorflow/lite/g3doc/guide/ops_select.md
@@ -53,7 +53,7 @@ partially supported by TensorFlow Lite, and one would like to avoid those
 limitations.
 
 The following example shows how to use this feature in the
-[`TFLiteConverter`](./convert/python_api.md) Python API.
+[`TFLiteConverter`](../convert/python_api.md) Python API.
 
 ```python
 import tensorflow as tf
diff --git a/tensorflow/lite/g3doc/performance/best_practices.md b/tensorflow/lite/g3doc/performance/best_practices.md
index 94436865e48..e4abb564b26 100644
--- a/tensorflow/lite/g3doc/performance/best_practices.md
+++ b/tensorflow/lite/g3doc/performance/best_practices.md
@@ -20,8 +20,9 @@ accuracy and latency tradeoffs for some common image classification models.
 
 One example of models optimized for mobile devices are
 [MobileNets](https://arxiv.org/abs/1704.04861), which are optimized for mobile
-vision applications. [Hosted models](../models/hosted.md) lists several other
-models that have been optimized specifically for mobile and embedded devices.
+vision applications. [Hosted models](../guide/hosted_models.md) lists several
+other models that have been optimized specifically for mobile and embedded
+devices.
 
 You can retrain the listed models on your own dataset by using transfer
 learning. Check out our transfer learning tutorial for

From aa99cf218c8bf13aeb15e64ec4c62ea14ecb5753 Mon Sep 17 00:00:00 2001
From: Terry Heo 
Date: Tue, 16 Jun 2020 01:27:48 -0700
Subject: [PATCH 0271/1390] Enable flex delegate on tensorflow.lite.Interpreter
 Python package

Usually, flex delegate is enabled by symbol override of AcquireFlexDelegate()
function. But this approach doesn't work well with shared library.

Since pywrap_tensorflow_internal.so is available for tensorflow PIP,
I've made the following changes to enable flex delegate.
- Included flex delegate module to the pywrap_tensorflow_internal.so.
  This file already contains most TF internal logic and having TFLite flex
  delegate impacts about 72K to the output.
- Added new function of TF_AcquireFlexDelegate() in the delegate module.
- Updated logic in AcquireFlexDelegate() of interpreter_builder.cc to check
  the availability of pywrap_tensorflow_internal.so and lookup the
  TF_AcquireFlexDelegate() symbol to enable flex delegate.

Also updated python/lite_flex_test.py since flex delegate is supported with
Python API

PiperOrigin-RevId: 316636275
Change-Id: I13a3246f27860ac0551fb04d81a84d4e82997ebc
---
 tensorflow/lite/delegates/flex/delegate.cc |  7 +++
 tensorflow/lite/interpreter_builder.cc     | 17 ++++++
 tensorflow/lite/python/BUILD               |  3 +-
 tensorflow/lite/python/lite_flex_test.py   | 61 +++++++++++++---------
 tensorflow/python/BUILD                    |  1 +
 5 files changed, 62 insertions(+), 27 deletions(-)

diff --git a/tensorflow/lite/delegates/flex/delegate.cc b/tensorflow/lite/delegates/flex/delegate.cc
index 4741bddc2f5..b8b0d4e6d01 100644
--- a/tensorflow/lite/delegates/flex/delegate.cc
+++ b/tensorflow/lite/delegates/flex/delegate.cc
@@ -136,3 +136,10 @@ TfLiteStatus FlexDelegate::CopyFromBufferHandle(
 }
 
 }  // namespace tflite
+
+// Exported C interface function which is used by AcquireFlexDelegate() at
+// interpreter_build.cc. To export the function name globally, the function name
+// must be matched with patterns in tf_version_script.lds
+extern "C" tflite::TfLiteDelegateUniquePtr TF_AcquireFlexDelegate() {
+  return tflite::AcquireFlexDelegate();
+}
diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index 43d81ef0770..d73b298e595 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/interpreter_builder.h"
 
+#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
+#include 
+#endif
 #include 
 #include 
 #include 
@@ -114,6 +117,20 @@ const char* kEmptyTensorName = "";
 // For flex delegate, see also the strong override in
 // lite/delegates/flex/delegate.cc.
 TFLITE_ATTRIBUTE_WEAK Interpreter::TfLiteDelegatePtr AcquireFlexDelegate() {
+#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
+  // If _pywrap_tensorflow_internal.so is available, use
+  // TF_AcquireFlexDelegate() to initialize flex delegate.
+  void* lib_tf_internal =
+      dlopen("_pywrap_tensorflow_internal.so", RTLD_NOW | RTLD_LOCAL);
+  if (lib_tf_internal) {
+    auto TF_AcquireFlexDelegate =
+        reinterpret_cast(
+            dlsym(lib_tf_internal, "TF_AcquireFlexDelegate"));
+    if (TF_AcquireFlexDelegate) {
+      return TF_AcquireFlexDelegate();
+    }
+  }
+#endif
   return Interpreter::TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
 }
 
diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD
index c1f37c81b7f..b0f605ed50d 100644
--- a/tensorflow/lite/python/BUILD
+++ b/tensorflow/lite/python/BUILD
@@ -194,8 +194,7 @@ py_test(
     python_version = "PY3",
     srcs_version = "PY2AND3",
     tags = [
-        # TODO(b/111881877): Enable in oss after resolving op registry issues.
-        "no_oss",
+        "no_mac",  # TODO(b/159077703): Enable Python API Flex support on MacOS.
         "no_windows",
     ],
     deps = [
diff --git a/tensorflow/lite/python/lite_flex_test.py b/tensorflow/lite/python/lite_flex_test.py
index 26bee206d27..ffc157c2128 100644
--- a/tensorflow/lite/python/lite_flex_test.py
+++ b/tensorflow/lite/python/lite_flex_test.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from absl.testing import parameterized
+import numpy as np
 
 from tensorflow.lite.python import lite
 from tensorflow.lite.python.interpreter import Interpreter
@@ -41,8 +42,7 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       ('DisableMlirConverter', False))  # disable mlir
   def testFlexMode(self, enable_mlir):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(
-          shape=[1, 16, 16, 3], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -54,19 +54,22 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
   def testDeprecatedFlags(self):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(
-          shape=[1, 16, 16, 3], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -83,14 +86,18 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
 
 class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
@@ -114,14 +121,18 @@ class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
     converter.experimental_new_converter = enable_mlir
     tflite_model = converter.convert()
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([4.0], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([24.0], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 87048ba9d40..343a95b85e9 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -6046,6 +6046,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/core/profiler/internal:print_model_analysis",
         "//tensorflow/core/profiler/internal/cpu:python_tracer",
         "//tensorflow/tools/graph_transforms:transform_graph_lib",
+        "//tensorflow/lite/delegates/flex:delegate",
         "//tensorflow/lite/toco/python:toco_python_api",
         "//tensorflow/python/eager:pywrap_tfe_lib",
         "//tensorflow/core/util/tensor_bundle",

From 8950c470bb11a9b94c0dd08d73156008dfac60c9 Mon Sep 17 00:00:00 2001
From: Pavithra Vijay 
Date: Tue, 16 Jun 2020 01:54:26 -0700
Subject: [PATCH 0272/1390] Remove automatic control dep wrapping from layers
 in v2.

PiperOrigin-RevId: 316638920
Change-Id: Iad14b1a4b0b14052f34784401b375a14b49a7641
---
 tensorflow/python/keras/engine/base_layer.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index a0ee25417c0..628e74db27a 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -40,7 +40,6 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import execute
 from tensorflow.python.eager import function
 from tensorflow.python.eager import monitoring
-from tensorflow.python.framework import auto_control_deps
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
@@ -1105,17 +1104,7 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
 
           try:
             with ops.enable_auto_cast_variables(self._compute_dtype_object):
-              # Add auto_control_deps in V2 when they are not already added by
-              # a `tf.function`.
-              if (ops.executing_eagerly_outside_functions() and
-                  not base_layer_utils.is_in_eager_or_tf_function()):
-                with auto_control_deps.AutomaticControlDependencies() as acd:
-                  outputs = call_fn(cast_inputs, *args, **kwargs)
-                  # Wrap Tensors in `outputs` in `tf.identity` to avoid
-                  # circular dependencies.
-                  outputs = base_layer_utils.mark_as_return(outputs, acd)
-              else:
-                outputs = call_fn(cast_inputs, *args, **kwargs)
+              outputs = call_fn(cast_inputs, *args, **kwargs)
 
           except errors.OperatorNotAllowedInGraphError as e:
             raise TypeError('You are attempting to use Python control '

From e2b5397f126ba9cbc76a840ea0a46331e0f10897 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 02:02:06 -0700
Subject: [PATCH 0273/1390] Update GraphDef version to 434.

PiperOrigin-RevId: 316639748
Change-Id: I2f62575a1ffdf72dbbafd5a2d6a10ae2a64d4b7c
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index e08b166df2c..8e3c66edfc2 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 433  // Updated: 2020/6/15
+#define TF_GRAPH_DEF_VERSION 434  // Updated: 2020/6/16
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From d2cba310e80fc545cb0f8075d32335c170d547f2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 02:02:10 -0700
Subject: [PATCH 0274/1390] compat: Update forward compatibility horizon to
 2020-06-16

PiperOrigin-RevId: 316639760
Change-Id: I5bfbc17f255457595771a2a4636abd59ee03feb1
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 33cd3404b4d..314acfdd38f 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 15)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 16)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 83f19c6a9e84fc6971ad0a7df5874603237a595f Mon Sep 17 00:00:00 2001
From: Yong Tang 
Date: Mon, 15 Jun 2020 17:16:56 +0000
Subject: [PATCH 0275/1390] Fix unknown output shape issue in autograph for
 tf.equal

This PR tries to address the issue raised in 40471 where
the output shape of an autograph consists of tf.equal
could not inference correctly. Specifically
`x.shape == [None, 10, 1]` and `y.shape == [None, 1, 4]`
only yield `shape == None` (should be `shape == [None, 10, 4]`).

The reason was that the shape inbference function for equal
didn't capture the cases where both x and y's dim are None.

This PR fixes the issue.

This PR fixes 40471.

Signed-off-by: Yong Tang 
---
 tensorflow/core/framework/common_shape_fns.cc      |  3 +++
 tensorflow/core/ops/math_ops_test.cc               |  2 +-
 .../python/autograph/operators/logical_test.py     | 14 ++++++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc
index 113adbdd432..7567db03c23 100644
--- a/tensorflow/core/framework/common_shape_fns.cc
+++ b/tensorflow/core/framework/common_shape_fns.cc
@@ -1936,6 +1936,7 @@ Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c,
       // in C++ op code, we must still assert that the unknown dim is either 1
       // or the same as the known dim.
       // - If either dimension is 1, the other dimension is the output.
+      // - If both are unknown then dimension is unknown
       if (c->Value(dim_x) > 1) {
         if (!incompatible_shape_error) {
           *out = c->UnknownShape();
@@ -1954,6 +1955,8 @@ Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c,
         dims.push_back(dim_x);
       } else if (dim_y.SameHandle(dim_x)) {
         dims.push_back(dim_x);
+      } else if (!c->ValueKnown(dim_x) && !c->ValueKnown(dim_y)) {
+        dims.push_back(c->UnknownDim());
       } else {
         if (!incompatible_shape_error) {
           *out = c->UnknownShape();
diff --git a/tensorflow/core/ops/math_ops_test.cc b/tensorflow/core/ops/math_ops_test.cc
index 5c69a2a7f1c..a2837d88bde 100644
--- a/tensorflow/core/ops/math_ops_test.cc
+++ b/tensorflow/core/ops/math_ops_test.cc
@@ -120,7 +120,7 @@ TEST(MathOpsTest, BroadcastBinaryOps_ShapeFn) {
     INFER_OK(op, "[1];[?]", "[d1_0]");
     INFER_OK(op, "[?];[2]", incompatible_shape_error ? "[d1_0]" : "?");
     INFER_OK(op, "[2];[?]", incompatible_shape_error ? "[d0_0]" : "?");
-    INFER_OK(op, "[?];[?]", incompatible_shape_error ? "[?]" : "?");
+    INFER_OK(op, "[?];[?]", "[?]");
     INFER_OK(op, "[];[?]", "[d1_0]");
     INFER_OK(op, "[?];[]", "[d0_0]");
 
diff --git a/tensorflow/python/autograph/operators/logical_test.py b/tensorflow/python/autograph/operators/logical_test.py
index e22f39932d1..0eab302a825 100644
--- a/tensorflow/python/autograph/operators/logical_test.py
+++ b/tensorflow/python/autograph/operators/logical_test.py
@@ -19,7 +19,9 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.autograph.operators import logical
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
 
@@ -83,6 +85,18 @@ class LogicalOperatorsTest(test.TestCase):
       t = logical.not_(self._tf_false())
       self.assertEqual(self.evaluate(t), True)
 
+  # Test case for GitHub issue 40471
+  def test_equal_output_shapes(self):
+
+    @def_function.function(input_signature=[
+        tensor_spec.TensorSpec([None, 10, 1]),
+        tensor_spec.TensorSpec([None, 1, 4])])
+    def f(x, y):
+      z = x == y
+      return z
+
+    self.assertAllEqual(f.get_concrete_function().output_shapes, [None, 10, 4])
+
 
 if __name__ == '__main__':
   test.main()

From d52f3465f56882ad169759a942448843d1b4b589 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 03:57:26 -0700
Subject: [PATCH 0276/1390] Remove automatic control dep wrapping from layers
 in v2.

PiperOrigin-RevId: 316652071
Change-Id: I90d3568fa727c8370de1f20e35742efbd9d615ac
---
 tensorflow/python/keras/engine/base_layer.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index 628e74db27a..a0ee25417c0 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -40,6 +40,7 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import execute
 from tensorflow.python.eager import function
 from tensorflow.python.eager import monitoring
+from tensorflow.python.framework import auto_control_deps
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
@@ -1104,7 +1105,17 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
 
           try:
             with ops.enable_auto_cast_variables(self._compute_dtype_object):
-              outputs = call_fn(cast_inputs, *args, **kwargs)
+              # Add auto_control_deps in V2 when they are not already added by
+              # a `tf.function`.
+              if (ops.executing_eagerly_outside_functions() and
+                  not base_layer_utils.is_in_eager_or_tf_function()):
+                with auto_control_deps.AutomaticControlDependencies() as acd:
+                  outputs = call_fn(cast_inputs, *args, **kwargs)
+                  # Wrap Tensors in `outputs` in `tf.identity` to avoid
+                  # circular dependencies.
+                  outputs = base_layer_utils.mark_as_return(outputs, acd)
+              else:
+                outputs = call_fn(cast_inputs, *args, **kwargs)
 
           except errors.OperatorNotAllowedInGraphError as e:
             raise TypeError('You are attempting to use Python control '

From a5ebf37c1d67340559b13be2b622d288af47368b Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 16 Jun 2020 04:20:01 -0700
Subject: [PATCH 0277/1390] Allow escaping return values in kernel lowering.

For simple cases where the launch only has function arguments and results as operands, the kernel signature will now be rewritten to have its operands ordered in the order of the function arguments followed by results.

This only works for simple functions without control flow or multiple launches and is intended for cases where the kernel is meant to be extracted. If the host-side is also used, signature rewriting should be disabled via the provided flag.

PiperOrigin-RevId: 316654295
Change-Id: Ie3e78b78bab2191610875acb79506e5986a821af
---
 .../xla/service/mlir_gpu/kernel_lowering.cc   | 81 +++++++++++++++----
 .../xla/service/mlir_gpu/kernel_lowering.h    |  2 +-
 2 files changed, 66 insertions(+), 17 deletions(-)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index f1e01bba27e..b0cbddcdb92 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "mlir/IR/OperationSupport.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "mlir/Transforms/BufferPlacement.h"  // from @llvm-project
@@ -278,15 +279,34 @@ struct MoveScalarComputationsIntoGpuLaunch
   }
 };
 
-// TODO(herhut): Make this a proper thing.
-struct FixKernelFunctionSignatures
-    : mlir::PassWrapper {
+// Sort the operands to the kernel for a deterministic order. First operands
+// that are defined by function arguments, followed by operands that are
+// returned from the function. This only works for simple functions without
+// control flow and can be used in cases where the kernel is extracted and used
+// independently of the host-side code.
+struct RewriteKernelSignature
+    : mlir::PassWrapper {
   void runOnFunction() override {
     mlir::FuncOp func = getFunction();
     mlir::ModuleOp module = func.getParentOfType();
     getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) {
       mlir::gpu::GPUFuncOp kernel =
           module.lookupSymbol(launchOp.kernel());
+
+      if (kernel.getNumFuncArguments() !=
+          func.getNumArguments() + func.getNumResults()) {
+        kernel.emitError()
+            << "number of kernel arguments does not match number"
+            << "of arguments and results of surrounding function";
+        signalPassFailure();
+        return;
+      }
+      if (func.getBlocks().size() != 1) {
+        func.emitError() << "surrounding function has more than one block";
+        signalPassFailure();
+        return;
+      }
+
       // Compute a map from function arguments to kernel function operands.
       mlir::BlockAndValueMapping func_to_kernel;
       for (mlir::BlockArgument arg : func.getArguments()) {
@@ -297,27 +317,54 @@ struct FixKernelFunctionSignatures
           }
         }
       }
+      // Also add function results that are computed by the launch.
+      mlir::Operation* returnOp = func.getBody().back().getTerminator();
+      for (mlir::Value result : returnOp->getOperands()) {
+        for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) {
+          if (launchOp.getKernelOperand(i) == result) {
+            func_to_kernel.map(result, kernel.getArgument(i));
+            break;
+          }
+        }
+      }
 
-      // Create a new kernel function with modified signature. We know that it
-      // will have the same signature as the original function, so just reuse it
-      // here.
+      // Create a new kernel function with modified signature. It will have the
+      // parameters and result types of the original funcion as its parameter
+      // type and otherwise will be void.
       auto gpu_module = kernel.getParentOfType();
       mlir::OpBuilder kernel_builder(gpu_module.body());
+      auto operand_types = llvm::to_vector<4>(llvm::concat(
+          func.getType().getInputs(), func.getType().getResults()));
       auto new_kernel = kernel_builder.create(
-          kernel.getLoc(), kernel.getName(), func.getType());
+          kernel.getLoc(), kernel.getName(),
+          kernel_builder.getFunctionType(operand_types, {}));
       new_kernel.setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(),
                          kernel_builder.getUnitAttr());
 
       // Create a map from old kernel argument to new one.
       mlir::BlockAndValueMapping old_kernel_to_new;
-      for (int i = 0, e = kernel.getNumFuncArguments(); i < e; ++i) {
+      for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
         mlir::Value func_arg = func.getArgument(i);
         mlir::Value new_kernel_arg = new_kernel.getArgument(i);
         mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(func_arg);
         if (!old_kernel_arg) {
           kernel.emitOpError()
               << "argument " << i
-              << "to kernel is not an argument to the containing function";
+              << " to containing function is not an argument to the kernel";
+          signalPassFailure();
+          return;
+        }
+        old_kernel_to_new.map(old_kernel_arg, new_kernel_arg);
+      }
+      for (int i = 0, e = returnOp->getNumOperands(); i < e; ++i) {
+        mlir::Value ret_op = returnOp->getOperand(i);
+        mlir::Value new_kernel_arg =
+            new_kernel.getArgument(func.getNumArguments() + i);
+        mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(ret_op);
+        if (!old_kernel_arg) {
+          kernel.emitOpError()
+              << "result " << i
+              << " of containing function is not an argument to the kernel";
           signalPassFailure();
           return;
         }
@@ -328,13 +375,16 @@ struct FixKernelFunctionSignatures
       kernel_builder.setInsertionPointToEnd(&new_kernel.body().front());
       kernel_builder.create(
           new_kernel.getLoc(), &*std::next(new_kernel.body().begin()));
-      // Now create a new launchOp calling the new kernel. We can just forward
-      // the arguments of the function to the launch, as we fixed the
-      // signature.
+      // Now create a new launchOp calling the new kernel. We need to forward
+      // the arguments of the surrounding function and operands to the return.
+      mlir::SmallVector new_operands;
+      new_operands.reserve(new_kernel.getNumFuncArguments());
+      new_operands.append(func.args_begin(), func.args_end());
+      new_operands.append(returnOp->operand_begin(), returnOp->operand_end());
       mlir::OpBuilder launch_builder(launchOp);
       launch_builder.create(
           launchOp.getLoc(), new_kernel, launchOp.getGridSizeOperandValues(),
-          launchOp.getBlockSizeOperandValues(), func.getArguments());
+          launchOp.getBlockSizeOperandValues(), new_operands);
       // Launch does not have results, so we can just erase it. And the kernel
       // also needs to go.
       launchOp.erase();
@@ -418,7 +468,6 @@ Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   // Transform LHLO operations to LinAlg.
   pm.addPass(::mlir::xla_lhlo::createLegalizeLhloToLinalgPass());
   // Fuse linalg operations.
-  // TODO(herhut): Make tiling conigurable.
   pm.addPass(::mlir::xla_lhlo::createLhloFuseLinalg(/*use_parallel_loops=*/true,
                                                     tiling_for_unrolling));
   // Legalize reduce operations directly to GPU dialect.
@@ -462,8 +511,8 @@ Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   pm.addPass(::mlir::createGpuKernelOutliningPass());
   // Make sure the kernel signature resembled the original function's
   // signature
-  if (options.fix_signature) {
-    pm.addPass(absl::make_unique());
+  if (options.rewrite_signature) {
+    pm.addPass(absl::make_unique());
   }
   if (failed(pm.run(module))) {
     return InternalError("Lowering to GPU kernels failed.");
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
index 7b5d5c35c05..77cf75b9e47 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
@@ -27,7 +27,7 @@ struct LowerLHLOToGPUOptions {
   llvm::ArrayRef tile_sizes = {16, 64};
   llvm::ArrayRef unroll_factors = {};
   bool collapse_parallel_loops = true;
-  bool fix_signature = true;
+  bool rewrite_signature = true;
 };
 
 Status LowerLHLOToGPU(mlir::ModuleOp module,

From fb7fe20ec732f4c13cf9a5e52697f0ae4e8e17b6 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 16 Jun 2020 05:20:07 -0700
Subject: [PATCH 0278/1390] Extend lowering from hlo to lhlo to also support
 buffer allocation with escaping result buffers. This is now a flag to the
 pass (defaults to the current preallocation behavior).

PiperOrigin-RevId: 316660810
Change-Id: I89e46b494d09acf2dbe14b300ee5b9df431ab09c
---
 .../mlir/xla/tests/hlo-legalize-to-lhlo.mlir  | 323 +++++++++---------
 .../xla/transforms/hlo_legalize_to_lhlo.cc    |  36 +-
 .../compiler/mlir/xla/transforms/passes.h     |  10 +-
 3 files changed, 203 insertions(+), 166 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/tests/hlo-legalize-to-lhlo.mlir b/tensorflow/compiler/mlir/xla/tests/hlo-legalize-to-lhlo.mlir
index 56429249d99..f3ce29f1bd2 100644
--- a/tensorflow/compiler/mlir/xla/tests/hlo-legalize-to-lhlo.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/hlo-legalize-to-lhlo.mlir
@@ -1,12 +1,13 @@
-// RUN: xla-opt -hlo-legalize-to-lhlo -buffer-placement -split-input-file %s -o - | FileCheck %s
+// RUN: xla-opt -hlo-legalize-to-lhlo -buffer-placement -split-input-file %s -o - | FileCheck --check-prefixes=PRE,BOTH %s
+// RUN: xla-opt -hlo-legalize-to-lhlo=results-escape-function=true -buffer-placement -split-input-file %s -o - | FileCheck --check-prefixes=ESC,BOTH %s
 
-// CHECK-LABEL: func @attrs
+// BOTH-LABEL: func @attrs
 func @attrs_copy(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.exponential"(%tensor_operand)
       {some_attr_1 = "exp.1", some_attr_2 = dense<1> : tensor<1xi64>}
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.exponential"(%{{.*}}, %{{.*}}) {some_attr_1 = "exp.1", some_attr_2 = dense<1> : tensor<1xi64>}
+  // BOTH: "xla_lhlo.exponential"(%{{.*}}, %{{.*}}) {some_attr_1 = "exp.1", some_attr_2 = dense<1> : tensor<1xi64>}
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
@@ -16,13 +17,16 @@ func @attrs_copy(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
 func @return_func(%arg0: tensor<4xf32>) -> tensor<4xf32> {
   return %arg0 : tensor<4xf32>
 }
-//      CHECK: (%[[ARG0:.*]]: [[TYPE:.*]], %[[RESULT:.*]]: [[TYPE]])
-// CHECK-NEXT: "xla_lhlo.copy"(%[[ARG0]], %[[RESULT]]) : ([[TYPE]], [[TYPE]]) -> ()
-// CHECK-NEXT: return
+//      PRE: (%[[ARG0:.*]]: [[TYPE:.*]], %[[RESULT:.*]]: [[TYPE]])
+// PRE-NEXT: "xla_lhlo.copy"(%[[ARG0]], %[[RESULT]]) : ([[TYPE]], [[TYPE]]) -> ()
+// PRE-NEXT: return
+//      ESC: (%[[ARG0:.*]]: [[TYPE:.*]]) -> [[TYPE]]
+//  ESC-NOT: "xla_lhlo.copy"
+// ESC-NEXT: return %[[ARG0]]
 
 // -----
 
-// CHECK-LABEL: func @func_op_long
+// BOTH-LABEL: func @func_op_long
 func @func_op_long(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
   %1 = xla_hlo.maximum %arg0, %arg1 : tensor<4xf32>
   %2 = xla_hlo.add %arg0, %1 : tensor<4xf32>
@@ -31,89 +35,91 @@ func @func_op_long(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
   %5 = xla_hlo.multiply %2, %4 : tensor<4xf32>
   return %5 : tensor<4xf32>
 }
-//      CHECK: (%[[NEW_ARG0:.*]]: memref<4xf32>, %[[NEW_ARG1:.*]]: memref<4xf32>, %[[RESULT:.*]]: memref<4xf32>)
-// CHECK-NEXT: %[[MAX_RESULT:.*]] = alloc() : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.maximum"(%[[NEW_ARG0]], %[[NEW_ARG1]], %[[MAX_RESULT]])
-// CHECK-NEXT: %[[ADD_RESULT:.*]] = alloc() : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.add"(%[[NEW_ARG0]], %[[MAX_RESULT]], %[[ADD_RESULT]])
-// CHECK-NEXT: dealloc %[[MAX_RESULT]] : memref<4xf32>
-// CHECK-NEXT: %[[MIN_RESULT:.*]] = alloc() : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.minimum"(%[[NEW_ARG0]], %[[NEW_ARG1]], %[[MIN_RESULT]])
-// CHECK-NEXT: %[[SUB_RESULT:.*]] = alloc() : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.subtract"(%[[NEW_ARG1]], %[[MIN_RESULT]], %[[SUB_RESULT]])
-// CHECK-NEXT: dealloc %[[MIN_RESULT]] : memref<4xf32>
-// CHECK-NEXT: %[[MUL_RESULT:.*]] = alloc() : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.multiply"(%[[ADD_RESULT]], %[[SUB_RESULT]], %[[MUL_RESULT]])
-// CHECK-NEXT: dealloc %[[SUB_RESULT]] : memref<4xf32>
-// CHECK-NEXT: dealloc %[[ADD_RESULT]] : memref<4xf32>
-// CHECK-NEXT: "xla_lhlo.copy"(%[[MUL_RESULT]], %[[RESULT]]) : (memref<4xf32>, memref<4xf32>) -> ()
-// CHECK-NEXT: dealloc %[[MUL_RESULT]] : memref<4xf32>
-// CHECK-NEXT: return
+//        PRE: (%[[NEW_ARG0:.*]]: memref<4xf32>, %[[NEW_ARG1:.*]]: memref<4xf32>, %[[RESULT:.*]]: memref<4xf32>)
+//        ESC: (%[[NEW_ARG0:.*]]: memref<4xf32>, %[[NEW_ARG1:.*]]: memref<4xf32>) -> memref<4xf32>
+//  BOTH-NEXT: %[[MAX_RESULT:.*]] = alloc() : memref<4xf32>
+//  BOTH-NEXT: "xla_lhlo.maximum"(%[[NEW_ARG0]], %[[NEW_ARG1]], %[[MAX_RESULT]])
+//  BOTH-NEXT: %[[ADD_RESULT:.*]] = alloc() : memref<4xf32>
+//  BOTH-NEXT: "xla_lhlo.add"(%[[NEW_ARG0]], %[[MAX_RESULT]], %[[ADD_RESULT]])
+//  BOTH-NEXT: dealloc %[[MAX_RESULT]] : memref<4xf32>
+//  BOTH-NEXT: %[[MIN_RESULT:.*]] = alloc() : memref<4xf32>
+//  BOTH-NEXT: "xla_lhlo.minimum"(%[[NEW_ARG0]], %[[NEW_ARG1]], %[[MIN_RESULT]])
+//  BOTH-NEXT: %[[SUB_RESULT:.*]] = alloc() : memref<4xf32>
+//  BOTH-NEXT: "xla_lhlo.subtract"(%[[NEW_ARG1]], %[[MIN_RESULT]], %[[SUB_RESULT]])
+//  BOTH-NEXT: dealloc %[[MIN_RESULT]] : memref<4xf32>
+//  BOTH-NEXT: %[[MUL_RESULT:.*]] = alloc() : memref<4xf32>
+//  BOTH-NEXT: "xla_lhlo.multiply"(%[[ADD_RESULT]], %[[SUB_RESULT]], %[[MUL_RESULT]])
+//  BOTH-NEXT: dealloc %[[SUB_RESULT]] : memref<4xf32>
+//  BOTH-NEXT: dealloc %[[ADD_RESULT]] : memref<4xf32>
+//   PRE-NEXT: "xla_lhlo.copy"(%[[MUL_RESULT]], %[[RESULT]]) : (memref<4xf32>, memref<4xf32>) -> ()
+//   PRE-NEXT: dealloc %[[MUL_RESULT]] : memref<4xf32>
+//   PRE-NEXT: return
+//   ESC-NEXT: return %[[MUL_RESULT]] : memref<4xf32>
 
 // -----
 
-// CHECK-LABEL: func @fusion
+// BOTH-LABEL: func @fusion
 func @fusion(%multiplier: memref<2x2xf32>, %summand_1: memref<2x2xf32>,
              %summand_2: memref<2x2xf32>, %result: memref<2x2xf32>) {
-  // CHECK: (%{{.*}}: {{.*}}, {{.*}}: {{.*}}, {{.*}}: {{.*}}, %[[RESULT:.*]]: {{.*}})
-  // CHECK-NEXT:  %[[ADD_RESULT:.*]] = alloc() : memref<2x2xf32>
+  // BOTH: (%{{.*}}: {{.*}}, {{.*}}: {{.*}}, {{.*}}: {{.*}}, %[[RESULT:.*]]: {{.*}})
+  // BOTH-NEXT:  %[[ADD_RESULT:.*]] = alloc() : memref<2x2xf32>
   %tensor_summand_1 = tensor_load %summand_1 : memref<2x2xf32>
   %tensor_summand_2 = tensor_load %summand_2 : memref<2x2xf32>
   %sum = "xla_hlo.add"(%tensor_summand_1, %tensor_summand_2)
       : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK-NEXT: "xla_lhlo.add"(%{{.*}}, %{{.*}}, %[[ADD_RESULT]])
-  // CHECK-NEXT:  %[[MUL_RESULT:.*]] = alloc() : memref<2x2xf32>
+  // BOTH-NEXT: "xla_lhlo.add"(%{{.*}}, %{{.*}}, %[[ADD_RESULT]])
+  // BOTH-NEXT:  %[[MUL_RESULT:.*]] = alloc() : memref<2x2xf32>
   %tensor_multiplier = tensor_load %multiplier : memref<2x2xf32>
   %tensor_result = "xla_hlo.multiply"(%sum, %tensor_multiplier)
       : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK-NEXT: "xla_lhlo.multiply"(%[[ADD_RESULT]], %{{.*}}, %[[MUL_RESULT]])
-  // CHECK-NEXT:  dealloc %[[ADD_RESULT]] : memref<2x2xf32>
-  // CHECK-NEXT: "xla_lhlo.copy"(%[[MUL_RESULT]], %[[RESULT]])
+  // BOTH-NEXT: "xla_lhlo.multiply"(%[[ADD_RESULT]], %{{.*}}, %[[MUL_RESULT]])
+  // BOTH-NEXT:  dealloc %[[ADD_RESULT]] : memref<2x2xf32>
+  // BOTH-NEXT: "xla_lhlo.copy"(%[[MUL_RESULT]], %[[RESULT]])
   tensor_store %tensor_result, %result : memref<2x2xf32>
-  // CHECK-NEXT:  dealloc %[[MUL_RESULT]] : memref<2x2xf32>
-  // CHECK-NEXT: "xla_lhlo.terminator"() : () -> ()
-  "xla_lhlo.terminator"() : () -> ()
+  // BOTH-NEXT:  dealloc %[[MUL_RESULT]] : memref<2x2xf32>
+  // BOTH-NEXT:  return
+  return
 }
 
 // -----
 
-// CHECK-LABEL: func @copy
+// BOTH-LABEL: func @copy
 func @copy(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.copy"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.copy"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.copy"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @exp
+// BOTH-LABEL: func @exp
 func @exp(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.exponential"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.exponential"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.exponential"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @log
+// BOTH-LABEL: func @log
 func @log(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.log"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.log"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.log"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @select
+// BOTH-LABEL: func @select
 func @select(%pred: memref<2x2xi1>, %lhs: memref<2x2xf32>,
              %rhs: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_pred = tensor_load %pred : memref<2x2xi1>
@@ -121,34 +127,34 @@ func @select(%pred: memref<2x2xi1>, %lhs: memref<2x2xf32>,
   %tensor_rhs = tensor_load %rhs : memref<2x2xf32>
   %tensor_result = "xla_hlo.select"(%tensor_pred, %tensor_lhs, %tensor_rhs)
       : (tensor<2x2xi1>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.select"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.select"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @compare
+// BOTH-LABEL: func @compare
 func @compare(%lhs: memref<2x2xf32>, %rhs: memref<2x2xf32>, %result: memref<2x2xi1>) {
   %tensor_lhs = tensor_load %lhs : memref<2x2xf32>
   %tensor_rhs = tensor_load %rhs : memref<2x2xf32>
   %tensor_result = "xla_hlo.compare"(%tensor_lhs, %tensor_rhs)
       {comparison_direction = "EQ"}
       : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xi1>
-  // CHECK: "xla_lhlo.compare"(%{{.*}}, %{{.*}}, %{{.*}}) {comparison_direction = "EQ"}
+  // BOTH: "xla_lhlo.compare"(%{{.*}}, %{{.*}}, %{{.*}}) {comparison_direction = "EQ"}
   tensor_store %tensor_result, %result : memref<2x2xi1>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @broadcast
+// BOTH-LABEL: func @broadcast
 func @broadcast(%operand: memref<5xf32>, %result: memref<10x5xf32>) {
   %tensor_operand = tensor_load %operand : memref<5xf32>
   %tensor_result = "xla_hlo.broadcast_in_dim"(%tensor_operand)
       {broadcast_dimensions = dense<1> : tensor<1xi64>}
         : (tensor<5xf32>) -> tensor<10x5xf32>
-  // CHECK: "xla_lhlo.broadcast_in_dim"(%{{.*}}, %{{.*}}) {broadcast_dimensions = dense<1> : tensor<1xi64>}
+  // BOTH: "xla_lhlo.broadcast_in_dim"(%{{.*}}, %{{.*}}) {broadcast_dimensions = dense<1> : tensor<1xi64>}
   tensor_store %tensor_result, %result : memref<10x5xf32>
   return
 }
@@ -157,55 +163,55 @@ func @broadcast(%operand: memref<5xf32>, %result: memref<10x5xf32>) {
 
 func @external_func() -> tensor<3xi64>
 
-// CHECK: #[[MAP:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
+// BOTH: #[[MAP:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
 
-// CHECK-LABEL: func @dyn_broadcast
+// BOTH-LABEL: func @dyn_broadcast
 func @dyn_broadcast(%operand: memref) {
-  // CHECK-SAME: (%[[OPERAND:.*]]: memref)
+  // BOTH-SAME: (%[[OPERAND:.*]]: memref)
   %tensor_operand = tensor_load %operand : memref
   %shape = call @external_func() : () -> tensor<3xi64>
   %tensor_result = "xla_hlo.dynamic_broadcast_in_dim"(%tensor_operand, %shape) {
     broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>
   } : (tensor, tensor<3xi64>) -> tensor
-  // CHECK: %[[SHAPE:.*]] = call @external_func()
-  // CHECK: %[[C0:.*]] = constant 0 : index
-  // CHECK: %[[EL0:.*]] = extract_element %[[SHAPE]][%[[C0]]] : tensor<3xi64>
-  // CHECK: %[[IC0:.*]]  = index_cast %[[EL0]] : i64 to index
-  // CHECK: %[[C1:.*]] = constant 1 : index
-  // CHECK: %[[EL1:.*]] = extract_element %[[SHAPE]][%[[C1]]] : tensor<3xi64>
-  // CHECK: %[[IC1:.*]]  = index_cast %[[EL1]] : i64 to index
-  // CHECK: %[[C2:.*]] = constant 2 : index
-  // CHECK: %[[EL2:.*]] = extract_element %[[SHAPE]][%[[C2]]] : tensor<3xi64>
-  // CHECK: %[[IC2:.*]]  = index_cast %[[EL2]] : i64 to index
-  // CHECK: %[[RESULT:.*]] = alloc(%[[IC0]], %[[IC1]], %[[IC2]])
+  // BOTH: %[[SHAPE:.*]] = call @external_func()
+  // BOTH: %[[C0:.*]] = constant 0 : index
+  // BOTH: %[[EL0:.*]] = extract_element %[[SHAPE]][%[[C0]]] : tensor<3xi64>
+  // BOTH: %[[IC0:.*]]  = index_cast %[[EL0]] : i64 to index
+  // BOTH: %[[C1:.*]] = constant 1 : index
+  // BOTH: %[[EL1:.*]] = extract_element %[[SHAPE]][%[[C1]]] : tensor<3xi64>
+  // BOTH: %[[IC1:.*]]  = index_cast %[[EL1]] : i64 to index
+  // BOTH: %[[C2:.*]] = constant 2 : index
+  // BOTH: %[[EL2:.*]] = extract_element %[[SHAPE]][%[[C2]]] : tensor<3xi64>
+  // BOTH: %[[IC2:.*]]  = index_cast %[[EL2]] : i64 to index
+  // BOTH: %[[RESULT:.*]] = alloc(%[[IC0]], %[[IC1]], %[[IC2]])
 
-  // CHECK: %[[C0_:.*]] = constant 0 : index
-  // CHECK: %[[C1_:.*]] = constant 1 : index
+  // BOTH: %[[C0_:.*]] = constant 0 : index
+  // BOTH: %[[C1_:.*]] = constant 1 : index
 
-  // CHECK: %[[C1__:.*]] = constant 1 : index
-  // CHECK: %[[EL1_:.*]] = extract_element %[[SHAPE]]{{\[}}%[[C1__]]] : tensor<3xi64>
-  // CHECK: %[[C0___:.*]] = constant 0 : index
-  // CHECK: %[[OPERAND_DIM_0:.*]] = dim %[[OPERAND]], %[[C0___]] : memref
-  // CHECK: %[[RESULT_DIM_1:.*]] = index_cast %[[EL1_]] : i64 to index
-  // CHECK: %[[EXPAND_0:.*]] = cmpi "slt", %[[OPERAND_DIM_0]], %[[RESULT_DIM_1]]
-  // CHECK: %[[STRIDE_0:.*]] = select %[[EXPAND_0]], %[[C0_]], %[[C1_]] : index
+  // BOTH: %[[C1__:.*]] = constant 1 : index
+  // BOTH: %[[EL1_:.*]] = extract_element %[[SHAPE]]{{\[}}%[[C1__]]] : tensor<3xi64>
+  // BOTH: %[[C0___:.*]] = constant 0 : index
+  // BOTH: %[[OPERAND_DIM_0:.*]] = dim %[[OPERAND]], %[[C0___]] : memref
+  // BOTH: %[[RESULT_DIM_1:.*]] = index_cast %[[EL1_]] : i64 to index
+  // BOTH: %[[EXPAND_0:.*]] = cmpi "slt", %[[OPERAND_DIM_0]], %[[RESULT_DIM_1]]
+  // BOTH: %[[STRIDE_0:.*]] = select %[[EXPAND_0]], %[[C0_]], %[[C1_]] : index
 
-  // CHECK: %[[C2_:.*]] = constant 2 : index
-  // CHECK: %[[EL2_:.*]] = extract_element %[[SHAPE]]{{\[}}%[[C2_]]] : tensor<3xi64>
-  // CHECK: %[[C1___:.*]] = constant 1 : index
-  // CHECK: %[[OPERAND_DIM_1:.*]] = dim %[[OPERAND]], %[[C1___]] : memref
-  // CHECK: %[[RESULT_DIM_2:.*]] = index_cast %[[EL2_]] : i64 to index
-  // CHECK: %[[EXPAND_1:.*]] = cmpi "slt", %[[OPERAND_DIM_1]], %[[RESULT_DIM_2]]
-  // CHECK: %[[STRIDE_1:.*]] = select %[[EXPAND_1]], %[[C0_]], %[[C1_]] : index
+  // BOTH: %[[C2_:.*]] = constant 2 : index
+  // BOTH: %[[EL2_:.*]] = extract_element %[[SHAPE]]{{\[}}%[[C2_]]] : tensor<3xi64>
+  // BOTH: %[[C1___:.*]] = constant 1 : index
+  // BOTH: %[[OPERAND_DIM_1:.*]] = dim %[[OPERAND]], %[[C1___]] : memref
+  // BOTH: %[[RESULT_DIM_2:.*]] = index_cast %[[EL2_]] : i64 to index
+  // BOTH: %[[EXPAND_1:.*]] = cmpi "slt", %[[OPERAND_DIM_1]], %[[RESULT_DIM_2]]
+  // BOTH: %[[STRIDE_1:.*]] = select %[[EXPAND_1]], %[[C0_]], %[[C1_]] : index
 
-  // CHECK: %[[TRANSFORMED_MEMREF:.*]] = xla_lhlo.dynamic_memref_cast
-  // CHECK-SAME: %[[OPERAND]](%[[RESULT_DIM_1]], %[[RESULT_DIM_2]])
-  // CHECK-SAME: {{\[}}%[[STRIDE_0]], %[[STRIDE_1]]]
-  // CHECK-SAME: : memref -> memref
+  // BOTH: %[[TRANSFORMED_MEMREF:.*]] = xla_lhlo.dynamic_memref_cast
+  // BOTH-SAME: %[[OPERAND]](%[[RESULT_DIM_1]], %[[RESULT_DIM_2]])
+  // BOTH-SAME: {{\[}}%[[STRIDE_0]], %[[STRIDE_1]]]
+  // BOTH-SAME: : memref -> memref
 
-  // CHECK: "xla_lhlo.broadcast_in_dim"(%[[TRANSFORMED_MEMREF]], %[[RESULT]]) {
-  // CHECK-SAME:   broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>
-  // CHECK-SAME: } : (memref, memref) -> ()
+  // BOTH: "xla_lhlo.broadcast_in_dim"(%[[TRANSFORMED_MEMREF]], %[[RESULT]]) {
+  // BOTH-SAME:   broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>
+  // BOTH-SAME: } : (memref, memref) -> ()
 
   // Do not store the value back to avoid the tensor-store being rewritten to
   // a copy into the pre-allocated argument.
@@ -214,7 +220,7 @@ func @dyn_broadcast(%operand: memref) {
 
 // -----
 
-// CHECK-LABEL: func @complex
+// BOTH-LABEL: func @complex
 func @complex(%real: memref<2x2xf32>,
               %imag: memref<2x2xf32>,
               %result: memref<2x2xcomplex>) {
@@ -222,164 +228,164 @@ func @complex(%real: memref<2x2xf32>,
   %tensor_imag = tensor_load %imag : memref<2x2xf32>
   %tensor_result = "xla_hlo.complex"(%tensor_real, %tensor_imag)
       : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xcomplex>
-  // CHECK: "xla_lhlo.complex"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.complex"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xcomplex>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @real
+// BOTH-LABEL: func @real
 func @real(%operand: memref<2x2xcomplex>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xcomplex>
   %tensor_result = "xla_hlo.real"(%tensor_operand)
       : (tensor<2x2xcomplex>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.real"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.real"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @imag
+// BOTH-LABEL: func @imag
 func @imag(%operand: memref<2x2xcomplex>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xcomplex>
   %tensor_result = "xla_hlo.imag"(%tensor_operand)
       : (tensor<2x2xcomplex>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.imag"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.imag"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @iota
+// BOTH-LABEL: func @iota
 func @iota(%result: memref<10xi32>) {
   %tensor_result = "xla_hlo.iota"()
       {iota_dimension = 0 : i64} : () -> tensor<10xi32>
-  // CHECK: "xla_lhlo.iota"(%{{.*}}) {iota_dimension = 0 : i64}
+  // BOTH: "xla_lhlo.iota"(%{{.*}}) {iota_dimension = 0 : i64}
   tensor_store %tensor_result, %result : memref<10xi32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @abs
+// BOTH-LABEL: func @abs
 func @abs(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.abs"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.abs"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.abs"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @ceil
+// BOTH-LABEL: func @ceil
 func @ceil(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.ceil"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.ceil"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.ceil"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @convert
+// BOTH-LABEL: func @convert
 func @convert(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.convert"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.copy"(%{{.*}}, %{{.*}})
-  // CHECK-NOT: tensor_store
+  // BOTH: "xla_lhlo.copy"(%{{.*}}, %{{.*}})
+  // BOTH-NOT: tensor_store
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @cos
+// BOTH-LABEL: func @cos
 func @cos(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.cosine"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.cosine"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.cosine"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @neg
+// BOTH-LABEL: func @neg
 func @neg(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.negate"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.negate"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.negate"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @rsqrt
+// BOTH-LABEL: func @rsqrt
 func @rsqrt(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.rsqrt"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.rsqrt"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.rsqrt"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @sign
+// BOTH-LABEL: func @sign
 func @sign(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.sign"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.sign"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.sign"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @sqrt
+// BOTH-LABEL: func @sqrt
 func @sqrt(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.sqrt"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.sqrt"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.sqrt"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @tanh
+// BOTH-LABEL: func @tanh
 func @tanh(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_operand = tensor_load %operand : memref<2x2xf32>
   %tensor_result = "xla_hlo.tanh"(%tensor_operand)
       : (tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.tanh"(%{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.tanh"(%{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @remainder
+// BOTH-LABEL: func @remainder
 func @remainder(%lhs: memref<2x2xf32>, %rhs: memref<2x2xf32>, %result: memref<2x2xf32>) {
   %tensor_lhs = tensor_load %lhs : memref<2x2xf32>
   %tensor_rhs = tensor_load %rhs : memref<2x2xf32>
   %tensor_result = "xla_hlo.remainder"(%tensor_lhs, %tensor_rhs)
       : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
-  // CHECK: "xla_lhlo.remainder"(%{{.*}}, %{{.*}}, %{{.*}})
+  // BOTH: "xla_lhlo.remainder"(%{{.*}}, %{{.*}}, %{{.*}})
   tensor_store %tensor_result, %result : memref<2x2xf32>
   return
 }
@@ -387,76 +393,79 @@ func @remainder(%lhs: memref<2x2xf32>, %rhs: memref<2x2xf32>, %result: memref<2x
 // -----
 
 // Dynamic shape binary element-wise operation.
-// CHECK-LABEL: func @add_dyn
+// BOTH-LABEL: func @add_dyn
 func @add_dyn(%lhs: tensor, %rhs: tensor) {
   %result = "xla_hlo.add"(%lhs, %rhs)
       : (tensor, tensor) -> tensor
-  // CHECK: %[[C0:.*]] = constant 0 : index
-  // CHECK: %[[DIM0:.*]] = dim %arg0, %[[C0]] : memref
-  // CHECK: %[[IC0:.*]] = index_cast %[[DIM0]] : index to i64
-  // CHECK: %[[C1:.*]] = constant 1 : index
-  // CHECK: %[[DIM1:.*]] = dim %arg0, %[[C1]] : memref
-  // CHECK: %[[IC1:.*]] = index_cast %[[DIM1]] : index to i64
-  // CHECK: %[[SHAPE:.*]] = tensor_from_elements(%[[IC0]], %[[IC1]]) : tensor<2xi64>
-  // CHECK: %[[C0_:.*]] = constant 0 : index
-  // CHECK: %[[EE0:.*]] = extract_element %[[SHAPE]][%[[C0_]]] : tensor<2xi64>
-  // CHECK: %[[ICS0:.*]] = index_cast %[[EE0]] : i64 to index
-  // CHECK: %[[C1_:.*]] = constant 1 : index
-  // CHECK: %[[EE1:.*]] = extract_element %[[SHAPE]][%[[C1_]]] : tensor<2xi64>
-  // CHECK: %[[ICS1:.*]] = index_cast %[[EE1]] : i64 to index
-  // CHECK: %[[RESULT:.*]] = alloc(%[[ICS0]], %[[ICS1]])
-  // CHECK: "xla_lhlo.add"(%arg0, %arg1, %[[RESULT]]) : (memref, memref, memref) -> ()
+  // BOTH: %[[C0:.*]] = constant 0 : index
+  // BOTH: %[[DIM0:.*]] = dim %arg0, %[[C0]] : memref
+  // BOTH: %[[IC0:.*]] = index_cast %[[DIM0]] : index to i64
+  // BOTH: %[[C1:.*]] = constant 1 : index
+  // BOTH: %[[DIM1:.*]] = dim %arg0, %[[C1]] : memref
+  // BOTH: %[[IC1:.*]] = index_cast %[[DIM1]] : index to i64
+  // BOTH: %[[SHAPE:.*]] = tensor_from_elements(%[[IC0]], %[[IC1]]) : tensor<2xi64>
+  // BOTH: %[[C0_:.*]] = constant 0 : index
+  // BOTH: %[[EE0:.*]] = extract_element %[[SHAPE]][%[[C0_]]] : tensor<2xi64>
+  // BOTH: %[[ICS0:.*]] = index_cast %[[EE0]] : i64 to index
+  // BOTH: %[[C1_:.*]] = constant 1 : index
+  // BOTH: %[[EE1:.*]] = extract_element %[[SHAPE]][%[[C1_]]] : tensor<2xi64>
+  // BOTH: %[[ICS1:.*]] = index_cast %[[EE1]] : i64 to index
+  // BOTH: %[[RESULT:.*]] = alloc(%[[ICS0]], %[[ICS1]])
+  // BOTH: "xla_lhlo.add"(%arg0, %arg1, %[[RESULT]]) : (memref, memref, memref) -> ()
   return
 }
 
 // -----
 
 // Dynamic shape unary element-wise operation.
-// CHECK-LABEL: func @tanh_dyn
+// BOTH-LABEL: func @tanh_dyn
 func @tanh_dyn(%arg0: tensor) {
   %result = "xla_hlo.tanh"(%arg0)
       : (tensor) -> tensor
-  // CHECK: %[[C0:.*]] = constant 0 : index
-  // CHECK: %[[DIM0:.*]] = dim %arg0, %[[C0]] : memref
-  // CHECK: %[[IC0:.*]] = index_cast %[[DIM0]] : index to i64
-  // CHECK: %[[C1:.*]] = constant 1 : index
-  // CHECK: %[[DIM1:.*]] = dim %arg0, %[[C1]] : memref
-  // CHECK: %[[IC1:.*]] = index_cast %[[DIM1]] : index to i64
-  // CHECK: %[[SHAPE:.*]] = tensor_from_elements(%[[IC0]], %[[IC1]]) : tensor<2xi64>
-  // CHECK: %[[C0_:.*]] = constant 0 : index
-  // CHECK: %[[EE0:.*]] = extract_element %[[SHAPE]][%[[C0_]]] : tensor<2xi64>
-  // CHECK: %[[ICS0:.*]] = index_cast %[[EE0]] : i64 to index
-  // CHECK: %[[C1_:.*]] = constant 1 : index
-  // CHECK: %[[EE1:.*]] = extract_element %[[SHAPE]][%[[C1_]]] : tensor<2xi64>
-  // CHECK: %[[ICS1:.*]] = index_cast %[[EE1]] : i64 to index
-  // CHECK: %[[RESULT:.*]] = alloc(%[[ICS0]], %[[ICS1]])
-  // CHECK: "xla_lhlo.tanh"(%arg0, %[[RESULT]]) : (memref, memref) -> ()
+  // BOTH: %[[C0:.*]] = constant 0 : index
+  // BOTH: %[[DIM0:.*]] = dim %arg0, %[[C0]] : memref
+  // BOTH: %[[IC0:.*]] = index_cast %[[DIM0]] : index to i64
+  // BOTH: %[[C1:.*]] = constant 1 : index
+  // BOTH: %[[DIM1:.*]] = dim %arg0, %[[C1]] : memref
+  // BOTH: %[[IC1:.*]] = index_cast %[[DIM1]] : index to i64
+  // BOTH: %[[SHAPE:.*]] = tensor_from_elements(%[[IC0]], %[[IC1]]) : tensor<2xi64>
+  // BOTH: %[[C0_:.*]] = constant 0 : index
+  // BOTH: %[[EE0:.*]] = extract_element %[[SHAPE]][%[[C0_]]] : tensor<2xi64>
+  // BOTH: %[[ICS0:.*]] = index_cast %[[EE0]] : i64 to index
+  // BOTH: %[[C1_:.*]] = constant 1 : index
+  // BOTH: %[[EE1:.*]] = extract_element %[[SHAPE]][%[[C1_]]] : tensor<2xi64>
+  // BOTH: %[[ICS1:.*]] = index_cast %[[EE1]] : i64 to index
+  // BOTH: %[[RESULT:.*]] = alloc(%[[ICS0]], %[[ICS1]])
+  // BOTH: "xla_lhlo.tanh"(%arg0, %[[RESULT]]) : (memref, memref) -> ()
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @dot
+// BOTH-LABEL: func @dot
 func @dot(%arg0: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
-// CHECK-SAME: (%[[ARG0:.*]]: [[TYPE:.*]],
-// CHECK-SAME:  %[[RESULT:.*]]: [[TYPE]])
-// CHECK: "xla_lhlo.dot"(%[[ARG0]], %[[ARG0]], %{{.*}}) : ([[TYPE]], [[TYPE]], [[TYPE]]) -> ()
+//  PRE-SAME: (%[[ARG0:.*]]: [[TYPE:.*]], %[[RESULT:.*]]: [[TYPE]])
+//  ESC-SAME: (%[[ARG0:.*]]: [[TYPE:.*]]) -> [[TYPE]]
+// BOTH-NEXT: %[[ALLOC:.*]] = alloc
+//      BOTH: "xla_lhlo.dot"(%[[ARG0]], %[[ARG0]], %[[ALLOC]]) : ([[TYPE]], [[TYPE]], [[TYPE]]) -> ()
   %dot = "xla_hlo.dot"(%arg0, %arg0)
           : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
+// PRE: "xla_lhlo.copy"(%[[ALLOC]], %[[RESULT]])
+// ESC: return %[[ALLOC]]
   return %dot : tensor<1024x1024xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @conv
+// BOTH-LABEL: func @conv
 func @conv(%input: tensor<3x5x5x3xf32>, %filter : tensor<2x2x3x4xf32>) -> tensor<3x5x5x4xf32> {
   %c0 = constant 0 : index
-  // CHECK: %[[OUT:.*]] = alloc() : memref<3x5x5x4xf32>
-  // CHECK: "xla_lhlo.convolution"(%{{.+}}, %{{.+}}, %[[OUT]])
-  // CHECK-SAME: padding = dense<[
-  // CHECK-SAME:                  [0, 1], [0, 1]]> : tensor<2x2xi64>
-  // CHECK-SAME: rhs_dilation = dense<[1, 2]>
-  // CHECK-SAME: window_strides = dense<[2, 1]>
+  // BOTH: %[[OUT:.*]] = alloc() : memref<3x5x5x4xf32>
+  // BOTH: "xla_lhlo.convolution"(%{{.+}}, %{{.+}}, %[[OUT]])
+  // BOTH-SAME: padding = dense<[
+  // BOTH-SAME:                  [0, 1], [0, 1]]> : tensor<2x2xi64>
+  // BOTH-SAME: rhs_dilation = dense<[1, 2]>
+  // BOTH-SAME: window_strides = dense<[2, 1]>
   %out = "xla_hlo.convolution"(%filter, %input) {
     batch_group_count = 1 : i64,
     dimension_numbers = {
diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
index f966ce765fc..1cfe0c12e20 100644
--- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
@@ -368,6 +368,15 @@ class HloToLhloTensorStoreOpConverter
 
 struct HloLegalizeToLhlo
     : public PassWrapper> {
+ public:
+  HloLegalizeToLhlo() = default;
+  HloLegalizeToLhlo(const HloLegalizeToLhlo& o) {
+    this->results_escape_function = o.results_escape_function.getValue();
+  }
+  explicit HloLegalizeToLhlo(bool results_escape_function) {
+    this->results_escape_function.setValue(results_escape_function);
+  }
+
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     auto& context = getContext();
@@ -398,10 +407,28 @@ struct HloLegalizeToLhlo
       OwningRewritePatternList patterns;
       populateHLOToLHLOConversionPattern(func.getContext(), &bufferAssignment,
                                          &converter, &patterns);
+      if (results_escape_function) {
+        populateWithBufferAssignmentOpConversionPatterns<
+            mlir::ReturnOp, mlir::ReturnOp, xla_lhlo::CopyOp,
+            /*allowMemrefFunctionResults=*/true>(&context, &bufferAssignment,
+                                                 &converter, &patterns);
+      } else {
+        populateWithBufferAssignmentOpConversionPatterns<
+            mlir::ReturnOp, mlir::ReturnOp, xla_lhlo::CopyOp,
+            /*allowMemrefFunctionResults=*/false>(&context, &bufferAssignment,
+                                                  &converter, &patterns);
+      }
       return WalkResult(
           applyPartialConversion(func, target, patterns, &converter));
     });
   }
+
+ private:
+  Option results_escape_function{
+      *this, "results-escape-function",
+      llvm::cl::desc(
+          "Allocate the results of functions within the functions body"),
+      llvm::cl::init(false)};
 };
 }  // namespace
 
@@ -446,14 +473,11 @@ void populateHLOToLHLOConversionPattern(
       HloToLhloTensorStoreOpConverter
   >(context, bufferAssignment, converter);
   // clang-format on
-  populateWithBufferAssignmentOpConversionPatterns<
-      mlir::ReturnOp, mlir::ReturnOp, xla_lhlo::CopyOp,
-      /*allowMemrefFunctionResults=*/false>(context, bufferAssignment,
-                                            converter, patterns);
 }
 
-std::unique_ptr> createLegalizeToLhloPass() {
-  return absl::make_unique();
+std::unique_ptr> createLegalizeToLhloPass(
+    bool results_escape_function) {
+  return absl::make_unique(results_escape_function);
 }
 
 static PassRegistration legalize_pass(
diff --git a/tensorflow/compiler/mlir/xla/transforms/passes.h b/tensorflow/compiler/mlir/xla/transforms/passes.h
index 9b9c799b2f0..f0c2d9b7372 100644
--- a/tensorflow/compiler/mlir/xla/transforms/passes.h
+++ b/tensorflow/compiler/mlir/xla/transforms/passes.h
@@ -59,9 +59,13 @@ std::unique_ptr> createLegalizeControlFlowPass();
 /// Lowers from HLO dialect to Standard dialect.
 std::unique_ptr> createLegalizeToStdPass();
 
-// Lowers from HLO dialect to LHLO dialect allocating/deallocating temporary
-// buffers if necessary.
-std::unique_ptr> createLegalizeToLhloPass();
+/// Lowers from HLO dialect to LHLO dialect allocating/deallocating temporary
+/// buffers if necessary. If `results_escape_functions` is set to true,
+/// allocated buffers for function results will be returned and escape the
+/// function. Otherwise, the signature is rewritten with extra arguments for the
+/// buffers that are to be used for results.
+std::unique_ptr> createLegalizeToLhloPass(
+    bool results_escape_functions = false);
 
 // Lowers from HLO dialect to Linalg dialect.
 std::unique_ptr> createLegalizeHloToLinalgPass();

From 2614035897d5c29fe54a3c45a605f11703db7097 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 05:39:27 -0700
Subject: [PATCH 0279/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/9b72b47ed633

PiperOrigin-RevId: 316662757
Change-Id: I94751da57f9731bd60aae6dd09186f3e173e5c89
---
 .../lhlo-legalize-select-and-scatter.mlir     | 30 ++++++++-----------
 .../lhlo-legalize-to-parallel-loops.mlir      | 10 ++-----
 2 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-select-and-scatter.mlir b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-select-and-scatter.mlir
index a8c8f5e73c5..9887860ca26 100644
--- a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-select-and-scatter.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-select-and-scatter.mlir
@@ -38,15 +38,15 @@ func @select_and_scatter(%arg: memref<112x112xf32>,
 // CHECK-SAME:   [[RESULT_BUF:%.*]]: memref<112x112xf32>) {
 
 // Constants.
-// CHECK:  [[C56:%.*]] = constant 56 : index
-// CHECK:  [[C1:%.*]] = constant 1 : index
-// CHECK:  [[C0_F32:%.*]] = constant 0.000000e+00 : f32
-// CHECK:  [[CFALSE:%.*]] = constant false
-// CHECK:  [[C3:%.*]] = constant 3 : index
-// CHECK:  [[C2:%.*]] = constant 2 : index
-// CHECK:  [[C0:%.*]] = constant 0 : index
-// CHECK:  [[C112:%.*]] = constant 112 : index
-// CHECK:  [[CTRUE:%.*]] = constant true
+// CHECK-DAG:  [[C56:%.*]] = constant 56 : index
+// CHECK-DAG:  [[C0:%.*]] = constant 0 : index
+// CHECK-DAG:  [[C1:%.*]] = constant 1 : index
+// CHECK-DAG:  [[C0_F32:%.*]] = constant 0.000000e+00 : f32
+// CHECK-DAG:  [[CFALSE:%.*]] = constant false
+// CHECK-DAG:  [[C3:%.*]] = constant 3 : index
+// CHECK-DAG:  [[C2:%.*]] = constant 2 : index
+// CHECK-DAG:  [[C112:%.*]] = constant 112 : index
+// CHECK-DAG:  [[CTRUE:%.*]] = constant true
 
 // Parallel loop to initialize the output buffer.
 // CHECK:    [[INIT:%.*]] = load [[INIT_BUF]][] : memref
@@ -80,23 +80,17 @@ func @select_and_scatter(%arg: memref<112x112xf32>,
 
 // Compute index I of the ARG buffer and check whether it is in padding area.
 // CHECK:  [[START_I:%.*]] = muli [[II]], [[C2]] : index
-// CHECK:  [[OFFSET_I:%.*]] = subi [[WIN_I]], [[C0]] : index
-// CHECK:  [[ARG_I:%.*]] = addi [[START_I]], [[OFFSET_I]] : index
+// CHECK:  [[ARG_I:%.*]] = addi [[START_I]], [[WIN_I]] : index
 // CHECK:  [[ARG_I_FITS:%.*]] = cmpi "ult", [[ARG_I]], [[C112]] : index
 
-// Update `INBOUNDS`, i.e. whether or not ARG indices are inside the boundaries
-// of the buffer or they are in the padding area.
-// CHECK:      [[INBOUNDS_0:%.*]] = and [[ARG_I_FITS]], [[CTRUE]] : i1
-
 // Compute index J of the ARG buffer and check whether it is in padding area.
 // CHECK:  [[START_J:%.*]] = muli [[JJ]], [[C2]] : index
-// CHECK:  [[OFFSET_J:%.*]] = subi [[WIN_J]], [[C0]] : index
-// CHECK:  [[ARG_J:%.*]] = addi [[START_J]], [[OFFSET_J]] : index
+// CHECK:  [[ARG_J:%.*]] = addi [[START_J]], [[WIN_J]] : index
 // CHECK:  [[ARG_J_FITS:%.*]] = cmpi "ult", [[ARG_J]], [[C112]] : index
 
 // Update `INBOUNDS`, i.e. whether or not ARG indices are inside the boundaries
 // of the buffer or they are in the padding area.
-// CHECK:  [[INBOUNDS_1:%.*]] = and [[INBOUNDS_0]], [[ARG_J_FITS]] : i1
+// CHECK:  [[INBOUNDS_1:%.*]] = and [[ARG_I_FITS]], [[ARG_J_FITS]] : i1
 
 // If ARG ivs are in the padding area, then 'select' function does not have to
 // be applied, current selected ivs (SEL_I, SEL_J) and value (SEL_VAL) are
diff --git a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-parallel-loops.mlir b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-parallel-loops.mlir
index 8c22c035edd..5127bcfcd8f 100644
--- a/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-parallel-loops.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/lhlo-legalize-to-parallel-loops.mlir
@@ -151,7 +151,6 @@ func @reduce_window(%arg: memref<112x112xf32>,
 // CHECK-SAME:      [[OPERAND_BUF:%.*]]: memref<112x112xf32>,
 // CHECK-SAME:      [[INIT_BUF:%.*]]: memref,
 // CHECK-SAME:      [[RESULT_BUF:%.*]]: memref<56x56xf32>) {
-// CHECK-DAG:  [[IN_BOUNDS:%.*]] = constant true
 // CHECK-DAG:  [[C0:%.*]] = constant 0 : index
 // CHECK-DAG:  [[C1:%.*]] = constant 1 : index
 // CHECK-DAG:  [[C2:%.*]] = constant 2 : index
@@ -167,16 +166,13 @@ func @reduce_window(%arg: memref<112x112xf32>,
 // CHECK-SAME:       init ([[INIT]]) -> f32 {
 
 // CHECK:          [[START_I:%.*]] = muli [[I]], [[C2]] : index
-// CHECK:          [[OFFSET_I:%.*]] = subi [[IW]], [[C0]] : index
-// CHECK:          [[INDEX_I:%.*]] = addi [[START_I]], [[OFFSET_I]] : index
+// CHECK:          [[INDEX_I:%.*]] = addi [[START_I]], [[IW]] : index
 // CHECK:          [[INDEX_I_FITS:%.*]] = cmpi "ult", [[INDEX_I]], [[C112]]
-// CHECK:          [[IN_BOUNDS_0:%.*]] = and [[INDEX_I_FITS]], [[IN_BOUNDS]]
 
 // CHECK:          [[START_J:%.*]] = muli [[J]], [[C2]] : index
-// CHECK:          [[OFFSET_J:%.*]] = subi [[JW]], [[C0]] : index
-// CHECK:          [[INDEX_J:%.*]] = addi [[START_J]], [[OFFSET_J]] : index
+// CHECK:          [[INDEX_J:%.*]] = addi [[START_J]], [[JW]] : index
 // CHECK:          [[INDEX_J_FITS:%.*]] = cmpi "ult", [[INDEX_J]], [[C112]]
-// CHECK:          [[IN_BOUNDS_1:%.*]] = and [[IN_BOUNDS_0]], [[INDEX_J_FITS]]
+// CHECK:          [[IN_BOUNDS_1:%.*]] = and [[INDEX_I_FITS]], [[INDEX_J_FITS]]
 
 // CHECK:          [[ELEM_TO_REDUCE:%.*]] = scf.if [[IN_BOUNDS_1]] -> (f32) {
 // CHECK:            [[OPERAND_ELEM:%.*]] =

From 55a1f169c3d2db146da385722213e8b4dfcfa6c3 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 16 Jun 2020 06:09:10 -0700
Subject: [PATCH 0280/1390] Fix same_shape constraint handling in cubin
 generator.

PiperOrigin-RevId: 316665956
Change-Id: Icd582db4e0c4c19b7f4b0c563a3323df0cda6836
---
 .../mlir/tools/kernel_gen/cubin_creator.cc    | 39 ++++++++++++-------
 tensorflow/core/kernels/cubin_headers/BUILD   |  1 +
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
index 30b60e8079f..b534b5a5604 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
@@ -125,7 +125,8 @@ Status LowerTfOpToLhloWithDynamicShapes(mlir::ModuleOp module) {
   pm.addNestedPass(
       absl::make_unique());
   pm.addNestedPass(absl::make_unique());
-  pm.addPass(mlir::xla_hlo::createLegalizeToLhloPass());
+  pm.addPass(mlir::xla_hlo::createLegalizeToLhloPass(
+      /*results_escape_functions=*/true));
   pm.addNestedPass(mlir::xla_lhlo::createLhloCopyRemovalPass());
 
   if (failed(pm.run(module))) {
@@ -148,7 +149,12 @@ struct PropagateStaticKnowledge
     // We do not change the signature so that we keep a somewhat stable ABI
     // that is easy to undertand by tools.
     mlir::LLVM::LLVMFuncOp func = getOperation();
+
+    // This only works if the function is local and we can rewrite it.
+    if (func.isExternal()) return;
+
     mlir::OpBuilder b(func.getBody());
+    // Steal the LLVM representation of the index type from the third argument.
     auto index_type = func.getArgument(3).getType();
     mlir::Value one = b.create(
         func.getLoc(), index_type, b.getIntegerAttr(b.getIndexType(), 1));
@@ -156,10 +162,21 @@ struct PropagateStaticKnowledge
         func.getLoc(), index_type, b.getIntegerAttr(b.getIndexType(), 0));
     uint32_t arg_pos = 0;
     std::vector positions;
-    for (mlir::Type arg_type : func_type.getInputs()) {
+    // Collect the agument and return types of the surrounding function.
+    auto arg_types = llvm::to_vector<4>(llvm::concat(
+        func_type.getInputs(), func_type.getResults()));
+    for (mlir::Type arg_type : arg_types) {
+      if (!arg_type.isa()) {
+        func.emitError() << "argument of surrounding func is not ranked memref";
+        signalPassFailure();
+        return;
+      }
       positions.push_back(arg_pos);
+      // Replace the offset with zero. Offset is argument number 3.
       func.getArgument(arg_pos + 2).replaceAllUsesWith(zero);
-      arg_pos += 3 + arg_type.cast().getRank() * 2;
+      // Forward over base_ptr, aligned_ptr, offset, size and stride arguments.
+      arg_pos += 3 + arg_type.cast().getRank() * 2;
+      // Replace the last stride with constant 1.
       func.getArgument(arg_pos - 1).replaceAllUsesWith(one);
     }
 
@@ -169,17 +186,17 @@ struct PropagateStaticKnowledge
     if (!same_shape.empty()) {
       auto first = same_shape.front();
       auto first_offset = positions.at(first);
-      mlir::ShapedType first_type =
-          func_type.getInput(first).cast();
+      auto first_type = arg_types[first].cast();
       uint32_t rank = first_type.getRank();
       for (auto same : same_shape.drop_front(1)) {
         uint32_t same_offset = positions.at(same);
-        auto same_type = func_type.getInput(same).cast();
+        auto same_type = arg_types[same].cast();
         if (same_type.getRank() != rank) {
           func.emitOpError() << "same shape constraints on arguments with "
                                 "non-matching shapes: #"
                              << first << " and #" << same;
           signalPassFailure();
+          continue;
         }
 
         for (uint32_t i = 0; i < 2 * rank; ++i) {
@@ -245,14 +262,8 @@ StatusOr> tensorflow::kernel_gen::GenerateCubinForTfCode(
     TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerLHLOToGPU(module.get(), options));
   }
   TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerKernelBodiesToNVVM(module.get()));
-  // TODO(b/156985522): Figure out why we get a segfault when generating Tanh
-  // with 'same_shape' containing {0, 1}. We would also get the crash if we
-  // unconditionally call PropagateStaticShapeKnowledgeToKernel while
-  // 'same_shape' is empty.
-  if (!same_shape.empty()) {
-    TF_RETURN_IF_ERROR(
-        PropagateStaticShapeKnowledgeToKernel(module.get(), same_shape));
-  }
+  TF_RETURN_IF_ERROR(
+      PropagateStaticShapeKnowledgeToKernel(module.get(), same_shape));
 
   mlir::OwningModuleRef kernel_module =
       xla::mlir_gpu::ExtractKernelModule(*module).ValueOrDie();
diff --git a/tensorflow/core/kernels/cubin_headers/BUILD b/tensorflow/core/kernels/cubin_headers/BUILD
index 49ab1b8a911..b8ba164fbc3 100644
--- a/tensorflow/core/kernels/cubin_headers/BUILD
+++ b/tensorflow/core/kernels/cubin_headers/BUILD
@@ -30,6 +30,7 @@ gen_kernel_library(
 
 gen_kernel_library(
     name = "tanh",
+    same_shape = "0,1",
     tile_size = "256",
     types = [
         "f16",

From 35fc2312816ffbb577843e2741568049639021e7 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev 
Date: Tue, 16 Jun 2020 06:15:38 -0700
Subject: [PATCH 0281/1390] Internal change

PiperOrigin-RevId: 316666669
Change-Id: I39909bfda1b991755cc4f7a69b00659fa947b8f4
---
 third_party/mlir/BUILD | 27 +++++----------------------
 1 file changed, 5 insertions(+), 22 deletions(-)

diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index 06e0ed8d4b4..cb0b2f9dc8e 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -707,14 +707,8 @@ gentbl(
 
 cc_library(
     name = "Shape",
-    srcs = glob(
-        [
-            "lib/Dialect/Shape/IR/*.cpp",
-        ],
-    ),
-    hdrs = glob([
-        "include/mlir/Dialect/Shape/IR/*.h",
-    ]),
+    srcs = glob(["lib/Dialect/Shape/IR/*.cpp"]),
+    hdrs = ["include/mlir/Dialect/Shape/IR/Shape.h"],
     includes = ["include"],
     deps = [
         ":CallOpInterfaces",
@@ -737,22 +731,16 @@ cc_library(
         "lib/Conversion/ShapeToStandard/*.cpp",
         "lib/Conversion/ShapeToStandard/*.h",
     ]) + ["lib/Conversion/PassDetail.h"],
-    hdrs = glob([
-        "include/mlir/Conversion/ShapeToStandard/*.h",
-    ]),
+    hdrs = ["include/mlir/Conversion/ShapeToStandard/ShapeToStandard.h"],
     includes = ["include"],
     deps = [
-        ":Affine",
         ":ConversionPassIncGen",
-        ":IR",
         ":Pass",
         ":SCFDialect",
         ":Shape",
         ":StandardOps",
         ":Support",
         ":Transforms",
-        "@llvm-project//llvm:Core",
-        "@llvm-project//llvm:Support",
     ],
 )
 
@@ -793,18 +781,15 @@ cc_library(
         "lib/Dialect/Shape/Transforms/*.cpp",
         "lib/Dialect/Shape/Transforms/*.h",
     ]),
-    hdrs = glob(["include/mlir/Dialect/Shape/Transforms/*.h"]),
+    hdrs = ["include/mlir/Dialect/Shape/Transforms/Passes.h"],
     includes = ["include"],
     deps = [
-        #":Analysis",
-        #":ControlFlowInterfaces",
         ":IR",
         ":Pass",
         ":Shape",
         ":ShapeTransformsPassIncGen",
         ":Support",
         ":Transforms",
-        #"@llvm-project//llvm:Support",
     ],
 )
 
@@ -820,9 +805,7 @@ cc_library(
     hdrs = glob([
         "include/mlir/Dialect/StandardOps/IR/*.h",
         "include/mlir/Dialect/StandardOps/EDSC/*.h",
-    ]) + [
-        "include/mlir/Transforms/InliningUtils.h",
-    ],
+    ]) + ["include/mlir/Transforms/InliningUtils.h"],
     includes = ["include"],
     deps = [
         ":CallOpInterfaces",

From b8b6bc58a1954353b6fbd49e6d7a0d1f6ad18072 Mon Sep 17 00:00:00 2001
From: Thomas Joerg 
Date: Tue, 16 Jun 2020 06:21:33 -0700
Subject: [PATCH 0282/1390] [XLA] Keep op metadata when replacing with bitcast.

PiperOrigin-RevId: 316667446
Change-Id: I3618b5eafba269eb087ff7607bdad1cb3873ffbc
---
 tensorflow/compiler/xla/service/algebraic_simplifier.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
index f88f08b9fa2..98e3229b062 100755
--- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc
+++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -573,6 +573,7 @@ void AlgebraicSimplifierVisitor::ReplaceWithBitcast(HloInstruction* instruction,
 
   auto bitcast = computation_->AddInstruction(
       HloInstruction::CreateBitcast(instruction->shape(), operand));
+  bitcast->set_metadata(instruction->metadata());
   TF_CHECK_OK(ReplaceInstruction(instruction, bitcast));
 }
 

From d9532e652600fcda0fee32b14cd62a4404a662fe Mon Sep 17 00:00:00 2001
From: Benjamin Kramer 
Date: Tue, 16 Jun 2020 06:27:29 -0700
Subject: [PATCH 0283/1390] Bump open source llvm revision to
 9b72b47ed63351ee5ceff4c44ccd9a71dc7dad27

PiperOrigin-RevId: 316668225
Change-Id: Ie49a7f46a51674c854e300629f4563c5be5e6448
---
 tensorflow/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 31daf2249a0..acba1598d1f 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "1a7f115dce22b2c09fdd4f7f79d24da5de6eaef8"
-    LLVM_SHA256 = "3ace55744a86211c9c837915b88c18e1ed3e5cd839aaeade6aa88b02bc86e47e"
+    LLVM_COMMIT = "9b72b47ed63351ee5ceff4c44ccd9a71dc7dad27"
+    LLVM_SHA256 = "03ce1e00901936e7259c6ee465773b7f231ca1724925460b71909868f5a61e11"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),

From 572442eb16bb989c1ae9665549c2c2d9daac7952 Mon Sep 17 00:00:00 2001
From: Peter Hawkins 
Date: Tue, 16 Jun 2020 06:56:05 -0700
Subject: [PATCH 0284/1390] [PJRT] Fix potential misuse of
 `PjRtBuffer::FromHostBuffer`.

Add a new `PjRtBuffer::HostBufferSemantics` enum that describes the possible contracts between caller and runtime.

* Change `FromHostBuffer(..., force_copy, ...)` to `FromHostBuffer(..., host_buffer_semantics, ...)`.

We were seeing some data races between modifications to a NumPy array and JAX on CPU, due to unintended buffer aliasing. This change allows clients to control whether they want zero-copy behavior or not.

PiperOrigin-RevId: 316672280
Change-Id: Ibee296305005e0aa306a2c0aacf4b35a3d6c3ac1
---
 tensorflow/compiler/xla/pjrt/cpu_device.cc    |   1 +
 .../compiler/xla/pjrt/gpu_multistream_test.cc |  17 +-
 .../compiler/xla/pjrt/interpreter_device.cc   |   1 +
 .../compiler/xla/pjrt/nvidia_gpu_device.cc    |   1 +
 tensorflow/compiler/xla/pjrt/pjrt_client.cc   | 162 +++++++++++++-----
 tensorflow/compiler/xla/pjrt/pjrt_client.h    |  49 ++++--
 tensorflow/compiler/xla/python/py_client.cc   |   9 +-
 tensorflow/compiler/xla/python/py_client.h    |   3 +-
 tensorflow/compiler/xla/python/xla.cc         |  11 +-
 tensorflow/compiler/xla/python/xla_client.py  |   1 +
 .../compiler/xla/python/xla_client_test.py    |   7 +-
 11 files changed, 195 insertions(+), 67 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/cpu_device.cc b/tensorflow/compiler/xla/pjrt/cpu_device.cc
index 75c3bfc1277..be70c16fc12 100644
--- a/tensorflow/compiler/xla/pjrt/cpu_device.cc
+++ b/tensorflow/compiler/xla/pjrt/cpu_device.cc
@@ -59,6 +59,7 @@ StatusOr> GetCpuClient(bool asynchronous) {
   return std::make_shared(
       kCpuPlatformName, client, std::move(devices), /*host_id=*/0,
       /*allocator=*/nullptr, /*host_memory_allocator=*/nullptr,
+      /*should_stage_host_to_device_transfers=*/false,
       /*gpu_run_options=*/nullptr);
 }
 
diff --git a/tensorflow/compiler/xla/pjrt/gpu_multistream_test.cc b/tensorflow/compiler/xla/pjrt/gpu_multistream_test.cc
index 2db7de3720d..d54be61fbb8 100644
--- a/tensorflow/compiler/xla/pjrt/gpu_multistream_test.cc
+++ b/tensorflow/compiler/xla/pjrt/gpu_multistream_test.cc
@@ -72,18 +72,21 @@ TEST(GpuMultiStream, Basics) {
     TF_ASSERT_OK_AND_ASSIGN(
         auto dummy_buffer,
         PjRtBuffer::FromHostBuffer(
-            dummy_inputs.data(), dummy_shape, /*force_copy=*/false,
+            dummy_inputs.data(), dummy_shape,
+            PjRtBuffer::HostBufferSemantics::kImmutableUntilTransferCompletes,
             /*buffer_reference=*/nullptr, client.get(), device));
     TF_ASSERT_OK_AND_ASSIGN(
         auto in_buffer0,
-        PjRtBuffer::FromHostBuffer(inputs.data(), shape, /*force_copy=*/false,
-                                   /*buffer_reference=*/nullptr, client.get(),
-                                   device));
+        PjRtBuffer::FromHostBuffer(
+            inputs.data(), shape,
+            PjRtBuffer::HostBufferSemantics::kImmutableUntilTransferCompletes,
+            /*buffer_reference=*/nullptr, client.get(), device));
     TF_ASSERT_OK_AND_ASSIGN(
         auto in_buffer1,
-        PjRtBuffer::FromHostBuffer(inputs.data(), shape, /*force_copy=*/false,
-                                   /*buffer_reference=*/nullptr, client.get(),
-                                   device));
+        PjRtBuffer::FromHostBuffer(
+            inputs.data(), shape,
+            PjRtBuffer::HostBufferSemantics::kImmutableUntilTransferCompletes,
+            /*buffer_reference=*/nullptr, client.get(), device));
     // The execution may be enqueued before the transfers complete, requiring
     // adequate device-side synchronization.
     ExecuteOptions options;
diff --git a/tensorflow/compiler/xla/pjrt/interpreter_device.cc b/tensorflow/compiler/xla/pjrt/interpreter_device.cc
index 63254d4aa70..f7138a8c181 100644
--- a/tensorflow/compiler/xla/pjrt/interpreter_device.cc
+++ b/tensorflow/compiler/xla/pjrt/interpreter_device.cc
@@ -53,6 +53,7 @@ StatusOr> GetInterpreterClient() {
   return std::make_shared(
       kInterpreterPlatformName, client, std::move(devices), /*host_id=*/0,
       /*allocator=*/nullptr, /*host_memory_allocator=*/nullptr,
+      /*should_stage_host_to_device_transfers=*/false,
       /*gpu_run_options=*/nullptr);
 }
 
diff --git a/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc b/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
index 4863e5e8165..de760af8fd9 100644
--- a/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
+++ b/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
@@ -316,6 +316,7 @@ StatusOr> GetNvidiaGpuClient(
       "gpu", xla_client, std::move(devices),
       /*node_id=*/node_id, std::move(allocator),
       std::move(host_memory_allocator),
+      /*should_stage_host_to_device_transfers=*/true,
       /*gpu_run_options=*/std::move(gpu_run_options));
   return pyclient;
 }
diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.cc b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
index c1b433845b2..ccb72b7ce30 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.cc
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
@@ -95,6 +95,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/mem.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
@@ -154,18 +155,35 @@ StatusOr DevicesToDeviceAssignment(
   return xla_assignment;
 }
 
+class CpuAllocator : public tensorflow::Allocator {
+ public:
+  CpuAllocator() = default;
+
+  string Name() override { return "cpu"; }
+
+  void* AllocateRaw(size_t alignment, size_t num_bytes) override {
+    return tensorflow::port::AlignedMalloc(num_bytes, alignment);
+  }
+  void DeallocateRaw(void* ptr) override {
+    return tensorflow::port::AlignedFree(ptr);
+  }
+};
+
 PjRtClient::PjRtClient(
     std::string platform_name, LocalClient* client,
     std::vector> devices, int host_id,
     std::unique_ptr allocator,
     std::unique_ptr host_memory_allocator,
+    bool should_stage_host_to_device_transfers,
     std::unique_ptr gpu_run_options)
     : platform_name_(std::move(platform_name)),
       client_(client),
+      host_memory_allocator_(std::move(host_memory_allocator)),
       devices_(std::move(devices)),
       host_id_(host_id),
       owned_allocator_(std::move(allocator)),
-      host_memory_allocator_(std::move(host_memory_allocator)),
+      should_stage_host_to_device_transfers_(
+          should_stage_host_to_device_transfers),
       gpu_run_options_(std::move(gpu_run_options)),
       h2d_transfer_pool_(tensorflow::Env::Default(), "py_xla_h2d_transfer",
                          client->device_count()) {
@@ -175,6 +193,10 @@ PjRtClient::PjRtClient(
     allocator_ = client_->backend().memory_allocator();
   }
 
+  if (!host_memory_allocator_) {
+    host_memory_allocator_ = std::make_unique();
+  }
+
   for (const std::unique_ptr& device : devices_) {
     CHECK(id_to_device_.insert({device->id(), device.get()}).second)
         << "Duplicate device id: " << device->id();
@@ -526,7 +548,8 @@ void PjRtBuffer::ScopedHold::AddToInput(
 
 /* static */
 StatusOr> PjRtBuffer::FromHostBuffer(
-    const void* data, const Shape& shape, bool force_copy,
+    const void* data, const Shape& shape,
+    HostBufferSemantics host_buffer_semantics,
     std::shared_ptr buffer_reference, PjRtClient* client,
     Device* device) {
   tensorflow::profiler::TraceMe traceme("PjRtBuffer::FromHostBuffer");
@@ -537,34 +560,63 @@ StatusOr> PjRtBuffer::FromHostBuffer(
   }
   TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device,
                       device->GetLocalDeviceState());
-
-  // If we are on the host platform and the input buffer is sufficiently
-  // aligned, we can simply point to the input array's data without any further
-  // copies. At the time of writing we require a 16-byte alignment because XLA
-  // may generate code which requires it.
-  if (!force_copy &&
-      ((absl::bit_cast(data) &
-        (cpu_function_runtime::kMinAlign - 1)) == 0) &&
-      local_device->executor()->platform()->id() == se::host::kHostPlatformId) {
-    std::function on_delete_callback =
-        [buffer_reference{std::move(buffer_reference)}]() {
-          // Frees buffer_reference.
-        };
-    se::DeviceMemoryBase buffer(const_cast(data),
-                                ShapeUtil::ByteSizeOf(shape));
-    absl::Span> definition_events;
-    auto device_buffer = std::make_shared(
-        /*allocator=*/nullptr, local_device->device_ordinal(),
-        std::initializer_list{buffer}, definition_events,
-        std::move(on_delete_callback));
-    return absl::make_unique(shape, shape, std::move(device_buffer),
-                                         client, device);
-  }
+  int64 size = ShapeUtil::ByteSizeOf(shape);
 
   TransferManager* transfer_manager =
       client->client()->backend().transfer_manager();
   TF_ASSIGN_OR_RETURN(Shape compact_shape,
                       transfer_manager->ChooseCompactLayoutForShape(shape));
+
+  // The CPU platform is special because the "host" and the "device" are in the
+  // same memory space. If the input shape is in the correct layout and we don't
+  // want to defer the copy onto a thread, we can use the following fast
+  // path.
+  bool is_cpu_platform =
+      local_device->executor()->platform()->id() == se::host::kHostPlatformId;
+  if (is_cpu_platform) {
+    // If we are on the host platform and the input buffer is sufficiently
+    // aligned, we can simply point to the input array's data without any
+    // further copies. At the time of writing we require a 16-byte alignment
+    // because XLA may generate code which requires it.
+    bool can_use_zero_copy =
+        host_buffer_semantics == HostBufferSemantics::kZeroCopy &&
+        ((absl::bit_cast(data) &
+          (cpu_function_runtime::kMinAlign - 1)) == 0);
+    if (shape.layout() == compact_shape.layout() &&
+        (host_buffer_semantics ==
+             HostBufferSemantics::kImmutableOnlyDuringCall ||
+         can_use_zero_copy)) {
+      std::function on_delete_callback;
+      se::DeviceMemoryBase buffer;
+      // If we are on the host platform and the input buffer is sufficiently
+      // aligned, we can simply point to the input array's data without any
+      // further copies. At the time of writing we require a 16-byte alignment
+      // because XLA may generate code which requires it.
+      if (can_use_zero_copy) {
+        on_delete_callback = [buffer_reference{std::move(buffer_reference)}]() {
+          // Frees buffer_reference.
+        };
+        buffer = se::DeviceMemoryBase(const_cast(data), size);
+      } else {
+        void* staging_buffer = client->host_memory_allocator()->AllocateRaw(
+            cpu_function_runtime::kMinAlign, size);
+        on_delete_callback = [staging_buffer, client]() {
+          client->host_memory_allocator()->DeallocateRaw(staging_buffer);
+        };
+        buffer = se::DeviceMemoryBase(staging_buffer, size);
+        std::memcpy(staging_buffer, data, size);
+      }
+      absl::Span>
+          definition_events;
+      auto device_buffer = std::make_shared(
+          /*allocator=*/nullptr, local_device->device_ordinal(),
+          std::initializer_list{buffer},
+          definition_events, std::move(on_delete_callback));
+      return absl::make_unique(
+          shape, shape, std::move(device_buffer), client, device);
+    }
+  }
+
   TF_ASSIGN_OR_RETURN(
       std::unique_ptr py_buffer,
       AllocateDestinationBuffer(compact_shape, device, local_device,
@@ -573,17 +625,41 @@ StatusOr> PjRtBuffer::FromHostBuffer(
   ScopedHold device_buffer(py_buffer->GetBufferWithUsageHold());
   CHECK(device_buffer.ok());
 
+  // If necessary, allocate a host-side buffer for staging host-to-device
+  // transfers. On GPU this is a buffer in pinned memory.
+  std::shared_ptr staging_buffer;
+  if (host_buffer_semantics == HostBufferSemantics::kImmutableOnlyDuringCall ||
+      client->should_stage_host_to_device_transfers()) {
+    void* ptr = client->host_memory_allocator()->AllocateRaw(
+        tensorflow::Allocator::kAllocatorAlignment, size);
+    staging_buffer = std::shared_ptr(ptr, [client](void* ptr) {
+      client->host_memory_allocator()->DeallocateRaw(ptr);
+    });
+  }
+
+  // Copy the buffer into a staging buffer before returning control to the
+  // caller if the caller only guaranteed that the buffer is valid for the
+  // duration of the call. Otherwise, we stage (if necessary) on a separate
+  // thread.
+  if (host_buffer_semantics == HostBufferSemantics::kImmutableOnlyDuringCall) {
+    std::memcpy(staging_buffer.get(), data, size);
+    buffer_reference.reset();
+    data = nullptr;
+  }
+
   // The host to device transfer is performed on a thread pool, mostly because
   // it includes linearization that may be slow. It is OK to capture the
   // py_buffer pointer because the py_buffer can't be deleted until all the
   // usage holds have gone away.
   // TODO(misard) assess if it would be preferable to introduce a heuristic to
   // put the transfer into the calling thread for small literals.
-  auto transfer_h2d = [client, transfer_manager, local_device,
-                       movable_device_buffer{device_buffer.ToClosure()}, data,
-                       shape, py_buffer{py_buffer.get()}, compact_shape,
+  auto transfer_h2d = [client, transfer_manager, local_device, data, size,
+                       movable_device_buffer{device_buffer.ToClosure()}, shape,
+                       py_buffer{py_buffer.get()}, compact_shape,
                        on_device_shape{py_buffer->on_device_shape()},
-                       buffer_reference{std::move(buffer_reference)}]() {
+                       staging_buffer{std::move(staging_buffer)},
+                       buffer_reference{std::move(buffer_reference)},
+                       host_buffer_semantics]() {
     ScopedHold device_buffer(movable_device_buffer);
     // This function uses TF_CHECK_OK and ValueOrDie() since we have no way
     // to report failures from a callback. However, the operations here are
@@ -593,20 +669,16 @@ StatusOr> PjRtBuffer::FromHostBuffer(
 
     ShapedBuffer buffer = device_buffer->AsShapedBuffer(
         compact_shape, on_device_shape, client->client()->platform());
-
-    std::shared_ptr staging_buffer;
-
     // If applicable on the backend, stage the transfer via host memory
     // allocated via the host_memory_allocator. On GPU, this is pinned
     // memory.
-    if (client->host_memory_allocator()) {
-      int64 size = ShapeUtil::ByteSizeOf(shape);
-      void* ptr = client->host_memory_allocator()->AllocateRaw(
-          tensorflow::Allocator::kAllocatorAlignment, size);
-      staging_buffer = std::shared_ptr(ptr, [client](void* ptr) {
-        client->host_memory_allocator()->DeallocateRaw(ptr);
-      });
-      std::memcpy(ptr, data, size);
+    if (staging_buffer) {
+      // If we didn't already copy the input buffer into the staging buffer,
+      // do so now.
+      if (host_buffer_semantics !=
+          HostBufferSemantics::kImmutableOnlyDuringCall) {
+        std::memcpy(staging_buffer.get(), data, size);
+      }
       BorrowingLiteral literal(static_cast(staging_buffer.get()),
                                shape);
       TF_CHECK_OK(transfer_manager->TransferLiteralToDeviceAsync(
@@ -626,9 +698,15 @@ StatusOr> PjRtBuffer::FromHostBuffer(
 
     local_device->ThenRelease(
         local_device->host_to_device_stream(),
-        std::make_pair(buffer_reference, std::move(staging_buffer)));
+        std::make_pair(std::move(buffer_reference), std::move(staging_buffer)));
   };
-  client->h2d_transfer_pool()->Schedule(transfer_h2d);
+  if (is_cpu_platform) {
+    // Using the h2d_transfer_pool would be a double thread hop; the code
+    // already defers its work onto a stream (= thread on CPU).
+    transfer_h2d();
+  } else {
+    client->h2d_transfer_pool()->Schedule(transfer_h2d);
+  }
   return py_buffer;
 }
 
diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.h b/tensorflow/compiler/xla/pjrt/pjrt_client.h
index c609abbf6fd..754eb19bec6 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.h
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.h
@@ -128,6 +128,7 @@ class PjRtClient {
       std::vector> devices, int host_id,
       std::unique_ptr allocator,
       std::unique_ptr host_memory_allocator,
+      bool should_stage_host_to_device_transfers,
       std::unique_ptr gpu_run_options);
   virtual ~PjRtClient() = default;
 
@@ -153,6 +154,9 @@ class PjRtClient {
   tensorflow::Allocator* host_memory_allocator() const {
     return host_memory_allocator_.get();
   }
+  bool should_stage_host_to_device_transfers() const {
+    return should_stage_host_to_device_transfers_;
+  }
 
   GpuExecutableRunOptions* gpu_run_options() const {
     return gpu_run_options_.get();
@@ -190,6 +194,9 @@ class PjRtClient {
   std::string platform_name_;
   LocalClient* client_;
 
+  // Allocator to be used for staging memory transfers to devices.
+  std::unique_ptr host_memory_allocator_;
+
   // Includes all devices, including non-local devices on multi-host platforms.
   std::vector> devices_;
   // Maps Device::id() to the corresponding Device. Includes all devices.
@@ -201,10 +208,10 @@ class PjRtClient {
   se::DeviceMemoryAllocator* allocator_;
   std::unique_ptr owned_allocator_;
 
-  // Allocator to be used for staging memory transfers to devices. Optional;
-  // only used on GPU where it is more efficient to copy buffers to and from the
-  // device via a staging area of pinned memory.
-  std::unique_ptr host_memory_allocator_;
+  // Should we always prefer to stage host-to-device transfers via memory
+  // allocated on host_memory_allocator_? True only on GPU, where we prefer to
+  // transfer via pinned memory.
+  bool should_stage_host_to_device_transfers_;
 
   std::unique_ptr gpu_run_options_;
 
@@ -396,13 +403,35 @@ class PjRtBuffer {
     StatusOr> buffer_or_;
   };
 
-  // If `force_copy` is true, forces a copy of the input buffer on CPU.
-  // Otherwise the library is free to alias the output buffer with `data`.
-  // `buffer_reference` is an optional shared pointer that should be kept alive
-  // by the runtime as long as the contents of `data` may still be accessed by
-  // the runtime (may be nullptr).
+  // Describes the semantics the caller to FromHostBuffer expects from the
+  // runtime, in a total order from most restrictive to least restrictive.
+  enum class HostBufferSemantics {
+    // The runtime may not hold references to `data` after the call to
+    // `FromHostBuffer` completes. The caller promises that `data` is immutable
+    // and will not be freed only for the duration of the FromHostBuffer call.
+    // `buffer_reference` will be freed by the time `FromHostBuffer` returns.
+    kImmutableOnlyDuringCall,
+
+    // The runtime may hold onto `data` after the call to `FromHostBuffer`
+    // returns while the runtime completes a transfer to the device. The caller
+    // promises not to mutate or free `data` until the transfer completes, at
+    // which point the runtime will release `buffer_reference`. It is also
+    // correct to wait on the host (directly or indirectly) for the buffer's
+    // definition event to complete.
+    kImmutableUntilTransferCompletes,
+
+    // The PjRtBuffer may alias `data` internally and the runtime may use the
+    // `data` contents as long as the buffer is alive.
+    // The caller promises to keep `data` alive and not to mutate its contents
+    // as long as the buffer is alive; to notify the caller that the buffer may
+    // be freed, the runtime will release its `buffer_reference` when the
+    // PjRtBuffer is freed. On non-CPU platforms this acts identically to
+    // kImmutableUntilTransferCompletes.
+    kZeroCopy,
+  };
   static StatusOr> FromHostBuffer(
-      const void* data, const Shape& shape, bool force_copy,
+      const void* data, const Shape& shape,
+      HostBufferSemantics host_buffer_semantics,
       std::shared_ptr buffer_reference, PjRtClient* client,
       Device* device);
 
diff --git a/tensorflow/compiler/xla/python/py_client.cc b/tensorflow/compiler/xla/python/py_client.cc
index a5779f3f8ee..bc7244cfc64 100644
--- a/tensorflow/compiler/xla/python/py_client.cc
+++ b/tensorflow/compiler/xla/python/py_client.cc
@@ -84,7 +84,8 @@ PyClient::GetDefaultDeviceAssignment1D(int num_replicas) {
 }
 
 StatusOr> PyClient::BufferFromPyal(
-    const pybind11::object& argument, Device* device, bool force_copy) {
+    const pybind11::object& argument, Device* device, bool force_copy,
+    PjRtBuffer::HostBufferSemantics host_buffer_semantics) {
   if (device == nullptr) {
     TF_RET_CHECK(!pjrt_client_->local_devices().empty());
     device = pjrt_client_->local_devices().front();
@@ -111,9 +112,9 @@ StatusOr> PyClient::BufferFromPyal(
   {
     py::gil_scoped_release gil_release;
     TF_ASSIGN_OR_RETURN(
-        buffer, PjRtBuffer::FromHostBuffer(c->buf_ptr, c->shape, force_copy,
-                                           std::move(py_buffer_ref),
-                                           pjrt_client_.get(), device));
+        buffer, PjRtBuffer::FromHostBuffer(
+                    c->buf_ptr, c->shape, host_buffer_semantics,
+                    std::move(py_buffer_ref), pjrt_client_.get(), device));
   }
   auto traceback = Traceback::Get();
   return std::make_unique(shared_from_this(), std::move(buffer),
diff --git a/tensorflow/compiler/xla/python/py_client.h b/tensorflow/compiler/xla/python/py_client.h
index 76419742c57..c94a206a926 100644
--- a/tensorflow/compiler/xla/python/py_client.h
+++ b/tensorflow/compiler/xla/python/py_client.h
@@ -120,7 +120,8 @@ class PyClient : public std::enable_shared_from_this {
   }
 
   StatusOr> BufferFromPyal(
-      const pybind11::object& argument, Device* device, bool force_copy);
+      const pybind11::object& argument, Device* device, bool force_copy,
+      PjRtBuffer::HostBufferSemantics host_buffer_semantics);
 
   StatusOr> Compile(
       const XlaComputation& computation, CompileOptions options);
diff --git a/tensorflow/compiler/xla/python/xla.cc b/tensorflow/compiler/xla/python/xla.cc
index 2c323783961..c0a440aa4bd 100644
--- a/tensorflow/compiler/xla/python/xla.cc
+++ b/tensorflow/compiler/xla/python/xla.cc
@@ -509,6 +509,13 @@ PYBIND11_MODULE(xla_extension, m) {
       .value("PLATFORM", GpuAllocatorConfig::Kind::kPlatform)
       .value("BFC", GpuAllocatorConfig::Kind::kBFC);
 
+  py::enum_(m, "HostBufferSemantics")
+      .value("IMMUTABLE_ONLY_DURING_CALL",
+             PjRtBuffer::HostBufferSemantics::kImmutableOnlyDuringCall)
+      .value("IMMUTABLE_UNTIL_TRANSFER_COMPLETES",
+             PjRtBuffer::HostBufferSemantics::kImmutableUntilTransferCompletes)
+      .value("ZERO_COPY", PjRtBuffer::HostBufferSemantics::kZeroCopy);
+
   py::class_> py_local_client(m, "Client");
   py_local_client.def_property_readonly("platform", &PyClient::platform_name)
       .def("device_count", &PyClient::device_count)
@@ -527,7 +534,9 @@ PYBIND11_MODULE(xla_extension, m) {
       .def("create_host_to_device_channel_handle",
            &PyClient::CreateHostToDeviceChannelHandle)
       .def("buffer_from_pyval", &PyClient::BufferFromPyal, py::arg("argument"),
-           py::arg("device") = nullptr, py::arg("force_copy") = false)
+           py::arg("device") = nullptr, py::arg("force_copy") = false,
+           py::arg("host_buffer_semantics") =
+               PjRtBuffer::HostBufferSemantics::kZeroCopy)
       .def("compile", &PyClient::Compile, py::arg("computation"),
            py::arg("compile_options") = CompileOptions())
       .def("heap_profile", &PyClient::HeapProfile);
diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py
index 8f176507215..38c55c6fe5d 100644
--- a/tensorflow/compiler/xla/python/xla_client.py
+++ b/tensorflow/compiler/xla/python/xla_client.py
@@ -304,6 +304,7 @@ def computation_count():
 Device = _xla.Device
 CompileOptions = _xla.CompileOptions
 
+HostBufferSemantics = _xla.HostBufferSemantics
 
 # An Executable is a C++ class that duck types with the following API:
 # class Executable(object):
diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py
index 6a316044734..49c57a27ac0 100644
--- a/tensorflow/compiler/xla/python/xla_client_test.py
+++ b/tensorflow/compiler/xla/python/xla_client_test.py
@@ -1986,7 +1986,8 @@ def TestFactory(xla_backend, cloud_tpu=False):
     def testRoundTrip(self, dtype, shape):
       x = np.array(np.random.rand(*shape) * 100, dtype=dtype)
       x_ptr = x.__array_interface__["data"][0]
-      buffer = self.backend.buffer_from_pyval(x)
+      buffer = self.backend.buffer_from_pyval(
+          x, host_buffer_semantics=xla_client.HostBufferSemantics.ZERO_COPY)
       y = np.array(buffer, copy=False)
       y_ptr = y.__array_interface__["data"][0]
       np.testing.assert_array_equal(x, y)
@@ -1995,7 +1996,9 @@ def TestFactory(xla_backend, cloud_tpu=False):
       self.assertTrue((x_ptr & 15) != 0 or x_ptr == y_ptr)
       self.assertEqual(y_ptr, buffer.unsafe_buffer_pointer())
 
-      buffer2 = self.backend.buffer_from_pyval(x, force_copy=True)
+      during_call = xla_client.HostBufferSemantics.IMMUTABLE_ONLY_DURING_CALL
+      buffer2 = self.backend.buffer_from_pyval(
+          x, host_buffer_semantics=during_call)
       z = np.array(buffer2, copy=False)
       self.assertNotEqual(x.__array_interface__["data"][0],
                           z.__array_interface__["data"][0])

From 2537f3d413a30ffd55187e795b256f29a1ce9d97 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 06:56:09 -0700
Subject: [PATCH 0285/1390] Enabled support for negative edge padding for
 XlaPad.

The underlying HLO pad already supports negative edge padding. We just removed
some checks preventing negative edge padding, and added a test.

PiperOrigin-RevId: 316672291
Change-Id: I38e55f86727d90f66aadb549a587b43df6588571
---
 tensorflow/compiler/tests/xla_ops_test.py      | 18 ++++++++++++++++++
 .../compiler/tf2xla/kernels/xla_pad_op.cc      |  8 --------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/tests/xla_ops_test.py b/tensorflow/compiler/tests/xla_ops_test.py
index f3e915daa67..35d36315464 100644
--- a/tensorflow/compiler/tests/xla_ops_test.py
+++ b/tensorflow/compiler/tests/xla_ops_test.py
@@ -211,6 +211,24 @@ class XlaOpsNumericalTest(xla_test.XLATestCase, parameterized.TestCase):
                [7, 7, 7, 7, 7], [7, 2, 3, 7, 7], [7, 7, 7, 7, 7]],
               dtype=dtype))
 
+  def testPadNegative(self):
+    for dtype in self.numeric_types:
+
+      def pad_fn(x):
+        return xla.pad(
+            x,
+            padding_value=7,
+            padding_low=[0, -1],
+            padding_high=[1, -2],
+            padding_interior=[1, 2])
+
+      self._assertOpOutputMatchesExpected(
+          pad_fn,
+          args=(np.arange(6, dtype=np.int32).astype(dtype).reshape([2, 3]),),
+          expected=np.array(
+              [[7, 7, 1, 7], [7, 7, 7, 7], [7, 7, 4, 7], [7, 7, 7, 7]],
+              dtype=dtype))
+
   @test_util.disable_mlir_bridge('Not supported yet')
   def testReduce(self):
     for dtype in set(self.numeric_types).intersection(
diff --git a/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc
index a3c2eef993c..d35101a771a 100644
--- a/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc
@@ -64,14 +64,6 @@ class XlaPadOp : public XlaOpKernel {
                     padding_interior.size(), " vs. ", rank, ")"));
 
     auto non_negative = [](int64 x) { return x >= 0; };
-    OP_REQUIRES(
-        context, absl::c_all_of(padding_low, non_negative),
-        errors::InvalidArgument("padding_low must be non-negative, got [",
-                                absl::StrJoin(padding_low, ","), "]"));
-    OP_REQUIRES(
-        context, absl::c_all_of(padding_high, non_negative),
-        errors::InvalidArgument("padding_high must be non-negative, got [",
-                                absl::StrJoin(padding_high, ","), "]"));
     OP_REQUIRES(
         context, absl::c_all_of(padding_interior, non_negative),
         errors::InvalidArgument("padding_interior must be non-negative, got [",

From 770251a7008c1d89e8141b37191e74008e01253d Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 16 Jun 2020 07:34:49 -0700
Subject: [PATCH 0286/1390] Support escaping result memrefs in
 lhlo_fuse_linalg.

So far, we have identified the root computation to fuse into by it writing into a function argument. Now writing into a buffer that is returned also qualifies.

PiperOrigin-RevId: 316677942
Change-Id: I7c3912419606555946c9111d12c4086d086d9456
---
 tensorflow/compiler/mlir/xla/BUILD            |  1 +
 .../mlir/xla/tests/lhlo-fuse-linalg.mlir      | 81 ++++++++++++++++---
 .../mlir/xla/transforms/lhlo_fuse_linalg.cc   | 16 +++-
 3 files changed, 82 insertions(+), 16 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/BUILD b/tensorflow/compiler/mlir/xla/BUILD
index 8f0f000b26a..43458aab2d3 100644
--- a/tensorflow/compiler/mlir/xla/BUILD
+++ b/tensorflow/compiler/mlir/xla/BUILD
@@ -377,6 +377,7 @@ cc_library(
         "@llvm-project//mlir:LinalgOps",
         "@llvm-project//mlir:LinalgTransforms",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:StandardOps",
         "@llvm-project//mlir:TransformUtils",
     ],
     alwayslink = 1,
diff --git a/tensorflow/compiler/mlir/xla/tests/lhlo-fuse-linalg.mlir b/tensorflow/compiler/mlir/xla/tests/lhlo-fuse-linalg.mlir
index 063487c00d8..b04c97f42d7 100644
--- a/tensorflow/compiler/mlir/xla/tests/lhlo-fuse-linalg.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/lhlo-fuse-linalg.mlir
@@ -1,13 +1,12 @@
-// RUN: xla-opt -lhlo-fuse-linalg %s -o - | FileCheck %s --dump-input=always
-// RUN: xla-opt -lhlo-fuse-linalg=tile-sizes=2,3 %s -o - | FileCheck %s -check-prefix=TILED
-// RUN: xla-opt -lhlo-fuse-linalg=use-parallel-loops %s -o - | FileCheck %s -check-prefix=PLOOP
-
+// RUN: xla-opt -lhlo-fuse-linalg %s -split-input-file | FileCheck %s --dump-input=always
+// RUN: xla-opt -lhlo-fuse-linalg=tile-sizes=2,3 %s -split-input-file | FileCheck %s -check-prefix=TILED
+// RUN: xla-opt -lhlo-fuse-linalg=use-parallel-loops %s -split-input-file | FileCheck %s -check-prefix=PLOOP
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #pointwise_2d_trait = {args_in = 2, args_out = 1, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
 func @fusion(%multiplier: memref<6x6xf32>, %summand_1: memref<6x6xf32>,
              %summand_2: memref<6x6xf32>, %result: memref<6x6xf32>) {
-  %temp_result = alloc() {temp = true} : memref<6x6xf32>
+  %temp_result = alloc() : memref<6x6xf32>
   linalg.generic #pointwise_2d_trait %summand_1, %summand_2, %temp_result {
   ^bb0(%summand_1_in: f32, %summand_2_in: f32, %temp_result_in: f32):
     %out = addf %summand_1_in, %summand_2_in : f32
@@ -19,7 +18,7 @@ func @fusion(%multiplier: memref<6x6xf32>, %summand_1: memref<6x6xf32>,
     linalg.yield %out : f32
   } : memref<6x6xf32>, memref<6x6xf32>, memref<6x6xf32>
   dealloc %temp_result : memref<6x6xf32>
-  "xla_lhlo.terminator"() : () -> ()
+  return
 }
 // CHECK-LABEL: func @fusion
 //       CHECK:  %[[C1:.*]] = constant 1
@@ -53,10 +52,12 @@ func @fusion(%multiplier: memref<6x6xf32>, %summand_1: memref<6x6xf32>,
 //       PLOOP:      linalg.generic
 //       PLOOP:        mulf
 
+// -----
+
 func @fusion_of_three(%arg0: memref<100x10xf32>,
                       %arg1: memref<100xf32>,
                       %arg2: memref<100x10xf32>) {
- %0 = alloc() {temp = true} : memref<100x10xf32>
+ %0 = alloc() : memref<100x10xf32>
  linalg.generic {
    args_in = 1 : i64,
    args_out = 1 : i64,
@@ -66,7 +67,7 @@ func @fusion_of_three(%arg0: memref<100x10xf32>,
      ^bb0(%arg3: f32, %arg4: f32): // no predecessors
        linalg.yield %arg3 : f32
      }: memref<100xf32>, memref<100x10xf32>
- %1 = alloc() {temp = true} : memref<100x10xf32>
+ %1 = alloc() : memref<100x10xf32>
  linalg.generic {
    args_in = 2 : i64,
    args_out = 1 : i64,
@@ -126,11 +127,13 @@ func @fusion_of_three(%arg0: memref<100x10xf32>,
 //       PLOOP:       linalg.generic
 //       PLOOP:         exp
 
-#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-#pointwise_4d_trait = {args_in = 2, args_out = 1, indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+#pointwise_4d_trait = {args_in = 2, args_out = 1, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
 func @fusion_4d(%multiplier: memref<6x6x6x6xf32>, %summand_1: memref<6x6x6x6xf32>,
              %summand_2: memref<6x6x6x6xf32>, %result: memref<6x6x6x6xf32>) {
-  %temp_result = alloc() {temp = true} : memref<6x6x6x6xf32>
+  %temp_result = alloc() : memref<6x6x6x6xf32>
   linalg.generic #pointwise_4d_trait %summand_1, %summand_2, %temp_result {
   ^bb0(%summand_1_in: f32, %summand_2_in: f32, %temp_result_in: f32):
     %out = addf %summand_1_in, %summand_2_in : f32
@@ -142,7 +145,7 @@ func @fusion_4d(%multiplier: memref<6x6x6x6xf32>, %summand_1: memref<6x6x6x6xf32
     linalg.yield %out : f32
   } : memref<6x6x6x6xf32>, memref<6x6x6x6xf32>, memref<6x6x6x6xf32>
   dealloc %temp_result : memref<6x6x6x6xf32>
-  "xla_lhlo.terminator"() : () -> ()
+  return
 }
 // CHECK-LABEL: func @fusion_4d
 //       CHECK:  %[[C1:.*]] = constant 1
@@ -177,3 +180,57 @@ func @fusion_4d(%multiplier: memref<6x6x6x6xf32>, %summand_1: memref<6x6x6x6xf32
 //       PLOOP:        addf
 //       PLOOP:      linalg.generic
 //       PLOOP:        mulf
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#pointwise_2d_trait = {args_in = 2, args_out = 1, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
+func @fusion(%multiplier: memref<6x6xf32>, %summand_1: memref<6x6xf32>,
+             %summand_2: memref<6x6xf32>) -> memref<6x6xf32> {
+  %temp_result = alloc() : memref<6x6xf32>
+  linalg.generic #pointwise_2d_trait %summand_1, %summand_2, %temp_result {
+  ^bb0(%summand_1_in: f32, %summand_2_in: f32, %temp_result_in: f32):
+    %out = addf %summand_1_in, %summand_2_in : f32
+    linalg.yield %out : f32
+  } : memref<6x6xf32>, memref<6x6xf32>, memref<6x6xf32>
+  %result = alloc() : memref<6x6xf32>
+  linalg.generic #pointwise_2d_trait %temp_result, %multiplier, %result {
+  ^bb0(%temp_result_in: f32, %multiplier_in: f32, %result_in: f32):
+    %out = mulf %temp_result_in, %multiplier_in : f32
+    linalg.yield %out : f32
+  } : memref<6x6xf32>, memref<6x6xf32>, memref<6x6xf32>
+  dealloc %temp_result : memref<6x6xf32>
+  return %result : memref<6x6xf32>
+}
+
+// CHECK-LABEL: func @fusion
+//       CHECK:  %[[C1:.*]] = constant 1
+//   CHECK-NOT:  linalg.generic
+//       CHECK:  scf.for {{.*}} step %[[C1]]
+//       CHECK:    scf.for {{.*}} step %[[C1]]
+//   CHECK-NOT:  scf.for
+//       CHECK:      linalg.generic
+//       CHECK:        addf
+//       CHECK:      linalg.generic
+//       CHECK:        mulf
+
+// TILED-LABEL: func @fusion
+//   TILED-DAG:  %[[C2:.*]] = constant 2
+//   TILED-DAG:  %[[C3:.*]] = constant 3
+//   TILED-NOT:  linalg.generic
+//       TILED:  scf.for {{.*}} step %[[C2]]
+//       TILED:    scf.for {{.*}} step %[[C3]]
+//   TILED-NOT:  scf.for
+//       TILED:      linalg.generic
+//       TILED:        addf
+//       TILED:      linalg.generic
+//       TILED:        mulf
+
+// PLOOP-LABEL: func @fusion
+//   PLOOP-NOT:  linalg.generic
+//       PLOOP:  scf.parallel
+//   PLOOP-NOT:  scf.parallel
+//       PLOOP:      linalg.generic
+//       PLOOP:        addf
+//       PLOOP:      linalg.generic
+//       PLOOP:        mulf
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc
index ddbb672c70a..e16ab571b4d 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "absl/memory/memory.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"  // from @llvm-project
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
 #include "mlir/Transforms/FoldUtils.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/xla/transforms/passes.h"
@@ -52,10 +53,17 @@ class LhloFuseLinalg : public PassWrapper {
     // The fusion in Linalg is currently possible only when the consumer op is
     // tiled. In order to greedily fuse the ops, we have to start from the tiled
     // root linalg ops, i.e. linalg ops that write to output buffers of the
-    // function.
-    llvm::SmallDenseSet func_args;
+    // function or are returned in case of escaping allocations.
+    llvm::SmallDenseSet result_buffers;
     for (auto func_arg : func.getArguments()) {
-      func_args.insert(func_arg);
+      result_buffers.insert(func_arg);
+    }
+    for (auto& block : func.getBlocks()) {
+      auto returnOp = mlir::dyn_cast(block.getTerminator());
+      if (!returnOp) continue;
+      for (auto operand : returnOp.getOperands()) {
+        result_buffers.insert(operand);
+      }
     }
     MLIRContext* ctx = func.getContext();
     OpBuilder b(func);
@@ -68,7 +76,7 @@ class LhloFuseLinalg : public PassWrapper {
       }
       auto op = cast(generic_op.getOperation());
       for (const Value result : op.getOutputBuffers()) {
-        if (!func_args.count(result)) continue;
+        if (!result_buffers.count(result)) continue;
         if (tileGenericOp(op, tile_sizes, &b)) {
           generic_op.erase();
           return;

From 426f62af5eb80e5f0c3b660451dac6f953b4ca0c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 07:36:58 -0700
Subject: [PATCH 0287/1390] Add an explicit distribution_strategy attribute to
 TrackableHandler. Since it's treated as a weight by Keras, it needs to pass
 DistStrat checks during compile.

PiperOrigin-RevId: 316678228
Change-Id: I132168f1ca3dd3729d7a499cef3564c5e04abb34
---
 tensorflow/python/keras/engine/base_layer_utils.py             | 1 +
 .../layers/preprocessing/index_lookup_distribution_test.py     | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/engine/base_layer_utils.py b/tensorflow/python/keras/engine/base_layer_utils.py
index cb2b9ed87f1..de67080af66 100644
--- a/tensorflow/python/keras/engine/base_layer_utils.py
+++ b/tensorflow/python/keras/engine/base_layer_utils.py
@@ -774,6 +774,7 @@ class TrackableWeightHandler(object):
     if not isinstance(trackable, tracking.Trackable):
       raise ValueError('%s is not a Trackable object.' % (trackable,))
     self._trackable = trackable
+    self._distribute_strategy = distribution_strategy_context.get_strategy()
 
     # TODO(b/141682913): Figure out why this is private and fix it.
     saveables = trackable._gather_saveables_for_checkpoint().values()  # pylint: disable=protected-access
diff --git a/tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py b/tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py
index 098e67f5f6b..c593cd41c85 100644
--- a/tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/index_lookup_distribution_test.py
@@ -44,7 +44,7 @@ def get_layer_class():
 @combinations.generate(
     combinations.combine(
         distribution=strategy_combinations.all_strategies,
-        mode=["eager", "graph"]))
+        mode=["eager"]))  # Eager-only, no graph: b/158793009
 class IndexLookupDistributionTest(
     keras_parameterized.TestCase,
     preprocessing_test_utils.PreprocessingLayerTest):
@@ -74,6 +74,7 @@ class IndexLookupDistributionTest(
       layer.adapt(vocab_dataset)
       int_data = layer(input_data)
       model = keras.Model(inputs=input_data, outputs=int_data)
+    model.compile(loss="mse")
     output_dataset = model.predict(input_dataset)
     self.assertAllEqual(expected_output, output_dataset)
 

From cb60e1c14b1470b41f60a905ced1176763dfc9ee Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 07:41:44 -0700
Subject: [PATCH 0288/1390] Introduces a new experimental package that: -
 Defines a schema for configuring delegates - Defines a C++ plugin mechanism
 using the schema, so that code can support   configuring arbitrary delegates
 without a build-time dependency

PiperOrigin-RevId: 316678829
Change-Id: I36ce8a6175b550d83dfe9cf1f237a04173fb8b16
---
 .../acceleration/configuration/BUILD          | 165 ++++++++++++++
 .../configuration/configuration.proto         | 208 ++++++++++++++++++
 .../configuration/delegate_registry.cc        |  60 +++++
 .../configuration/delegate_registry.h         |  95 ++++++++
 .../acceleration/configuration/gpu_plugin.cc  |  62 ++++++
 .../configuration/hexagon_plugin.cc           |  73 ++++++
 .../configuration/nnapi_plugin.cc             |  93 ++++++++
 .../configuration/nnapi_plugin_test.cc        | 175 +++++++++++++++
 .../configuration/proto_to_flatbuffer.cc      |  58 +++++
 .../configuration/proto_to_flatbuffer.h       |  32 +++
 10 files changed, 1021 insertions(+)
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/BUILD
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/configuration.proto
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
 create mode 100644 tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h

diff --git a/tensorflow/lite/experimental/acceleration/configuration/BUILD b/tensorflow/lite/experimental/acceleration/configuration/BUILD
new file mode 100644
index 00000000000..1bfd2494fe1
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/BUILD
@@ -0,0 +1,165 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_java_library", "flatc_path")
+
+package(
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+genrule(
+    name = "configuration_schema",
+    srcs = ["configuration.proto"],
+    outs = ["configuration.fbs"],
+    # We rename the namespace since otherwise the proto classes and flatbuffer
+    # classes would have the same names.
+    cmd = """
+    $(location {}) --proto -o $(@D) $(location :configuration.proto)
+    perl -p -i -e 's/tflite.proto/tflite/' $(@D)/configuration.fbs
+    """.format(flatc_path),
+    tools = [
+        flatc_path,
+    ],
+)
+
+genrule(
+    name = "configuration_fbs_contents_cc",
+    srcs = ["configuration.fbs"],
+    outs = ["configuration_fbs_contents-inl.h"],
+    cmd = """
+      echo 'constexpr char configuration_fbs_contents[] = R"Delimiter(' > $(@)
+      cat < $(<) >> $(@)
+      echo ')Delimiter";' >> $(@)
+    """,
+)
+
+proto_library(
+    name = "configuration_proto",
+    srcs = [
+        "configuration.proto",
+    ],
+)
+
+cc_proto_library(
+    name = "configuration_cc_proto",
+    deps = [":configuration_proto"],
+)
+
+java_lite_proto_library(
+    name = "configuration_java_proto_lite",
+    deps = [":configuration_proto"],
+)
+
+flatbuffer_cc_library(
+    name = "configuration_fbs",
+    srcs = [":configuration.fbs"],
+)
+
+flatbuffer_java_library(
+    name = "configuration_fbs_java",
+    srcs = [":configuration.fbs"],
+)
+
+cc_library(
+    name = "proto_to_flatbuffer",
+    srcs = [
+        "configuration_fbs_contents-inl.h",
+        "proto_to_flatbuffer.cc",
+    ],
+    hdrs = ["proto_to_flatbuffer.h"],
+    deps = [
+        ":configuration_cc_proto",
+        ":configuration_fbs",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/lite:minimal_logging",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "delegate_registry",
+    srcs = ["delegate_registry.cc"],
+    hdrs = ["delegate_registry.h"],
+    deps = [
+        ":configuration_fbs",
+        "//tensorflow/lite/c:common",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+cc_library(
+    name = "nnapi_plugin",
+    srcs = ["nnapi_plugin.cc"],
+    deps = [
+        ":configuration_fbs",
+        ":delegate_registry",
+        "//tensorflow/lite/delegates/nnapi:nnapi_delegate",
+        "@com_google_absl//absl/memory",
+    ],
+    alwayslink = 1,  # For registration to always run.
+)
+
+cc_test(
+    name = "nnapi_plugin_test",
+    srcs = ["nnapi_plugin_test.cc"],
+    tags = [
+        "no_mac",
+        "no_windows",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":configuration_fbs",
+        ":delegate_registry",
+        ":nnapi_plugin",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/delegates/nnapi:nnapi_delegate",
+        "//tensorflow/lite/delegates/nnapi:nnapi_delegate_mock_test",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "hexagon_plugin",
+    srcs = ["hexagon_plugin.cc"],
+    deps = [
+        ":configuration_fbs",
+        ":delegate_registry",
+        "@com_google_absl//absl/memory",
+    ] + select({
+        "//tensorflow:android": [
+            "//tensorflow/lite/delegates/hexagon:hexagon_delegate",
+        ],
+        "//conditions:default": [],
+    }),
+    alwayslink = 1,  # For registration to always run.
+)
+
+cc_library(
+    name = "gpu_plugin",
+    srcs = ["gpu_plugin.cc"],
+    deps = [
+        ":configuration_fbs",
+        ":delegate_registry",
+        "//tensorflow/lite/delegates/gpu:delegate",
+        "@com_google_absl//absl/memory",
+    ],
+    alwayslink = 1,  # For registration to always run.
+)
diff --git a/tensorflow/lite/experimental/acceleration/configuration/configuration.proto b/tensorflow/lite/experimental/acceleration/configuration/configuration.proto
new file mode 100644
index 00000000000..e1c49f02856
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/configuration.proto
@@ -0,0 +1,208 @@
+// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+//
+// 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
+//
+//     http://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.
+
+// This schema defines how to configure TFLite for delegation. These
+// definitions can be used in multiple ways: as output of a compatibility list,
+// in benchmarking tools and to decouple delegate instantiation from code.
+//
+// The schema is work-in-progress, covering the most broadly used delegates and
+// options.
+
+syntax = "proto2";
+
+package tflite.proto;
+
+// ExecutionPreference is used to match accelerators against the preferences of
+// the current application or usecase. Some of the values here can appear both
+// in the compatibility list and as input, some only as input.
+//
+// These are separate from NNAPIExecutionPreference - the compatibility list
+// design doesn't assume a one-to-one mapping between which usecases
+// compatibility list entries have been developed for and what settings are used
+// for NNAPI.
+enum ExecutionPreference {
+  // Match any selected preference. Whitelist (semantically - value is same as
+  // on input).
+  ANY = 0;
+  // Match low latency preference. Both compatibility list and input.
+  LOW_LATENCY = 1;
+  // Math low power preference. Both compatibility list and input.
+  LOW_POWER = 2;
+  // Never accelerate. Can be used for input to compatibility list or for
+  // standalone Acceleration configuration.
+  FORCE_CPU = 3;
+}
+
+// TFLite delegate to use.
+enum Delegate {
+  NONE = 0;
+  NNAPI = 1;
+  GPU = 2;
+  HEXAGON = 3;
+  XNNPACK = 4;
+  // TODO(b/157893534): Support exposing edgetpu tflite delegate creation
+  // options.
+  EDGETPU = 5;
+}
+
+enum NNAPIExecutionPreference {
+  // Undefined.
+  UNDEFINED = 0;
+  // Prefer executing in a way that minimizes battery drain.
+  NNAPI_LOW_POWER = 1;
+  // Prefer returning a single answer as fast as possible, even if this causes
+  // more power consumption.
+  NNAPI_FAST_SINGLE_ANSWER = 2;
+  // Prefer maximizing the throughput of successive frames, for example when
+  // processing successive frames coming from the camera.
+  NNAPI_SUSTAINED_SPEED = 3;
+}
+
+// One possible acceleration configuration.
+message ComputeSettings {
+  // Which preference to use this accelerator for.
+  optional ExecutionPreference preference = 1;
+  // How to configure TFLite
+  optional TFLiteSettings tflite_settings = 2;
+  // Identifiers to use for instrumentation and telemetry.
+  optional string model_namespace_for_statistics = 3;
+  optional string model_identifier_for_statistics = 4;
+}
+
+// NNAPI delegate settings.
+message NNAPISettings {
+  // Which instance (NNAPI accelerator) to use. One driver may provide several
+  // accelerators (though a driver may also hide several back-ends behind one
+  // name, at the choice of the driver vendor).
+  // Note that driver introspection is only available in Android Q and later.
+  optional string accelerator_name = 1;
+
+  // NNAPI model compilation caching settings to be passed to
+  // tflite::StatefulNnApiDelegate
+  optional string cache_directory = 2;
+  optional string model_token = 3;
+
+  // NNAPI execution preference to pass. See
+  // https://developer.android.com/ndk/reference/group/neural-networks.html
+  optional NNAPIExecutionPreference execution_preference = 4;
+
+  // Number of instances to cache for the same model (for input size
+  // changes). This is mandatory for getting reasonable performance in that
+  // case.
+  optional int32 no_of_nnapi_instances_to_cache = 5;
+
+  // Whether to automatically fall back to TFLite CPU path.
+  optional FallbackSettings fallback_settings = 6;
+
+  // Whether to allow use of NNAPI CPU (nnapi-reference accelerator) on Android
+  // 10+ when an accelerator name is not specified. The NNAPI CPU typically
+  // performs less well than the TfLite built-in kernels; but allowing allows a
+  // model to be partially accelerated which may be a win.
+  optional bool allow_nnapi_cpu_on_android_10_plus = 7;
+}
+
+// Which GPU backend to select. Default behaviour on Android is to try OpenCL
+// and if it's not available fall back to OpenGL.
+enum GPUBackend {
+  UNSET = 0;
+  OPENCL = 1;
+  OPENGL = 2;
+  // Not yet supported.
+  // VULKAN = 3;
+  // METAL = 4;
+}
+
+// GPU Delegate settings.
+//
+// See
+// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/gpu/delegate.h
+message GPUSettings {
+  optional bool is_precision_loss_allowed = 1;
+  optional bool enable_quantized_inference = 2 [default = true];
+  optional GPUBackend force_backend = 3;
+  // TODO(b/152019007): add remaining options.
+}
+
+// Hexagon Delegate settings.
+//
+// See
+// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/hexagon/hexagon_delegate.h
+message HexagonSettings {
+  optional int32 debug_level = 1;
+  optional int32 powersave_level = 2;
+  optional bool print_graph_profile = 3;
+  optional bool print_graph_debug = 4;
+}
+
+// XNNPack Delegate settings.
+//
+// See
+// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h
+message XNNPackSettings {
+  optional int32 num_threads = 1;
+}
+
+message CPUSettings {
+  optional int32 num_threads = 1;
+}
+
+// How to configure TFLite.
+message TFLiteSettings {
+  // Which delegate to use.
+  optional Delegate delegate = 1;
+
+  // How to configure the chosen delegate.
+  // (In principle we would like to use 'oneof', but flatc turns that into an
+  // nested anonymous table rather than a union. See
+  // https://github.com/google/flatbuffers/issues/4628).
+  optional NNAPISettings nnapi_settings = 2;
+  optional GPUSettings gpu_settings = 3;
+  optional HexagonSettings hexagon_settings = 4;
+  optional XNNPackSettings xnnpack_settings = 5;
+
+  // How to configure CPU execution.
+  optional CPUSettings cpu_settings = 6;
+
+  // Shared delegation settings.
+  optional int32 max_delegated_partitions = 7;
+}
+
+// Whether to automatically fallback to TFLite CPU path on delegation errors.
+//
+// Typically fallback is enabled in production use but disabled in tests and
+// benchmarks to ensure they test the intended path.
+message FallbackSettings {
+  // Whether to allow automatically falling back to TfLite CPU path on
+  // compilation failure. Default is not allowing automatic fallback.
+  //
+  // This is useful in naive production usecases where the caller would prefer
+  // for the model to run even if it's not accelerated. More advanced users will
+  // implement fallback themselves; e.g., by using a different model on CPU.
+  //
+  // Note that compilation errors may occur either at initial
+  // ModifyGraphWithDelegate() time, or when calling AllocateTensors() after
+  // resizing.
+  optional bool allow_automatic_fallback_on_compilation_error = 7;
+  // Whether to allow automatically falling back to TfLite CPU path on
+  // execution error. Default is not allowing automatic fallback.
+  //
+  // Experimental, use with care (only when you have complete control over the
+  // client code).
+  //
+  // The caveat above for compilation error holds.  Additionally, execution-time
+  // errors are harder to handle automatically as they require invalidating the
+  // TfLite interpreter which most client code has not been designed to deal
+  // with.
+  optional bool allow_automatic_fallback_on_execution_error = 8;
+}
diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
new file mode 100644
index 00000000000..b8d80342d5f
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.cc
@@ -0,0 +1,60 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
+
+#include "absl/synchronization/mutex.h"
+
+namespace tflite {
+namespace delegates {
+
+void DelegatePluginRegistry::RegisterImpl(
+    const std::string& name,
+    std::function<
+        std::unique_ptr(const TFLiteSettings&)>
+        creator_function) {
+  absl::MutexLock lock(&mutex_);
+  factories_[name] = creator_function;
+}
+
+std::unique_ptr DelegatePluginRegistry::CreateImpl(
+    const std::string& name, const TFLiteSettings& settings) {
+  absl::MutexLock lock(&mutex_);
+  auto it = factories_.find(name);
+  if (it != factories_.end()) {
+    return it->second(settings);
+  } else {
+    return nullptr;
+  }
+}
+
+DelegatePluginRegistry* DelegatePluginRegistry::GetSingleton() {
+  static auto* instance = new DelegatePluginRegistry();
+  return instance;
+}
+
+std::unique_ptr DelegatePluginRegistry::CreateByName(
+    const std::string& name, const TFLiteSettings& settings) {
+  auto* const instance = DelegatePluginRegistry::GetSingleton();
+  return instance->CreateImpl(name, settings);
+}
+
+DelegatePluginRegistry::Register::Register(const std::string& name,
+                                           CreatorFunction creator_function) {
+  auto* const instance = DelegatePluginRegistry::GetSingleton();
+  instance->RegisterImpl(name, creator_function);
+}
+
+}  // namespace delegates
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
new file mode 100644
index 00000000000..c86759dcc3f
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h
@@ -0,0 +1,95 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
+
+#include 
+#include 
+
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+
+// Defines an interface for TFLite delegate plugins.
+//
+// The acceleration library aims to support all TFLite delegates based on
+// configuration expressed as data (flatbuffers). However, consumers tend to
+// care about size and also use a subset of delegates. Hence we don't want to
+// statically build against all delegates.
+//
+// This interface allows plugins to handle specific delegates.
+//
+// Goal of this interface is not to abstract away all the differences between
+// delegates. The goal is only to avoid static linking.
+//
+// Note to implementers: this interface may change if new delegates don't fit
+// into the same design.
+namespace tflite {
+namespace delegates {
+
+// Same w/ Interpreter::TfLiteDelegatePtr to avoid pulling
+// tensorflow/lite/interpreter.h dependency
+using TfLiteDelegatePtr =
+    std::unique_ptr;
+
+class DelegatePluginInterface {
+ public:
+  virtual TfLiteDelegatePtr Create() = 0;
+  virtual int GetDelegateErrno(TfLiteDelegate* from_delegate) = 0;
+  virtual ~DelegatePluginInterface() = default;
+};
+
+// A stripped-down registry that allows delegate plugins to be created by name.
+//
+// Limitations:
+// - Doesn't allow deregistration.
+// - Doesn't check for duplication registration.
+//
+class DelegatePluginRegistry {
+ public:
+  typedef std::function(
+      const TFLiteSettings&)>
+      CreatorFunction;
+  // Returns a DelegatePluginInterface registered with `name` or nullptr if no
+  // matching plugin found.
+  // TFLiteSettings is per-plugin, so that the corresponding delegate options
+  // data lifetime is maintained.
+  static std::unique_ptr CreateByName(
+      const std::string& name, const TFLiteSettings& settings);
+
+  // Struct to be statically allocated for registration.
+  struct Register {
+    Register(const std::string& name, CreatorFunction creator_function);
+  };
+
+ private:
+  void RegisterImpl(const std::string& name, CreatorFunction creator_function);
+  std::unique_ptr CreateImpl(
+      const std::string& name, const TFLiteSettings& settings);
+  static DelegatePluginRegistry* GetSingleton();
+  std::unordered_map factories_;
+  absl::Mutex mutex_;
+};
+
+}  // namespace delegates
+}  // namespace tflite
+
+#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f) \
+  static auto* g_delegate_plugin_##name##_ =                     \
+      new DelegatePluginRegistry::Register(#name, f);
+#define TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(name, f) \
+  TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION_VNAME(name, f);
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_DELEGATE_REGISTRY_H_
diff --git a/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
new file mode 100644
index 00000000000..25b8171c5ea
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/gpu_plugin.cc
@@ -0,0 +1,62 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/delegates/gpu/delegate.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
+
+namespace tflite {
+namespace delegates {
+class GpuPlugin : public DelegatePluginInterface {
+ public:
+  TfLiteDelegatePtr Create() override {
+    return TfLiteDelegatePtr(TfLiteGpuDelegateV2Create(&options_),
+                             TfLiteGpuDelegateV2Delete);
+  }
+  int GetDelegateErrno(TfLiteDelegate* from_delegate) override { return 0; }
+  static std::unique_ptr New(
+      const TFLiteSettings& acceleration) {
+    return absl::make_unique(acceleration);
+  }
+  explicit GpuPlugin(const TFLiteSettings& tflite_settings)
+      : options_(TfLiteGpuDelegateOptionsV2Default()) {
+    const auto* gpu_settings = tflite_settings.gpu_settings();
+    if (gpu_settings) {
+      options_.inference_priority1 =
+          gpu_settings->is_precision_loss_allowed()
+              ? TFLITE_GPU_INFERENCE_PRIORITY_MIN_LATENCY
+              : TFLITE_GPU_INFERENCE_PRIORITY_MAX_PRECISION;
+      if (gpu_settings->enable_quantized_inference()) {
+        options_.experimental_flags |=
+            TFLITE_GPU_EXPERIMENTAL_FLAGS_ENABLE_QUANT;
+      }
+      if (gpu_settings->force_backend() == GPUBackend_OPENCL) {
+        options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_CL_ONLY;
+      } else if (gpu_settings->force_backend() == GPUBackend_OPENGL) {
+        options_.experimental_flags |= TFLITE_GPU_EXPERIMENTAL_FLAGS_GL_ONLY;
+      }
+    }
+  }
+
+ private:
+  TfLiteGpuDelegateOptionsV2 options_;
+};
+
+TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(GpuPlugin, GpuPlugin::New);
+
+}  // namespace delegates
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
new file mode 100644
index 00000000000..7f2674604b0
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/hexagon_plugin.cc
@@ -0,0 +1,73 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
+
+#if defined(__ARM_ARCH)
+#include "tensorflow/lite/delegates/hexagon/hexagon_delegate.h"
+#endif
+
+namespace tflite {
+namespace delegates {
+class HexagonPlugin : public DelegatePluginInterface {
+ public:
+  TfLiteDelegatePtr Create() override {
+#if defined(__ARM_ARCH)
+    TfLiteHexagonInit();
+    auto* delegate_ptr = TfLiteHexagonDelegateCreate(&options_);
+    TfLiteDelegatePtr delegate(delegate_ptr, [](TfLiteDelegate* delegate) {
+      TfLiteHexagonDelegateDelete(delegate);
+      TfLiteHexagonTearDown();
+    });
+    return delegate;
+#else   // !defined(__ARM_ARCH)
+    return TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
+#endif  // defined(__ARM_ARCH)
+  }
+  int GetDelegateErrno(TfLiteDelegate* /* from_delegate */) override {
+    return 0;
+  }
+  static std::unique_ptr New(
+      const TFLiteSettings& tflite_settings) {
+    return absl::make_unique(tflite_settings);
+  }
+  explicit HexagonPlugin(const TFLiteSettings& tflite_settings) {
+    const HexagonSettings* settings = tflite_settings.hexagon_settings();
+#if defined(__ARM_ARCH)
+    options_ = TfLiteHexagonDelegateOptions({0});
+    if (settings) {
+      options_.debug_level = settings->debug_level();
+      options_.powersave_level = settings->powersave_level();
+      options_.print_graph_profile = settings->print_graph_profile();
+      options_.print_graph_debug = settings->print_graph_debug();
+    }
+#else
+    (void)settings;
+#endif
+  }
+
+ private:
+#if defined(__ARM_ARCH)
+  TfLiteHexagonDelegateOptions options_;
+#endif
+};
+
+TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(HexagonPlugin, HexagonPlugin::New);
+
+}  // namespace delegates
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
new file mode 100644
index 00000000000..7301983a815
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin.cc
@@ -0,0 +1,93 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
+
+namespace tflite {
+namespace delegates {
+
+inline tflite::StatefulNnApiDelegate::Options::ExecutionPreference
+ConvertExecutionPrefence(
+    NNAPIExecutionPreference from_compatibility_preference) {
+  using TflitePreference =
+      tflite::StatefulNnApiDelegate::Options::ExecutionPreference;
+  switch (from_compatibility_preference) {
+    case NNAPIExecutionPreference_NNAPI_LOW_POWER:
+      return TflitePreference::kLowPower;
+    case NNAPIExecutionPreference_NNAPI_FAST_SINGLE_ANSWER:
+      return TflitePreference::kFastSingleAnswer;
+    case NNAPIExecutionPreference_NNAPI_SUSTAINED_SPEED:
+      return TflitePreference::kSustainedSpeed;
+    default:
+      return TflitePreference::kUndefined;
+  }
+}
+
+class NnapiPlugin : public DelegatePluginInterface {
+ public:
+  TfLiteDelegatePtr Create() override {
+    auto nnapi_delegate =
+        absl::make_unique(options_);
+    return TfLiteDelegatePtr(
+        nnapi_delegate.release(), [](TfLiteDelegate* delegate) {
+          delete reinterpret_cast(delegate);
+        });
+  }
+  int GetDelegateErrno(TfLiteDelegate* from_delegate) override {
+    auto nnapi_delegate =
+        reinterpret_cast(from_delegate);
+    return nnapi_delegate->GetNnApiErrno();
+  }
+  static std::unique_ptr New(
+      const TFLiteSettings& tflite_settings) {
+    return absl::make_unique(tflite_settings);
+  }
+  explicit NnapiPlugin(const TFLiteSettings& tflite_settings) {
+    const NNAPISettings* nnapi_settings = tflite_settings.nnapi_settings();
+    if (!nnapi_settings) return;
+    if (nnapi_settings->accelerator_name() &&
+        nnapi_settings->accelerator_name()->Length() != 0) {
+      accelerator_ = nnapi_settings->accelerator_name()->str();
+      options_.accelerator_name = accelerator_.c_str();
+    }
+    if (nnapi_settings->cache_directory() &&
+        nnapi_settings->cache_directory()->Length() != 0) {
+      cache_dir_ = nnapi_settings->cache_directory()->str();
+      options_.cache_dir = cache_dir_.c_str();
+    }
+    if (nnapi_settings->model_token() &&
+        nnapi_settings->model_token()->Length() != 0) {
+      model_token_ = nnapi_settings->model_token()->str();
+      options_.model_token = model_token_.c_str();
+    }
+    options_.execution_preference =
+        ConvertExecutionPrefence(nnapi_settings->execution_preference());
+    options_.disallow_nnapi_cpu =
+        !nnapi_settings->allow_nnapi_cpu_on_android_10_plus();
+  }
+
+ private:
+  std::string accelerator_, cache_dir_, model_token_;
+  tflite::StatefulNnApiDelegate::Options options_;
+};
+
+TFLITE_REGISTER_DELEGATE_FACTORY_FUNCTION(NnapiPlugin, NnapiPlugin::New);
+
+}  // namespace delegates
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
new file mode 100644
index 00000000000..4f9f5dd08c1
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/nnapi_plugin_test.cc
@@ -0,0 +1,175 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+
+#include 
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
+#include "tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/delegate_registry.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+// Tests for checking that the NNAPI Delegate plugin correctly handles all the
+// options from the flatbuffer.
+//
+// Checking done at NNAPI call level, as that is where we have a mockable
+// layer.
+namespace tflite {
+namespace {
+
+using delegate::nnapi::NnApiMock;
+
+class SingleAddOpModel : tflite::SingleOpModel {
+ public:
+  void Build() {
+    int input = AddInput({tflite::TensorType_FLOAT32, {1, 2, 2}});
+    int constant = AddConstInput({tflite::TensorType_FLOAT32, {1, 2, 2}},
+                                 {1.0f, 1.0f, 1.0f, 1.0f});
+    AddOutput({tflite::TensorType_FLOAT32, {}});
+
+    SetBuiltinOp(tflite::BuiltinOperator_ADD, tflite::BuiltinOptions_AddOptions,
+                 tflite::CreateAddOptions(builder_).Union());
+    BuildInterpreter({GetShape(input), GetShape(constant)});
+  }
+
+  tflite::Interpreter* Interpreter() const { return interpreter_.get(); }
+};
+
+class NNAPIPluginTest : public ::testing::Test {
+ protected:
+  NNAPIPluginTest() : delegate_(nullptr, [](TfLiteDelegate*) {}) {}
+  void SetUp() override {
+    nnapi_ = const_cast(NnApiImplementation());
+    nnapi_mock_ = absl::make_unique(nnapi_);
+    nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
+        [](const ANeuralNetworksModel* model,
+           const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
+           bool* supportedOps) -> int {
+      supportedOps[0] = true;
+      return 0;
+    };
+    model_.Build();
+  }
+  template 
+  void CheckExecutionPreference() {
+    // Note - this uses a template since the NNAPI functions are C function
+    // pointers rather than lambdas so can't capture variables.
+    nnapi_->ANeuralNetworksCompilation_setPreference =
+        [](ANeuralNetworksCompilation* compilation, int32_t preference) {
+          return preference - output;
+        };
+    CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0, input));
+    // Since delegation succeeds, the model becomes immutable and hence can't
+    // reuse it.
+    SingleAddOpModel model;
+    model.Build();
+    EXPECT_EQ(model.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+              kTfLiteOk)
+        << " given input: " << input << " expected output: " << output;
+  }
+
+  void CreateDelegate(flatbuffers::Offset settings) {
+    settings_ = flatbuffers::GetTemporaryPointer(
+        fbb_, CreateTFLiteSettings(fbb_, tflite::Delegate_NNAPI, settings));
+
+    plugin_ = delegates::DelegatePluginRegistry::CreateByName("NnapiPlugin",
+                                                              *settings_);
+    delegate_ = plugin_->Create();
+  }
+
+  NnApi* nnapi_;
+  std::unique_ptr nnapi_mock_;
+  SingleAddOpModel model_;
+  flatbuffers::FlatBufferBuilder fbb_;
+  const TFLiteSettings* settings_ = nullptr;
+  delegates::TfLiteDelegatePtr delegate_;
+  std::unique_ptr plugin_;
+};
+
+TEST_F(NNAPIPluginTest, PassesAcceleratorName) {
+  // Fails with non-existent "foo".
+  CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("foo")));
+  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+            kTfLiteDelegateError);
+
+  // Succeeds with "test-device" supported by the mock.
+  CreateDelegate(CreateNNAPISettings(fbb_, fbb_.CreateString("test-device")));
+  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+            kTfLiteOk);
+}
+
+TEST_F(NNAPIPluginTest, PassesExecutionPreference) {
+  CheckExecutionPreference();
+  CheckExecutionPreference();
+  CheckExecutionPreference();
+  CheckExecutionPreference();
+}
+
+TEST_F(NNAPIPluginTest, PassesCachingParameters) {
+  nnapi_->ANeuralNetworksCompilation_setCaching =
+      [](ANeuralNetworksCompilation* compilation, const char* cacheDir,
+         const uint8_t* token) -> int {
+    if (std::string(cacheDir) != "d") return 1;
+    // Token is hashed with other bits, just check that it's not empty.
+    if (std::string(reinterpret_cast(token)).empty()) return 2;
+    return 0;
+  };
+  CreateDelegate(CreateNNAPISettings(fbb_, 0, fbb_.CreateString("d"),
+                                     fbb_.CreateString("t")));
+  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+            kTfLiteOk);
+}
+
+TEST_F(NNAPIPluginTest, PassesFalseNNAPICpuFlag) {
+  CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0,
+                                     NNAPIExecutionPreference_UNDEFINED, 0, 0,
+                                     /* allow CPU */ false));
+  nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
+      [](const ANeuralNetworksModel* model,
+         const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
+         bool* supportedOps) -> int {
+    supportedOps[0] = true;
+    // Since no CPU, should only pass one device.
+    return numDevices - 1;
+  };
+  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+            kTfLiteOk);
+}
+
+TEST_F(NNAPIPluginTest, PassesTrueNNAPICpuFlag) {
+  CreateDelegate(CreateNNAPISettings(fbb_, 0, 0, 0,
+                                     NNAPIExecutionPreference_UNDEFINED, 0, 0,
+                                     /* allow CPU */ true));
+  nnapi_->ANeuralNetworksModel_getSupportedOperationsForDevices =
+      [](const ANeuralNetworksModel* model,
+         const ANeuralNetworksDevice* const* devices, uint32_t numDevices,
+         bool* supportedOps) -> int {
+    supportedOps[0] = true;
+    // With CPU allowed, should pass two devices.
+    return numDevices - 2;
+  };
+  EXPECT_EQ(model_.Interpreter()->ModifyGraphWithDelegate(delegate_.get()),
+            kTfLiteOk);
+}
+
+}  // namespace
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
new file mode 100644
index 00000000000..709bb70ca70
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.cc
@@ -0,0 +1,58 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h"
+
+#include 
+
+#include "flatbuffers/idl.h"  // from @flatbuffers
+#include "flatbuffers/util.h"  // from @flatbuffers
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+
+namespace {
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_fbs_contents-inl.h"
+}
+
+const ComputeSettings* ConvertFromProto(
+    flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings) {
+  std::string json;
+  tensorflow::protobuf::util::JsonPrintOptions options;
+  options.preserve_proto_field_names = true;
+  options.always_print_primitive_fields = true;  // For catching problems.
+  auto status = tensorflow::protobuf::util::MessageToJsonString(proto_settings,
+                                                                &json, options);
+  if (!status.ok()) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to convert to Json: %s",
+                    status.ToString().c_str());
+    return nullptr;
+  }
+  if (!parser->Parse(configuration_fbs_contents)) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse schema: %s",
+                    parser->error_.c_str());
+    return nullptr;
+  }
+  parser->SetRootType("tflite.ComputeSettings");
+  if (!parser->Parse(json.c_str())) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse json: %s",
+                    parser->error_.c_str());
+    return nullptr;
+  }
+  return flatbuffers::GetRoot(
+      parser->builder_.GetBufferPointer());
+}
+
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h
new file mode 100644
index 00000000000..3b69e8465a5
--- /dev/null
+++ b/tensorflow/lite/experimental/acceleration/configuration/proto_to_flatbuffer.h
@@ -0,0 +1,32 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_
+
+#include "flatbuffers/idl.h"  // from @flatbuffers
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration.pb.h"
+#include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
+
+namespace tflite {
+
+// Converts the protobuf version ComputeSettings to the flatbuffer version, via
+// json. The parser is used for state - the returned pointer is valid only as
+// long as the parser is kept alive and unmutated.
+const ComputeSettings* ConvertFromProto(
+    flatbuffers::Parser* parser, const proto::ComputeSettings& proto_settings);
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_CONFIGURATION_PROTO_TO_FLATBUFFER_H_

From bedc750c7d06c79ca61b53c6eeaba727ba7e886a Mon Sep 17 00:00:00 2001
From: Meteorix 
Date: Tue, 16 Jun 2020 23:38:53 +0800
Subject: [PATCH 0289/1390] fix LaunchDepthwiseConvBackpropFilterOp

---
 tensorflow/core/kernels/depthwise_conv_op_gpu.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/depthwise_conv_op_gpu.h b/tensorflow/core/kernels/depthwise_conv_op_gpu.h
index f712d9c48c0..f69878bea89 100644
--- a/tensorflow/core/kernels/depthwise_conv_op_gpu.h
+++ b/tensorflow/core/kernels/depthwise_conv_op_gpu.h
@@ -1763,7 +1763,7 @@ void LaunchDepthwiseConvBackpropFilterOp::operator()(
   int num_filter_backprop =
       args.filter_rows * args.filter_cols * args.out_depth;
   se::DeviceMemoryBase filter_bp_ptr(filter_backprop, num_filter_backprop);
-  stream->ThenMemset32(&filter_bp_ptr, 0, num_filter_backprop * sizeof(T));
+  stream->ThenMemZero(&filter_bp_ptr, num_filter_backprop * sizeof(T));
 
   if (args.filter_rows == 3 && args.filter_cols == 3) {
     OP_REQUIRES_OK(

From 878ac5ae83120cd5f2d013ee3dd08b8d1cd14040 Mon Sep 17 00:00:00 2001
From: Thomas O'Malley 
Date: Tue, 16 Jun 2020 08:54:22 -0700
Subject: [PATCH 0290/1390] Add fused/non-fused inference/training overhead
 benchmarks for BatchNormalization.

PiperOrigin-RevId: 316690878
Change-Id: I36a0c8595b973657ae3cb8f95c11ba797cc4dcab
---
 .../benchmark/eager_microbenchmarks_test.py   | 36 +++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/keras/benchmark/eager_microbenchmarks_test.py b/tensorflow/python/keras/benchmark/eager_microbenchmarks_test.py
index 9c755286ee0..b4d0837c326 100644
--- a/tensorflow/python/keras/benchmark/eager_microbenchmarks_test.py
+++ b/tensorflow/python/keras/benchmark/eager_microbenchmarks_test.py
@@ -246,10 +246,40 @@ class MicroBenchmarksBase(test.Benchmark):
 
     self._run(fn, 10000)
 
-  def benchmark_layers_normalization_batch_normalization_overhead(self):
+  def benchmark_layers_batch_norm_fused_inf(self):
 
-    layer = normalization.BatchNormalization()
-    x = array_ops.ones((1, 1))
+    layer = normalization.BatchNormalization(fused=True)
+    x = array_ops.ones((1, 1, 1, 1))
+
+    def fn():
+      layer(x)
+
+    self._run(fn, 10000)
+
+  def benchmark_layers_batch_norm_fused_train(self):
+
+    layer = normalization.BatchNormalization(fused=True)
+    x = array_ops.ones((1, 1, 1, 1))
+
+    def fn():
+      layer(x, training=True)
+
+    self._run(fn, 10000)
+
+  def benchmark_layers_batch_norm_nonfused_inf(self):
+
+    layer = normalization.BatchNormalization(fused=False)
+    x = array_ops.ones((1, 1, 1, 1))
+
+    def fn():
+      layer(x)
+
+    self._run(fn, 10000)
+
+  def benchmark_layers_batch_norm_nonfused_train(self):
+
+    layer = normalization.BatchNormalization(fused=False)
+    x = array_ops.ones((1, 1, 1, 1))
 
     def fn():
       layer(x, training=True)

From 8a081946857e0951fd206daf69d159df8cdca34b Mon Sep 17 00:00:00 2001
From: Andy Ly 
Date: Tue, 16 Jun 2020 08:55:28 -0700
Subject: [PATCH 0291/1390] Add verifier for tf.BatchToSpace.

The verifier checks for input, output and crops for whether the op is valid or not. If the contents of crops can be determined, its values will be used directly in checking the input and output shapes.

PiperOrigin-RevId: 316691063
Change-Id: Idd495d2102604e267cdaa5f45a21c0bfa3dcbcb0
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    |   4 +
 .../compiler/mlir/tensorflow/ir/tf_ops.cc     | 143 ++++++++++++++++++
 .../mlir/tensorflow/tests/tf-ops.mlir         | 137 +++++++++++++++++
 3 files changed, 284 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index c51e9f4d026..7db3539fcef 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -820,6 +820,10 @@ followed by cropping along the `height` and `width` dimensions.
 
   TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
   TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;
+
+  let verifier = [{
+    return Verify(*this);
+  }];
 }
 
 def TF_BatchToSpaceNDOp : TF_Op<"BatchToSpaceND", [NoSideEffect]> {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
index f4e5dc05eb0..8410929a19f 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
@@ -695,6 +695,149 @@ void BatchMatMulV2Op::getCanonicalizationPatterns(
   results.insert(context);
 }
 
+//===----------------------------------------------------------------------===//
+// BatchToSpaceOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult Verify(BatchToSpaceOp op) {
+  // Op already has a constraint that block_size >= 2.
+  int64_t block_size = op.block_size().getSExtValue();
+
+  llvm::SmallVector input_shape(4, ShapedType::kDynamicSize);
+  auto input_type = op.input().getType().cast();
+  if (input_type.hasRank()) {
+    if (input_type.getRank() != 4)
+      return op.emitOpError()
+             << "requires input to be a 4D tensor, but got " << input_type;
+
+    int64_t input_batch = input_type.getDimSize(0);
+    if (input_batch != ShapedType::kDynamicSize &&
+        input_batch % (block_size * block_size) != 0) {
+      return op.emitOpError()
+             << "requires input batch (dimension 0) to be evenly divisible "
+                "by (block_size * block_size), but got input batch "
+             << input_batch << " and block_size " << block_size;
+    }
+
+    input_shape.assign(input_type.getShape().begin(),
+                       input_type.getShape().end());
+  }
+
+  auto crops_type = op.crops().getType().cast();
+  if (crops_type.hasRank()) {
+    if (crops_type.getRank() != 2)
+      return op.emitOpError()
+             << "requires crops to be a 2D tensor, but got " << crops_type;
+
+    auto dim_of_size = [&](int64_t dim, int64_t size) {
+      if (crops_type.isDynamicDim(dim)) return true;
+      return crops_type.getDimSize(dim) == size;
+    };
+    if (!dim_of_size(0, 2) || !dim_of_size(1, 2))
+      return op.emitOpError()
+             << "requires crops to be a tensor<2x2>, but got " << crops_type;
+  }
+
+  DenseIntElementsAttr crops_attr;
+  // Crops are defined as [[crop_top, crop_bottom], [crop_left, crop_right]],
+  // and flattened as [crop_top, crop_bottom, crop_left, crop_right]
+  llvm::SmallVector crops_values;
+  if (matchPattern(op.crops(), m_Constant(&crops_attr))) {
+    assert(crops_attr.getNumElements() == 4 &&
+           "tf.BatchToSpace crops must have 4 elements");
+
+    auto crops_range = crops_attr.getIntValues();
+    for (const auto &crops_value : crops_range) {
+      int64_t crops_value_int = crops_value.getSExtValue();
+      if (crops_value_int < 0)
+        return op.emitOpError()
+               << "requires all crop values to be nonnegative, but got "
+               << crops_attr;
+
+      crops_values.push_back(crops_value_int);
+    }
+  }
+
+  auto output_type = op.output().getType().cast();
+  if (output_type.hasRank()) {
+    if (output_type.getRank() != 4)
+      return op.emitOpError()
+             << "requires output to be a 4D tensor, but got " << output_type;
+
+    auto static_dims = [](int64_t dim_a, int64_t dim_b) {
+      return dim_a != ShapedType::kDynamicSize &&
+             dim_b != ShapedType::kDynamicSize;
+    };
+
+    auto output_shape = output_type.getShape();
+
+    // output batch = input batch / (block_size * block_size).
+    int64_t input_batch = input_shape[0];
+    int64_t output_batch = output_shape[0];
+    if (static_dims(input_batch, output_batch) &&
+        (output_batch * block_size * block_size) != input_batch)
+      return op.emitOpError()
+             << "requires output batch (dimension 0) to be equal to input "
+                "batch (dimension 0) / (block_size * block_size), but got "
+                "output batch "
+             << output_batch << ", input batch " << input_batch
+             << ", and block_size " << block_size;
+
+    auto check_spatial_dim = [&](int64_t spatial_dim_index,
+                                 llvm::StringRef dim_name,
+                                 llvm::StringRef crop_a_name,
+                                 llvm::StringRef crop_b_name) -> LogicalResult {
+      int64_t input_dim = input_shape[spatial_dim_index];
+      int64_t output_dim = output_shape[spatial_dim_index];
+      if (!static_dims(input_dim, output_dim)) return success();
+
+      int64_t input_dim_pad = input_dim * block_size;
+      // If crops are unknown, the maximum output spatial dim size is input
+      // spatial dim size * block_size, as crops can be minimum 0.
+      if (crops_values.empty() && output_dim > input_dim * block_size)
+        return op.emitOpError()
+               << "requires output " << dim_name << " (dimension "
+               << spatial_dim_index << ") to be less than or equal to input "
+               << dim_name << " (dimension " << spatial_dim_index
+               << ") * block_size, but got output " << dim_name << " "
+               << output_dim << ", input " << dim_name << " " << input_dim
+               << ", and block_size " << block_size;
+
+      if (!crops_values.empty()) {
+        // output spatial dim = input spatial dim * block_size - crops.
+        int64_t crop_a = crops_values[2 * (spatial_dim_index - 1)];
+        int64_t crop_b = crops_values[2 * (spatial_dim_index - 1) + 1];
+        if (output_dim != input_dim_pad - crop_a - crop_b)
+          return op.emitOpError()
+                 << "requires output " << dim_name << " (dimension "
+                 << spatial_dim_index << ") to be equal to input " << dim_name
+                 << " (dimension " << spatial_dim_index << ") * block_size - "
+                 << crop_a_name << " - " << crop_b_name << ", but got output "
+                 << dim_name << " " << output_dim << ", input " << dim_name
+                 << " " << input_dim << ", " << crop_a_name << " " << crop_a
+                 << ", " << crop_b_name << " " << crop_b << ", and block_size "
+                 << block_size;
+      }
+
+      return success();
+    };
+
+    if (failed(check_spatial_dim(1, "height", "crop_top", "crop_bottom")) ||
+        failed(check_spatial_dim(2, "width", "crop_left", "crop_right")))
+      return failure();
+
+    int64_t input_depth = input_shape[3];
+    int64_t output_depth = output_shape[3];
+    if (static_dims(input_depth, output_depth) && output_depth != input_depth)
+      return op.emitOpError()
+             << "requires output depth (dimension 3) to be equal to input "
+                "depth (dimension 3), but got output depth "
+             << output_depth << " and input depth " << input_depth;
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // BiasAddOp
 //===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
index 058585b41d7..4ba2e83300b 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir
@@ -2872,4 +2872,141 @@ func @testSendTPUEmbeddingGradients(%x: tensor<512x256xf32>) {
   return
 }
 
+// -----
 
+//===--------------------------------------------------------------------===//
+//  tf.BatchToSpace
+//===--------------------------------------------------------------------===//
+
+func @testBatchToSpaceDynamic(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
+  return
+}
+
+func @testBatchToSpaceRankedInput(%arg0: tensor, %arg1: tensor<*xi32>) {
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor, tensor<*xi32>) -> tensor<*xf32>
+  return
+}
+
+func @testBatchToSpaceRankedCrops(%arg0: tensor<*xf32>, %arg1: tensor) {
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor) -> tensor<*xf32>
+  return
+}
+
+func @testBatchToSpaceRankedOutput(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor
+  return
+}
+
+func @testBatchToSpaceStatic(%arg0: tensor<36x8x8x8xf32>) {
+  %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 3 : i64} : (tensor<36x8x8x8xf32>, tensor<2x2xi32>) -> tensor<4x21x17x8xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidInputRank(%arg0: tensor<8xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires input to be a 4D tensor, but got 'tensor<8xf32>'}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<8xf32>, tensor<*xi32>) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidInputBatch(%arg0: tensor<2x4x6x8xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires input batch (dimension 0) to be evenly divisible by (block_size * block_size), but got input batch 2 and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<2x4x6x8xf32>, tensor<*xi32>) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsRank(%arg0: tensor<*xf32>, %arg1: tensor) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a 2D tensor, but got 'tensor'}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsFirstDim(%arg0: tensor<*xf32>, %arg1: tensor<3x?xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a tensor<2x2>, but got 'tensor<3x?xi32>'}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<3x?xi32>) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidCropsSecondDim(%arg0: tensor<*xf32>, %arg1: tensor) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires crops to be a tensor<2x2>, but got 'tensor'}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceBadCropValues(%arg0: tensor<*xf32>) {
+  %crops = "tf.Const"() {value = dense<[[-1, -2], [-3, -4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  // expected-error @+1 {{'tf.BatchToSpace' op requires all crop values to be nonnegative, but got dense<[[-1, -2], [-3, -4]]> : tensor<2x2xi32>}}
+  %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<*xf32>, tensor<2x2xi32>) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputRank(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output to be a 4D tensor, but got 'tensor<8xf32>'}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<8xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputBatch(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output batch (dimension 0) to be equal to input batch (dimension 0) / (block_size * block_size), but got output batch 8, input batch 16, and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<8x8x8x3xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputHeight(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output height (dimension 1) to be less than or equal to input height (dimension 1) * block_size, but got output height 17, input height 8, and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<4x17x8x3xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputHeightCrops(%arg0: tensor<16x8x8x3xf32>) {
+  %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output height (dimension 1) to be equal to input height (dimension 1) * block_size - crop_top - crop_bottom, but got output height 8, input height 8, crop_top 1, crop_bottom 2, and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<2x2xi32>) -> tensor<4x8x9x3xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputWidth(%arg0: tensor<16x4x4x3xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output width (dimension 2) to be less than or equal to input width (dimension 2) * block_size, but got output width 9, input width 4, and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x4x4x3xf32>, tensor<*xi32>) -> tensor<4x4x9x3xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputWidthCrops(%arg0: tensor<16x8x8x3xf32>) {
+  %crops = "tf.Const"() {value = dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output width (dimension 2) to be equal to input width (dimension 2) * block_size - crop_left - crop_right, but got output width 8, input width 8, crop_left 3, crop_right 4, and block_size 2}}
+  %0 = "tf.BatchToSpace"(%arg0, %crops) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<2x2xi32>) -> tensor<4x13x8x3xf32>
+  return
+}
+
+// -----
+
+func @testBatchToSpaceInvalidOutputDepth(%arg0: tensor<16x8x8x3xf32>, %arg1: tensor<*xi32>) {
+  // expected-error @+1 {{'tf.BatchToSpace' op requires output depth (dimension 3) to be equal to input depth (dimension 3), but got output depth 8 and input depth 3}}
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 2 : i64} : (tensor<16x8x8x3xf32>, tensor<*xi32>) -> tensor<4x8x8x8xf32>
+  return
+}

From 68e3ce2fe773f1731f40ce69da7585ebd3257b3d Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Tue, 16 Jun 2020 23:19:28 +0700
Subject: [PATCH 0292/1390] Add GetTempFileName

---
 tensorflow/c/env.cc              | 10 ++++++++++
 tensorflow/c/env.h               |  6 ++++++
 tensorflow/core/platform/path.cc |  2 ++
 3 files changed, 18 insertions(+)

diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc
index 1c35ff9001d..3d490d95e66 100644
--- a/tensorflow/c/env.cc
+++ b/tensorflow/c/env.cc
@@ -146,6 +146,16 @@ TF_StringStream* TF_GetLocalTempDirectories() {
   return list;
 }
 
+void TF_GetTempFileName(const char* extension, std::string* name,
+                        TF_Status* status) {
+  *name = ::tensorflow::Env::Default()->GetTempFilename(extension);
+  if (*name.length() == 0) {
+    TF_SetStatus(status, TF_INTERNAL, "Can not create temp file name");
+  } else {
+    TF_SetStatus(status, TF_OK, "");
+  }
+}
+
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) {
   return ::tensorflow::Env::Default()->NowNanos();
 }
diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h
index 2a763730bc3..b50d0fdec03 100644
--- a/tensorflow/c/env.h
+++ b/tensorflow/c/env.h
@@ -152,6 +152,12 @@ TF_CAPI_EXPORT extern TF_StringStream* TF_GetChildren(const char* filename,
 // The caller is responsible for freeing the list (see TF_StringStreamDone).
 TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void);
 
+// Creates a temporary file name with an extension.
+// The caller is responsible for freeing the returned pointer.
+TF_CAPI_EXPORT extern void TF_GetTempFileName(const char* extension,
+                                              std::string* name,
+                                              TF_Status* status);
+
 // Returns the number of nanoseconds since the Unix epoch.
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void);
 
diff --git a/tensorflow/core/platform/path.cc b/tensorflow/core/platform/path.cc
index 1e88328aace..f9442ccba0f 100644
--- a/tensorflow/core/platform/path.cc
+++ b/tensorflow/core/platform/path.cc
@@ -327,6 +327,8 @@ string GetTempFilename(const string& extension) {
   }
   LOG(FATAL) << "No temp directory found.";
 #endif
+  // Return an empty string to indicate that we can not create temp file name.
+  return "";
 }
 
 bool GetTestUndeclaredOutputsDir(string* dir) {

From 12ec80d2395da75c18db5b8371c68f07889dba11 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Tue, 16 Jun 2020 09:18:02 -0700
Subject: [PATCH 0293/1390] Move the module test wrt to Keras private API to
 keras test.

PiperOrigin-RevId: 316695477
Change-Id: I4201f62e97bb1acb768b980a51d26bf0c76853ab
---
 tensorflow/python/keras/layers/BUILD          | 12 ++++++
 tensorflow/python/keras/layers/layers_test.py | 40 +++++++++++++++++++
 tensorflow/tools/api/tests/module_test.py     | 12 ------
 3 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 tensorflow/python/keras/layers/layers_test.py

diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD
index 91bace936d3..f67838d591a 100644
--- a/tensorflow/python/keras/layers/BUILD
+++ b/tensorflow/python/keras/layers/BUILD
@@ -886,3 +886,15 @@ tf_py_test(
         "@absl_py//absl/testing:parameterized",
     ],
 )
+
+tf_py_test(
+    name = "layers_test",
+    size = "small",
+    srcs = ["layers_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":layers",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:tf2",
+    ],
+)
diff --git a/tensorflow/python/keras/layers/layers_test.py b/tensorflow/python/keras/layers/layers_test.py
new file mode 100644
index 00000000000..35ba029dd83
--- /dev/null
+++ b/tensorflow/python/keras/layers/layers_test.py
@@ -0,0 +1,40 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+# pylint: disable=g-classes-have-attributes
+"""Tests for layers.__init__."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python import tf2
+from tensorflow.python.keras import layers
+from tensorflow.python.platform import test
+
+
+class LayersTest(test.TestCase):
+
+  def test_keras_private_symbol(self):
+    normalization_parent = layers.Normalization.__module__.split('.')[-1]
+    if tf2.enabled():
+      self.assertEqual('normalization', normalization_parent)
+      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
+    else:
+      self.assertEqual('normalization_v1', normalization_parent)
+      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
+
+
+if __name__ == '__main__':
+  test.main()
diff --git a/tensorflow/tools/api/tests/module_test.py b/tensorflow/tools/api/tests/module_test.py
index aa8a224d00b..5397278f5f3 100644
--- a/tensorflow/tools/api/tests/module_test.py
+++ b/tensorflow/tools/api/tests/module_test.py
@@ -79,18 +79,6 @@ class ModuleTest(test.TestCase):
       tf.compat.v1.summary.FileWriter
     # pylint: enable=pointless-statement
 
-  def testInternalKerasImport(self):
-    # pylint: disable=g-import-not-at-top
-    from tensorflow.python.keras import layers
-    normalization_parent = layers.Normalization.__module__.split('.')[-1]
-    if tf._major_api_version == 2:
-      self.assertEqual('normalization', normalization_parent)
-      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    else:
-      self.assertEqual('normalization_v1', normalization_parent)
-      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    # pylint: enable=g-import-not-at-top
-
 
 if __name__ == '__main__':
   test.main()

From cf599cade160e40d9f7051afebe89672f60e7137 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Tue, 16 Jun 2020 09:30:51 -0700
Subject: [PATCH 0294/1390] Move the serialization_test to keras/tests

PiperOrigin-RevId: 316697481
Change-Id: I918b29a976b166662acf9045f87512aef485441b
---
 tensorflow/python/keras/tests/BUILD           | 15 +++++
 .../keras/tests/serialization_util_test.py    | 67 +++++++++++++++++++
 tensorflow/python/util/serialization_test.py  | 36 ----------
 3 files changed, 82 insertions(+), 36 deletions(-)
 create mode 100644 tensorflow/python/keras/tests/serialization_util_test.py

diff --git a/tensorflow/python/keras/tests/BUILD b/tensorflow/python/keras/tests/BUILD
index 36af32184e6..b01fc3ca903 100644
--- a/tensorflow/python/keras/tests/BUILD
+++ b/tensorflow/python/keras/tests/BUILD
@@ -359,6 +359,21 @@ cuda_py_test(
     ],
 )
 
+tf_py_test(
+    name = "serialization_util_test",
+    size = "small",
+    srcs = ["serialization_util_test.py"],
+    python_version = "PY3",
+    deps = [
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:constant_op",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:util",
+        "//tensorflow/python/keras/engine",
+        "//tensorflow/python/keras/layers:core",
+    ],
+)
+
 tf_py_test(
     name = "temporal_sample_weights_correctness_test",
     srcs = ["temporal_sample_weights_correctness_test.py"],
diff --git a/tensorflow/python/keras/tests/serialization_util_test.py b/tensorflow/python/keras/tests/serialization_util_test.py
new file mode 100644
index 00000000000..058bdaec56c
--- /dev/null
+++ b/tensorflow/python/keras/tests/serialization_util_test.py
@@ -0,0 +1,67 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for serialization functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.keras.engine import input_layer
+from tensorflow.python.keras.engine import sequential
+from tensorflow.python.keras.engine import training
+from tensorflow.python.keras.layers import core
+from tensorflow.python.platform import test
+from tensorflow.python.util import serialization
+
+
+class SerializationTests(test.TestCase):
+
+  def test_serialize_dense(self):
+    dense = core.Dense(3)
+    dense(constant_op.constant([[4.]]))
+    round_trip = json.loads(json.dumps(
+        dense, default=serialization.get_json_type))
+    self.assertEqual(3, round_trip["config"]["units"])
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialize_sequential(self):
+    model = sequential.Sequential()
+    model.add(core.Dense(4))
+    model.add(core.Dense(5))
+    model(constant_op.constant([[1.]]))
+    sequential_round_trip = json.loads(
+        json.dumps(model, default=serialization.get_json_type))
+    self.assertEqual(
+        # Note that `config['layers'][0]` will be an InputLayer in V2
+        # (but not in V1)
+        5, sequential_round_trip["config"]["layers"][-1]["config"]["units"])
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_serialize_model(self):
+    x = input_layer.Input(shape=[3])
+    y = core.Dense(10)(x)
+    model = training.Model(x, y)
+    model(constant_op.constant([[1., 1., 1.]]))
+    model_round_trip = json.loads(
+        json.dumps(model, default=serialization.get_json_type))
+    self.assertEqual(
+        10, model_round_trip["config"]["layers"][1]["config"]["units"])
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/python/util/serialization_test.py b/tensorflow/python/util/serialization_test.py
index a66dd11ba99..fea2e0feb6e 100644
--- a/tensorflow/python/util/serialization_test.py
+++ b/tensorflow/python/util/serialization_test.py
@@ -20,26 +20,13 @@ from __future__ import print_function
 
 import json
 
-from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import tensor_shape
-from tensorflow.python.framework import test_util
-from tensorflow.python.keras.engine import input_layer
-from tensorflow.python.keras.engine import sequential
-from tensorflow.python.keras.engine import training
-from tensorflow.python.keras.layers import core
 from tensorflow.python.platform import test
 from tensorflow.python.util import serialization
 
 
 class SerializationTests(test.TestCase):
 
-  def test_serialize_dense(self):
-    dense = core.Dense(3)
-    dense(constant_op.constant([[4.]]))
-    round_trip = json.loads(json.dumps(
-        dense, default=serialization.get_json_type))
-    self.assertEqual(3, round_trip["config"]["units"])
-
   def test_serialize_shape(self):
     round_trip = json.loads(json.dumps(
         tensor_shape.TensorShape([None, 2, 3]),
@@ -47,29 +34,6 @@ class SerializationTests(test.TestCase):
     self.assertIs(round_trip[0], None)
     self.assertEqual(round_trip[1], 2)
 
-  @test_util.run_in_graph_and_eager_modes
-  def test_serialize_sequential(self):
-    model = sequential.Sequential()
-    model.add(core.Dense(4))
-    model.add(core.Dense(5))
-    model(constant_op.constant([[1.]]))
-    sequential_round_trip = json.loads(
-        json.dumps(model, default=serialization.get_json_type))
-    self.assertEqual(
-        # Note that `config['layers'][0]` will be an InputLayer in V2
-        # (but not in V1)
-        5, sequential_round_trip["config"]["layers"][-1]["config"]["units"])
-
-  @test_util.run_in_graph_and_eager_modes
-  def test_serialize_model(self):
-    x = input_layer.Input(shape=[3])
-    y = core.Dense(10)(x)
-    model = training.Model(x, y)
-    model(constant_op.constant([[1., 1., 1.]]))
-    model_round_trip = json.loads(
-        json.dumps(model, default=serialization.get_json_type))
-    self.assertEqual(
-        10, model_round_trip["config"]["layers"][1]["config"]["units"])
 
 if __name__ == "__main__":
   test.main()

From e178928ea82c020c234cabf0ea672ada29a04292 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Tue, 16 Jun 2020 09:33:22 -0700
Subject: [PATCH 0295/1390] Rollback of breaking
 aa99cf218c8bf13aeb15e64ec4c62ea14ecb5753

PiperOrigin-RevId: 316697856
Change-Id: I1ea9a700a1b6d6d4bf824b3d810e4a61127da912
---
 tensorflow/lite/delegates/flex/delegate.cc |  7 ---
 tensorflow/lite/interpreter_builder.cc     | 17 ------
 tensorflow/lite/python/BUILD               |  3 +-
 tensorflow/lite/python/lite_flex_test.py   | 61 +++++++++-------------
 tensorflow/python/BUILD                    |  1 -
 5 files changed, 27 insertions(+), 62 deletions(-)

diff --git a/tensorflow/lite/delegates/flex/delegate.cc b/tensorflow/lite/delegates/flex/delegate.cc
index b8b0d4e6d01..4741bddc2f5 100644
--- a/tensorflow/lite/delegates/flex/delegate.cc
+++ b/tensorflow/lite/delegates/flex/delegate.cc
@@ -136,10 +136,3 @@ TfLiteStatus FlexDelegate::CopyFromBufferHandle(
 }
 
 }  // namespace tflite
-
-// Exported C interface function which is used by AcquireFlexDelegate() at
-// interpreter_build.cc. To export the function name globally, the function name
-// must be matched with patterns in tf_version_script.lds
-extern "C" tflite::TfLiteDelegateUniquePtr TF_AcquireFlexDelegate() {
-  return tflite::AcquireFlexDelegate();
-}
diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index d73b298e595..43d81ef0770 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -14,9 +14,6 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/interpreter_builder.h"
 
-#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
-#include 
-#endif
 #include 
 #include 
 #include 
@@ -117,20 +114,6 @@ const char* kEmptyTensorName = "";
 // For flex delegate, see also the strong override in
 // lite/delegates/flex/delegate.cc.
 TFLITE_ATTRIBUTE_WEAK Interpreter::TfLiteDelegatePtr AcquireFlexDelegate() {
-#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
-  // If _pywrap_tensorflow_internal.so is available, use
-  // TF_AcquireFlexDelegate() to initialize flex delegate.
-  void* lib_tf_internal =
-      dlopen("_pywrap_tensorflow_internal.so", RTLD_NOW | RTLD_LOCAL);
-  if (lib_tf_internal) {
-    auto TF_AcquireFlexDelegate =
-        reinterpret_cast(
-            dlsym(lib_tf_internal, "TF_AcquireFlexDelegate"));
-    if (TF_AcquireFlexDelegate) {
-      return TF_AcquireFlexDelegate();
-    }
-  }
-#endif
   return Interpreter::TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
 }
 
diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD
index b0f605ed50d..c1f37c81b7f 100644
--- a/tensorflow/lite/python/BUILD
+++ b/tensorflow/lite/python/BUILD
@@ -194,7 +194,8 @@ py_test(
     python_version = "PY3",
     srcs_version = "PY2AND3",
     tags = [
-        "no_mac",  # TODO(b/159077703): Enable Python API Flex support on MacOS.
+        # TODO(b/111881877): Enable in oss after resolving op registry issues.
+        "no_oss",
         "no_windows",
     ],
     deps = [
diff --git a/tensorflow/lite/python/lite_flex_test.py b/tensorflow/lite/python/lite_flex_test.py
index ffc157c2128..26bee206d27 100644
--- a/tensorflow/lite/python/lite_flex_test.py
+++ b/tensorflow/lite/python/lite_flex_test.py
@@ -19,7 +19,6 @@ from __future__ import division
 from __future__ import print_function
 
 from absl.testing import parameterized
-import numpy as np
 
 from tensorflow.lite.python import lite
 from tensorflow.lite.python.interpreter import Interpreter
@@ -42,7 +41,8 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       ('DisableMlirConverter', False))  # disable mlir
   def testFlexMode(self, enable_mlir):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(
+          shape=[1, 16, 16, 3], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -54,22 +54,19 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Check the model works with TensorFlow ops.
+    # Ensures the model contains TensorFlow ops.
+    # TODO(nupurgarg): Check values once there is a Python delegate interface.
     interpreter = Interpreter(model_content=tflite_model)
-    interpreter.allocate_tensors()
-    input_details = interpreter.get_input_details()
-    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
-    interpreter.set_tensor(input_details[0]['index'], test_input)
-    interpreter.invoke()
-
-    output_details = interpreter.get_output_details()
-    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
-    output_data = interpreter.get_tensor(output_details[0]['index'])
-    self.assertTrue((expected_output == output_data).all())
+    with self.assertRaises(RuntimeError) as error:
+      interpreter.allocate_tensors()
+    self.assertIn(
+        'Regular TensorFlow ops are not supported by this interpreter.',
+        str(error.exception))
 
   def testDeprecatedFlags(self):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(
+          shape=[1, 16, 16, 3], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -86,18 +83,14 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Check the model works with TensorFlow ops.
+    # Ensures the model contains TensorFlow ops.
+    # TODO(nupurgarg): Check values once there is a Python delegate interface.
     interpreter = Interpreter(model_content=tflite_model)
-    interpreter.allocate_tensors()
-    input_details = interpreter.get_input_details()
-    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
-    interpreter.set_tensor(input_details[0]['index'], test_input)
-    interpreter.invoke()
-
-    output_details = interpreter.get_output_details()
-    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
-    output_data = interpreter.get_tensor(output_details[0]['index'])
-    self.assertTrue((expected_output == output_data).all())
+    with self.assertRaises(RuntimeError) as error:
+      interpreter.allocate_tensors()
+    self.assertIn(
+        'Regular TensorFlow ops are not supported by this interpreter.',
+        str(error.exception))
 
 
 class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
@@ -121,18 +114,14 @@ class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
     converter.experimental_new_converter = enable_mlir
     tflite_model = converter.convert()
 
-    # Check the model works with TensorFlow ops.
+    # Ensures the model contains TensorFlow ops.
+    # TODO(nupurgarg): Check values once there is a Python delegate interface.
     interpreter = Interpreter(model_content=tflite_model)
-    interpreter.allocate_tensors()
-    input_details = interpreter.get_input_details()
-    test_input = np.array([4.0], dtype=np.float32)
-    interpreter.set_tensor(input_details[0]['index'], test_input)
-    interpreter.invoke()
-
-    output_details = interpreter.get_output_details()
-    expected_output = np.array([24.0], dtype=np.float32)
-    output_data = interpreter.get_tensor(output_details[0]['index'])
-    self.assertTrue((expected_output == output_data).all())
+    with self.assertRaises(RuntimeError) as error:
+      interpreter.allocate_tensors()
+    self.assertIn(
+        'Regular TensorFlow ops are not supported by this interpreter.',
+        str(error.exception))
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 343a95b85e9..87048ba9d40 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -6046,7 +6046,6 @@ pywrap_tensorflow_macro(
         "//tensorflow/core/profiler/internal:print_model_analysis",
         "//tensorflow/core/profiler/internal/cpu:python_tracer",
         "//tensorflow/tools/graph_transforms:transform_graph_lib",
-        "//tensorflow/lite/delegates/flex:delegate",
         "//tensorflow/lite/toco/python:toco_python_api",
         "//tensorflow/python/eager:pywrap_tfe_lib",
         "//tensorflow/core/util/tensor_bundle",

From 9ee6864c2c1acc356fcd50c56d1d1a6ae38392a7 Mon Sep 17 00:00:00 2001
From: Sachin Joglekar 
Date: Tue, 16 Jun 2020 09:56:08 -0700
Subject: [PATCH 0296/1390] Minor doc update to use TfLiteIntArray better

PiperOrigin-RevId: 316701561
Change-Id: If14f772f0053eb4ecb0bc225cc11e93e0f07adb0
---
 tensorflow/lite/g3doc/performance/delegates.md | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/g3doc/performance/delegates.md b/tensorflow/lite/g3doc/performance/delegates.md
index 1d11a8eb391..760e7273fc4 100644
--- a/tensorflow/lite/g3doc/performance/delegates.md
+++ b/tensorflow/lite/g3doc/performance/delegates.md
@@ -85,6 +85,8 @@ To see it in code, let's define a delegate and call it `MyDelegate`, which can
 execute Conv2D and Mean operations faster.
 
 ```c++
+#include "tensorflow/lite/util.h"
+
 // This is where the execution of the operations or whole graph happens.
 // The class below has an empty implementation just as a guideline
 // on the structure.
@@ -156,8 +158,7 @@ TfLiteRegistration GetMyDelegateNodeRegistration() {
 TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
   // Claim all nodes that can be evaluated by the delegate and ask the
   // framework to update the graph with delegate kernel instead.
-  // Reserve 1 element, since we need first element to be size.
-  std::vector supported_nodes(1);
+  std::vector supported_nodes;
   TfLiteIntArray* plan;
   TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
   TfLiteNode* node;
@@ -169,17 +170,19 @@ TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
       supported_nodes.push_back(node_index);
     }
   }
-  // Set first element to the number of nodes to replace.
-  supported_nodes[0] = supported_nodes.size() - 1;
   TfLiteRegistration my_delegate_kernel_registration =
       GetMyDelegateNodeRegistration();
 
   // This call split the graphs into subgraphs, for subgraphs that can be
   // handled by the delegate, it will replace it with a
   // 'my_delegate_kernel_registration'
-  return context->ReplaceNodeSubsetsWithDelegateKernels(
+  TfLiteIntArray* supported_nodes_int_array =
+      ::tflite::ConvertVectorToTfLiteIntArray(supported_nodes);
+  auto status = context->ReplaceNodeSubsetsWithDelegateKernels(
       context, my_delegate_kernel_registration,
-      reinterpret_cast(supported_nodes.data()), delegate);
+      supported_nodes_int_array, delegate);
+  TfLiteIntArrayFree(supported_nodes_int_array);
+  return status
 }
 
 void FreeBufferHandle(TfLiteContext* context, TfLiteDelegate* delegate,

From f6b8e93a5f614c962db73d4409f2b89026fbf4da Mon Sep 17 00:00:00 2001
From: Mingming Liu 
Date: Tue, 16 Jun 2020 10:09:39 -0700
Subject: [PATCH 0297/1390] A refactor to add helper functions {Split, Concat},
 which do Tensor dtype deduction automatically.

PiperOrigin-RevId: 316703807
Change-Id: I234bea9fce3cf3b2cb352be12246ee9f4e8c405a
---
 tensorflow/core/kernels/batch_kernels.cc | 73 ++++++++++++++----------
 1 file changed, 44 insertions(+), 29 deletions(-)

diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc
index e94aef641e3..6449a399573 100644
--- a/tensorflow/core/kernels/batch_kernels.cc
+++ b/tensorflow/core/kernels/batch_kernels.cc
@@ -100,9 +100,9 @@ typedef Eigen::SyclDevice SYCLDevice;
 #endif  // TENSORFLOW_USE_SYCL
 
 // Concatenates 'inputs' into a single tensor along the zeroth dimension.
-// Requires that all elements of 'inputs' have element type T. Writes to the
-// op's output at position 'output_index', using 'context' for the allocation to
-// ensure proper device placement.
+// Requires that all elements of 'inputs' have element type T. Writes to
+// 'output' using 'context' for the allocation to ensure proper device
+// placement.
 template 
 Status Concat(OpKernelContext* context, const gtl::ArraySlice inputs,
               Tensor* output) {
@@ -157,6 +157,25 @@ Status Concat(OpKernelContext* context, const gtl::ArraySlice inputs,
   return Status::OK();
 }
 
+// Same as 'Concat' above, but handles Tensor dtype deduction automatically.
+Status Concat(OpKernelContext* context, const gtl::ArraySlice inputs,
+              Tensor* output) {
+  const DataType type = inputs[0].dtype();
+  Status concat_status;
+  switch (type) {
+#define CASE(type)                                         \
+  case DataTypeToEnum::value:                        \
+    concat_status = Concat(context, inputs, output); \
+    break;
+    TF_CALL_ALL_TYPES(CASE);
+#undef CASE
+    default:
+      concat_status = errors::InvalidArgument("Unsupported data type: ", type);
+      break;
+  }
+  return concat_status;
+}
+
 // The Split*() functions split 'input' with element type T into 'sizes.size()'
 // tensors along the zeroth dimension, with the ith split having zeroth-
 // dimension size 'sizes[i]'. They allocate the output tensors using 'context',
@@ -268,6 +287,25 @@ Status Split(OpKernelContext* context, const Tensor& input,
   return SplitCPU(context, input, sizes, outputs);
 }
 
+// Same as 'Split' above, but handles Tensor dtype automatically.
+Status Split(OpKernelContext* context, const Tensor& input,
+             const gtl::ArraySlice sizes, std::vector* outputs) {
+  const DataType type = input.dtype();
+  Status split_status;
+  switch (type) {
+#define CASE(type)                                              \
+  case DataTypeToEnum::value:                             \
+    split_status = Split(context, input, sizes, outputs); \
+    break;
+    TF_CALL_ALL_TYPES(CASE);
+#undef CASE
+    default:
+      split_status = errors::InvalidArgument("Unsupported data type: ", type);
+      break;
+  }
+  return split_status;
+}
+
 // A class encapsulating the state and logic for batching tensors.
 class BatchResource : public ResourceBase {
  public:
@@ -449,22 +487,9 @@ class BatchResource : public ResourceBase {
         }
       }
 
-      const DataType type = to_concatenate[0].dtype();
-      Status concat_status;
       Tensor concatenated_tensor;
-      switch (type) {
-#define CASE(type)                                                   \
-  case DataTypeToEnum::value:                                  \
-    concat_status =                                                  \
-        Concat(context, to_concatenate, &concatenated_tensor); \
-    break;
-        TF_CALL_ALL_TYPES(CASE);
-#undef CASE
-        default:
-          concat_status =
-              errors::InvalidArgument("Unsupported data type: ", type);
-          break;
-      }
+      Status concat_status =
+          Concat(context, to_concatenate, &concatenated_tensor);
       TF_RETURN_IF_ERROR(concat_status);
       concatenated_tensors->push_back(concatenated_tensor);
     }
@@ -1001,17 +1026,7 @@ class UnbatchResource : public ResourceBase {
         batch_keys.push_back(batch_indices(i, 0));
       }
 
-      const DataType type = data_t.dtype();
-      switch (type) {
-#define CASE(type)                                                          \
-  case DataTypeToEnum::value:                                         \
-    TF_RETURN_IF_ERROR(Split(context, data_t, sizes, &split_inputs)); \
-    break;
-        TF_CALL_ALL_TYPES(CASE);
-#undef CASE
-        default:
-          return errors::InvalidArgument("Unsupported data type: ", type);
-      }
+      TF_RETURN_IF_ERROR(Split(context, data_t, sizes, &split_inputs));
     }
 
     // Critical section.

From 57c533a79af76f65af04f4ea5d0b30e777a31b66 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 10:24:58 -0700
Subject: [PATCH 0298/1390] Winograd back conversion updated to new style.

PiperOrigin-RevId: 316707136
Change-Id: I710af1fd5748f2d8ca371e870e150d79a39af0ce
---
 .../lite/delegates/gpu/cl/kernels/winograd.cc | 131 +++++++-----------
 .../lite/delegates/gpu/cl/kernels/winograd.h  |   3 -
 2 files changed, 53 insertions(+), 81 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
index d71513e4de4..2dcb72637ec 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.cc
@@ -213,20 +213,8 @@ std::string GetWinograd4x4To36Code(
   return c;
 }
 
-std::string GetWinograd36To4x4Code(
-    const OperationDef& op_def, const LinearStorage& at_arr,
-    const LinearStorage& biases,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHSBPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHSBPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
-
-  const std::string batch_id = op_def.IsBatchSupported() ? "batch_id" : "";
+std::string GetWinograd36To4x4Code(const OperationDef& op_def,
+                                   Arguments* args) {
   std::string c = GetCommonDefines(op_def.precision);
 
   switch (op_def.precision) {
@@ -243,6 +231,15 @@ std::string GetWinograd36To4x4Code(
                                   ? DataType::FLOAT16
                                   : DataType::FLOAT32;
 
+  std::string cl_type = accum_type == DataType::FLOAT16 ? "half" : "float";
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  src_desc->SetStateVar("ACCUM_FLT", cl_type);
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+  args->AddInt("tiles_x");
+
   auto at_mat = AtMatrixForWinograd4x4To6x6();
   c += "constant ACCUM_FLT At[24] = {\n";
   for (int y = 0; y < 4; ++y) {
@@ -255,30 +252,21 @@ std::string GetWinograd36To4x4Code(
   c += "};\n";
 
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ) + ",\n";
-  c += at_arr.GetDeclaration() + ",\n";
-  c += biases.GetDeclaration();
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,                              \n";
-  c += "    int4 dst_size,                              \n";
-  c += "    int tiles_x                                 \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int tile_id = get_global_id(0);\n";
   c += "  int DST_Y = get_global_id(1);\n";
   c += "  int DST_Z = get_global_id(2);\n";
-  c += "  int tile_x = (tile_id % tiles_x) * 4;\n";
-  c += "  int tile_y = (tile_id / tiles_x) * 4 + DST_Y;\n";
-  c += "  if (tile_x >= dst_size.x || tile_y >= dst_size.y || DST_Z >= "
-       "dst_size.z) {\n";
+  c += "  int tile_x = (tile_id % args.tiles_x) * 4;\n";
+  c += "  int tile_y = (tile_id / args.tiles_x) * 4 + DST_Y;\n";
+
+  c += "  if (tile_x >= args.dst_tensor.Width() || tile_y >= "
+       "args.dst_tensor.Height() || DST_Z >= args.dst_tensor.Slices()) {\n";
   c += "    return; \n";
   c += "  }\n";
   c += "  ACCUM_FLT4 I0, I1, I2, I3, I4, I5;\n";
   c += "  ACCUM_FLT at_ar[6];\n";
-  c += "  ACCUM_FLT4 t00 = TO_ACCUM_TYPE(" +
-       at_arr.ReadLinearFLT4("DST_Y * 2 + 0") + ");\n";
-  c += "  ACCUM_FLT4 t01 = TO_ACCUM_TYPE(" +
-       at_arr.ReadLinearFLT4("DST_Y * 2 + 1") + ");\n";
+  c += "  ACCUM_FLT4 t00 = TO_ACCUM_TYPE(args.at.Read(DST_Y * 2 + 0));\n";
+  c += "  ACCUM_FLT4 t01 = TO_ACCUM_TYPE(args.at.Read(DST_Y * 2 + 1));\n";
   c += "  at_ar[0] = t00.x;\n";
   c += "  at_ar[1] = t00.y;\n";
   c += "  at_ar[2] = t00.z;\n";
@@ -290,10 +278,8 @@ std::string GetWinograd36To4x4Code(
   for (int x = 0; x < 6; ++x) {
     const std::string yc = std::to_string(x);
     const std::string src = "src" + std::to_string(x);
-    c += "    ACCUM_FLT4 " + src + " = " +
-         src_tensor.ReadAsTypeWHSB(accum_type, "tile_id", yc, "DST_Z",
-                                   batch_id) +
-         ";\n";
+    c += "    ACCUM_FLT4 " + src +
+         " = args.src_tensor.Read(tile_id, " + yc + ", DST_Z);\n";
     c += "    I" + std::to_string(x) + " = at * " + src + ";\n";
   }
   c += "  }\n";
@@ -303,46 +289,35 @@ std::string GetWinograd36To4x4Code(
     for (int x = 0; x < 6; ++x) {
       const std::string yc = std::to_string(y * 6 + x);
       const std::string src = "src" + std::to_string(x);
-      c += "    ACCUM_FLT4 " + src + " = " +
-           src_tensor.ReadAsTypeWHSB(accum_type, "tile_id", yc, "DST_Z",
-                                     batch_id) +
-           ";\n";
+      c += "    ACCUM_FLT4 " + src +
+           " = args.src_tensor.Read(tile_id, " + yc + ", DST_Z);\n";
       c += "    I" + std::to_string(x) + " += at * " + src + ";\n";
     }
     c += "  }\n";
   }
   c += "  ACCUM_FLT4 t0 = I1 + I2;\n";
   c += "  ACCUM_FLT4 t1 = I3 + I4;\n";
-  c += "  FLT4 bias_val = " + biases.ReadLinearFLT4("DST_Z") + ";\n";
+  c += "  FLT4 bias_val = args.biases.Read(DST_Z);\n";
   c += "  {\n";
-  const LinkingContext context{"r0", "tile_x", "tile_y", "DST_Z"};
   c += "    FLT4 r0 = TO_FLT4(I0 + t0 + t1) + bias_val;\n";
-  c += PostProcess(linked_operations, context);
-  c += "    " +
-       dst_tensor.WriteWHSB("r0", "tile_x", "tile_y", "DST_Z", batch_id);
+  c += "    args.dst_tensor.Write(r0, tile_x, tile_y, DST_Z);\n";
   c += "    tile_x++;\n";
   c += "  }\n";
   c += "  ACCUM_FLT4 t2 = I1 - I2;\n";
   c += "  ACCUM_FLT4 t3 = I3 - I4;\n";
-  c += "  if (tile_x < dst_size.x) {\n";
+  c += "  if (tile_x < args.dst_tensor.Width()) {\n";
   c += "    FLT4 r0 = TO_FLT4(t2 * At[7] + t3 * At[9]) + bias_val;\n";
-  c += PostProcess(linked_operations, context);
-  c += "    " +
-       dst_tensor.WriteWHSB("r0", "tile_x", "tile_y", "DST_Z", batch_id);
+  c += "    args.dst_tensor.Write(r0, tile_x, tile_y, DST_Z);\n";
   c += "    tile_x++;\n";
   c += "  }\n";
-  c += "  if (tile_x < dst_size.x) {\n";
+  c += "  if (tile_x < args.dst_tensor.Width()) {\n";
   c += "    FLT4 r0 = TO_FLT4(t0 * At[13] + t1 * At[15]) + bias_val;\n";
-  c += PostProcess(linked_operations, context);
-  c += "    " +
-       dst_tensor.WriteWHSB("r0", "tile_x", "tile_y", "DST_Z", batch_id);
+  c += "    args.dst_tensor.Write(r0, tile_x, tile_y, DST_Z);\n";
   c += "    tile_x++;\n";
   c += "  }\n";
-  c += "  if (tile_x < dst_size.x) {\n";
+  c += "  if (tile_x < args.dst_tensor.Width()) {\n";
   c += "    FLT4 r0 = TO_FLT4(t2 * At[19] + t3 * At[21] + I5) + bias_val;\n";
-  c += PostProcess(linked_operations, context);
-  c += "    " +
-       dst_tensor.WriteWHSB("r0", "tile_x", "tile_y", "DST_Z", batch_id);
+  c += "    args.dst_tensor.Write(r0, tile_x, tile_y, DST_Z);\n";
   c += "    tile_x++;\n";
   c += "  }\n";
   c += "}\n";
@@ -406,7 +381,6 @@ absl::Status Winograd4x4To36::UploadBt(CLContext* context) {
   LinearStorageCreateInfo create_info;
   create_info.storage_type = LinearStorageType::TEXTURE_2D;
   create_info.data_type = definition_.GetDataType();
-  create_info.name = "bt_arr";
 
   LinearStorage lt;
   RETURN_IF_ERROR(CreateLinearStorage(create_info, bt_aligned, context, <));
@@ -473,15 +447,11 @@ absl::Status CreateWinograd4x4To36(const CreationContext& creation_context,
 
 Winograd36To4x4::Winograd36To4x4(Winograd36To4x4&& operation)
     : GPUOperation(std::move(operation)),
-      at_(std::move(operation.at_)),
-      biases_(std::move(operation.biases_)),
       kernel_(std::move(operation.kernel_)),
       work_group_size_(operation.work_group_size_) {}
 
 Winograd36To4x4& Winograd36To4x4::operator=(Winograd36To4x4&& operation) {
   if (this != &operation) {
-    at_ = std::move(operation.at_);
-    biases_ = std::move(operation.biases_);
     kernel_ = std::move(operation.kernel_);
     std::swap(work_group_size_, operation.work_group_size_);
     GPUOperation::operator=(std::move(operation));
@@ -495,8 +465,13 @@ absl::Status Winograd36To4x4::Compile(const CreationContext& creation_context) {
       creation_context.device->IsPowerVR()) {
     options.push_back(CompilerOptions::POWERVR_FP16);
   }
-  const auto code =
-      GetWinograd36To4x4Code(definition_, at_, biases_, linked_operations_);
+  std::string code = GetWinograd36To4x4Code(definition_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   RETURN_IF_ERROR(creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", options, *creation_context.context,
       *creation_context.device, &kernel_));
@@ -520,8 +495,11 @@ absl::Status Winograd36To4x4::UploadAt(CLContext* context) {
   LinearStorageCreateInfo create_info;
   create_info.storage_type = LinearStorageType::TEXTURE_2D;
   create_info.data_type = definition_.GetDataType();
-  create_info.name = "at_arr";
-  return CreateLinearStorage(create_info, at_aligned, context, &at_);
+  LinearStorage lt;
+  RETURN_IF_ERROR(CreateLinearStorage(create_info, at_aligned, context, <));
+  args_.AddObject("at", AccessType::READ,
+                  absl::make_unique(std::move(lt)));
+  return absl::OkStatus();
 }
 
 int3 Winograd36To4x4::SelectBestWorkGroup() {
@@ -532,17 +510,12 @@ int3 Winograd36To4x4::SelectBestWorkGroup() {
 }
 
 absl::Status Winograd36To4x4::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(at_.GetMemoryPtr()));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(biases_.GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWHSB()));
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
   const int tiles_x = DivideRoundUp(dst_[0]->Width(), 4);
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(tiles_x));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetInt("tiles_x", tiles_x));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Winograd36To4x4::GetGridSize() const {
@@ -580,9 +553,11 @@ absl::Status CreateWinograd36To4x4(
   LinearStorageCreateInfo create_info;
   create_info.storage_type = LinearStorageType::TEXTURE_2D;
   create_info.data_type = definition.GetDataType();
-  create_info.name = "biases";
-  RETURN_IF_ERROR(CreateLinearStorage(
-      create_info, biases, creation_context.context, &result->biases_));
+  LinearStorage lt;
+  RETURN_IF_ERROR(
+      CreateLinearStorage(create_info, biases, creation_context.context, <));
+  result->args_.AddObject("biases", AccessType::READ,
+                          absl::make_unique(std::move(lt)));
   return result->UploadAt(creation_context.context);
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.h b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.h
index ec8fe22ea11..84ebd87042d 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/winograd.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/winograd.h
@@ -99,9 +99,6 @@ class Winograd36To4x4 : public GPUOperation {
   absl::Status BindArguments();
   int3 GetGridSize() const;
 
-  LinearStorage at_;
-  LinearStorage biases_;
-
   CLKernel kernel_;
   int3 work_group_size_ = int3(128, 1, 1);
 };

From 7fb05db5b0d4caad58ecb536b1ed80393dabc4a8 Mon Sep 17 00:00:00 2001
From: Thomas Joerg 
Date: Tue, 16 Jun 2020 10:27:26 -0700
Subject: [PATCH 0299/1390] Add op metadata to bitcasts inserted by
 ReductionLayoutNormalizer.

PiperOrigin-RevId: 316707766
Change-Id: I11aa5e9c32e86d63981398605c50a73ac9149303
---
 .../compiler/xla/service/gpu/reduction_layout_normalizer.cc      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.cc b/tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.cc
index 295ccebd442..64a1403a3f9 100644
--- a/tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.cc
+++ b/tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.cc
@@ -101,6 +101,7 @@ class EnforceMinorToMajorReduceOpVisitor : public DfsHloRewriteVisitor {
         new_reduce_shape_layout);
     HloInstruction *canonical_reduce_input = reduce->parent()->AddInstruction(
         HloInstruction::CreateBitcast(new_operand_shape, operand));
+    canonical_reduce_input->set_metadata(reduce->metadata());
 
     VLOG(5) << "Reduction input: " << canonical_reduce_input->ToString();
     std::unique_ptr new_reduce = HloInstruction::CreateReduce(

From c42314ef70f9e093e53cb85d0dac19535bd677ae Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 10:29:33 -0700
Subject: [PATCH 0300/1390] fix flaky test multi_process_runner_test on
 tensorflow msan TAP. allow more margin of errors.

PiperOrigin-RevId: 316708284
Change-Id: I3ab912bd7a546bbedf5ed7825fa4b345e92880d8
---
 tensorflow/python/distribute/multi_process_runner_test.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/multi_process_runner_test.py b/tensorflow/python/distribute/multi_process_runner_test.py
index f7158d1fdb4..aeba43b6b7c 100644
--- a/tensorflow/python/distribute/multi_process_runner_test.py
+++ b/tensorflow/python/distribute/multi_process_runner_test.py
@@ -314,8 +314,12 @@ class MultiProcessRunnerTest(test.TestCase):
           timeout=5)
 
     list_to_assert = cm.exception.mpr_result.stdout
+    # We should see 5 iterations from worker and ps, however sometime on TAP
+    # due to CPU throttling and slugginess of msan/asan build, this became
+    # flaky. Therefore we allow more margin of errors to only check the first
+    # 3 iterations.
     for job in ['worker', 'ps']:
-      for iteration in range(0, 5):
+      for iteration in range(0, 3):
         self.assertTrue(
             any('(logging) {}-0, i: {}'.format(job, iteration) in line
                 for line in list_to_assert))

From 08e445e37f27c2228c499ed261e343aae15551cd Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Tue, 16 Jun 2020 10:51:13 -0700
Subject: [PATCH 0301/1390] [TF/XLA] [NFC] Simplify
 XlaComputationLaunchContext::PopulateInputs

Try to make the logic more transparent

PiperOrigin-RevId: 316713452
Change-Id: I41e8d691e6cab9c7a6b5bc40d50d660fcbe05906
---
 tensorflow/compiler/jit/xla_launch_util.cc | 43 +++++++++-------------
 1 file changed, 17 insertions(+), 26 deletions(-)

diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 209220938ed..ec5a372875c 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -198,50 +198,41 @@ void XlaComputationLaunchContext::PopulateInputs(
     OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel,
     const std::map& variables,
     int missing_ctx_input_prefix) {
-  se::Stream* stream =
-      ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr;
   // Build ShapedBuffers that point directly to the Tensor buffers.
   arg_ptrs_ = std::vector(kernel->xla_input_shapes.size());
 
-  // Pass remaining parameters.
-  const Tensor* t;
+  xla::TransferManager* transfer_manager =
+      client_->backend().transfer_manager();
   for (int i = 0; i < kernel->xla_input_shapes.size(); ++i) {
     int arg_num = kernel->input_mapping[i];
-    DCHECK_GE(arg_num, missing_ctx_input_prefix);
+    CHECK_GE(arg_num, missing_ctx_input_prefix);
     const xla::Shape& shape = kernel->xla_input_shapes[i];
-    if (variables.count(arg_num)) {
-      t = &(variables.at(arg_num).value);
-      CHECK(t);
-    } else {
-      t = &(ctx->input(arg_num - missing_ctx_input_prefix));
-    }
+    const Tensor* t = variables.count(arg_num)
+                          ? &(variables.at(arg_num).value)
+                          : &(ctx->input(arg_num - missing_ctx_input_prefix));
+    CHECK(t);
 
     if (use_multiple_streams_) {
-      CHECK(stream) << "Must have a stream available when using XLA tensors!";
+      CHECK(ctx->op_device_context() && ctx->op_device_context()->stream())
+          << "Must have a stream available when using XLA tensors!";
       XlaTensor* xla_tensor = XlaTensor::FromTensor(t);
       CHECK(xla_tensor);
-      xla_tensor->WaitForDefinitionEventOnStream(stream);
+      xla_tensor->WaitForDefinitionEventOnStream(
+          ctx->op_device_context()->stream());
     }
 
-    const xla::Shape on_device_shape =
-        client_->backend().transfer_manager()->HostShapeToDeviceShape(shape);
-    if (on_device_shape.IsTuple()) {
-      const XlaTensor* xla_tensor = XlaTensor::FromTensor(t);
-      CHECK(xla_tensor && xla_tensor->has_shaped_buffer());
-      arg_ptrs_[i] = const_cast(&xla_tensor->shaped_buffer());
-    } else {
-      CHECK(xla::Shape::Equal().MinorToMajorOnlyInLayout()(shape,
-                                                           on_device_shape))
-          << "On-device shape "
-          << xla::ShapeUtil::HumanStringWithLayout(on_device_shape)
-          << " not the same as on-host shape "
-          << xla::ShapeUtil::HumanStringWithLayout(shape);
+    if (xla::Shape::Equal().MinorToMajorOnlyInLayout()(
+            shape, transfer_manager->HostShapeToDeviceShape(shape))) {
       se::DeviceMemoryBase dmem = XlaTensor::DeviceMemoryFromTensor(*t);
       arg_buffers_.emplace_back(
           /*on_host_shape=*/shape, /*on_device_shape=*/shape,
           client_->platform(), client_->default_device_ordinal());
       arg_buffers_.back().set_buffer(dmem, /*index=*/{});
       arg_ptrs_[i] = &arg_buffers_.back();
+    } else {
+      const XlaTensor* xla_tensor = XlaTensor::FromTensor(t);
+      CHECK(xla_tensor && xla_tensor->has_shaped_buffer());
+      arg_ptrs_[i] = const_cast(&xla_tensor->shaped_buffer());
     }
   }
 }

From e6e8d48f8b4c17a19e968ccff3c88b084d7053e3 Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Tue, 16 Jun 2020 10:55:37 -0700
Subject: [PATCH 0302/1390] Improve numerics of Log1p for XLA:CPU.

PiperOrigin-RevId: 316714497
Change-Id: I20bf0148a850451077e435190034418e7eb9b38c
---
 .../compiler/tests/special_math_test.py       | 75 +++++++++++++++++++
 tensorflow/compiler/tests/unary_ops_test.py   | 14 ++--
 .../xla/service/elemental_ir_emitter.cc       | 47 +++++++++++-
 .../xla/service/elemental_ir_emitter.h        |  4 +
 4 files changed, 132 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/tests/special_math_test.py b/tensorflow/compiler/tests/special_math_test.py
index 3efaa6434be..246ab2a1641 100644
--- a/tensorflow/compiler/tests/special_math_test.py
+++ b/tensorflow/compiler/tests/special_math_test.py
@@ -61,6 +61,81 @@ def implicit_reparameterization_grad(a, x):
   return -gen_math_ops.igamma_grad_a(a, x) / prob
 
 
+@def_function.function(experimental_compile=True)
+def _log1p(x):
+  return math_ops.log1p(x)
+
+
+class Log1pTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      if six.PY2:
+        answer = int(entropy.encode('hex'), 16)
+      else:
+        answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(Log1pTest, self).setUp()
+
+  def adjust_tolerance_for_tpu(self, dtype, rtol, atol):
+    if self.device not in ['TPU']:
+      return rtol, atol
+
+    if dtype == np.float32:
+      return 4e-4, 0.
+    return 1e-10, 0.
+
+  def _test_range(self, low, high, dtype, rtol, atol, is_negative=False):
+    # Test values near zero.
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    x = np.exp(np.random.uniform(
+        low=low, high=high, size=[NUM_SAMPLES])).astype(dtype)
+    if is_negative:
+      x = -x
+    expected_values = np.log1p(x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = _log1p(x)
+      actual = sess.run(actual)
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testSmallX(self, dtype, rtol, atol):
+    self._test_range(-40., -20., dtype, rtol, atol, is_negative=False)
+    self._test_range(-40., -20., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 1e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeTwentyExponent(self, dtype, rtol, atol):
+    self._test_range(-20., -10., dtype, rtol, atol, is_negative=False)
+    self._test_range(-20., -10., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 1e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeTenExponent(self, dtype, rtol, atol):
+    self._test_range(-10., -5., dtype, rtol, atol, is_negative=False)
+    self._test_range(-10., -5., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeFiveExponent(self, dtype, rtol, atol):
+    self._test_range(-5., -1., dtype, rtol, atol, is_negative=False)
+    self._test_range(-5., -1., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 4e-7, 0.),
+                            (np.float64, 3e-14, 0.))
+  def testXGreaterThanOneTenth(self, dtype, rtol, atol):
+    self._test_range(-1., 0., dtype, rtol, atol, is_negative=False)
+    self._test_range(-1., 0., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 2e-15, 0.))
+  def testXGreaterThanOne(self, dtype, rtol, atol):
+    self._test_range(0., 3., dtype, rtol, atol, is_negative=False)
+
+
 class IgammaTest(xla_test.XLATestCase, parameterized.TestCase):
 
   def setUp(self):
diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index 5b3f9465513..162693a9eb1 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -292,13 +292,17 @@ class UnaryOpsTest(xla_test.XLATestCase):
           np.array([[1, 2]], dtype=dtype),
           expected=np.array([[0.540297, -0.41614]], dtype=dtype))
 
+      # Confirm that log1p will remain precise across a range of small values.
       self._assertOpOutputMatchesExpected(
           math_ops.log1p,
-          np.array([[1e-14, 1e-15, 0.6]], dtype=dtype),
-          expected=np.log1p(np.array([[1e-14, 1e-15, 0.6]],
-                                     dtype=dtype)).astype(dtype),
-          rtol=1e-4,
-          atol=1e-6)
+          np.array([[1e-14, 1e-15, 0.6, 2] + [x * 1e-5 for x in range(1, 20)]],
+                   dtype=dtype),
+          expected=np.log1p(
+              np.array(
+                  [[1e-14, 1e-15, 0.6, 2] + [x * 1e-5 for x in range(1, 20)]],
+                  dtype=dtype)).astype(dtype),
+          rtol=1e-15 if dtype == np.float64 else 1e-4,
+          atol=1e-15 if dtype == np.float64 else 1e-4)
 
       self._assertOpOutputMatchesExpected(
           math_ops.rint,
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
index 8cb660de46c..e4097b0c06f 100644
--- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
@@ -1336,9 +1336,40 @@ StatusOr ElementalIrEmitter::EmitLog1p(PrimitiveType prim_type,
   // When x is large, the naive evaluation of ln(x + 1) is more
   // accurate than the Taylor series.
   TF_ASSIGN_OR_RETURN(auto for_large_x, EmitLog(prim_type, FAdd(x, one)));
-  // The Taylor series for ln(x+1) is x - x^2/2 - x^3/3 + ….
-  auto for_small_x = FMul(FAdd(FMul(negative_half, x), one), x);
-  const auto kAntilogarithmIsSmallThreshold = 1e-4;
+  // When x is small, (defined to be less than sqrt(2) / 2), use a rational
+  // approximation. The approximation below is based on one from the Cephes
+  // Mathematical Library.
+  //
+  // sqrt(2) - 1.
+  const auto kAntilogarithmIsSmallThreshold = 0.41421356237309504880;
+
+  static const std::array kDenominatorCoeffs{
+      1.,
+      1.5062909083469192043167E1,
+      8.3047565967967209469434E1,
+      2.2176239823732856465394E2,
+      3.0909872225312059774938E2,
+      2.1642788614495947685003E2,
+      6.0118660497603843919306E1,
+  };
+
+  static const std::array kNumeratorCoeffs{
+      4.5270000862445199635215E-5, 4.9854102823193375972212E-1,
+      6.5787325942061044846969E0,  2.9911919328553073277375E1,
+      6.0949667980987787057556E1,  5.7112963590585538103336E1,
+      2.0039553499201281259648E1,
+  };
+
+  auto x_squared = FMul(x, x);
+  TF_ASSIGN_OR_RETURN(auto denominator,
+                      EvaluatePolynomial(type, x, kDenominatorCoeffs));
+  TF_ASSIGN_OR_RETURN(auto numerator,
+                      EvaluatePolynomial(type, x, kNumeratorCoeffs));
+  auto for_small_x = FDiv(numerator, denominator);
+  for_small_x = FMul(FMul(x, x_squared), for_small_x);
+  for_small_x = FAdd(FMul(negative_half, x_squared), for_small_x);
+  for_small_x = FAdd(x, for_small_x);
+
   auto abs_x =
       llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::fabs, {value}, {type}, b_);
   auto x_is_small = FCmpOLT(
@@ -2699,4 +2730,14 @@ StatusOr ElementalIrEmitter::EmitElementalReduce(
   }
 }
 
+// Evaluate polynomial using Horner's method.
+StatusOr ElementalIrEmitter::EvaluatePolynomial(
+    llvm::Type* type, llvm::Value* x, absl::Span coefficients) {
+  llvm::Value* poly = llvm::ConstantFP::get(type, 0.0);
+  for (const double c : coefficients) {
+    poly = FAdd(FMul(poly, x), llvm::ConstantFP::get(type, c));
+  }
+  return poly;
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/elemental_ir_emitter.h
index 06a9d7b194c..e39d2dd99ec 100644
--- a/tensorflow/compiler/xla/service/elemental_ir_emitter.h
+++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.h
@@ -258,6 +258,10 @@ class ElementalIrEmitter : public IrBuilderMixin {
   StatusOr EmitComplexPower(const HloInstruction* op,
                                           llvm::Value* a, llvm::Value* b,
                                           llvm::Value* c, llvm::Value* d);
+
+  // Evaluates a polynomial using Horner's method.
+  StatusOr EvaluatePolynomial(
+      llvm::Type* type, llvm::Value* x, absl::Span coefficients);
 };
 
 }  // namespace xla

From 70075e5acc70783336c220bc484909ceadfdd9ac Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 10:56:01 -0700
Subject: [PATCH 0303/1390] StridedSlice converted to new style.

PiperOrigin-RevId: 316714604
Change-Id: I5bc49b5d49704b46b5b4f3169fd679b91f60d235
---
 .../delegates/gpu/cl/kernels/strided_slice.cc | 102 +++++++++---------
 1 file changed, 52 insertions(+), 50 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc b/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
index 19f1b185d3c..2cf65f24447 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
@@ -25,72 +25,67 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetStridedSliceCode(
-    const OperationDef& op_def, bool alignedx4,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHSBPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHSBPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
+std::string GetStridedSliceCode(const OperationDef& op_def, bool alignedx4,
+                                Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+  args->AddInt("offset_x");
+  args->AddInt("offset_y");
+  args->AddInt("offset_z");
+  args->AddInt("offset_b");
+  args->AddInt("stride_x");
+  args->AddInt("stride_y");
+  args->AddInt("stride_z");
+  args->AddInt("stride_b");
 
   const std::string dst_batch = op_def.IsBatchSupported() ? "B" : "";
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 offset,            \n";
-  c += "    int4 stride,            \n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size              \n";
-  c += ") {\n";
+  c += "$0) {\n";
   if (op_def.IsBatchSupported()) {
     c += "  int linear_id = get_global_id(0);\n";
-    c += "  int X = linear_id / dst_size.w;\n";
-    c += "  int B = linear_id % dst_size.w;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
   } else {
     c += "  int X = get_global_id(0);\n";
   }
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) { \n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
   c += "    return; \n";
   c += "  } \n";
-  c += "  int s_x = X * stride.x + offset.x;\n";
-  c += "  int s_y = Y * stride.y + offset.y;\n";
+  c += "  int s_x = X * args.stride_x + args.offset_x;\n";
+  c += "  int s_y = Y * args.stride_y + args.offset_y;\n";
   if (op_def.IsBatchSupported()) {
-    c += "  int s_b = B * stride.w + offset.w;\n";
+    c += "  int s_b = B * args.stride_b + args.offset_b;\n";
+    c += "  args.src_tensor.SetBatchRef(s_b);\n";
   }
   const std::string src_batch = op_def.IsBatchSupported() ? "s_b" : "";
   if (alignedx4) {
-    c += "  int s_z = Z + offset.z;\n";
-    c += "  FLT4 result = " +
-         src_tensor.ReadWHSB("s_x", "s_y", "s_z", src_batch) + ";\n";
+    c += "  int s_z = Z + args.offset_z;\n";
+    c += "  FLT4 result = args.src_tensor.Read(s_x, s_y, s_z);\n";
   } else {
     c += "  FLT4 result;\n";
     const std::string postfixes[] = {"x", "y", "z", "w"};
     for (int i = 0; i < 4; ++i) {
       c += "  {\n";
       const std::string channel = "(Z * 4 + " + std::to_string(i) + ")";
-      c += "    int s_ch = " + channel + " * stride.z + offset.z;\n";
-      c += "    int s_z = min(s_ch >> 2, src_size.z - 1);\n";
+      c += "    int s_ch = " + channel + " * args.stride_z + args.offset_z;\n";
+      c += "    int s_z = min(s_ch >> 2, args.src_tensor.Slices() - 1);\n";
       c += "    int s_z_rem = s_ch & 3;\n";
-      c += "    FLT4 t = " +
-           src_tensor.ReadWHSB("s_x", "s_y", "s_z", src_batch) + ";\n";
+      c += "    FLT4 t = args.src_tensor.Read(s_x, s_y, s_z);\n";
       c += "    FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
       c += "    result." + postfixes[i] + " = t_ar[s_z_rem];\n";
       c += "  }\n";
     }
   }
-  std::string x_3dcoord =
-      op_def.IsBatchSupported() ? "X * dst_size.w + B" : "X";
-  const LinkingContext context{"result", x_3dcoord, "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHSB("result", "X", "Y", "Z", dst_batch);
+  c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
   c += "}\n";
   return c;
 }
@@ -167,27 +162,34 @@ StridedSlice& StridedSlice::operator=(StridedSlice&& operation) {
 }
 
 absl::Status StridedSlice::Compile(const CreationContext& creation_context) {
-  const auto code = GetStridedSliceCode(definition_, Is4Aligned(attributes_),
-                                        linked_operations_);
+  std::string code =
+      GetStridedSliceCode(definition_, Is4Aligned(attributes_), &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status StridedSlice::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
   int4 offset = GetOffset(attributes_, src_[0]->Width(), src_[0]->Height(),
                           src_[0]->Channels(), src_[0]->Batch());
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(offset));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int4(attributes_.strides.w, attributes_.strides.h,
-                                attributes_.strides.c, attributes_.strides.b)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWHSB()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetInt("offset_x", offset.x));
+  RETURN_IF_ERROR(args_.SetInt("offset_y", offset.y));
+  RETURN_IF_ERROR(args_.SetInt("offset_z", offset.z));
+  RETURN_IF_ERROR(args_.SetInt("offset_b", offset.w));
+  RETURN_IF_ERROR(args_.SetInt("stride_x", attributes_.strides.w));
+  RETURN_IF_ERROR(args_.SetInt("stride_y", attributes_.strides.h));
+  RETURN_IF_ERROR(args_.SetInt("stride_z", attributes_.strides.c));
+  RETURN_IF_ERROR(args_.SetInt("stride_b", attributes_.strides.b));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 StridedSlice::GetGridSize() const {

From 9a0838f66ad078e7ae089eaf87447cc00e5d65c5 Mon Sep 17 00:00:00 2001
From: Andy Ly 
Date: Tue, 16 Jun 2020 10:59:42 -0700
Subject: [PATCH 0304/1390] Add canonicalization for tf.BatchToSpace to more
 generalized tf.BatchToSpaceND.

This canonicalization is restricted to ranked inputs and crops, so in the case where an invalid shaped input or crops is passed to tf.BatchToSpace, a similar error will occur instead of potentially succeeding with the more generalized tf.BatchToSpaceND.

PiperOrigin-RevId: 316715564
Change-Id: I41b8cedf5ac770772e4c943fee8049af9e34418b
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    |  2 ++
 .../compiler/mlir/tensorflow/ir/tf_ops.cc     |  5 ++++
 .../mlir/tensorflow/tests/canonicalize.mlir   | 24 +++++++++++++++++++
 .../tensorflow/transforms/canonicalize.td     | 17 +++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index 7db3539fcef..d403462e6a6 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -824,6 +824,8 @@ followed by cropping along the `height` and `width` dimensions.
   let verifier = [{
     return Verify(*this);
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def TF_BatchToSpaceNDOp : TF_Op<"BatchToSpaceND", [NoSideEffect]> {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
index 8410929a19f..6d8c5af297d 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc
@@ -838,6 +838,11 @@ static LogicalResult Verify(BatchToSpaceOp op) {
   return success();
 }
 
+void BatchToSpaceOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert(context);
+}
+
 //===----------------------------------------------------------------------===//
 // BiasAddOp
 //===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir
index a127168157f..542c5b3c166 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir
@@ -586,3 +586,27 @@ func @sub(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
   %0 = "tf.Sub"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
   return %0 : tensor<*xf32>
 }
+
+// CHECK-LABEL: testBatchToSpaceToBatchToSpaceND
+// CHECK-SAME: ([[INPUT:%.*]]: tensor, [[CROPS:%.*]]: tensor)
+func @testBatchToSpaceToBatchToSpaceND(%arg0: tensor, %arg1: tensor) -> tensor<*xf32> {
+  // CHECK: [[BLOCK_SHAPE:%.*]] = "tf.Const"() {value = dense<8> : tensor<2xi64>}
+  // CHECK: [[BATCH_TO_SHAPE_ND:%.*]] = "tf.BatchToSpaceND"([[INPUT]], [[BLOCK_SHAPE]], [[CROPS]])
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 8 : i64} : (tensor, tensor) -> tensor<*xf32>
+  // CHECK: return [[BATCH_TO_SHAPE_ND]]
+  return %0 : tensor<*xf32>
+}
+
+// CHECK-LABEL: testBatchToSpaceDynamicInput
+func @testBatchToSpaceDynamicInput(%arg0: tensor<*xf32>, %arg1: tensor) -> tensor<*xf32> {
+  // CHECK-NOT: "tf.BatchToSpaceND"
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 8 : i64} : (tensor<*xf32>, tensor) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
+
+// CHECK-LABEL: testBatchToSpaceDynamicCrops
+func @testBatchToSpaceDynamicCrops(%arg0: tensor, %arg1: tensor<*xi32>) -> tensor<*xf32> {
+  // CHECK-NOT: "tf.BatchToSpaceND"
+  %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 8 : i64} : (tensor, tensor<*xi32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/canonicalize.td b/tensorflow/compiler/mlir/tensorflow/transforms/canonicalize.td
index 041d1c6cdaf..edd36a7b4c7 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/canonicalize.td
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/canonicalize.td
@@ -27,6 +27,8 @@ def SingleResultAndOperandHaveSameType : Constraint<
 
 def IsRank2Tensor : Type, "Rank 2 tensor">;
 
+def IsRank4Tensor : Type, "Rank 4 tensor">;
+
 // Checks if all the users is ReadVariableOp.
 def HasOnlyReadVariableOpUsers : Constraint<
   CPred<"llvm::all_of($0.getUsers(), [](mlir::OpOperand op) { "
@@ -65,6 +67,21 @@ def BatchMatMulV2ToMatMul : Pat<(TF_BatchMatMulV2Op $x, $y, $adj_x, $adj_y),
                               (TF_MatMulOp $x, $y, $adj_x, $adj_y),
                               [(IsRank2Tensor $x), (IsRank2Tensor $y)]>;
 
+//===----------------------------------------------------------------------===//
+// BatchToSpace op patterns.
+//===----------------------------------------------------------------------===//
+
+def BatchToSpaceBlockSizeToBlockShape : NativeCodeCall<
+  "DenseElementsAttr::get(RankedTensorType::get({2}, $_builder.getI64Type()), "
+  "ArrayRef{$0.getValue(), $0.getValue()})">;
+
+def BatchToSpaceToBatchToSpaceND :
+    Pat<(TF_BatchToSpaceOp $input, $crops, $block_size),
+        (TF_BatchToSpaceNDOp $input,
+         (TF_ConstOp (BatchToSpaceBlockSizeToBlockShape $block_size)),
+         $crops),
+        [(IsRank4Tensor $input), (IsRank2Tensor $crops)]>;
+
 //===----------------------------------------------------------------------===//
 // BiasAddV1 op patterns.
 //===----------------------------------------------------------------------===//

From 30357d1d9b5599b08a43eba918d0349ce4a29acc Mon Sep 17 00:00:00 2001
From: Sachin Joglekar 
Date: Tue, 16 Jun 2020 11:08:54 -0700
Subject: [PATCH 0305/1390] Improve floating-point NMS tests to use smaller
 error thresholds

PiperOrigin-RevId: 316717898
Change-Id: Iab097dcf4ac3feca17c6d54ad84a2437341d0bb3
---
 .../kernels/detection_postprocess_test.cc     | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/tensorflow/lite/kernels/detection_postprocess_test.cc b/tensorflow/lite/kernels/detection_postprocess_test.cc
index cf0d3ba2f3d..b9c42e75f21 100644
--- a/tensorflow/lite/kernels/detection_postprocess_test.cc
+++ b/tensorflow/lite/kernels/detection_postprocess_test.cc
@@ -174,17 +174,17 @@ TEST(DetectionPostprocessOpTest, FloatTest) {
   std::vector output_shape2 = m.GetOutputShape2();
   EXPECT_THAT(output_shape2, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput2(),
-              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-4)));
   // detection_scores
   std::vector output_shape3 = m.GetOutputShape3();
   EXPECT_THAT(output_shape3, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput3(),
-              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-4)));
   // num_detections
   std::vector output_shape4 = m.GetOutputShape4();
   EXPECT_THAT(output_shape4, ElementsAre(1));
   EXPECT_THAT(m.GetOutput4(),
-              ElementsAreArray(ArrayFloatNear({3.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({3.0}, 1e-4)));
 }
 
 TEST(DetectionPostprocessOpTest, QuantizedTest) {
@@ -385,17 +385,17 @@ TEST(DetectionPostprocessOpTest, FloatTestFastNMS) {
   std::vector output_shape2 = m.GetOutputShape2();
   EXPECT_THAT(output_shape2, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput2(),
-              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-4)));
   // detection_scores
   std::vector output_shape3 = m.GetOutputShape3();
   EXPECT_THAT(output_shape3, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput3(),
-              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-4)));
   // num_detections
   std::vector output_shape4 = m.GetOutputShape4();
   EXPECT_THAT(output_shape4, ElementsAre(1));
   EXPECT_THAT(m.GetOutput4(),
-              ElementsAreArray(ArrayFloatNear({3.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({3.0}, 1e-4)));
 }
 
 TEST(DetectionPostprocessOpTest, QuantizedTestFastNMS) {
@@ -492,22 +492,22 @@ TEST(DetectionPostprocessOpTest, FloatTestRegularNMS) {
   EXPECT_THAT(m.GetOutput1(),
               ElementsAreArray(ArrayFloatNear({0.0, 10.0, 1.0, 11.0, 0.0, 10.0,
                                                1.0, 11.0, 0.0, 0.0, 0.0, 0.0},
-                                              3e-1)));
+                                              3e-4)));
   // detection_classes
   std::vector output_shape2 = m.GetOutputShape2();
   EXPECT_THAT(output_shape2, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput2(),
-              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-4)));
   // detection_scores
   std::vector output_shape3 = m.GetOutputShape3();
   EXPECT_THAT(output_shape3, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput3(),
-              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({0.95, 0.93, 0.0}, 1e-4)));
   // num_detections
   std::vector output_shape4 = m.GetOutputShape4();
   EXPECT_THAT(output_shape4, ElementsAre(1));
   EXPECT_THAT(m.GetOutput4(),
-              ElementsAreArray(ArrayFloatNear({2.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({2.0}, 1e-4)));
 }
 
 TEST(DetectionPostprocessOpTest, QuantizedTestRegularNMS) {
@@ -666,17 +666,17 @@ TEST(DetectionPostprocessOpTest, FloatTestwithBackgroundClassAndKeypoints) {
   std::vector output_shape2 = m.GetOutputShape2();
   EXPECT_THAT(output_shape2, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput2(),
-              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-4)));
   // detection_scores
   std::vector output_shape3 = m.GetOutputShape3();
   EXPECT_THAT(output_shape3, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput3(),
-              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-4)));
   // num_detections
   std::vector output_shape4 = m.GetOutputShape4();
   EXPECT_THAT(output_shape4, ElementsAre(1));
   EXPECT_THAT(m.GetOutput4(),
-              ElementsAreArray(ArrayFloatNear({3.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({3.0}, 1e-4)));
 }
 
 TEST(DetectionPostprocessOpTest,
@@ -780,17 +780,17 @@ TEST(DetectionPostprocessOpTest, FloatTestwithNoBackgroundClassAndKeypoints) {
   std::vector output_shape2 = m.GetOutputShape2();
   EXPECT_THAT(output_shape2, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput2(),
-              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-4)));
   // detection_scores
   std::vector output_shape3 = m.GetOutputShape3();
   EXPECT_THAT(output_shape3, ElementsAre(1, 3));
   EXPECT_THAT(m.GetOutput3(),
-              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-4)));
   // num_detections
   std::vector output_shape4 = m.GetOutputShape4();
   EXPECT_THAT(output_shape4, ElementsAre(1));
   EXPECT_THAT(m.GetOutput4(),
-              ElementsAreArray(ArrayFloatNear({3.0}, 1e-1)));
+              ElementsAreArray(ArrayFloatNear({3.0}, 1e-4)));
 }
 }  // namespace
 }  // namespace custom

From 421e64c0c6d54a877059f3743c0cbcaf625c51ee Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 11:10:26 -0700
Subject: [PATCH 0306/1390] Build fixes for CLion.

This change allows targets that have TF dependencies to build in the IDE without extraneous errors in the output.

PiperOrigin-RevId: 316718211
Change-Id: Ia2c8b87dbb7b755b65f8c1d390f13eb70c8cf4ee
---
 tensorflow/lite/toco/python/BUILD | 1 +
 tensorflow/python/BUILD           | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/toco/python/BUILD b/tensorflow/lite/toco/python/BUILD
index 7dfa714d1d6..bada1016d26 100644
--- a/tensorflow/lite/toco/python/BUILD
+++ b/tensorflow/lite/toco/python/BUILD
@@ -35,6 +35,7 @@ cc_library(
     ],
     deps = [
         "@com_google_protobuf//:protobuf_headers",
+        "//third_party/python_runtime:headers",  # build_cleaner: keep; DNR: b/35864863
         "//tensorflow/core:lib",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/core/api",
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 87048ba9d40..d141b719aef 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1058,11 +1058,12 @@ cc_library(
         ":safe_ptr",
         "//tensorflow/c:tensor_interface",
         "//tensorflow/c:tf_tensor_internal",
+        "//tensorflow/c/eager:c_api_internal",
         "//tensorflow/c/eager:tfe_context_internal",
         "//tensorflow/c/eager:tfe_tensorhandle_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
-        "//third_party/python_runtime:headers",
+        "//third_party/python_runtime:headers",  # build_cleaner: keep; DNR: b/35864863
     ],
 )
 

From 2d50164bdb08f5c02ffe9dba2c597fcded1f07a1 Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Tue, 16 Jun 2020 11:19:06 -0700
Subject: [PATCH 0307/1390] [TF/XLA] [NFC] Simplify
 XlaComputationLaunchContext::PopulateOutputs

Reduce the nesting level, extract a function for gathering VariableInfo.

PiperOrigin-RevId: 316720004
Change-Id: I49982058d9f7efbc2dcbb2b180c1fc95193cfa39
---
 tensorflow/compiler/jit/xla_launch_util.cc | 234 +++++++++++----------
 tensorflow/compiler/jit/xla_launch_util.h  |   9 +-
 2 files changed, 133 insertions(+), 110 deletions(-)

diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index ec5a372875c..25eed134e35 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -195,18 +195,20 @@ XlaComputationLaunchContext::XlaComputationLaunchContext(
 }
 
 void XlaComputationLaunchContext::PopulateInputs(
-    OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel,
+    OpKernelContext* ctx,
+    const XlaCompiler::CompilationResult* compilation_result,
     const std::map& variables,
     int missing_ctx_input_prefix) {
   // Build ShapedBuffers that point directly to the Tensor buffers.
-  arg_ptrs_ = std::vector(kernel->xla_input_shapes.size());
+  arg_ptrs_ =
+      std::vector(compilation_result->xla_input_shapes.size());
 
   xla::TransferManager* transfer_manager =
       client_->backend().transfer_manager();
-  for (int i = 0; i < kernel->xla_input_shapes.size(); ++i) {
-    int arg_num = kernel->input_mapping[i];
+  for (int i = 0; i < compilation_result->xla_input_shapes.size(); ++i) {
+    int arg_num = compilation_result->input_mapping[i];
     CHECK_GE(arg_num, missing_ctx_input_prefix);
-    const xla::Shape& shape = kernel->xla_input_shapes[i];
+    const xla::Shape& shape = compilation_result->xla_input_shapes[i];
     const Tensor* t = variables.count(arg_num)
                           ? &(variables.at(arg_num).value)
                           : &(ctx->input(arg_num - missing_ctx_input_prefix));
@@ -361,13 +363,94 @@ static Status SetBufferForResourceVarTensorUnderAllocateXlaTensors(
   return Status::OK();
 }
 
+// Sets output `output_num` for `ctx` provided it is known at a compile time.
+static Status SetOutputForConstant(
+    OpKernelContext* ctx, se::Stream* stream,
+    const XlaCompiler::CompilationResult* compilation_result, int output_num) {
+  CHECK(compilation_result->outputs[output_num].is_constant);
+  // Output is a constant.
+  const Tensor& const_tensor =
+      compilation_result->outputs[output_num].constant_value;
+  Tensor* output_tensor;
+  const size_t total_bytes = const_tensor.TotalBytes();
+  if (stream && total_bytes > 0) {
+    // Copy host -> device. (Empty tensors don't have backing buffers.)
+    // Manually allocate memory using an XlaTensorBuffer so we can allocate
+    // as much memory as the device requires (as given by
+    // GetByteSizeRequirement). This avoids XlaTransferManager having to
+    // reallocate the device buffer later.
+    VLOG(1) << "Constant output tensor on device";
+
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_output(output_num, const_tensor.shape(), &output_tensor));
+    Device* device = dynamic_cast(ctx->device());
+    if (device == nullptr) {
+      return errors::Internal("DeviceBase was not a Device.");
+    }
+    ctx->op_device_context()->CopyCPUTensorToDevice(
+        &const_tensor, device, output_tensor,
+        [&](Status status) { TF_CHECK_OK(status); });
+
+    if (device->device_type() == DEVICE_GPU) {
+      // The GPUDeviceContext enqueues the host->device transfer in a
+      // separate stream from the main compute stream. We must ensure the
+      // compute stream is synchronized with the host->device transfer
+      // stream now otherwise we will create a race condition.
+      auto* gpu_device_context =
+          static_cast(ctx->op_device_context());
+      gpu_device_context->stream()->ThenWaitFor(
+          gpu_device_context->host_to_device_stream());
+    }
+  } else {
+    // No copy required.
+    ctx->set_output(output_num, const_tensor);
+    output_tensor = ctx->mutable_output(output_num);
+  }
+  if (XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor)) {
+    xla_tensor->set_host_tensor(const_tensor);
+  }
+  return Status::OK();
+}
+
+// Creates a list of updates resource variables.
+static xla::StatusOr> GatherVariableInfo(
+    OpKernelContext* ctx,
+    const XlaCompiler::CompilationResult* compilation_result,
+    int missing_ctx_input_prefix) {
+  std::vector variable_infos;
+  variable_infos.reserve(compilation_result->resource_updates.size());
+
+  for (int i = 0; i < compilation_result->resource_updates.size(); ++i) {
+    const XlaCompiler::ResourceUpdate& write =
+        compilation_result->resource_updates[i];
+    int actual_input_index = write.input_index - missing_ctx_input_prefix;
+    if (actual_input_index < 0 || actual_input_index >= ctx->num_inputs()) {
+      return errors::Internal("Invalid input index for variable write.");
+    }
+
+    // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor,
+    // not a Tensor.
+    Var* variable = nullptr;
+    TF_RETURN_IF_ERROR(LookupOrCreateResource(
+        ctx, HandleFromInput(ctx, actual_input_index), &variable,
+        [&write](Var** ptr) {
+          *ptr = new Var(write.type);
+          return Status::OK();
+        }));
+    variable_infos.emplace_back(actual_input_index, variable);
+  }
+  return variable_infos;
+}
+
 Status XlaComputationLaunchContext::PopulateOutputs(
-    OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel,
+    OpKernelContext* ctx,
+    const XlaCompiler::CompilationResult* compilation_result,
     ScopedShapedBuffer output, int missing_ctx_input_prefix,
     const xla::HloInputOutputAliasConfig& input_output_alias,
     const std::map& resource_var_snapshots) {
   se::Stream* stream =
       ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr;
+  Allocator* allocator = ctx->device()->GetAllocator({});
 
   // Computation output should always be a tuple.
   if (VLOG_IS_ON(2)) {
@@ -375,7 +458,7 @@ Status XlaComputationLaunchContext::PopulateOutputs(
     VLOG(2) << "Result tuple shape (on device): "
             << output.on_device_shape().DebugString();
   }
-  CHECK_EQ(ctx->num_outputs(), kernel->outputs.size());
+  CHECK_EQ(ctx->num_outputs(), compilation_result->outputs.size());
 
   // If the on-host-shape isn't a tuple, create a new single-element tuple
   // buffer with a nullptr root index table. This allows the code below to treat
@@ -404,82 +487,41 @@ Status XlaComputationLaunchContext::PopulateOutputs(
   // Copy XLA results to the OpOutputList.
   int output_num = 0;
   for (int i = 0; i < ctx->num_outputs(); ++i) {
-    Allocator* allocator = ctx->device()->GetAllocator({});
-    if (kernel->outputs[i].is_constant) {
-      // Output is a constant.
-      const Tensor& const_tensor = kernel->outputs[i].constant_value;
-      Tensor* output_tensor;
-      const size_t total_bytes = const_tensor.TotalBytes();
-      if (stream && total_bytes > 0) {
-        // Copy host -> device. (Empty tensors don't have backing buffers.)
-        // Manually allocate memory using an XlaTensorBuffer so we can allocate
-        // as much memory as the device requires (as given by
-        // GetByteSizeRequirement). This avoids XlaTransferManager having to
-        // reallocate the device buffer later.
-        VLOG(1) << "Constant output tensor on device";
+    const TensorShape& shape = compilation_result->outputs[i].shape;
+    const DataType& type = compilation_result->outputs[i].type;
+    VLOG(2) << "Retval " << i << " shape " << shape.DebugString() << " type "
+            << DataTypeString(type);
+    if (type == DT_VARIANT) {
+      return errors::Unimplemented(
+          "Support for TensorList crossing the XLA/TF boundary "
+          "is not implemented");
+    }
 
-        TF_RETURN_IF_ERROR(
-            ctx->allocate_output(i, const_tensor.shape(), &output_tensor));
-
-        Device* device = dynamic_cast(ctx->device());
-        if (device == nullptr) {
-          return errors::Internal("DeviceBase was not a Device.");
-        }
-        ctx->op_device_context()->CopyCPUTensorToDevice(
-            &const_tensor, device, output_tensor,
-            [&](Status status) { TF_CHECK_OK(status); });
-
-        if (device->device_type() == DEVICE_GPU) {
-          // The GPUDeviceContext enqueues the host->device transfer in a
-          // separate stream from the main compute stream. We must ensure the
-          // compute stream is synchronized with the host->device transfer
-          // stream now otherwise we will create a race condition.
-          auto* gpu_device_context =
-              static_cast(ctx->op_device_context());
-          gpu_device_context->stream()->ThenWaitFor(
-              gpu_device_context->host_to_device_stream());
-        }
-      } else {
-        // No copy required.
-        ctx->set_output(i, const_tensor);
-        output_tensor = ctx->mutable_output(i);
-      }
-      if (XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor)) {
-        xla_tensor->set_host_tensor(const_tensor);
-      }
+    if (compilation_result->outputs[i].is_constant) {
+      TF_RETURN_IF_ERROR(
+          SetOutputForConstant(ctx, stream, compilation_result, i));
+    } else if (type == DT_RESOURCE) {
+      int input_index =
+          compilation_result->outputs[i].input_index - missing_ctx_input_prefix;
+      TF_RET_CHECK(input_index >= 0 && input_index < ctx->num_inputs())
+          << "Invalid input for outputs " << i << ": " << input_index;
+      ctx->set_output(i, ctx->input(input_index));
     } else {
-      const TensorShape& shape = kernel->outputs[i].shape;
-      const DataType& type = kernel->outputs[i].type;
-      VLOG(2) << "Retval " << i << " shape " << shape.DebugString() << " type "
-              << DataTypeString(type);
-      if (type == DT_RESOURCE) {
-        int input_index =
-            kernel->outputs[i].input_index - missing_ctx_input_prefix;
-        TF_RET_CHECK(input_index >= 0 && input_index < ctx->num_inputs())
-            << "Invalid input for outputs " << i << ": " << input_index;
-        ctx->set_output(i, ctx->input(input_index));
-      } else {
-        if (allocate_xla_tensors_) {
-          TF_RETURN_IF_ERROR(SetBufferForTensorUnderAllocateXlaTensors(
-              input_output_alias, output_num, ctx, i, shape, &output,
-              definition_event, stream, use_multiple_streams_));
-        } else {
-          if (type == DT_VARIANT) {
-            return errors::Unimplemented(
-                "Support for TensorList crossing the XLA/TF boundary "
-                "is not implemented");
-          }
+      if (allocate_xla_tensors_) {
+        TF_RETURN_IF_ERROR(SetBufferForTensorUnderAllocateXlaTensors(
+            input_output_alias, output_num, ctx, i, shape, &output,
+            definition_event, stream, use_multiple_streams_));
 
-          se::DeviceMemoryBase buffer = output.buffer({output_num});
-          Tensor output_tensor = GetOrCreateTensorForOutput(
-              output_num, ctx, missing_ctx_input_prefix, input_output_alias,
-              kernel->input_mapping, resource_var_snapshots,
-              ctx->expected_output_dtype(i), shape, buffer, allocator);
-          output.set_buffer(se::OwningDeviceMemory(), {output_num});
-          ctx->set_output(i, output_tensor);
-        }
-        ++output_num;
+      } else {
+        se::DeviceMemoryBase buffer = output.buffer({output_num});
+        Tensor output_tensor = GetOrCreateTensorForOutput(
+            output_num, ctx, missing_ctx_input_prefix, input_output_alias,
+            compilation_result->input_mapping, resource_var_snapshots,
+            ctx->expected_output_dtype(i), shape, buffer, allocator);
+        output.set_buffer(se::OwningDeviceMemory(), {output_num});
+        ctx->set_output(i, output_tensor);
       }
+      ++output_num;
     }
 
     if (VLOG_IS_ON(3)) {
@@ -489,34 +531,14 @@ Status XlaComputationLaunchContext::PopulateOutputs(
 
   // Apply variable updates, if any.
   VLOG(2) << "Applying variable updates";
-  std::vector variable_infos;
-  variable_infos.reserve(kernel->resource_updates.size());
-
-  for (int i = 0; i < kernel->resource_updates.size(); ++i) {
-    const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i];
-    int actual_input_index = write.input_index - missing_ctx_input_prefix;
-    if (actual_input_index < 0 || actual_input_index >= ctx->num_inputs()) {
-      return errors::Internal("Invalid input index for variable write.");
-    }
-
-    // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor,
-    // not a Tensor.
-    Var* variable = nullptr;
-    TF_RETURN_IF_ERROR(LookupOrCreateResource(
-        ctx, HandleFromInput(ctx, actual_input_index), &variable,
-        [&write](Var** ptr) {
-          *ptr = new Var(write.type);
-          return Status::OK();
-        }));
-    variable_infos.emplace_back(actual_input_index, variable);
-  }
-
+  TF_ASSIGN_OR_RETURN(
+      std::vector variable_infos,
+      GatherVariableInfo(ctx, compilation_result, missing_ctx_input_prefix));
   TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos)));
 
-  for (int i = 0; i < kernel->resource_updates.size(); ++i) {
-    Allocator* allocator = ctx->device()->GetAllocator({});
-    const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i];
-
+  for (int i = 0; i < compilation_result->resource_updates.size(); ++i) {
+    const XlaCompiler::ResourceUpdate& write =
+        compilation_result->resource_updates[i];
     if (variable_infos[i].var()->tensor()->dtype() != write.type) {
       return errors::Internal("Mismatched type in variable write");
     }
@@ -530,7 +552,7 @@ Status XlaComputationLaunchContext::PopulateOutputs(
       output.set_buffer(se::OwningDeviceMemory(), {output_num});
       Tensor output_tensor = GetOrCreateTensorForOutput(
           output_num, ctx, missing_ctx_input_prefix, input_output_alias,
-          kernel->input_mapping, resource_var_snapshots, write.type,
+          compilation_result->input_mapping, resource_var_snapshots, write.type,
           write.shape, buffer, allocator);
       *variable_infos[i].var()->tensor() = output_tensor;
       variable_infos[i].var()->is_initialized |= write.modified;
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index cf68dcb7dd6..9a7f20cb310 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -136,7 +136,7 @@ class XlaComputationLaunchContext {
   // input_mapping must be greater than or equal to `missing_ctx_input_prefix`
   // (in other words, no inputs actually required by the kernel can be missing).
   void PopulateInputs(OpKernelContext* ctx,
-                      const XlaCompiler::CompilationResult* kernel,
+                      const XlaCompiler::CompilationResult* compilation_result,
                       const std::map& variables,
                       int missing_ctx_input_prefix);
 
@@ -148,10 +148,11 @@ class XlaComputationLaunchContext {
   // See jit/resource_operation_safety_analysis for details.
   //
   //
-  // Assumes that the first `missing_ctx_input_prefix` inputs to the kernel are
-  // missing and adjusts input indices accordingly.
+  // Assumes that the first `missing_ctx_input_prefix` inputs to the
+  // compilation_result are missing and adjusts input indices accordingly.
   Status PopulateOutputs(
-      OpKernelContext* ctx, const XlaCompiler::CompilationResult* kernel,
+      OpKernelContext* ctx,
+      const XlaCompiler::CompilationResult* compilation_result,
       xla::ScopedShapedBuffer output, int missing_ctx_input_prefix,
       const xla::HloInputOutputAliasConfig& input_output_alias,
       const std::map& resource_var_snapshots);

From ed557008d681b6bd612f0ea6b1aa056c4cfa744b Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Tue, 16 Jun 2020 11:19:41 -0700
Subject: [PATCH 0308/1390] All LSTM implementations: Rename cell_scratch to
 cell_gate_scratch, and cell_bias_ptr to cell_gate_bias_ptr to better reflect
 what those arrays are.

Do note this is not the same thing as the LSTM cell "state", but a layer/gate that calculates the update. The cell state depends on the input, forget, and cell gates; these arrays are the output and the bias for the last gate.

PiperOrigin-RevId: 316720132
Change-Id: I71c370dabd27f776987e061b9393022c775589c9
---
 tensorflow/lite/kernels/lstm_eval.cc          | 156 +++++++++---------
 .../calibration/builtin_logging_ops/lstm.cc   |  61 +++----
 2 files changed, 114 insertions(+), 103 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index b285ed1030f..b4d43414d89 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -212,13 +212,13 @@ inline void LstmStepFloat(
     const float* cell_layer_norm_coefficients_ptr,
     const float* output_layer_norm_coefficients_ptr,
     const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
-    const float* cell_bias_ptr, const float* output_gate_bias_ptr,
+    const float* cell_gate_bias_ptr, const float* output_gate_bias_ptr,
     const float* projection_weights_ptr, const float* projection_bias_ptr,
     const TfLiteLSTMParams* params, int n_batch, int n_cell, int n_input,
     int n_aux_input, int n_output, int output_batch_leading_dim,
     float* output_state_ptr, float* cell_state_ptr, float* input_gate_scratch,
-    float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch,
-    float* output_ptr) {
+    float* forget_gate_scratch, float* cell_gate_scratch,
+    float* output_gate_scratch, float* output_ptr) {
   ruy::profiler::ScopeLabel label("LstmStepFloat");
   // Since we have already checked that weights are all there or none, we can
   // check the existence of only one to the get the condition.
@@ -233,7 +233,7 @@ inline void LstmStepFloat(
       std::fill_n(input_gate_scratch, n_cell * n_batch, 0.0f);
     }
     std::fill_n(forget_gate_scratch, n_cell * n_batch, 0.0f);
-    std::fill_n(cell_scratch, n_cell * n_batch, 0.0f);
+    std::fill_n(cell_gate_scratch, n_cell * n_batch, 0.0f);
     std::fill_n(output_gate_scratch, n_cell * n_batch, 0.0f);
   } else {
     if (!use_cifg) {
@@ -242,8 +242,8 @@ inline void LstmStepFloat(
     }
     tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch,
                                           forget_gate_scratch);
-    tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch,
-                                          cell_scratch);
+    tensor_utils::VectorBatchVectorAssign(cell_gate_bias_ptr, n_cell, n_batch,
+                                          cell_gate_scratch);
     tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch,
                                           output_gate_scratch);
   }
@@ -262,7 +262,7 @@ inline void LstmStepFloat(
         forget_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         input_to_cell_weights_ptr, n_cell, n_input, input_ptr, n_batch,
-        cell_scratch);
+        cell_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         input_to_output_weights_ptr, n_cell, n_input, input_ptr, n_batch,
         output_gate_scratch);
@@ -283,7 +283,7 @@ inline void LstmStepFloat(
         n_batch, forget_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_cell_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
-        n_batch, cell_scratch);
+        n_batch, cell_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_output_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
         n_batch, output_gate_scratch);
@@ -300,7 +300,7 @@ inline void LstmStepFloat(
       n_batch, forget_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
       recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr,
-      n_batch, cell_scratch);
+      n_batch, cell_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
       recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr,
       n_batch, output_gate_scratch);
@@ -347,24 +347,26 @@ inline void LstmStepFloat(
   tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr,
                                          n_batch * n_cell, cell_state_ptr);
   if (use_layer_norm) {
-    tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell,
-                                          n_batch);
+    tensor_utils::MeanStddevNormalization(cell_gate_scratch, cell_gate_scratch,
+                                          n_cell, n_batch);
     tensor_utils::VectorBatchVectorCwiseProduct(
-        cell_layer_norm_coefficients_ptr, n_cell, cell_scratch, n_batch,
-        cell_scratch);
-    tensor_utils::VectorBatchVectorAdd(cell_bias_ptr, n_cell, n_batch,
-                                       cell_scratch);
+        cell_layer_norm_coefficients_ptr, n_cell, cell_gate_scratch, n_batch,
+        cell_gate_scratch);
+    tensor_utils::VectorBatchVectorAdd(cell_gate_bias_ptr, n_cell, n_batch,
+                                       cell_gate_scratch);
   }
-  tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell,
-                                        params->activation, cell_scratch);
+  tensor_utils::ApplyActivationToVector(cell_gate_scratch, n_batch * n_cell,
+                                        params->activation, cell_gate_scratch);
   if (use_cifg) {
     tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell,
                              forget_gate_scratch);
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, forget_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   } else {
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, input_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   }
   if (params->cell_clip > 0.0) {
     tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell,
@@ -389,8 +391,8 @@ inline void LstmStepFloat(
   tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell,
                                      output_gate_scratch);
   tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell,
-                                        params->activation, cell_scratch);
-  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch,
+                                        params->activation, cell_gate_scratch);
+  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_gate_scratch,
                                          n_batch * n_cell, output_gate_scratch);
 
   const bool use_projection_weight = (projection_weights_ptr != nullptr);
@@ -525,19 +527,19 @@ inline void LstmStepHybrid(
     const float* cell_layer_norm_coefficients_ptr,
     const float* output_layer_norm_coefficients_ptr,
     const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
-    const float* cell_bias_ptr, const float* output_gate_bias_ptr,
+    const float* cell_gate_bias_ptr, const float* output_gate_bias_ptr,
     const int8_t* projection_weights_ptr, float projection_weights_scale,
     const float* projection_bias_ptr, const TfLiteLSTMParams* params,
     int n_batch, int n_cell, int n_input, int n_aux_input, int n_output,
     int output_batch_leading_dim, float* input_gate_scratch,
-    float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch,
-    float* scaling_factors, float* scaling_factors_scratch,
-    float* recovered_cell_weights, int8_t* quantized_input_ptr,
-    int8_t* quantized_aux_input_ptr, int8_t* quantized_output_state_ptr,
-    int8_t* quantized_cell_state_ptr, float* output_state_ptr,
-    float* cell_state_ptr, int32_t* accum_scratch_ptr, float* output_ptr,
-    int32_t* zero_points, int32_t* row_sums, int row_sums_size,
-    bool* compute_row_sums, bool asymmetric_quantize_inputs,
+    float* forget_gate_scratch, float* cell_gate_scratch,
+    float* output_gate_scratch, float* scaling_factors,
+    float* scaling_factors_scratch, float* recovered_cell_weights,
+    int8_t* quantized_input_ptr, int8_t* quantized_aux_input_ptr,
+    int8_t* quantized_output_state_ptr, int8_t* quantized_cell_state_ptr,
+    float* output_state_ptr, float* cell_state_ptr, int32_t* accum_scratch_ptr,
+    float* output_ptr, int32_t* zero_points, int32_t* row_sums,
+    int row_sums_size, bool* compute_row_sums, bool asymmetric_quantize_inputs,
     CpuBackendContext* context) {
   ruy::profiler::ScopeLabel label("LstmStepHybrid");
   // Since we have already checked that weights are all there or none, we
@@ -553,7 +555,7 @@ inline void LstmStepHybrid(
       std::fill_n(input_gate_scratch, n_cell * n_batch, 0.0f);
     }
     std::fill_n(forget_gate_scratch, n_cell * n_batch, 0.0f);
-    std::fill_n(cell_scratch, n_cell * n_batch, 0.0f);
+    std::fill_n(cell_gate_scratch, n_cell * n_batch, 0.0f);
     std::fill_n(output_gate_scratch, n_cell * n_batch, 0.0f);
   } else {
     if (!use_cifg) {
@@ -562,8 +564,8 @@ inline void LstmStepHybrid(
     }
     tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch,
                                           forget_gate_scratch);
-    tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch,
-                                          cell_scratch);
+    tensor_utils::VectorBatchVectorAssign(cell_gate_bias_ptr, n_cell, n_batch,
+                                          cell_gate_scratch);
     tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch,
                                           output_gate_scratch);
   }
@@ -657,7 +659,8 @@ inline void LstmStepHybrid(
 
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         input_to_cell_weights_ptr, n_cell, n_input, quantized_input_ptr,
-        input_to_cell_weights_scale, scaling_factors, n_batch, cell_scratch,
+        input_to_cell_weights_scale, scaling_factors, n_batch,
+        cell_gate_scratch,
         /*per_channel_scale=*/nullptr, zero_points, accum_scratch_ptr,
         input_to_cell_row_sums, compute_row_sums, scaling_factors_scratch,
         context);
@@ -699,9 +702,10 @@ inline void LstmStepHybrid(
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_cell_weights_ptr, n_cell, n_aux_input,
         quantized_aux_input_ptr, aux_input_to_cell_weights_scale,
-        scaling_factors, n_batch, cell_scratch, /*per_channel_scale=*/nullptr,
-        zero_points, accum_scratch_ptr, aux_input_to_cell_row_sums,
-        compute_row_sums, scaling_factors_scratch, context);
+        scaling_factors, n_batch, cell_gate_scratch,
+        /*per_channel_scale=*/nullptr, zero_points, accum_scratch_ptr,
+        aux_input_to_cell_row_sums, compute_row_sums, scaling_factors_scratch,
+        context);
 
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_output_weights_ptr, n_cell, n_aux_input,
@@ -739,9 +743,10 @@ inline void LstmStepHybrid(
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         recurrent_to_cell_weights_ptr, n_cell, n_output,
         quantized_output_state_ptr, recurrent_to_cell_weights_scale,
-        scaling_factors, n_batch, cell_scratch, /*per_channel_scale=*/nullptr,
-        zero_points, accum_scratch_ptr, recurrent_to_cell_row_sums,
-        compute_row_sums, scaling_factors_scratch, context);
+        scaling_factors, n_batch, cell_gate_scratch,
+        /*per_channel_scale=*/nullptr, zero_points, accum_scratch_ptr,
+        recurrent_to_cell_row_sums, compute_row_sums, scaling_factors_scratch,
+        context);
 
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         recurrent_to_output_weights_ptr, n_cell, n_output,
@@ -800,24 +805,26 @@ inline void LstmStepHybrid(
   tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr,
                                          n_batch * n_cell, cell_state_ptr);
   if (use_layer_norm) {
-    tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell,
-                                          n_batch);
+    tensor_utils::MeanStddevNormalization(cell_gate_scratch, cell_gate_scratch,
+                                          n_cell, n_batch);
     tensor_utils::VectorBatchVectorCwiseProduct(
-        cell_layer_norm_coefficients_ptr, n_cell, cell_scratch, n_batch,
-        cell_scratch);
-    tensor_utils::VectorBatchVectorAdd(cell_bias_ptr, n_cell, n_batch,
-                                       cell_scratch);
+        cell_layer_norm_coefficients_ptr, n_cell, cell_gate_scratch, n_batch,
+        cell_gate_scratch);
+    tensor_utils::VectorBatchVectorAdd(cell_gate_bias_ptr, n_cell, n_batch,
+                                       cell_gate_scratch);
   }
-  tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell,
-                                        params->activation, cell_scratch);
+  tensor_utils::ApplyActivationToVector(cell_gate_scratch, n_batch * n_cell,
+                                        params->activation, cell_gate_scratch);
   if (use_cifg) {
     tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell,
                              forget_gate_scratch);
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, forget_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   } else {
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, input_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   }
   if (params->cell_clip > 0.0) {
     tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell,
@@ -845,8 +852,8 @@ inline void LstmStepHybrid(
   tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell,
                                      output_gate_scratch);
   tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell,
-                                        params->activation, cell_scratch);
-  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch,
+                                        params->activation, cell_gate_scratch);
+  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_gate_scratch,
                                          n_batch * n_cell, output_gate_scratch);
 
   const bool use_projection_weight = (projection_weights_ptr != nullptr);
@@ -940,7 +947,7 @@ inline void LstmStepHybrid(
 // Gate biases of size 'n_cell':
 //   input_bias_ptr                 - optional
 //   forget_bias_ptr
-//   cell_bias_ptr
+//   cell_gate_bias_ptr
 //   output_bias_ptr
 //
 // Layer norm coefficients of size 'n_cell', representing diagonal matrices.
@@ -1028,7 +1035,7 @@ inline void LstmStepInteger(
     const int16_t* layer_norm_output_weight_ptr,
     int32_t layer_norm_output_scale_a, int32_t layer_norm_output_scale_b,
     const int32_t* input_bias_ptr, const int32_t* forget_bias_ptr,
-    const int32_t* cell_bias_ptr, const int32_t* output_bias_ptr,
+    const int32_t* cell_gate_bias_ptr, const int32_t* output_bias_ptr,
     int16_t quantized_cell_clip, int8_t quantized_proj_clip, int32_t cell_scale,
     int32_t input_variance_guard, int32_t forget_variance_guard,
     int32_t cell_variance_guard, int32_t output_variance_guard,
@@ -1115,7 +1122,7 @@ inline void LstmStepInteger(
 
   if (use_layer_norm) {
     tensor_utils::ApplyLayerNorm(scratch_2_ptr, layer_norm_cell_weight_ptr,
-                                 cell_bias_ptr, layer_norm_cell_scale_a,
+                                 cell_gate_bias_ptr, layer_norm_cell_scale_a,
                                  layer_norm_cell_scale_b, cell_variance_guard,
                                  n_batch, n_cell, scratch_2_ptr);
   }
@@ -1266,7 +1273,7 @@ inline void LstmStepInteger(
 // Gate biases of size 'n_cell':
 //   input_bias_ptr                 - optional
 //   forget_bias_ptr
-//   cell_bias_ptr
+//   cell_gate_bias_ptr
 //   output_bias_ptr
 //
 // Layer norm coefficients of size 'n_cell', representing diagonal matrices.
@@ -1355,7 +1362,7 @@ void LstmStepInteger(
     const int16_t* layer_norm_output_weight_ptr,
     int32_t layer_norm_output_scale_a, int32_t layer_norm_output_scale_b,
     const int32_t* input_bias_ptr, const int32_t* forget_bias_ptr,
-    const int32_t* cell_bias_ptr, const int32_t* output_bias_ptr,
+    const int32_t* cell_gate_bias_ptr, const int32_t* output_bias_ptr,
     const int32_t* proj_bias_ptr, const TfLiteLSTMParams* params,
     const int32_t* intermediate_scale_a, const int32_t* intermediate_scale_b,
     const int32_t* intermediate_zp, int32 quantized_cell_clip,
@@ -1413,7 +1420,7 @@ void LstmStepInteger(
   // Update gate with layer norm.
   tensor_utils::ApplyLayerNormFloat(
       scratch3, layer_norm_cell_weight_ptr, layer_norm_cell_scale_a,
-      layer_norm_cell_scale_b, cell_bias_ptr, n_batch, n_cell, scratch3);
+      layer_norm_cell_scale_b, cell_gate_bias_ptr, n_batch, n_cell, scratch3);
 
   // Update gate tanh.
   tensor_utils::ApplyTanhFloat(scratch3, n_batch, n_cell, -12, scratch3);
@@ -1538,16 +1545,16 @@ TfLiteStatus EvalFloat(
   // Index the scratch buffers pointers to the global scratch buffer.
   float* scratch_buffer_ptr = GetTensorData(scratch_buffer);
   float* input_gate_scratch = nullptr;
-  float* cell_scratch = nullptr;
+  float* cell_gate_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
   float* output_gate_scratch = nullptr;
   if (use_cifg) {
-    cell_scratch = scratch_buffer_ptr;
+    cell_gate_scratch = scratch_buffer_ptr;
     forget_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
   } else {
     input_gate_scratch = scratch_buffer_ptr;
-    cell_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    cell_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     forget_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 3 * n_cell * n_batch;
   }
@@ -1599,7 +1606,8 @@ TfLiteStatus EvalFloat(
           n_input, aux_input_size, n_output, output_batch_leading_dim,
           GetTensorData(activation_state),
           GetTensorData(cell_state), input_gate_scratch,
-          forget_gate_scratch, cell_scratch, output_gate_scratch, output_ptr);
+          forget_gate_scratch, cell_gate_scratch, output_gate_scratch,
+          output_ptr);
     }
   } else {
     for (int b = 0; b < n_batch; b++) {
@@ -1628,7 +1636,7 @@ TfLiteStatus EvalFloat(
         float* input_gate_scratch_ptr =
             input_gate_scratch ? input_gate_scratch + b * n_cell : nullptr;
         float* forget_gate_scratch_ptr = forget_gate_scratch + b * n_cell;
-        float* cell_scratch_ptr = cell_scratch + b * n_cell;
+        float* cell_gate_scratch_ptr = cell_gate_scratch + b * n_cell;
         float* output_gate_scratch_ptr = output_gate_scratch + b * n_cell;
 
         LstmStepFloat(
@@ -1659,8 +1667,8 @@ TfLiteStatus EvalFloat(
             GetTensorData(projection_bias), params, /*n_batch=*/1,
             n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
             activation_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
-            forget_gate_scratch_ptr, cell_scratch_ptr, output_gate_scratch_ptr,
-            output_ptr);
+            forget_gate_scratch_ptr, cell_gate_scratch_ptr,
+            output_gate_scratch_ptr, output_ptr);
       }
     }
   }
@@ -1723,16 +1731,16 @@ TfLiteStatus EvalHybrid(
 
   float* scratch_buffer_ptr = GetTensorData(scratch_buffer);
   float* input_gate_scratch = nullptr;
-  float* cell_scratch = nullptr;
+  float* cell_gate_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
   float* output_gate_scratch = nullptr;
   if (use_cifg) {
-    cell_scratch = scratch_buffer_ptr;
+    cell_gate_scratch = scratch_buffer_ptr;
     forget_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
   } else {
     input_gate_scratch = scratch_buffer_ptr;
-    cell_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    cell_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     forget_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 3 * n_cell * n_batch;
   }
@@ -1805,7 +1813,7 @@ TfLiteStatus EvalHybrid(
           GetTensorScale(projection_weights),
           GetTensorData(projection_bias), params, n_batch, n_cell,
           n_input, aux_input_size, n_output, output_batch_leading_dim,
-          input_gate_scratch, forget_gate_scratch, cell_scratch,
+          input_gate_scratch, forget_gate_scratch, cell_gate_scratch,
           output_gate_scratch, GetTensorData(scaling_factors),
           GetTensorData(prod_scaling_factors),
           GetTensorData(recovered_cell_weights),
@@ -1845,7 +1853,7 @@ TfLiteStatus EvalHybrid(
         float* input_gate_scratch_ptr =
             input_gate_scratch ? input_gate_scratch + b * n_cell : nullptr;
         float* forget_gate_scratch_ptr = forget_gate_scratch + b * n_cell;
-        float* cell_scratch_ptr = cell_scratch + b * n_cell;
+        float* cell_gate_scratch_ptr = cell_gate_scratch + b * n_cell;
         float* output_gate_scratch_ptr = output_gate_scratch + b * n_cell;
 
         LstmStepHybrid(
@@ -1892,8 +1900,8 @@ TfLiteStatus EvalHybrid(
             GetTensorData(projection_bias), params,
             /*n_batch=*/1, n_cell, n_input, aux_input_size, n_output,
             output_batch_leading_dim, input_gate_scratch_ptr,
-            forget_gate_scratch_ptr, cell_scratch_ptr, output_gate_scratch_ptr,
-            GetTensorData(scaling_factors),
+            forget_gate_scratch_ptr, cell_gate_scratch_ptr,
+            output_gate_scratch_ptr, GetTensorData(scaling_factors),
             GetTensorData(prod_scaling_factors),
             GetTensorData(recovered_cell_weights),
             GetTensorData(input_quantized),
@@ -2119,7 +2127,7 @@ TfLiteStatus EvalInteger8x8_8(
       GetTensorData(output_layer_norm_coefficients);
   const int32_t* input_bias_ptr = GetTensorData(input_gate_bias);
   const int32_t* forget_bias_ptr = GetTensorData(forget_gate_bias);
-  const int32_t* cell_bias_ptr = GetTensorData(cell_bias);
+  const int32_t* cell_gate_bias_ptr = GetTensorData(cell_bias);
   const int32_t* output_bias_ptr = GetTensorData(output_gate_bias);
   const int32_t* proj_bias_ptr = GetTensorData(projection_bias);
   int16_t* cell_ptr = GetTensorData(cell_state);
@@ -2206,7 +2214,7 @@ TfLiteStatus EvalInteger8x8_8(
         integer_lstm_param->layer_norm_output_scale_a,
         integer_lstm_param->layer_norm_output_scale_b,
 
-        input_bias_ptr, forget_bias_ptr, cell_bias_ptr, output_bias_ptr,
+        input_bias_ptr, forget_bias_ptr, cell_gate_bias_ptr, output_bias_ptr,
         proj_bias_ptr,
 
         params, integer_lstm_param->intermediate_scale_a,
diff --git a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
index b58900c0bc6..0d4c614511d 100644
--- a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
+++ b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
@@ -58,13 +58,13 @@ inline void LstmStepWithAuxInput(
     const float* cell_layer_norm_coefficients_ptr,
     const float* output_layer_norm_coefficients_ptr,
     const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
-    const float* cell_bias_ptr, const float* output_gate_bias_ptr,
+    const float* cell_gate_bias_ptr, const float* output_gate_bias_ptr,
     const float* projection_weights_ptr, const float* projection_bias_ptr,
     const TfLiteLSTMParams* params, int n_batch, int n_cell, int n_input,
     int n_aux_input, int n_output, int output_batch_leading_dim,
     float* output_state_ptr, float* cell_state_ptr, float* input_gate_scratch,
-    float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch,
-    float* output_ptr, Logger* logger,
+    float* forget_gate_scratch, float* cell_gate_scratch,
+    float* output_gate_scratch, float* output_ptr, Logger* logger,
     const std::vector& intermediate_tensor_indexes,
     ErrorReporter* error_reporter) {
   // Since we have already checked that weights are all there or none, we can
@@ -80,7 +80,7 @@ inline void LstmStepWithAuxInput(
       std::fill_n(input_gate_scratch, n_cell * n_batch, 0.0f);
     }
     std::fill_n(forget_gate_scratch, n_cell * n_batch, 0.0f);
-    std::fill_n(cell_scratch, n_cell * n_batch, 0.0f);
+    std::fill_n(cell_gate_scratch, n_cell * n_batch, 0.0f);
     std::fill_n(output_gate_scratch, n_cell * n_batch, 0.0f);
   } else {
     if (!use_cifg) {
@@ -89,8 +89,8 @@ inline void LstmStepWithAuxInput(
     }
     tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch,
                                           forget_gate_scratch);
-    tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch,
-                                          cell_scratch);
+    tensor_utils::VectorBatchVectorAssign(cell_gate_bias_ptr, n_cell, n_batch,
+                                          cell_gate_scratch);
     tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch,
                                           output_gate_scratch);
   }
@@ -107,7 +107,7 @@ inline void LstmStepWithAuxInput(
       forget_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(input_to_cell_weights_ptr,
                                                     n_cell, n_input, input_ptr,
-                                                    n_batch, cell_scratch);
+                                                    n_batch, cell_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
       input_to_output_weights_ptr, n_cell, n_input, input_ptr, n_batch,
       output_gate_scratch);
@@ -125,7 +125,7 @@ inline void LstmStepWithAuxInput(
         n_batch, forget_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_cell_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
-        n_batch, cell_scratch);
+        n_batch, cell_gate_scratch);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         aux_input_to_output_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
         n_batch, output_gate_scratch);
@@ -142,7 +142,7 @@ inline void LstmStepWithAuxInput(
       n_batch, forget_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
       recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr,
-      n_batch, cell_scratch);
+      n_batch, cell_gate_scratch);
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
       recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr,
       n_batch, output_gate_scratch);
@@ -193,26 +193,28 @@ inline void LstmStepWithAuxInput(
   tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr,
                                          n_batch * n_cell, cell_state_ptr);
   if (use_layer_norm) {
-    logger->LogTensorValue(intermediate_tensor_indexes[2], cell_scratch,
+    logger->LogTensorValue(intermediate_tensor_indexes[2], cell_gate_scratch,
                            n_cell * n_batch, error_reporter);
-    tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell,
-                                          n_batch);
+    tensor_utils::MeanStddevNormalization(cell_gate_scratch, cell_gate_scratch,
+                                          n_cell, n_batch);
     tensor_utils::VectorBatchVectorCwiseProduct(
-        cell_layer_norm_coefficients_ptr, n_cell, cell_scratch, n_batch,
-        cell_scratch);
-    tensor_utils::VectorBatchVectorAdd(cell_bias_ptr, n_cell, n_batch,
-                                       cell_scratch);
+        cell_layer_norm_coefficients_ptr, n_cell, cell_gate_scratch, n_batch,
+        cell_gate_scratch);
+    tensor_utils::VectorBatchVectorAdd(cell_gate_bias_ptr, n_cell, n_batch,
+                                       cell_gate_scratch);
   }
-  tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell,
-                                        params->activation, cell_scratch);
+  tensor_utils::ApplyActivationToVector(cell_gate_scratch, n_batch * n_cell,
+                                        params->activation, cell_gate_scratch);
   if (use_cifg) {
     tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell,
                              forget_gate_scratch);
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, forget_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   } else {
     tensor_utils::VectorVectorCwiseProductAccumulate(
-        cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr);
+        cell_gate_scratch, input_gate_scratch, n_batch * n_cell,
+        cell_state_ptr);
   }
   if (params->cell_clip > 0.0) {
     tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell,
@@ -239,8 +241,8 @@ inline void LstmStepWithAuxInput(
   tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell,
                                      output_gate_scratch);
   tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell,
-                                        params->activation, cell_scratch);
-  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch,
+                                        params->activation, cell_gate_scratch);
+  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_gate_scratch,
                                          n_batch * n_cell, output_gate_scratch);
 
   logger->LogTensorValue(intermediate_tensor_indexes[4], output_gate_scratch,
@@ -329,16 +331,16 @@ TfLiteStatus EvalFloat(
   // Index the scratch buffers pointers to the global scratch buffer.
   float* scratch_buffer_ptr = GetTensorData(scratch_buffer);
   float* input_gate_scratch = nullptr;
-  float* cell_scratch = nullptr;
+  float* cell_gate_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
   float* output_gate_scratch = nullptr;
   if (use_cifg) {
-    cell_scratch = scratch_buffer_ptr;
+    cell_gate_scratch = scratch_buffer_ptr;
     forget_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
   } else {
     input_gate_scratch = scratch_buffer_ptr;
-    cell_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    cell_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
     forget_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
     output_gate_scratch = scratch_buffer_ptr + 3 * n_cell * n_batch;
   }
@@ -390,7 +392,7 @@ TfLiteStatus EvalFloat(
           n_input, aux_input_size, n_output, output_batch_leading_dim,
           GetTensorData(activation_state),
           GetTensorData(cell_state), input_gate_scratch,
-          forget_gate_scratch, cell_scratch, output_gate_scratch,
+          forget_gate_scratch, cell_gate_scratch, output_gate_scratch,
           output_ptr_time, logger, intermediate_tensor_indexes, error_reporter);
     }
   } else {
@@ -420,7 +422,7 @@ TfLiteStatus EvalFloat(
         float* input_gate_scratch_ptr =
             input_gate_scratch ? input_gate_scratch + b * n_cell : nullptr;
         float* forget_gate_scratch_ptr = forget_gate_scratch + b * n_cell;
-        float* cell_scratch_ptr = cell_scratch + b * n_cell;
+        float* cell_gate_scratch_ptr = cell_gate_scratch + b * n_cell;
         float* output_gate_scratch_ptr = output_gate_scratch + b * n_cell;
 
         LstmStepWithAuxInput(
@@ -451,8 +453,9 @@ TfLiteStatus EvalFloat(
             GetTensorData(projection_bias), params, /*n_batch=*/1,
             n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
             activation_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
-            forget_gate_scratch_ptr, cell_scratch_ptr, output_gate_scratch_ptr,
-            output_ptr, logger, intermediate_tensor_indexes, error_reporter);
+            forget_gate_scratch_ptr, cell_gate_scratch_ptr,
+            output_gate_scratch_ptr, output_ptr, logger,
+            intermediate_tensor_indexes, error_reporter);
       }
     }
   }

From 430b00361b76827c055c63fb6398a520b25ed770 Mon Sep 17 00:00:00 2001
From: Sachin Joglekar 
Date: Tue, 16 Jun 2020 11:20:59 -0700
Subject: [PATCH 0309/1390] Audit and improve TfLiteType checks in kernels

PiperOrigin-RevId: 316720436
Change-Id: I2032e799ee6afa533b932385c2a70f7621f4ac1b
---
 tensorflow/lite/c/common.h                    |  1 +
 tensorflow/lite/kernels/activations.cc        | 16 ++--
 tensorflow/lite/kernels/add.cc                |  2 +-
 tensorflow/lite/kernels/add_n.cc              |  2 +-
 tensorflow/lite/kernels/audio_spectrogram.cc  |  4 +-
 tensorflow/lite/kernels/basic_rnn.cc          |  9 +-
 tensorflow/lite/kernels/batch_matmul.cc       |  2 +-
 .../kernels/bidirectional_sequence_lstm.cc    | 54 ++++++------
 .../kernels/bidirectional_sequence_rnn.cc     |  2 +-
 tensorflow/lite/kernels/ceil.cc               |  2 +-
 tensorflow/lite/kernels/concatenation.cc      |  2 +-
 tensorflow/lite/kernels/conv.cc               | 16 ++--
 tensorflow/lite/kernels/depth_to_space.cc     |  2 +-
 tensorflow/lite/kernels/depthwise_conv.cc     | 14 ++--
 tensorflow/lite/kernels/div.cc                |  2 +-
 tensorflow/lite/kernels/elementwise.cc        |  4 +-
 .../lite/kernels/embedding_lookup_sparse.cc   |  2 +-
 tensorflow/lite/kernels/floor.cc              |  2 +-
 tensorflow/lite/kernels/floor_div.cc          |  2 +-
 tensorflow/lite/kernels/fully_connected.cc    | 16 ++--
 tensorflow/lite/kernels/if.cc                 |  2 +-
 tensorflow/lite/kernels/l2norm.cc             |  6 +-
 .../lite/kernels/local_response_norm.cc       |  4 +-
 tensorflow/lite/kernels/logical.cc            |  2 +-
 tensorflow/lite/kernels/lstm.cc               | 82 +++++++++----------
 tensorflow/lite/kernels/maximum_minimum.cc    |  3 +-
 tensorflow/lite/kernels/mfcc.cc               |  6 +-
 tensorflow/lite/kernels/mul.cc                |  2 +-
 tensorflow/lite/kernels/one_hot.cc            |  9 +-
 tensorflow/lite/kernels/pack.cc               |  4 +-
 tensorflow/lite/kernels/pad.cc                | 12 +--
 tensorflow/lite/kernels/pooling.cc            | 12 +--
 tensorflow/lite/kernels/pow.cc                |  5 +-
 tensorflow/lite/kernels/range.cc              |  4 +-
 tensorflow/lite/kernels/read_variable.cc      |  2 +-
 tensorflow/lite/kernels/reduce.cc             |  2 +-
 .../lite/kernels/resize_nearest_neighbor.cc   |  8 +-
 tensorflow/lite/kernels/reverse.cc            |  2 +-
 tensorflow/lite/kernels/reverse_sequence.cc   |  2 +-
 tensorflow/lite/kernels/round.cc              |  2 +-
 tensorflow/lite/kernels/select.cc             |  4 +-
 tensorflow/lite/kernels/skip_gram.cc          |  6 +-
 tensorflow/lite/kernels/space_to_batch_nd.cc  |  3 +-
 tensorflow/lite/kernels/space_to_depth.cc     |  2 +-
 tensorflow/lite/kernels/sparse_to_dense.cc    | 14 ++--
 tensorflow/lite/kernels/squared_difference.cc |  2 +-
 tensorflow/lite/kernels/strided_slice.cc      | 14 ++--
 tensorflow/lite/kernels/sub.cc                |  6 +-
 tensorflow/lite/kernels/tile.cc               |  2 +-
 tensorflow/lite/kernels/topk_v2.cc            | 11 ++-
 tensorflow/lite/kernels/transpose.cc          |  9 +-
 tensorflow/lite/kernels/transpose_conv.cc     | 16 ++--
 .../kernels/unidirectional_sequence_lstm.cc   | 22 ++---
 .../kernels/unidirectional_sequence_rnn.cc    |  9 +-
 tensorflow/lite/kernels/unpack.cc             |  2 +-
 tensorflow/lite/kernels/while.cc              |  4 +-
 .../micro/kernels/arc_mli/fully_connected.cc  |  2 +-
 tensorflow/lite/micro/kernels/ceil.cc         |  4 +-
 .../lite/micro/kernels/circular_buffer.cc     |  4 +-
 .../micro/kernels/cmsis-nn/fully_connected.cc |  2 +-
 tensorflow/lite/micro/kernels/cmsis-nn/mul.cc |  2 +-
 tensorflow/lite/micro/kernels/elementwise.cc  |  4 +-
 tensorflow/lite/micro/kernels/floor.cc        |  2 +-
 .../lite/micro/kernels/fully_connected.cc     |  2 +-
 tensorflow/lite/micro/kernels/l2norm.cc       |  6 +-
 tensorflow/lite/micro/kernels/logistic.cc     |  2 +-
 tensorflow/lite/micro/kernels/mul.cc          |  2 +-
 tensorflow/lite/micro/kernels/reshape.cc      |  2 +-
 tensorflow/lite/micro/kernels/round.cc        |  4 +-
 tensorflow/lite/micro/kernels/svdf.cc         |  4 +-
 tensorflow/lite/micro/kernels/tanh.cc         |  2 +-
 .../lite/micro/kernels/xtensa_hifi/floor.cc   |  2 +-
 .../lite/micro/kernels/xtensa_hifi/svdf.cc    |  4 +-
 .../kernels/xtensa_hifimini_legacy/svdf.cc    |  2 +-
 .../benchmark/experimental/c/c_api_types.h    |  1 +
 75 files changed, 258 insertions(+), 248 deletions(-)

diff --git a/tensorflow/lite/c/common.h b/tensorflow/lite/c/common.h
index ab769fec249..15823784d12 100644
--- a/tensorflow/lite/c/common.h
+++ b/tensorflow/lite/c/common.h
@@ -205,6 +205,7 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
 // the current function, while also reporting the location of the error.
 // `a` and `b` may be evaluated more than once, so no side effects or
 // extremely expensive computations should be done.
+// NOTE: Use TF_LITE_ENSURE_TYPES_EQ if comparing TfLiteTypes.
 #define TF_LITE_ENSURE_EQ(context, a, b)                                   \
   do {                                                                     \
     if ((a) != (b)) {                                                      \
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 2b2428f3f92..7ad33973b38 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -254,7 +254,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   return context->ResizeTensor(context, output,
                                TfLiteIntArrayCopy(input->dims));
@@ -274,7 +274,7 @@ TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (input->type == kTfLiteInt8 || input->type == kTfLiteUInt8) {
     double real_multiplier = input->params.scale / output->params.scale;
@@ -355,7 +355,7 @@ TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   LeakyReluOpData* data = reinterpret_cast(node->user_data);
 
@@ -384,7 +384,7 @@ TfLiteStatus TanhPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (kernel_type == kFixedPointOptimized) {
     if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
@@ -469,7 +469,7 @@ TfLiteStatus SigmoidPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (kernel_type == kFixedPointOptimized) {
     if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
@@ -569,7 +569,7 @@ TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
                                 input->type == kTfLiteUInt8 ||
                                 input->type == kTfLiteInt16);
   } else {
-    TF_LITE_ENSURE_EQ(context, input->type, output->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   }
 
   TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
@@ -632,7 +632,7 @@ TfLiteStatus LogSoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
     TF_LITE_ENSURE_EQ(context, output->params.scale, 16.0 / 256);
@@ -671,7 +671,7 @@ TfLiteStatus PreluPrepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* alpha = GetInput(context, node, 1);
   PreluOpData* data = reinterpret_cast(node->user_data);
 
-  TF_LITE_ENSURE_EQ(context, input->type, alpha->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, alpha->type);
 
   output->type = input->type;
 
diff --git a/tensorflow/lite/kernels/add.cc b/tensorflow/lite/kernels/add.cc
index 279f6aa12ce..d6e5db90a97 100644
--- a/tensorflow/lite/kernels/add.cc
+++ b/tensorflow/lite/kernels/add.cc
@@ -90,7 +90,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
   output->type = input2->type;
 
   const bool requires_broadcast = !HaveSameShapes(input1, input2);
diff --git a/tensorflow/lite/kernels/add_n.cc b/tensorflow/lite/kernels/add_n.cc
index 7b4d52c5272..e933c5bbd66 100644
--- a/tensorflow/lite/kernels/add_n.cc
+++ b/tensorflow/lite/kernels/add_n.cc
@@ -41,7 +41,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   for (int i = kInputTensor1 + 1; i < num_inputs; ++i) {
     const TfLiteTensor* input = GetInput(context, node, i);
     TF_LITE_ENSURE(context, HaveSameShapes(input1, input));
-    TF_LITE_ENSURE_EQ(context, input1->type, input->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input->type);
   }
 
   // Use the first input node's dimension to be the dimension of the output
diff --git a/tensorflow/lite/kernels/audio_spectrogram.cc b/tensorflow/lite/kernels/audio_spectrogram.cc
index 29c9eeef3d0..8132130f4ab 100644
--- a/tensorflow/lite/kernels/audio_spectrogram.cc
+++ b/tensorflow/lite/kernels/audio_spectrogram.cc
@@ -81,8 +81,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 2);
 
-  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   TF_LITE_ENSURE(context, params->spectrogram->Initialize(params->window_size,
                                                           params->stride));
diff --git a/tensorflow/lite/kernels/basic_rnn.cc b/tensorflow/lite/kernels/basic_rnn.cc
index 920e8cd223a..c2e503d6462 100644
--- a/tensorflow/lite/kernels/basic_rnn.cc
+++ b/tensorflow/lite/kernels/basic_rnn.cc
@@ -79,8 +79,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                     bias->dims->data[0]);
   TF_LITE_ENSURE_EQ(context, recurrent_weights->dims->data[1],
                     bias->dims->data[0]);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, input_weights->type, recurrent_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_weights->type,
+                          recurrent_weights->type);
   TF_LITE_ENSURE_EQ(context, NumDimensions(hidden_state), 2);
   TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[0], batch_size);
   TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[1], num_units);
@@ -288,8 +289,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
                         accum_scratch, row_sums, &op_data->compute_row_sums);
     }
     default:
-      context->ReportError(context, "Type %d not currently supported.",
-                           input_weights->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
+                         TfLiteTypeGetName(input_weights->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/batch_matmul.cc b/tensorflow/lite/kernels/batch_matmul.cc
index d2115f96e1c..8bc23c9c94a 100644
--- a/tensorflow/lite/kernels/batch_matmul.cc
+++ b/tensorflow/lite/kernels/batch_matmul.cc
@@ -282,7 +282,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* rhs_data = GetInput(context, node, kInputRHSTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, lhs_data->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, lhs_data->type, kTfLiteFloat32);
   TF_LITE_ENSURE(context, rhs_data->type == kTfLiteFloat32 ||
                               rhs_data->type == kTfLiteInt8);
   // Support dimensions between 2 and 4, inclusive.
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
index a984ff5124f..439fc94afad 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
@@ -203,8 +203,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[0], n_cell);
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[1], n_input);
-    TF_LITE_ENSURE_EQ(context, input_to_input_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, input_to_input_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* input_to_cell_weights =
@@ -212,16 +212,16 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[1], n_input);
-  TF_LITE_ENSURE_EQ(context, input_to_cell_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_to_cell_weights->type,
+                          input_to_forget_weights->type);
 
   const TfLiteTensor* input_to_output_weights =
       GetInput(context, node, input_to_output_weights_tensor);
   TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, input_to_output_weights->dims->data[1], n_input);
-  TF_LITE_ENSURE_EQ(context, input_to_output_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_to_output_weights->type,
+                          input_to_forget_weights->type);
 
   const TfLiteTensor* recurrent_to_input_weights =
       GetOptionalInputTensor(context, node, recurrent_to_input_weights_tensor);
@@ -231,8 +231,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
                       n_cell);
     TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[1],
                       n_output);
-    TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_input_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* recurrent_to_forget_weights =
@@ -242,8 +242,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
                     n_cell);
   TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->data[1],
                     n_output);
-  TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_forget_weights->type,
+                          input_to_forget_weights->type);
 
   const TfLiteTensor* recurrent_to_cell_weights =
       GetInput(context, node, recurrent_to_cell_weights_tensor);
@@ -251,8 +251,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[1],
                     n_output);
-  TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_cell_weights->type,
+                          input_to_forget_weights->type);
 
   // We make sure the input-gate's parameters are either both present (regular
   // LSTM) or not at all (CIFG-LSTM).
@@ -268,8 +268,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   if (cell_to_input_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(context, cell_to_input_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_to_input_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* cell_to_forget_weights =
@@ -277,8 +277,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   if (cell_to_forget_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_to_forget_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* cell_to_output_weights =
@@ -286,8 +286,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   if (cell_to_output_weights != nullptr) {
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(context, cell_to_output_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_to_output_weights->type,
+                            input_to_forget_weights->type);
   }
 
   // Making sure the peephole weights are there all or none.
@@ -309,14 +309,14 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   } else {
     TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(context, input_gate_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, input_gate_bias->type, kTfLiteFloat32);
   }
 
   const TfLiteTensor* forget_gate_bias =
       GetInput(context, node, forget_gate_bias_tensor);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell);
-  TF_LITE_ENSURE_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
 
   const TfLiteTensor* cell_bias =
       GetInput(context, node, cell_gate_bias_tensor);
@@ -328,7 +328,7 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
       GetInput(context, node, output_gate_bias_tensor);
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->data[0], n_cell);
-  TF_LITE_ENSURE_EQ(context, output_gate_bias->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output_gate_bias->type, kTfLiteFloat32);
 
   const TfLiteTensor* projection_weights =
       GetOptionalInputTensor(context, node, projection_weights_tensor);
@@ -336,8 +336,8 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[0], n_output);
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[1], n_cell);
-    TF_LITE_ENSURE_EQ(context, projection_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, projection_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* projection_bias =
@@ -345,7 +345,7 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   if (projection_bias != nullptr) {
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->data[0], n_output);
-    TF_LITE_ENSURE_EQ(context, projection_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, projection_bias->type, kTfLiteFloat32);
   }
 
   // Making sure the projection tensors are consistent:
@@ -410,7 +410,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Inferring batch size, number of outputs and sequence length and
   // number of cells from the input tensors.
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
   const bool time_major = params->time_major;
   const int max_time = time_major ? input->dims->data[0] : input->dims->data[1];
@@ -1140,8 +1140,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteOk;
     }
     default:
-      context->ReportError(context, "Type %d is not currently supported.",
-                           fw_input_to_output_weights->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s is not currently supported.",
+                         TfLiteTypeGetName(fw_input_to_output_weights->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc b/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
index abaf6df9fa8..bc88740b6ed 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
@@ -129,7 +129,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Check all the parameters of tensor match within themselves and match the
   // input configuration.
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
 
   TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
   const bool time_major = params->time_major;
diff --git a/tensorflow/lite/kernels/ceil.cc b/tensorflow/lite/kernels/ceil.cc
index 9914dbe09ce..d8c6eaad7a4 100644
--- a/tensorflow/lite/kernels/ceil.cc
+++ b/tensorflow/lite/kernels/ceil.cc
@@ -32,7 +32,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   output->type = input->type;
   TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims);
   return context->ResizeTensor(context, output, output_size);
diff --git a/tensorflow/lite/kernels/concatenation.cc b/tensorflow/lite/kernels/concatenation.cc
index 61748e5ce58..5d5f06ba013 100644
--- a/tensorflow/lite/kernels/concatenation.cc
+++ b/tensorflow/lite/kernels/concatenation.cc
@@ -81,7 +81,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   }
 
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, output->type, input_type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input_type);
 
   if (input_type == kTfLiteInt8) {
     // Make sure there is no re-scaling needed for Int8 quantized kernel. This
diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc
index fa6caff5baa..88765b2f9c4 100644
--- a/tensorflow/lite/kernels/conv.cc
+++ b/tensorflow/lite/kernels/conv.cc
@@ -320,7 +320,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
   TF_LITE_ENSURE(context,
                  input_type == kTfLiteFloat32 || input_type == kTfLiteUInt8 ||
                      input_type == kTfLiteInt8 || input_type == kTfLiteInt16);
-  TF_LITE_ENSURE_EQ(context, output->type, input_type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input_type);
 
   const TfLiteTensor* bias = nullptr;
 
@@ -331,15 +331,15 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
   if (has_bias) {
     bias = GetInput(context, node, 2);
     if (input_type == kTfLiteUInt8 || input_type == kTfLiteInt8) {
-      TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt32);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt32);
       TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
     } else if (input_type == kTfLiteInt16) {
-      TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt64);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt64);
       TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
       TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0);
       TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
     } else {
-      TF_LITE_ENSURE_EQ(context, bias->type, input_type);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, input_type);
     }
     TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 0));
   }
@@ -984,8 +984,8 @@ TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) {
           context, node, params, data, input, filter, bias, output, im2col);
       break;
     default:
-      context->ReportError(context, "Type %s currently not supported.",
-                           TfLiteTypeGetName(input->type));
+      TF_LITE_KERNEL_LOG(context, "Type %s currently not supported.",
+                         TfLiteTypeGetName(input->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
@@ -1005,8 +1005,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
     case kTfLiteInt16:
       return EvalImpl(context, node);
     default:
-      context->ReportError(context, "Type %d not currently supported.",
-                           input->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
+                         TfLiteTypeGetName(input->type));
       return kTfLiteError;
   }
 }
diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc
index 8a81ea932bf..1637ad4350f 100644
--- a/tensorflow/lite/kernels/depth_to_space.cc
+++ b/tensorflow/lite/kernels/depth_to_space.cc
@@ -55,7 +55,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                  data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8 ||
                      data_type == kTfLiteInt8 || data_type == kTfLiteInt32 ||
                      data_type == kTfLiteInt64);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
   const int input_height = input->dims->data[1];
diff --git a/tensorflow/lite/kernels/depthwise_conv.cc b/tensorflow/lite/kernels/depthwise_conv.cc
index 1897d14a065..961a987cf02 100644
--- a/tensorflow/lite/kernels/depthwise_conv.cc
+++ b/tensorflow/lite/kernels/depthwise_conv.cc
@@ -122,7 +122,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context,
                  data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8 ||
                      data_type == kTfLiteInt8 || data_type == kTfLiteInt16);
-  TF_LITE_ENSURE_EQ(context, output->type, data_type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, data_type);
   if (!is_hybrid) {
     TF_LITE_ENSURE(context,
                    filter->type == data_type || data_type == kTfLiteInt16);
@@ -134,15 +134,15 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   if (hasBias) {
     bias = GetInput(context, node, kBiasTensor);
     if (data_type == kTfLiteUInt8 || data_type == kTfLiteInt8) {
-      TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt32);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt32);
       TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
     } else if (data_type == kTfLiteInt16) {
-      TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt64);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt64);
       TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
       TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0);
       TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
     } else {
-      TF_LITE_ENSURE_EQ(context, bias->type, data_type);
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, data_type);
     }
     TF_LITE_ENSURE_EQ(context, NumDimensions(bias), 1);
     TF_LITE_ENSURE_EQ(context, SizeOfDimension(filter, 3),
@@ -520,9 +520,9 @@ TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) {
         return EvalHybridPerChannel(context, node, params, data,
                                                  input, filter, bias, output);
       } else {
-        context->ReportError(
-            context, "Type %d with filter type %d not currently supported.",
-            input->type, filter->type);
+        TF_LITE_KERNEL_LOG(
+            context, "Type %s with filter type %s not currently supported.",
+            TfLiteTypeGetName(input->type), TfLiteTypeGetName(filter->type));
         return kTfLiteError;
       }
       break;
diff --git a/tensorflow/lite/kernels/div.cc b/tensorflow/lite/kernels/div.cc
index cdd02277ec9..c9eb1db531a 100644
--- a/tensorflow/lite/kernels/div.cc
+++ b/tensorflow/lite/kernels/div.cc
@@ -78,7 +78,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
   output->type = input2->type;
 
   data->requires_broadcast = !HaveSameShapes(input1, input2);
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 95b791be3f2..1b91244af33 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -45,7 +45,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   if (!IsSupportedType(input->type)) {
     context->ReportError(context, "Current data type %d is not supported.",
                          input->type);
@@ -60,7 +60,7 @@ inline TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node,
                              T func(T), TfLiteType expected_type) {
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, expected_type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, expected_type);
   const int64_t num_elements = NumElements(input);
   const T* in_data = GetTensorData(input);
   T* out_data = GetTensorData(output);
diff --git a/tensorflow/lite/kernels/embedding_lookup_sparse.cc b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
index 92574817e3b..745b5090094 100644
--- a/tensorflow/lite/kernels/embedding_lookup_sparse.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_sparse.cc
@@ -109,7 +109,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Mark the output as a dynamic tensor.
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
   output->allocation_type = kTfLiteDynamic;
 
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/floor.cc b/tensorflow/lite/kernels/floor.cc
index 2e341218700..d629b48d1a6 100644
--- a/tensorflow/lite/kernels/floor.cc
+++ b/tensorflow/lite/kernels/floor.cc
@@ -39,7 +39,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   output->type = input->type;
   TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims);
   return context->ResizeTensor(context, output, output_size);
diff --git a/tensorflow/lite/kernels/floor_div.cc b/tensorflow/lite/kernels/floor_div.cc
index 5677dc4d9b7..24682fdebe1 100644
--- a/tensorflow/lite/kernels/floor_div.cc
+++ b/tensorflow/lite/kernels/floor_div.cc
@@ -68,7 +68,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   const TfLiteType type = input1->type;
   switch (type) {
diff --git a/tensorflow/lite/kernels/fully_connected.cc b/tensorflow/lite/kernels/fully_connected.cc
index a1893878232..8b7a7832dbb 100644
--- a/tensorflow/lite/kernels/fully_connected.cc
+++ b/tensorflow/lite/kernels/fully_connected.cc
@@ -101,13 +101,13 @@ inline TfLiteStatus CheckTypes(TfLiteContext* context,
 
   if (is_quantized) {
     if (is_shuffled) {
-      TF_LITE_ENSURE_EQ(context, input->type, kTfLiteUInt8);
-      TF_LITE_ENSURE_EQ(context, filter->type, kTfLiteUInt8);
-      TF_LITE_ENSURE_EQ(context, output->type, kTfLiteInt16);
+      TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteUInt8);
+      TF_LITE_ENSURE_TYPES_EQ(context, filter->type, kTfLiteUInt8);
+      TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt16);
       TF_LITE_ENSURE_EQ(context, is_optional_bias_int, true);
     } else if (is_hybrid) {
-      TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-      TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
       TF_LITE_ENSURE_EQ(context, is_optional_bias_float, true);
     } else {
       TF_LITE_ENSURE(context, input->type == kTfLiteUInt8 ||
@@ -120,9 +120,9 @@ inline TfLiteStatus CheckTypes(TfLiteContext* context,
     }
   } else {
     // Only float32 is supported currently
-    TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-    TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
-    TF_LITE_ENSURE_EQ(context, filter->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, filter->type, kTfLiteFloat32);
     TF_LITE_ENSURE_EQ(context, is_optional_bias_float, true);
   }
 
diff --git a/tensorflow/lite/kernels/if.cc b/tensorflow/lite/kernels/if.cc
index d3f92a92b08..4c39a07bf8b 100644
--- a/tensorflow/lite/kernels/if.cc
+++ b/tensorflow/lite/kernels/if.cc
@@ -88,7 +88,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                             input->dims->data + input->dims->size);
       subgraph->ResizeInputTensor(i, dims);
       TfLiteTensor* subgraph_input = subgraph->tensor(subgraph->inputs()[i]);
-      TF_LITE_ENSURE_EQ(context, input->type, subgraph_input->type);
+      TF_LITE_ENSURE_TYPES_EQ(context, input->type, subgraph_input->type);
     }
     // Note: The `Prepare` function is responsible to run `AllocateTensors` on
     // both subgraphs. It's intentionally not to break out of the loop when
diff --git a/tensorflow/lite/kernels/l2norm.cc b/tensorflow/lite/kernels/l2norm.cc
index a7fb35ed594..857ef62a155 100644
--- a/tensorflow/lite/kernels/l2norm.cc
+++ b/tensorflow/lite/kernels/l2norm.cc
@@ -52,7 +52,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context, output->type == kTfLiteFloat32 ||
                               output->type == kTfLiteUInt8 ||
                               output->type == kTfLiteInt8);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8) {
     TF_LITE_ENSURE_EQ(context, output->params.scale, (1. / 128.));
@@ -133,8 +133,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
                                            depth, GetTensorData(input),
                                            GetTensorData(output));
   } else {
-    context->ReportError(context, "Output type is %d, requires float.",
-                         output->type);
+    TF_LITE_KERNEL_LOG(context, "Output type is %s, requires float.",
+                       TfLiteTypeGetName(output->type));
     return kTfLiteError;
   }
 
diff --git a/tensorflow/lite/kernels/local_response_norm.cc b/tensorflow/lite/kernels/local_response_norm.cc
index f4b996c45a1..ed964365920 100644
--- a/tensorflow/lite/kernels/local_response_norm.cc
+++ b/tensorflow/lite/kernels/local_response_norm.cc
@@ -44,8 +44,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
 
-  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
   output_size->data[0] = input->dims->data[0];
diff --git a/tensorflow/lite/kernels/logical.cc b/tensorflow/lite/kernels/logical.cc
index ec650dd4210..a703f3f5358 100644
--- a/tensorflow/lite/kernels/logical.cc
+++ b/tensorflow/lite/kernels/logical.cc
@@ -58,7 +58,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   const TfLiteType type = input1->type;
   if (type != kTfLiteBool) {
diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index e022bfb85ba..74caafbd0c7 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -762,8 +762,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[0], n_cell);
     TF_LITE_ENSURE_EQ(context, input_to_input_weights->dims->data[1], n_input);
-    TF_LITE_ENSURE_EQ(context, input_to_input_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, input_to_input_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* input_to_cell_weights =
@@ -771,8 +771,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->size, 2);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, input_to_cell_weights->dims->data[1], n_input);
-  TF_LITE_ENSURE_EQ(context, input_to_cell_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_to_cell_weights->type,
+                          input_to_forget_weights->type);
 
   const TfLiteTensor* recurrent_to_input_weights =
       GetOptionalInputTensor(context, node, kRecurrentToInputWeightsTensor);
@@ -782,8 +782,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                       n_cell);
     TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->dims->data[1],
                       n_output);
-    TF_LITE_ENSURE_EQ(context, recurrent_to_input_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_input_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* recurrent_to_forget_weights =
@@ -793,8 +793,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                     n_cell);
   TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->dims->data[1],
                     n_output);
-  TF_LITE_ENSURE_EQ(context, recurrent_to_forget_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_forget_weights->type,
+                          input_to_forget_weights->type);
 
   const TfLiteTensor* recurrent_to_cell_weights =
       GetInput(context, node, kRecurrentToCellWeightsTensor);
@@ -802,8 +802,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[0], n_cell);
   TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->dims->data[1],
                     n_output);
-  TF_LITE_ENSURE_EQ(context, recurrent_to_cell_weights->type,
-                    input_to_forget_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, recurrent_to_cell_weights->type,
+                          input_to_forget_weights->type);
 
   // We make sure the input-gate's parameters are either both present (regular
   // LSTM) or not at all (CIFG-LSTM).
@@ -819,7 +819,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   if (cell_to_input_weights) {
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_input_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(
+    TF_LITE_ENSURE_TYPES_EQ(
         context, cell_to_input_weights->type,
         is_integer ? kTfLiteInt16 : input_to_forget_weights->type);
   }
@@ -829,7 +829,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   if (cell_to_forget_weights) {
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_forget_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(
+    TF_LITE_ENSURE_TYPES_EQ(
         context, cell_to_forget_weights->type,
         is_integer ? kTfLiteInt16 : input_to_forget_weights->type);
   }
@@ -839,7 +839,7 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   if (cell_to_output_weights) {
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_to_output_weights->dims->data[0], n_cell);
-    TF_LITE_ENSURE_EQ(
+    TF_LITE_ENSURE_TYPES_EQ(
         context, cell_to_output_weights->type,
         is_integer ? kTfLiteInt16 : input_to_forget_weights->type);
   }
@@ -863,9 +863,9 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, input_gate_bias->dims->data[0], n_cell);
     if (is_integer) {
-      TF_LITE_ENSURE_EQ(context, input_gate_bias->type, kTfLiteInt32);
+      TF_LITE_ENSURE_TYPES_EQ(context, input_gate_bias->type, kTfLiteInt32);
     } else {
-      TF_LITE_ENSURE_EQ(context, input_gate_bias->type, kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, input_gate_bias->type, kTfLiteFloat32);
     }
   }
 
@@ -874,18 +874,18 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell);
   if (is_integer) {
-    TF_LITE_ENSURE_EQ(context, forget_gate_bias->type, kTfLiteInt32);
+    TF_LITE_ENSURE_TYPES_EQ(context, forget_gate_bias->type, kTfLiteInt32);
   } else {
-    TF_LITE_ENSURE_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
   }
 
   const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
   TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell);
   if (is_integer) {
-    TF_LITE_ENSURE_EQ(context, cell_bias->type, kTfLiteInt32);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_bias->type, kTfLiteInt32);
   } else {
-    TF_LITE_ENSURE_EQ(context, cell_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_bias->type, kTfLiteFloat32);
   }
 
   const TfLiteTensor* output_gate_bias =
@@ -893,9 +893,9 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, output_gate_bias->dims->data[0], n_cell);
   if (is_integer) {
-    TF_LITE_ENSURE_EQ(context, output_gate_bias->type, kTfLiteInt32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output_gate_bias->type, kTfLiteInt32);
   } else {
-    TF_LITE_ENSURE_EQ(context, output_gate_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output_gate_bias->type, kTfLiteFloat32);
   }
 
   const TfLiteTensor* projection_weights =
@@ -904,8 +904,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->size, 2);
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[0], n_output);
     TF_LITE_ENSURE_EQ(context, projection_weights->dims->data[1], n_cell);
-    TF_LITE_ENSURE_EQ(context, projection_weights->type,
-                      input_to_forget_weights->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, projection_weights->type,
+                            input_to_forget_weights->type);
   }
 
   const TfLiteTensor* projection_bias =
@@ -914,9 +914,9 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, projection_bias->dims->data[0], n_output);
     if (is_integer) {
-      TF_LITE_ENSURE_EQ(context, projection_bias->type, kTfLiteInt32);
+      TF_LITE_ENSURE_TYPES_EQ(context, projection_bias->type, kTfLiteInt32);
     } else {
-      TF_LITE_ENSURE_EQ(context, projection_bias->type, kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, projection_bias->type, kTfLiteFloat32);
     }
   }
 
@@ -940,11 +940,11 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
       TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->dims->data[0],
                         n_cell);
       if (is_integer) {
-        TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->type,
-                          kTfLiteInt16);
+        TF_LITE_ENSURE_TYPES_EQ(context, input_layer_norm_coefficients->type,
+                                kTfLiteInt16);
       } else {
-        TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->type,
-                          kTfLiteFloat32);
+        TF_LITE_ENSURE_TYPES_EQ(context, input_layer_norm_coefficients->type,
+                                kTfLiteFloat32);
       }
     }
 
@@ -955,11 +955,11 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->dims->data[0],
                       n_cell);
     if (is_integer) {
-      TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->type,
-                        kTfLiteInt16);
+      TF_LITE_ENSURE_TYPES_EQ(context, forget_layer_norm_coefficients->type,
+                              kTfLiteInt16);
     } else {
-      TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->type,
-                        kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, forget_layer_norm_coefficients->type,
+                              kTfLiteFloat32);
     }
 
     const TfLiteTensor* cell_layer_norm_coefficients =
@@ -969,11 +969,11 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->dims->data[0],
                       n_cell);
     if (is_integer) {
-      TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->type,
-                        kTfLiteInt16);
+      TF_LITE_ENSURE_TYPES_EQ(context, cell_layer_norm_coefficients->type,
+                              kTfLiteInt16);
     } else {
-      TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->type,
-                        kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, cell_layer_norm_coefficients->type,
+                              kTfLiteFloat32);
     }
 
     const TfLiteTensor* output_layer_norm_coefficients = GetOptionalInputTensor(
@@ -983,11 +983,11 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->dims->data[0],
                       n_cell);
     if (is_integer) {
-      TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->type,
-                        kTfLiteInt16);
+      TF_LITE_ENSURE_TYPES_EQ(context, output_layer_norm_coefficients->type,
+                              kTfLiteInt16);
     } else {
-      TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->type,
-                        kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, output_layer_norm_coefficients->type,
+                              kTfLiteFloat32);
     }
   }
 
diff --git a/tensorflow/lite/kernels/maximum_minimum.cc b/tensorflow/lite/kernels/maximum_minimum.cc
index ae1920e53db..777e51442f1 100644
--- a/tensorflow/lite/kernels/maximum_minimum.cc
+++ b/tensorflow/lite/kernels/maximum_minimum.cc
@@ -57,7 +57,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
   OpContext op_context(context, node);
-  TF_LITE_ENSURE_EQ(context, op_context.input1->type, op_context.input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input1->type,
+                          op_context.input2->type);
   op_context.output->type = op_context.input1->type;
 
   bool requires_broadcast =
diff --git a/tensorflow/lite/kernels/mfcc.cc b/tensorflow/lite/kernels/mfcc.cc
index 5fe5b948a87..a3bf5baafaa 100644
--- a/tensorflow/lite/kernels/mfcc.cc
+++ b/tensorflow/lite/kernels/mfcc.cc
@@ -80,9 +80,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumDimensions(input_wav), 3);
   TF_LITE_ENSURE_EQ(context, NumElements(input_rate), 1);
 
-  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, input_wav->type, output->type);
-  TF_LITE_ENSURE_EQ(context, input_rate->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_wav->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_rate->type, kTfLiteInt32);
 
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
   output_size->data[0] = input_wav->dims->data[0];
diff --git a/tensorflow/lite/kernels/mul.cc b/tensorflow/lite/kernels/mul.cc
index 0ab378e278d..5c34c9c7199 100644
--- a/tensorflow/lite/kernels/mul.cc
+++ b/tensorflow/lite/kernels/mul.cc
@@ -79,7 +79,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   const bool requires_broadcast = !HaveSameShapes(input1, input2);
 
diff --git a/tensorflow/lite/kernels/one_hot.cc b/tensorflow/lite/kernels/one_hot.cc
index 76d53c6396f..f7b4e8e7e19 100644
--- a/tensorflow/lite/kernels/one_hot.cc
+++ b/tensorflow/lite/kernels/one_hot.cc
@@ -136,8 +136,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       op_context.output->type = op_context.dtype;
       break;
     default:
-      context->ReportError(context, "Unknown output data type: %d",
-                           op_context.dtype);
+      TF_LITE_KERNEL_LOG(context, "Unknown output data type: %s",
+                         TfLiteTypeGetName(op_context.dtype));
       return kTfLiteError;
   }
 
@@ -148,8 +148,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumElements(op_context.depth), 1);
   TF_LITE_ENSURE_EQ(context, NumElements(op_context.on_value), 1);
   TF_LITE_ENSURE_EQ(context, NumElements(op_context.off_value), 1);
-  TF_LITE_ENSURE_EQ(context, op_context.on_value->type, op_context.dtype);
-  TF_LITE_ENSURE_EQ(context, op_context.off_value->type, op_context.dtype);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.on_value->type, op_context.dtype);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.off_value->type,
+                          op_context.dtype);
 
   if (!IsConstantTensor(op_context.depth)) {
     SetTensorToDynamic(op_context.output);
diff --git a/tensorflow/lite/kernels/pack.cc b/tensorflow/lite/kernels/pack.cc
index fc7a87692c4..90a87b0c8c7 100644
--- a/tensorflow/lite/kernels/pack.cc
+++ b/tensorflow/lite/kernels/pack.cc
@@ -57,7 +57,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   for (int i = 1; i < data->values_count; ++i) {
     const TfLiteTensor* input = GetInput(context, node, i);
     TF_LITE_ENSURE(context, HaveSameShapes(input0, input));
-    TF_LITE_ENSURE_EQ(context, input0->type, input->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, input0->type, input->type);
   }
 
   // Resize output. rank R will become rank R + 1
@@ -73,7 +73,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   }
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-  TF_LITE_ENSURE_EQ(context, output->type, input0->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input0->type);
 
   // Guarantee input/output quantization params match as we do not support
   // packing quantized tensors.
diff --git a/tensorflow/lite/kernels/pad.cc b/tensorflow/lite/kernels/pad.cc
index 2239511b60a..4d9df6b89ab 100644
--- a/tensorflow/lite/kernels/pad.cc
+++ b/tensorflow/lite/kernels/pad.cc
@@ -111,10 +111,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
   PadContext op_context(context, node);
-  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type,
+                          op_context.output->type);
   if (op_context.constant_values != nullptr) {
-    TF_LITE_ENSURE_EQ(context, op_context.input->type,
-                      op_context.constant_values->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type,
+                            op_context.constant_values->type);
   }
 
   // TODO(nupurgarg): Current implementations rely on the inputs being <= 4D.
@@ -268,9 +269,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       }
     } break;
     default:
-      context->ReportError(context,
-                           "Type %d is currently not supported by Pad.",
-                           op_context.input->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s is currently not supported by Pad.",
+                         TfLiteTypeGetName(op_context.input->type));
       return kTfLiteError;
   }
 #undef TF_LITE_PAD
diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc
index 1dc5cbb6199..a1380080a1e 100644
--- a/tensorflow/lite/kernels/pooling.cc
+++ b/tensorflow/lite/kernels/pooling.cc
@@ -74,7 +74,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, 0);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   int batches = input->dims->data[0];
   int height = input->dims->data[1];
@@ -98,7 +98,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
     }
     if (pool_type == kL2) {
       // We currently don't have a quantized implementation of L2Pool
-      TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
     }
   }
 
@@ -387,8 +387,8 @@ TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
                                              output);
       break;
     default:
-      context->ReportError(context, "Type %d not currently supported.",
-                           input->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
+                         TfLiteTypeGetName(input->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
@@ -418,8 +418,8 @@ TfLiteStatus MaxEval(TfLiteContext* context, TfLiteNode* node) {
                                          output);
       break;
     default:
-      context->ReportError(context, "Type %d not currently supported.",
-                           input->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
+                         TfLiteTypeGetName(input->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/pow.cc b/tensorflow/lite/kernels/pow.cc
index a76c77a3f9f..7f995929ec7 100644
--- a/tensorflow/lite/kernels/pow.cc
+++ b/tensorflow/lite/kernels/pow.cc
@@ -58,11 +58,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   const TfLiteType type = input1->type;
   if (type != kTfLiteInt32 && type != kTfLiteFloat32) {
-    context->ReportError(context, "Unsupported data type %d.", type);
+    TF_LITE_KERNEL_LOG(context, "Unsupported data type %s.",
+                       TfLiteTypeGetName(type));
     return kTfLiteError;
   }
   output->type = type;
diff --git a/tensorflow/lite/kernels/range.cc b/tensorflow/lite/kernels/range.cc
index 55cc543d745..fe67d055ded 100644
--- a/tensorflow/lite/kernels/range.cc
+++ b/tensorflow/lite/kernels/range.cc
@@ -100,8 +100,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     return kTfLiteError;
   }
 
-  TF_LITE_ENSURE_EQ(context, limit->type, dtype);
-  TF_LITE_ENSURE_EQ(context, delta->type, dtype);
+  TF_LITE_ENSURE_TYPES_EQ(context, limit->type, dtype);
+  TF_LITE_ENSURE_TYPES_EQ(context, delta->type, dtype);
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   output->type = dtype;
diff --git a/tensorflow/lite/kernels/read_variable.cc b/tensorflow/lite/kernels/read_variable.cc
index ad6e8d43858..78b6a136be4 100644
--- a/tensorflow/lite/kernels/read_variable.cc
+++ b/tensorflow/lite/kernels/read_variable.cc
@@ -58,7 +58,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* variable_tensor = variable->GetTensor();
   TfLiteTensor* output = GetOutput(context, node, kOutputValue);
 
-  TF_LITE_ENSURE_EQ(context, variable_tensor->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, variable_tensor->type, output->type);
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(
                    context, output, TfLiteIntArrayCopy(variable_tensor->dims)));
diff --git a/tensorflow/lite/kernels/reduce.cc b/tensorflow/lite/kernels/reduce.cc
index af42b2a369c..6107b01cd46 100644
--- a/tensorflow/lite/kernels/reduce.cc
+++ b/tensorflow/lite/kernels/reduce.cc
@@ -235,7 +235,7 @@ TfLiteStatus PrepareSimple(TfLiteContext* context, TfLiteNode* node) {
 TfLiteStatus PrepareAny(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   const TfLiteTensor* input = GetInput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteBool);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteBool);
   return PrepareSimple(context, node);
 }
 
diff --git a/tensorflow/lite/kernels/resize_nearest_neighbor.cc b/tensorflow/lite/kernels/resize_nearest_neighbor.cc
index fff45ac13cc..13c54c4f906 100644
--- a/tensorflow/lite/kernels/resize_nearest_neighbor.cc
+++ b/tensorflow/lite/kernels/resize_nearest_neighbor.cc
@@ -68,7 +68,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // and the size being 1D tensor with exactly 2 elements.
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
   TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1);
-  TF_LITE_ENSURE_EQ(context, size->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, size->type, kTfLiteInt32);
   TF_LITE_ENSURE_EQ(context, size->dims->data[0], 2);
 
   output->type = input->type;
@@ -122,9 +122,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
         GetTensorShape(size), GetTensorData(size),
         GetTensorShape(output), GetTensorData(output));
   } else {
-    context->ReportError(context,
-                         "Output type is %d, requires float, uint8 or int8.",
-                         output->type);
+    TF_LITE_KERNEL_LOG(context,
+                       "Output type is %s, requires float, uint8 or int8.",
+                       TfLiteTypeGetName(output->type));
     return kTfLiteError;
   }
 
diff --git a/tensorflow/lite/kernels/reverse.cc b/tensorflow/lite/kernels/reverse.cc
index 9ce845b4b7b..34cc92da5d8 100644
--- a/tensorflow/lite/kernels/reverse.cc
+++ b/tensorflow/lite/kernels/reverse.cc
@@ -61,7 +61,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
-  TF_LITE_ENSURE_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
 
   return context->ResizeTensor(context, output, output_shape);
 }
diff --git a/tensorflow/lite/kernels/reverse_sequence.cc b/tensorflow/lite/kernels/reverse_sequence.cc
index 7390876d39b..b36b1f803ca 100644
--- a/tensorflow/lite/kernels/reverse_sequence.cc
+++ b/tensorflow/lite/kernels/reverse_sequence.cc
@@ -58,7 +58,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
-  TF_LITE_ENSURE_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
 
   return context->ResizeTensor(context, output, output_shape);
 }
diff --git a/tensorflow/lite/kernels/round.cc b/tensorflow/lite/kernels/round.cc
index 341d2880705..72c793c1152 100644
--- a/tensorflow/lite/kernels/round.cc
+++ b/tensorflow/lite/kernels/round.cc
@@ -34,7 +34,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   output->type = input->type;
   TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims);
   return context->ResizeTensor(context, output, output_size);
diff --git a/tensorflow/lite/kernels/select.cc b/tensorflow/lite/kernels/select.cc
index 281425253c5..62c8ddbaa97 100644
--- a/tensorflow/lite/kernels/select.cc
+++ b/tensorflow/lite/kernels/select.cc
@@ -66,8 +66,8 @@ TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   // Input must be bool.
-  TF_LITE_ENSURE(context, input_condition->type == kTfLiteBool);
-  TF_LITE_ENSURE_EQ(context, input_x->type, input_y->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_condition->type, kTfLiteBool);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_x->type, input_y->type);
   output->type = input_x->type;
 
   bool same_shape = HaveSameShapes(input_condition, input_x) &&
diff --git a/tensorflow/lite/kernels/skip_gram.cc b/tensorflow/lite/kernels/skip_gram.cc
index 8348a25bba7..f81d152bb70 100644
--- a/tensorflow/lite/kernels/skip_gram.cc
+++ b/tensorflow/lite/kernels/skip_gram.cc
@@ -48,8 +48,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  TF_LITE_ENSURE_EQ(context, GetInput(context, node, 0)->type, kTfLiteString);
-  TF_LITE_ENSURE_EQ(context, GetOutput(context, node, 0)->type, kTfLiteString);
+  TF_LITE_ENSURE_TYPES_EQ(context, GetInput(context, node, 0)->type,
+                          kTfLiteString);
+  TF_LITE_ENSURE_TYPES_EQ(context, GetOutput(context, node, 0)->type,
+                          kTfLiteString);
   return kTfLiteOk;
 }
 
diff --git a/tensorflow/lite/kernels/space_to_batch_nd.cc b/tensorflow/lite/kernels/space_to_batch_nd.cc
index 7fc58e7ee6b..0d537e2d189 100644
--- a/tensorflow/lite/kernels/space_to_batch_nd.cc
+++ b/tensorflow/lite/kernels/space_to_batch_nd.cc
@@ -100,7 +100,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                  NumDimensions(op_context.input) >= kInputMinDimensionNum);
   TF_LITE_ENSURE(context,
                  NumDimensions(op_context.input) <= kInputMaxDimensionNum);
-  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type,
+                          op_context.output->type);
 
   if (!IsConstantTensor(op_context.block_shape) ||
       !IsConstantTensor(op_context.paddings)) {
diff --git a/tensorflow/lite/kernels/space_to_depth.cc b/tensorflow/lite/kernels/space_to_depth.cc
index e4c7efaaf99..ac001d903a4 100644
--- a/tensorflow/lite/kernels/space_to_depth.cc
+++ b/tensorflow/lite/kernels/space_to_depth.cc
@@ -55,7 +55,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                  data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8 ||
                      data_type == kTfLiteInt8 || data_type == kTfLiteInt32 ||
                      data_type == kTfLiteInt64);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const int block_size = params->block_size;
   const int input_height = input->dims->data[1];
diff --git a/tensorflow/lite/kernels/sparse_to_dense.cc b/tensorflow/lite/kernels/sparse_to_dense.cc
index bdf0f4e703a..4aea0f491bc 100644
--- a/tensorflow/lite/kernels/sparse_to_dense.cc
+++ b/tensorflow/lite/kernels/sparse_to_dense.cc
@@ -172,7 +172,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                               values->type == kTfLiteInt8 ||
                               values->type == kTfLiteUInt8 ||
                               values->type == kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, values->type, default_value->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, values->type, default_value->type);
 
   // Ensure dimensions match.
   TF_LITE_ENSURE_OK(
@@ -229,10 +229,10 @@ TfLiteStatus EvalForIndexType(TfLiteContext* context, TfLiteNode* node,
       return SparseToDenseImpl(context, node);
     }
     default:
-      context->ReportError(
+      TF_LITE_KERNEL_LOG(
           context,
-          "Indice type %d is currently not supported by sparse to dense.",
-          indices->type);
+          "Indice type %s is currently not supported by sparse to dense.",
+          TfLiteTypeGetName(indices->type));
       return kTfLiteError;
   }
 }
@@ -253,10 +253,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
     case kTfLiteUInt8:
       return EvalForIndexType(context, node, indices);
     default:
-      context->ReportError(
+      TF_LITE_KERNEL_LOG(
           context,
-          "Value type %d is currently not supported by sparse to dense.",
-          values->type);
+          "Value type %s is currently not supported by sparse to dense.",
+          TfLiteTypeGetName(values->type));
       return kTfLiteError;
   }
 }
diff --git a/tensorflow/lite/kernels/squared_difference.cc b/tensorflow/lite/kernels/squared_difference.cc
index e17ff8e3191..ff09995845e 100644
--- a/tensorflow/lite/kernels/squared_difference.cc
+++ b/tensorflow/lite/kernels/squared_difference.cc
@@ -64,7 +64,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
   output->type = input2->type;
 
   data->requires_broadcast = !HaveSameShapes(input1, input2);
diff --git a/tensorflow/lite/kernels/strided_slice.cc b/tensorflow/lite/kernels/strided_slice.cc
index 50c2255e526..83221cd4a3d 100644
--- a/tensorflow/lite/kernels/strided_slice.cc
+++ b/tensorflow/lite/kernels/strided_slice.cc
@@ -145,9 +145,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
   // Only INT32 begin/end/strides are supported
   // TODO(soroosh) add support for INT64
-  TF_LITE_ENSURE_EQ(context, op_context.begin->type, kTfLiteInt32);
-  TF_LITE_ENSURE_EQ(context, op_context.end->type, kTfLiteInt32);
-  TF_LITE_ENSURE_EQ(context, op_context.strides->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.begin->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.end->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.strides->type, kTfLiteInt32);
   TF_LITE_ENSURE_MSG(context, op_context.dims <= 5,
                      "StridedSlice op only supports 1D-5D input arrays.");
 
@@ -223,10 +223,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       }
       break;
     default:
-      context->ReportError(context,
-                           "Type %d is currently not supported "
-                           "by StridedSlice.",
-                           op_context.input->type);
+      TF_LITE_KERNEL_LOG(context,
+                         "Type %s is currently not supported "
+                         "by StridedSlice.",
+                         TfLiteTypeGetName(op_context.input->type));
       return kTfLiteError;
   }
 #undef TF_LITE_STRIDED_SLICE
diff --git a/tensorflow/lite/kernels/sub.cc b/tensorflow/lite/kernels/sub.cc
index aa628fa5408..83b2714135d 100644
--- a/tensorflow/lite/kernels/sub.cc
+++ b/tensorflow/lite/kernels/sub.cc
@@ -206,7 +206,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
   output->type = input2->type;
 
   data->requires_broadcast = !HaveSameShapes(input1, input2);
@@ -287,8 +287,8 @@ void EvalSub(TfLiteContext* context, TfLiteNode* node, TfLiteSubParams* params,
                                       input2, requires_broadcast, output);
       break;
     default:
-      TF_LITE_KERNEL_LOG(context, "output type %d is not supported.",
-                         output->type);
+      TF_LITE_KERNEL_LOG(context, "output type %s is not supported.",
+                         TfLiteTypeGetName(output->type));
   }
 }
 
diff --git a/tensorflow/lite/kernels/tile.cc b/tensorflow/lite/kernels/tile.cc
index 884456fcbf2..08d246203ae 100644
--- a/tensorflow/lite/kernels/tile.cc
+++ b/tensorflow/lite/kernels/tile.cc
@@ -211,7 +211,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   const TfLiteTensor* multipliers = GetInput(context, node, kInputMultipliers);
   // Only int32 and int64 multipliers type is supported.
diff --git a/tensorflow/lite/kernels/topk_v2.cc b/tensorflow/lite/kernels/topk_v2.cc
index 6a5bd392086..3fb241356e1 100644
--- a/tensorflow/lite/kernels/topk_v2.cc
+++ b/tensorflow/lite/kernels/topk_v2.cc
@@ -37,7 +37,7 @@ namespace {
 TfLiteStatus ResizeOutput(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* top_k = GetInput(context, node, kInputTopK);
   // INT32 number of top results is supported.
-  TF_LITE_ENSURE_EQ(context, top_k->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, top_k->type, kTfLiteInt32);
   // Check that the tensor contains only one value.
   TF_LITE_ENSURE_EQ(context, NumElements(top_k), 1);
   const int32 k = *GetTensorData(top_k);
@@ -197,10 +197,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output_values = GetOutput(context, node, kOutputValues);
-  TF_LITE_ENSURE_EQ(context, input->type, output_values->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output_values->type);
 
   const TfLiteTensor* top_k = GetInput(context, node, kInputTopK);
-  TF_LITE_ENSURE_EQ(context, top_k->type, kTfLiteInt32);
+  TF_LITE_ENSURE_TYPES_EQ(context, top_k->type, kTfLiteInt32);
 
   // Set output dynamic if the input is not const.
   if (IsConstantTensor(top_k)) {
@@ -252,9 +252,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
            output_values->data.i64);
       break;
     default:
-      context->ReportError(context,
-                           "Type %d is currently not supported by TopK.",
-                           output_values->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s is currently not supported by TopK.",
+                         TfLiteTypeGetName(output_values->type));
       return kTfLiteError;
   }
 
diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc
index 27f5cf6f065..3a6d1b1f1ed 100644
--- a/tensorflow/lite/kernels/transpose.cc
+++ b/tensorflow/lite/kernels/transpose.cc
@@ -77,7 +77,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Ensure validity of input tensor.
   TF_LITE_ENSURE_MSG(context, NumDimensions(op_context.input) <= 5,
                      "Transpose op only supports 1D-5D input arrays.");
-  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, op_context.input->type,
+                          op_context.output->type);
 
   if (!IsConstantTensor(op_context.perm)) {
     SetTensorToDynamic(op_context.output);
@@ -144,9 +145,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       }
       break;
     default:
-      context->ReportError(context,
-                           "Type %d is currently not supported by Transpose.",
-                           op_context.input->type);
+      TF_LITE_KERNEL_LOG(context,
+                         "Type %s is currently not supported by Transpose.",
+                         TfLiteTypeGetName(op_context.input->type));
       return kTfLiteError;
   }
 #undef TF_LITE_TRANSPOSE
diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc
index 33e122ba037..07dc4bbac53 100644
--- a/tensorflow/lite/kernels/transpose_conv.cc
+++ b/tensorflow/lite/kernels/transpose_conv.cc
@@ -111,8 +111,8 @@ TfLiteStatus ResizeTensor(TfLiteContext* context,
                           TfLiteTensor* tensor_to_resize) {
   // Currently only support int32 for output shape.
   if (shape_tensor->type != kTfLiteInt32) {
-    context->ReportError(context, "Output shape is %d, not int32.",
-                         shape_tensor->type);
+    TF_LITE_KERNEL_LOG(context, "Output shape is %s, not int32.",
+                       TfLiteTypeGetName(shape_tensor->type));
     return kTfLiteError;
   }
 
@@ -176,8 +176,8 @@ TfLiteStatus ResizeCol2ImTensor(TfLiteContext* context,
                                 const TfLiteTensor* input,
                                 TfLiteTensor* col2im) {
   if (output_shape->type != kTfLiteInt32) {
-    context->ReportError(context, "col2im shape is %d, not int32.",
-                         output_shape->type);
+    TF_LITE_KERNEL_LOG(context, "col2im shape is %s, not int32.",
+                       TfLiteTypeGetName(output_shape->type));
     return kTfLiteError;
   }
   TF_LITE_ENSURE_EQ(context, NumElements(output_shape), 4);
@@ -274,7 +274,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     bias = GetOptionalInputTensor(context, node, kBiasTensor);
     if (bias) {
       if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
-        TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt32);
+        TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt32);
         if (input->type == kTfLiteInt8) {
           TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
         }
@@ -282,7 +282,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
         TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt64);
         TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
       } else {
-        TF_LITE_ENSURE_EQ(context, bias->type, input->type);
+        TF_LITE_ENSURE_TYPES_EQ(context, bias->type, input->type);
       }
       TF_LITE_ENSURE_EQ(context, NumElements(bias),
                         SizeOfDimension(weights, 0));
@@ -294,9 +294,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0);
     TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
   } else {
-    TF_LITE_ENSURE_EQ(context, weights->type, input->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, weights->type, input->type);
   }
-  TF_LITE_ENSURE_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
   // Ensure that weights and inputs have the same channel dimension.
   // Note: TOCO will reorder weights in the following format: OHWI.
   TF_LITE_ENSURE_EQ(context, SizeOfDimension(input, 3),
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
index b8b9396f436..95864196f18 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
@@ -223,8 +223,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
       TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->dims->size, 1);
       TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->dims->data[0],
                         n_cell);
-      TF_LITE_ENSURE_EQ(context, input_layer_norm_coefficients->type,
-                        kTfLiteFloat32);
+      TF_LITE_ENSURE_TYPES_EQ(context, input_layer_norm_coefficients->type,
+                              kTfLiteFloat32);
     }
 
     const TfLiteTensor* forget_layer_norm_coefficients =
@@ -233,8 +233,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->dims->data[0],
                       n_cell);
-    TF_LITE_ENSURE_EQ(context, forget_layer_norm_coefficients->type,
-                      kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, forget_layer_norm_coefficients->type,
+                            kTfLiteFloat32);
 
     const TfLiteTensor* cell_layer_norm_coefficients =
         GetInput(context, node, lstm::full::kCellLayerNormCoefficientsTensor);
@@ -242,8 +242,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->dims->data[0],
                       n_cell);
-    TF_LITE_ENSURE_EQ(context, cell_layer_norm_coefficients->type,
-                      kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_layer_norm_coefficients->type,
+                            kTfLiteFloat32);
 
     const TfLiteTensor* output_layer_norm_coefficients =
         GetInput(context, node, lstm::full::kOutputLayerNormCoefficientsTensor);
@@ -251,8 +251,8 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->dims->size, 1);
     TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->dims->data[0],
                       n_cell);
-    TF_LITE_ENSURE_EQ(context, output_layer_norm_coefficients->type,
-                      kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output_layer_norm_coefficients->type,
+                            kTfLiteFloat32);
   }
 
   return kTfLiteOk;
@@ -290,7 +290,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Inferring batch size, number of outputs and sequence length and
   // number of cells from the input tensors.
   const TfLiteTensor* input = GetInput(context, node, lstm::full::kInputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   TF_LITE_ENSURE(context, input->dims->size > 1);
   const auto* params =
       reinterpret_cast(
@@ -659,8 +659,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           CpuBackendContext::GetFromContext(context));
     }
     default:
-      context->ReportError(context, "Type %d is not currently supported.",
-                           input_to_output_weights->type);
+      TF_LITE_KERNEL_LOG(context, "Type %s is not currently supported.",
+                         TfLiteTypeGetName(input_to_output_weights->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_rnn.cc b/tensorflow/lite/kernels/unidirectional_sequence_rnn.cc
index 7ed67c1614d..350ca293cbf 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_rnn.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_rnn.cc
@@ -85,8 +85,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                     bias->dims->data[0]);
   TF_LITE_ENSURE_EQ(context, recurrent_weights->dims->data[1],
                     bias->dims->data[0]);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, input_weights->type, recurrent_weights->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_weights->type,
+                          recurrent_weights->type);
   TF_LITE_ENSURE_EQ(context, NumDimensions(hidden_state), 2);
   TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[0], batch_size);
   TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[1], num_units);
@@ -364,8 +365,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
                         accum_scratch, row_sums, &op_data->compute_row_sums);
     }
     default:
-      context->ReportError(context, "Type %d not currently supported.",
-                           input_weights->type);
+      TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.",
+                         TfLiteTypeGetName(input_weights->type));
       return kTfLiteError;
   }
   return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/unpack.cc b/tensorflow/lite/kernels/unpack.cc
index 8d307acb268..a41556ed079 100644
--- a/tensorflow/lite/kernels/unpack.cc
+++ b/tensorflow/lite/kernels/unpack.cc
@@ -68,7 +68,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   for (int i = 0; i < data->num; ++i) {
     TfLiteIntArray* copied_output_shape = TfLiteIntArrayCopy(output_shape);
     TfLiteTensor* output = GetOutput(context, node, i);
-    TF_LITE_ENSURE_EQ(context, output->type, input->type);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
     // Guarantee input/output quantization params match as we do not support
     // rescaling of unpacked quantized tensors.
     TF_LITE_ENSURE_EQ(context, input->params.zero_point,
diff --git a/tensorflow/lite/kernels/while.cc b/tensorflow/lite/kernels/while.cc
index 99d6d2cc1c8..b50cdff9974 100644
--- a/tensorflow/lite/kernels/while.cc
+++ b/tensorflow/lite/kernels/while.cc
@@ -90,7 +90,7 @@ TfLiteStatus CopyTensorsData(TfLiteContext* context, Subgraph* src_subgraph,
 TfLiteStatus CheckCondOutput(TfLiteContext* context,
                              const TfLiteTensor* cond_output) {
   // The condition output must be a single boolean value.
-  TF_LITE_ENSURE_EQ(context, cond_output->type, kTfLiteBool);
+  TF_LITE_ENSURE_TYPES_EQ(context, cond_output->type, kTfLiteBool);
   if (cond_output->dims->size == 0) {
     // It's okay if it's a 0D scalar.
     return kTfLiteOk;
@@ -179,7 +179,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
           body_subgraph->tensor(body_subgraph->inputs()[i]);
       TfLiteTensor* body_output =
           body_subgraph->tensor(body_subgraph->outputs()[i]);
-      TF_LITE_ENSURE_EQ(context, body_input->type, body_output->type);
+      TF_LITE_ENSURE_TYPES_EQ(context, body_input->type, body_output->type);
 
       // TODO(ycling): Support dynamic sized body subgraph.
       TF_LITE_ENSURE(context, !IsDynamicTensor(body_output));
diff --git a/tensorflow/lite/micro/kernels/arc_mli/fully_connected.cc b/tensorflow/lite/micro/kernels/arc_mli/fully_connected.cc
index fe077c99fac..8247875a7ac 100644
--- a/tensorflow/lite/micro/kernels/arc_mli/fully_connected.cc
+++ b/tensorflow/lite/micro/kernels/arc_mli/fully_connected.cc
@@ -111,7 +111,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   TF_LITE_ENSURE(context, data != nullptr);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   TF_LITE_ENSURE_MSG(context, input->type == filter->type,
                      "Hybrid models are not supported on TFLite Micro.");
 
diff --git a/tensorflow/lite/micro/kernels/ceil.cc b/tensorflow/lite/micro/kernels/ceil.cc
index 89831a767fe..ace038aaac5 100644
--- a/tensorflow/lite/micro/kernels/ceil.cc
+++ b/tensorflow/lite/micro/kernels/ceil.cc
@@ -32,8 +32,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
   TF_LITE_ENSURE_EQ(context, output->bytes, input->bytes);
   TF_LITE_ENSURE_EQ(context, output->dims->size, input->dims->size);
   for (int i = 0; i < output->dims->size; ++i) {
diff --git a/tensorflow/lite/micro/kernels/circular_buffer.cc b/tensorflow/lite/micro/kernels/circular_buffer.cc
index 590bdbe00f8..f588d64dcd5 100644
--- a/tensorflow/lite/micro/kernels/circular_buffer.cc
+++ b/tensorflow/lite/micro/kernels/circular_buffer.cc
@@ -89,10 +89,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, 1, input->dims->data[2]);
   TF_LITE_ENSURE_EQ(context, output->dims->data[3], input->dims->data[3]);
 
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   // The circular buffer custom operator currently only supports int8.
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteInt8);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt8);
 
   // TODO(b/132070898): Use statically slotted OpData structures until a
   // scratch memory API is ready.
diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
index 6a9eb882fdf..d8827b36d06 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
@@ -81,7 +81,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   TF_LITE_ENSURE_MSG(context, input->type == filter->type,
                      "Hybrid models are not supported on TFLite Micro.");
 #if defined(__ARM_FEATURE_DSP)
diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc b/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc
index 814daa526d2..a0e7af6d2d5 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/mul.cc
@@ -48,7 +48,7 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
diff --git a/tensorflow/lite/micro/kernels/elementwise.cc b/tensorflow/lite/micro/kernels/elementwise.cc
index b69d260a826..aa97907df24 100644
--- a/tensorflow/lite/micro/kernels/elementwise.cc
+++ b/tensorflow/lite/micro/kernels/elementwise.cc
@@ -40,7 +40,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   if (!IsSupportedType(input->type)) {
     TF_LITE_KERNEL_LOG(context, "Input data type %s (%d) is not supported.",
                        TfLiteTypeGetName(input->type), input->type);
@@ -54,7 +54,7 @@ inline TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node,
                              T func(T), TfLiteType expected_type) {
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
-  TF_LITE_ENSURE_EQ(context, input->type, expected_type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, expected_type);
   const int64_t num_elements = NumElements(input);
   const T* in_data = GetTensorData(input);
   T* out_data = GetTensorData(output);
diff --git a/tensorflow/lite/micro/kernels/floor.cc b/tensorflow/lite/micro/kernels/floor.cc
index 435934fe39f..d8134e96cd6 100644
--- a/tensorflow/lite/micro/kernels/floor.cc
+++ b/tensorflow/lite/micro/kernels/floor.cc
@@ -29,7 +29,7 @@ constexpr int kOutputTensor = 0;
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   reference_ops::Floor(GetTensorShape(input), GetTensorData(input),
                        GetTensorShape(output), GetTensorData(output));
diff --git a/tensorflow/lite/micro/kernels/fully_connected.cc b/tensorflow/lite/micro/kernels/fully_connected.cc
index bd949e6f552..8478b13d90e 100644
--- a/tensorflow/lite/micro/kernels/fully_connected.cc
+++ b/tensorflow/lite/micro/kernels/fully_connected.cc
@@ -93,7 +93,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   TF_LITE_ENSURE_MSG(context, input->type == filter->type,
                      "Hybrid models are not supported on TFLite Micro.");
 
diff --git a/tensorflow/lite/micro/kernels/l2norm.cc b/tensorflow/lite/micro/kernels/l2norm.cc
index 4dd71fe1c4b..050f9d1e184 100644
--- a/tensorflow/lite/micro/kernels/l2norm.cc
+++ b/tensorflow/lite/micro/kernels/l2norm.cc
@@ -48,7 +48,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE(context, output->type == kTfLiteFloat32 ||
                               output->type == kTfLiteUInt8 ||
                               output->type == kTfLiteInt8);
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
 
   if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8) {
     TF_LITE_ENSURE_EQ(context, output->params.scale, (1. / 128.));
@@ -118,8 +118,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
                                            depth, GetTensorData(input),
                                            GetTensorData(output));
   } else {
-    TF_LITE_KERNEL_LOG(context, "Output type is %d, requires float.",
-                         output->type);
+    TF_LITE_KERNEL_LOG(context, "Output type is %s, requires float.",
+                       TfLiteTypeGetName(output->type));
     return kTfLiteError;
   }
 
diff --git a/tensorflow/lite/micro/kernels/logistic.cc b/tensorflow/lite/micro/kernels/logistic.cc
index cc360c58bd9..cb1140e0839 100644
--- a/tensorflow/lite/micro/kernels/logistic.cc
+++ b/tensorflow/lite/micro/kernels/logistic.cc
@@ -44,7 +44,7 @@ TfLiteStatus CalculateArithmeticOpData(TfLiteContext* context, TfLiteNode* node,
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   if (input->type == kTfLiteInt8) {
     TF_LITE_ENSURE_EQ(context, output->params.zero_point,
                       std::numeric_limits::min());
diff --git a/tensorflow/lite/micro/kernels/mul.cc b/tensorflow/lite/micro/kernels/mul.cc
index fb47728a1a4..82b01b11baf 100644
--- a/tensorflow/lite/micro/kernels/mul.cc
+++ b/tensorflow/lite/micro/kernels/mul.cc
@@ -48,7 +48,7 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
-  TF_LITE_ENSURE_EQ(context, input1->type, input2->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
 
   if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
diff --git a/tensorflow/lite/micro/kernels/reshape.cc b/tensorflow/lite/micro/kernels/reshape.cc
index 407682a6ff4..36601b1a43d 100644
--- a/tensorflow/lite/micro/kernels/reshape.cc
+++ b/tensorflow/lite/micro/kernels/reshape.cc
@@ -61,7 +61,7 @@ TfLiteStatus ReshapeOutput(TfLiteContext* context, TfLiteNode* node) {
     num_output_elements *= output_shape->data[stretch_dim];
   }
 
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements);
   return kTfLiteOk;
 }
diff --git a/tensorflow/lite/micro/kernels/round.cc b/tensorflow/lite/micro/kernels/round.cc
index b88c9fe0581..dc93817729b 100644
--- a/tensorflow/lite/micro/kernels/round.cc
+++ b/tensorflow/lite/micro/kernels/round.cc
@@ -32,8 +32,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
-  TF_LITE_ENSURE_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
   TF_LITE_ENSURE_EQ(context, output->bytes, input->bytes);
   TF_LITE_ENSURE_EQ(context, output->dims->size, input->dims->size);
   for (int i = 0; i < output->dims->size; ++i) {
diff --git a/tensorflow/lite/micro/kernels/svdf.cc b/tensorflow/lite/micro/kernels/svdf.cc
index ba7cb05da57..717301e2261 100644
--- a/tensorflow/lite/micro/kernels/svdf.cc
+++ b/tensorflow/lite/micro/kernels/svdf.cc
@@ -419,7 +419,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt32);
     }
 
-    TF_LITE_ENSURE_EQ(context, output->type, kTfLiteInt8);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt8);
 
     const auto* input_params =
         reinterpret_cast(input->quantization.params);
@@ -467,7 +467,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     if (bias != nullptr) {
       TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteFloat32);
     }
-    TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
 
     TFLITE_DCHECK(node->user_data != nullptr);
     OpData* data = static_cast(node->user_data);
diff --git a/tensorflow/lite/micro/kernels/tanh.cc b/tensorflow/lite/micro/kernels/tanh.cc
index 9ee5b74bde4..d978c7a1308 100644
--- a/tensorflow/lite/micro/kernels/tanh.cc
+++ b/tensorflow/lite/micro/kernels/tanh.cc
@@ -44,7 +44,7 @@ TfLiteStatus CalculateArithmeticOpData(TfLiteContext* context, TfLiteNode* node,
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_EQ(context, input->type, output->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   if (input->type == kTfLiteInt8) {
     TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
 
diff --git a/tensorflow/lite/micro/kernels/xtensa_hifi/floor.cc b/tensorflow/lite/micro/kernels/xtensa_hifi/floor.cc
index 7ea2c2c906e..e1507b7ab27 100644
--- a/tensorflow/lite/micro/kernels/xtensa_hifi/floor.cc
+++ b/tensorflow/lite/micro/kernels/xtensa_hifi/floor.cc
@@ -51,7 +51,7 @@ constexpr int kOutputTensor = 0;
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   int err;
   const float* inp_data_ptr;
diff --git a/tensorflow/lite/micro/kernels/xtensa_hifi/svdf.cc b/tensorflow/lite/micro/kernels/xtensa_hifi/svdf.cc
index 92e5a476197..26a5a4e87d3 100644
--- a/tensorflow/lite/micro/kernels/xtensa_hifi/svdf.cc
+++ b/tensorflow/lite/micro/kernels/xtensa_hifi/svdf.cc
@@ -432,7 +432,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // EvalIntegerSVDF().
 
     // Validate output tensor:
-    TF_LITE_ENSURE_EQ(context, output->type, kTfLiteInt8);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt8);
   } else {
     TF_LITE_ENSURE_EQ(context, node->inputs->size, 5);
 
@@ -457,7 +457,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // TODO(b/132070898): Use input tensor as variable until scratch tensor
     // allocation has been implemented.
     // TF_LITE_ENSURE_EQ(context, node->temporaries->size, 1);
-    TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
   }
 
   return kTfLiteOk;
diff --git a/tensorflow/lite/micro/kernels/xtensa_hifimini_legacy/svdf.cc b/tensorflow/lite/micro/kernels/xtensa_hifimini_legacy/svdf.cc
index 02bb72976dd..760e1290e73 100644
--- a/tensorflow/lite/micro/kernels/xtensa_hifimini_legacy/svdf.cc
+++ b/tensorflow/lite/micro/kernels/xtensa_hifimini_legacy/svdf.cc
@@ -348,7 +348,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, activation_state->type, kTfLiteInt16);
 
   // Validate output tensor:
-  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteInt8);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt8);
 
   // Calculate effective scales.
   auto* input_params =
diff --git a/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h b/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
index ab769fec249..15823784d12 100644
--- a/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
+++ b/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
@@ -205,6 +205,7 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
 // the current function, while also reporting the location of the error.
 // `a` and `b` may be evaluated more than once, so no side effects or
 // extremely expensive computations should be done.
+// NOTE: Use TF_LITE_ENSURE_TYPES_EQ if comparing TfLiteTypes.
 #define TF_LITE_ENSURE_EQ(context, a, b)                                   \
   do {                                                                     \
     if ((a) != (b)) {                                                      \

From 2bba7c464ecd2a1baa234cd72bac5f6ddc436ff5 Mon Sep 17 00:00:00 2001
From: Tim Shen 
Date: Tue, 16 Jun 2020 11:31:59 -0700
Subject: [PATCH 0310/1390] [MLIR] Add more LHLO ops.

PiperOrigin-RevId: 316722878
Change-Id: I996646ef895cdb5ce2f68c6cb128ac55c443adb1
---
 tensorflow/compiler/mlir/xla/ir/hlo_ops.td    |  10 +-
 .../compiler/mlir/xla/ir/hlo_ops_base.td      |   9 ++
 tensorflow/compiler/mlir/xla/ir/lhlo_ops.td   |  57 +++++++++
 .../compiler/mlir/xla/tests/lhlo_ops.mlir     | 121 ++++++++++++++++++
 4 files changed, 188 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
index f4947f2aadb..b1745c73fbf 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
@@ -1112,21 +1112,13 @@ def HLO_DynamicReshapeOp: HLO_Op<"dynamic_reshape", [NoSideEffect]> {
   let hasCustomHLOConverter = 1;
 }
 
-def ScatterDimensionNumbers : StructAttr<"ScatterDimensionNumbers", HLO_Dialect,
-      [StructFieldAttr<"update_window_dims", I64ElementsAttr>,
-      StructFieldAttr<"inserted_window_dims", I64ElementsAttr>,
-      StructFieldAttr<"scatter_dims_to_operand_dims", I64ElementsAttr>,
-      StructFieldAttr<"index_vector_dim", I64Attr>]> {
-  let description = "Structure of dimension information for scatter";
-}
-
 def HLO_ScatterOp: HLO_Op<"scatter", [RecursiveSideEffects]>,
       BASE_HLO_ScatterOp {
   let arguments = (ins
     HLO_Tensor:$operand,
     HLO_Tensor:$scatter_indices,
     HLO_Tensor:$updates,
-    ScatterDimensionNumbers:$scatter_dimension_numbers,
+    ScatterDimensionNumbers:$scatter_dimension_numbers,
     DefaultValuedAttr:$indices_are_sorted,
     DefaultValuedAttr:$unique_indices
   );
diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
index 64d374692a8..b0975d9ab03 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops_base.td
@@ -1098,6 +1098,15 @@ class BASE_HLO_ReshapeOp {
   }];
 }
 
+class ScatterDimensionNumbers : StructAttr<
+    "ScatterDimensionNumbers", dialect, [
+      StructFieldAttr<"update_window_dims", I64ElementsAttr>,
+      StructFieldAttr<"inserted_window_dims", I64ElementsAttr>,
+      StructFieldAttr<"scatter_dims_to_operand_dims", I64ElementsAttr>,
+      StructFieldAttr<"index_vector_dim", I64Attr>]> {
+  let description = "Structure of dimension information for scatter";
+}
+
 class BASE_HLO_ScatterOp {
   string summary = "Scatter operator";
 
diff --git a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td
index b3ba4afa97a..aed7c83570e 100644
--- a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td
+++ b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td
@@ -471,6 +471,12 @@ def LHLO_BatchNormTrainingOp : LHLO_Op<"batch_norm_training", []>,
   );
 }
 
+// TODO(timshen): add a custom verifier.
+def LHLO_BitcastOp: LHLO_Op<"bitcast", []> {
+  let arguments = (ins Arg:$input,
+                       Arg:$output);
+}
+
 def LHLO_BroadcastOp : LHLO_Op<"broadcast",
       []>, BASE_HLO_BroadcastOp {
   let arguments = (ins
@@ -578,6 +584,19 @@ def LHLO_ReshapeOp: LHLO_Op<"reshape", []>, BASE_HLO_ReshapeOp {
   );
 }
 
+def LHLO_ScatterOp: LHLO_Op<"scatter", []>, BASE_HLO_ScatterOp {
+  let arguments = (ins
+    Arg:$operand,
+    Arg:$scatter_indices,
+    Arg:$updates,
+    Arg:$output,
+    ScatterDimensionNumbers:$scatter_dimension_numbers,
+    DefaultValuedAttr:$indices_are_sorted,
+    DefaultValuedAttr:$unique_indices
+  );
+
+  let regions = (region SizedRegion<1>:$update_computation);
+}
 
 def LHLO_SelectOp: LHLO_Op<"select", []>, BASE_HLO_SelectOp {
   let arguments = (ins
@@ -712,6 +731,44 @@ def LHLO_TriangularSolveOp: LHLO_Op<"triangular_solve", [SameOperandsElementType
   );
 }
 
+// TODO(timshen): add a custom verifier.
+def LHLO_MapOp: LHLO_Op<"map", [SameOperandsShape]>, BASE_HLO_MapOp {
+  let arguments = (ins
+    Arg, "", [MemRead]>:$operands,
+    Arg:$output,
+    I64ElementsAttr:$dimensions
+  );
+  let regions = (region SizedRegion<1>:$computation);
+}
+
+def LHLO_RngGetAndUpdateStateOp: LHLO_Op<"rng_get_and_update_state", []> {
+  let arguments = (ins
+    Arg, "", [MemRead, MemWrite]>:$state,
+    I64Attr:$delta
+  );
+}
+
+// TODO(timshen): add a custom verifier.
+def LHLO_SortOp: LHLO_Op<"sort", []>, BASE_HLO_SortOp {
+  let arguments = (ins
+    Arg, "", [MemRead]>:$operands,
+    LHLO_BufferOrTuple:$output,
+    DefaultValuedAttr:$dimension,
+    DefaultValuedAttr:$is_stable
+  );
+
+  let regions = (region SizedRegion<1>:$comparator);
+}
+
+def LHLO_TupleSelectOp: LHLO_Op<"tuple_select", [SameOperandsShape]> {
+  let arguments = (ins
+    Arg:$pred,
+    Arg:$on_true,
+    Arg:$on_false,
+    Arg:$output
+  );
+}
+
 //===----------------------------------------------------------------------===//
 // Late operations
 //===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/xla/tests/lhlo_ops.mlir b/tensorflow/compiler/mlir/xla/tests/lhlo_ops.mlir
index 6747b1e536e..0ed8b36466e 100644
--- a/tensorflow/compiler/mlir/xla/tests/lhlo_ops.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/lhlo_ops.mlir
@@ -863,3 +863,124 @@ func @while_memrefs(%arg0: memref, %arg_out: memref) -> () {
   ) : (memref, memref) -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @bitcast_memrefs
+func @bitcast_memrefs(%arg0: memref<1xf64>, %arg_out: memref<2xi32>) -> () {
+  "xla_lhlo.bitcast"(%arg0, %arg_out) : (memref<1xf64>, memref<2xi32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @scatter_memrefs
+func @scatter_memrefs(%input: memref<200x100x300xf32>, %indices: memref<10x2xi32>,
+                      %updates: memref<10x300xf32>, %arg_out: memref<200x100x300xf32>) -> () {
+  "xla_lhlo.scatter" (%input, %indices, %updates, %arg_out) ({
+  ^bb0(%lhs: tensor, %rhs: tensor): // no predecessors
+    %add = xla_hlo.add %lhs, %rhs : tensor
+    "xla_hlo.return"(%add) : (tensor) -> ()
+  }) {
+    scatter_dimension_numbers = {
+      update_window_dims = dense<[1]> : tensor<1xi64>,
+      inserted_window_dims = dense<[0, 1]> : tensor<2xi64>,
+      scatter_dims_to_operand_dims = dense<[0, 1]> : tensor<2xi64>,
+      index_vector_dim = 1 : i64
+    },
+    indices_are_sorted = true,
+    unique_indices = true
+  } : (memref<200x100x300xf32>, memref<10x2xi32>, memref<10x300xf32>, memref<200x100x300xf32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @map_memrefs
+func @map_memrefs(%arg0: memref<20xf32>, %arg1: memref<20xf32>, %arg_out: memref<20xf32>) -> () {
+  "xla_lhlo.map"(%arg0, %arg1, %arg_out) ({
+    ^bb0(%a: tensor, %b: tensor):
+    %c = xla_hlo.add %a, %b : tensor
+    "xla_hlo.return"(%c) : (tensor) -> ()
+  }) {dimensions = dense<0> : tensor<1xi64>} : (memref<20xf32>, memref<20xf32>, memref<20xf32>) -> ()
+  return
+}
+
+// -----
+
+func @map_memrefs(%arg0: memref<20xf32>, %arg1: memref<20xf32>, %arg_out: memref<10xf32>) -> () {
+  // expected-error@+1{{requires the same shape for all operands}}
+  "xla_lhlo.map"(%arg0, %arg1, %arg_out) ({
+    ^bb0(%a: tensor, %b: tensor):
+    %c = xla_hlo.add %a, %b : tensor
+    "xla_hlo.return"(%c) : (tensor) -> ()
+  }) {dimensions = dense<0> : tensor<1xi64>} : (memref<20xf32>, memref<20xf32>, memref<10xf32>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @rng_get_and_update_state_memrefs
+func @rng_get_and_update_state_memrefs(%state: memref<1xui64>) -> () {
+  "xla_lhlo.rng_get_and_update_state"(%state) { delta = 1 : i64 } : (memref<1xui64>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @sort_memrefs
+func @sort_memrefs(%arg0: memref<16x16xf32>, %arg1: memref<16x16xf16>,
+                   %arg_out: tuple, memref<16x16xf16>>) -> () {
+  "xla_lhlo.sort"(%arg0, %arg1, %arg_out) ( {
+  ^bb0(%a: tensor, %b: tensor, %c: tensor, %d: tensor):
+    %7 = "xla_hlo.compare"(%a, %b) {comparison_direction = "GT"} : (tensor, tensor) -> tensor
+    "xla_hlo.return"(%7) : (tensor) -> ()
+  }) {dimension = 1 : i64, is_stable = true} : (memref<16x16xf32>, memref<16x16xf16>, tuple, memref<16x16xf16>>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @sort_memrefs
+func @sort_memrefs(%arg0: memref<16x16xf32>, %arg1: memref<16x16xf16>,
+                   %arg_out: tuple, memref<16x16xf16>>) -> () {
+  "xla_lhlo.sort"(%arg0, %arg1, %arg_out) ( {
+  ^bb0(%a: tensor, %b: tensor, %c: tensor, %d: tensor):
+    %7 = "xla_hlo.compare"(%a, %b) {comparison_direction = "GT"} : (tensor, tensor) -> tensor
+    "xla_hlo.return"(%7) : (tensor) -> ()
+  }) {dimension = 1 : i64} : (memref<16x16xf32>, memref<16x16xf16>, tuple, memref<16x16xf16>>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @sort_memrefs
+func @sort_memrefs(%arg0: memref<16x16xf32>, %arg1: memref<16x16xf16>,
+                   %arg_out: tuple, memref<16x16xf16>>) -> () {
+  "xla_lhlo.sort"(%arg0, %arg1, %arg_out) ( {
+  ^bb0(%a: tensor, %b: tensor, %c: tensor, %d: tensor):
+    %7 = "xla_hlo.compare"(%a, %b) {comparison_direction = "GT"} : (tensor, tensor) -> tensor
+    "xla_hlo.return"(%7) : (tensor) -> ()
+  }) : (memref<16x16xf32>, memref<16x16xf16>, tuple, memref<16x16xf16>>) -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @tuple_select_memrefs
+func @tuple_select_memrefs(%pred: memref<20xi1>, %true_values: memref<20xf32>,
+                           %false_values: memref<20xf32>, %arg_out: memref<20xf32>) -> () {
+  "xla_lhlo.tuple_select"(%pred, %true_values, %false_values, %arg_out)
+      : (memref<20xi1>, memref<20xf32>, memref<20xf32>, memref<20xf32>) -> ()
+  return
+}
+
+// -----
+
+func @tuple_select_memrefs(%pred: memref<10xi1>, %true_values: memref<20xf32>,
+                           %false_values: memref<20xf32>, %arg_out: memref<20xf32>) -> () {
+  // expected-error@+1{{requires the same shape for all operands}}
+  "xla_lhlo.tuple_select"(%pred, %true_values, %false_values, %arg_out)
+      : (memref<10xi1>, memref<20xf32>, memref<20xf32>, memref<20xf32>) -> ()
+  return
+}

From dd3e910e21e3c22e2f4c8172d48386642527bbc4 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 11:57:20 -0700
Subject: [PATCH 0311/1390] Changed address space of default samplers.

PiperOrigin-RevId: 316728298
Change-Id: Ib74d1bbbb5561a29aa569e2960a3776eceb9f237
---
 tensorflow/lite/delegates/gpu/cl/kernels/util.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
index 73fc0e744bb..3161a73a18f 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
@@ -119,13 +119,13 @@ std::string GetCommonDefines(CalculationsPrecision precision) {
   }
 
   result +=
-      "const sampler_t smp_edge = CLK_NORMALIZED_COORDS_FALSE | "
+      "__constant sampler_t smp_edge = CLK_NORMALIZED_COORDS_FALSE | "
       "CLK_ADDRESS_CLAMP_TO_EDGE | CLK_FILTER_NEAREST;\n";
   result +=
-      "const sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | "
+      "__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | "
       "CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;\n";
   result +=
-      "const sampler_t smp_zero = CLK_NORMALIZED_COORDS_FALSE | "
+      "__constant sampler_t smp_zero = CLK_NORMALIZED_COORDS_FALSE | "
       "CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n";
 
   return result;

From 266387ecb45ca06a2bdf117a119cb65b9ab5dcd9 Mon Sep 17 00:00:00 2001
From: nammbash 
Date: Tue, 16 Jun 2020 12:02:14 -0700
Subject: [PATCH 0312/1390] create a single oneDNN string

---
 tensorflow/core/platform/cpu_feature_guard.cc | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/platform/cpu_feature_guard.cc b/tensorflow/core/platform/cpu_feature_guard.cc
index a020d3fd70e..13116de49dc 100644
--- a/tensorflow/core/platform/cpu_feature_guard.cc
+++ b/tensorflow/core/platform/cpu_feature_guard.cc
@@ -138,16 +138,9 @@ void InfoAboutUnusedCPUFeatures() {
     CheckIfFeatureUnused(CPUFeature::FMA, "FMA", missing_instructions);
 #endif  // __FMA__
 #endif  // else of if defined(_MSC_VER) && !defined(__clang__)
-
-    string intel_library_official_name(
-        "Intel(R) oneAPI Deep Neural Network Library (oneDNN) ");
-#ifndef INTEL_MKL
-    intel_library_official_name = "oneAPI Deep Neural Network Library (oneDNN) ";
-#endif
-
-        if (!missing_instructions.empty()) {
+    if (!missing_instructions.empty()) {
       LOG(INFO) << "This TensorFlow binary is optimized with "
-                << intel_library_official_name
+                << "oneAPI Deep Neural Network Library (oneDNN)"
                 << "to use the following CPU instructions in performance-"
                 << "critical operations: " << missing_instructions << std::endl
                 << "To enable them in other operations, rebuild TensorFlow "

From aee694363c354881ee659cd6a70fafc5a3d30507 Mon Sep 17 00:00:00 2001
From: Kibeom Kim 
Date: Tue, 16 Jun 2020 12:22:48 -0700
Subject: [PATCH 0313/1390] Support TFRT async config.

PiperOrigin-RevId: 316733436
Change-Id: I0eef2279b9c77d6084c1300a8d38a987a2cee065
---
 tensorflow/c/eager/c_api.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc
index e71073ec79f..fdc91675f8b 100644
--- a/tensorflow/c/eager/c_api.cc
+++ b/tensorflow/c/eager/c_api.cc
@@ -713,8 +713,8 @@ TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) {
     status->status = tfrt::ListOpHandlerChains(
         opts->session_options.options, &op_handler_chains, &device_attributes);
     if (!status->status.ok()) return nullptr;
-    return tensorflow::wrap(
-        new tfrt::ContextInterface(op_handler_chains, device_attributes));
+    return tensorflow::wrap(new tfrt::ContextInterface(
+        op_handler_chains, device_attributes, opts->async));
 #else
     status->status = tensorflow::errors::Unimplemented("TFRT is not supported");
     return nullptr;

From 1a03ea7e61d54135047d50c942af346d50b5af8b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 12:31:05 -0700
Subject: [PATCH 0314/1390] PR #39767: Wider vector for FP16 RELU Grad on GPUs

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/39767

This PR uses wider vector (8 FP16 values) for loading and storing in FP16 ReluGrad kernel to improve performance on Nvidia Ampere GPUs.

For older GPUs, the performance is expected to be unchanged.

fyi @nluehr
Copybara import of the project:

--
ab809024a4a5b0887c360b3e5542c149f4a5f14d by Kaixi Hou :

Enable wider vector for reluGrad...

PiperOrigin-RevId: 316735014
Change-Id: Ic4d93a211e52844f9804ed6c2c4a0346052ceb1e
---
 tensorflow/core/kernels/relu_op_gpu.cu.cc | 87 ++---------------------
 1 file changed, 7 insertions(+), 80 deletions(-)

diff --git a/tensorflow/core/kernels/relu_op_gpu.cu.cc b/tensorflow/core/kernels/relu_op_gpu.cu.cc
index 568bed57c70..27fd5f64249 100644
--- a/tensorflow/core/kernels/relu_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/relu_op_gpu.cu.cc
@@ -35,7 +35,6 @@ namespace tensorflow {
 
 typedef Eigen::GpuDevice GPUDevice;
 
-static constexpr int VectorSizeElements = 8;
 namespace functor {
 
 // This kernel computes ReluGrad by processing one half2, two fp16, at a time.
@@ -94,65 +93,6 @@ __global__ void ReluGradHalfKernel(const Eigen::half* __restrict__ gradient,
   }
 }
 
-__global__ void ReluGradHalfKernelVector(
-    const Eigen::half* __restrict__ gradient,
-    const Eigen::half* __restrict__ feature, Eigen::half* __restrict__ backprop,
-    int32 count) {
-  int32 half8_count = count / VectorSizeElements;
-  int32 index = blockIdx.x * blockDim.x + threadIdx.x;
-
-  if (index < half8_count) {
-    // Cast to xx_h8 for vector load and store.
-    float4 gradient_h8 = reinterpret_cast(gradient)[index];
-    float4 feature_h8 = reinterpret_cast(feature)[index];
-    float4* p_backprop_h8 = reinterpret_cast(backprop) + index;
-
-    half2* gradient_h2 = reinterpret_cast(&gradient_h8);
-    half2* feature_h2 = reinterpret_cast(&feature_h8);
-    float4 backprop_h8;
-    half2* p_backprop_h2 = reinterpret_cast(&backprop_h8);
-
-    // Fast path, when half2 primitives are available.
-#if __CUDA_ARCH__ >= 530
-    const half2 kZeroH2 = __float2half2_rn(0.f);
-#endif
-    for (int i = 0; i < VectorSizeElements / 2; i++) {
-#if __CUDA_ARCH__ >= 530
-      // mask = (feature > 0)
-      half2 mask_h2 = __hgt2(feature_h2[i], kZeroH2);
-      // backprop = mask * gradient
-      half2 backprop_h2 = __hmul2(mask_h2, gradient_h2[i]);
-#else
-      // Fall back: convert half2 to float2 for processing.
-      float2 feature_f2 = __half22float2(feature_h2[i]);
-      float2 gradient_f2 = __half22float2(gradient_h2[i]);
-      float2 backprop_f2 = make_float2((feature_f2.x > 0) ? gradient_f2.x : 0,
-                                       (feature_f2.y > 0) ? gradient_f2.y : 0);
-      // Convert back to half2.
-      half2 backprop_h2 = __float22half2_rn(backprop_f2);
-#endif
-      p_backprop_h2[i] = backprop_h2;
-    }
-    // Write back the result.
-    *p_backprop_h8 = backprop_h8;
-  }
-
-  int remaining_count = (count % VectorSizeElements);
-
-  if (index < remaining_count) {
-    // Use first threads to process the remaining elements.
-    Eigen::half grad_h = gradient[half8_count * VectorSizeElements + index];
-    Eigen::half feature_h = feature[half8_count * VectorSizeElements + index];
-
-    float grad_f = static_cast(grad_h);
-    float feature_f = static_cast(feature_h);
-    float backprop_f = (feature_f > 0) ? grad_f : 0;
-
-    Eigen::half backprop_h(backprop_f);
-    backprop[half8_count * VectorSizeElements + index] = backprop_h;
-  }
-}
-
 template 
 struct ReluGrad {
   // Computes ReluGrad backprop.
@@ -168,28 +108,15 @@ struct ReluGrad {
     // NOTE: When the activation is exactly zero, we do not propagate the
     // associated gradient value. This allows the output of the Relu to be used,
     // as well as its input.
-    auto gradient_ptr = reinterpret_cast(gradient.data());
-    auto feature_ptr = reinterpret_cast(feature.data());
-    auto backprop_ptr = reinterpret_cast(backprop.data());
-    bool aligned = gradient_ptr % 16 == 0 && feature_ptr % 16 == 0 &&
-                   backprop_ptr % 16 == 0;
     int32 count = gradient.size();
-    constexpr int32 kThreadInBlock = 512;
     if (count == 0) return;
-    if (aligned) {
-      int32 half8_count = Eigen::divup(count, VectorSizeElements);
-      int32 kBlock = Eigen::divup(half8_count, kThreadInBlock);
-      TF_CHECK_OK(GpuLaunchKernel(
-          ReluGradHalfKernelVector, kBlock, kThreadInBlock, 0, d.stream(),
-          gradient.data(), feature.data(), backprop.data(), count));
-    } else {
-      int32 half2_count = Eigen::divup(count, 2);
-      GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize(
-          half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock);
-      TF_CHECK_OK(GpuLaunchKernel(
-          ReluGradHalfKernel, config.block_count, config.thread_per_block, 0,
-          d.stream(), gradient.data(), feature.data(), backprop.data(), count));
-    }
+    int32 half2_count = Eigen::divup(count, 2);
+    constexpr int32 kThreadInBlock = 512;
+    GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize(
+        half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock);
+    TF_CHECK_OK(GpuLaunchKernel(
+        ReluGradHalfKernel, config.block_count, config.thread_per_block, 0,
+        d.stream(), gradient.data(), feature.data(), backprop.data(), count));
   }
 };
 

From d97234dc564b2db7543b7f69517e5beefb89dc0e Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Tue, 16 Jun 2020 12:34:27 -0700
Subject: [PATCH 0315/1390] More compatibility fixes for typing.Generic:  *
 types.new_class is required in some distributions  * avoid calling
 `isinstance` on some function objects in python 3.6  * account for some
 strange zombie pointer issue on windows Required for #40132.

PiperOrigin-RevId: 316735720
Change-Id: I1b08ef5f18c77c9343d587562f50632336b684d5
---
 tensorflow/python/framework/test_util.py |  2 +-
 tensorflow/python/util/tf_should_use.py  | 18 +++++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index 2967bb3de84..a46bb7c9bda 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -736,7 +736,7 @@ def assert_no_new_tensors(f):
         return isinstance(obj,
                           (ops.Tensor, variables.Variable,
                            tensor_shape.Dimension, tensor_shape.TensorShape))
-      except ReferenceError:
+      except (ReferenceError, AttributeError):
         # If the object no longer exists, we don't care about it.
         return False
 
diff --git a/tensorflow/python/util/tf_should_use.py b/tensorflow/python/util/tf_should_use.py
index 1671b078fa3..41c3220f5ca 100644
--- a/tensorflow/python/util/tf_should_use.py
+++ b/tensorflow/python/util/tf_should_use.py
@@ -21,15 +21,12 @@ import copy
 import sys
 import textwrap
 import traceback
-
-import six  # pylint: disable=unused-import
-
+import types
 
 from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util import tf_decorator
-# pylint: enable=g-bad-import-order,g-import-not-at-top
 
 
 class _TFShouldUseHelper(object):
@@ -154,7 +151,18 @@ def _get_wrapper(x, tf_should_use_helper):
   tx = copy.deepcopy(type_x)
   # Prefer using __orig_bases__, which preserve generic type arguments.
   bases = getattr(tx, '__orig_bases__', tx.__bases__)
-  copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
+
+  # Use types.new_class when available, which is preferred over plain type in
+  # some distributions.
+  if sys.version_info >= (3, 5):
+    def set_body(ns):
+      ns.update(tx.__dict__)
+      return ns
+
+    copy_tx = types.new_class(tx.__name__, bases, exec_body=set_body)
+  else:
+    copy_tx = type(tx.__name__, bases, dict(tx.__dict__))
+
   copy_tx.__init__ = _new__init__
   copy_tx.__getattribute__ = _new__getattribute__
   copy_tx.mark_used = _new_mark_used

From e0266dbf39deac09315b764524835299b513926c Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Tue, 16 Jun 2020 12:48:20 -0700
Subject: [PATCH 0316/1390] Use `static_cast` instead of C-style casts.

PiperOrigin-RevId: 316738458
Change-Id: I54f2f2f43d31606246475df0eae8d20e673aee6b
---
 tensorflow/core/platform/types.h | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/platform/types.h b/tensorflow/core/platform/types.h
index 5f26dabda55..b2fefcaa960 100644
--- a/tensorflow/core/platform/types.h
+++ b/tensorflow/core/platform/types.h
@@ -37,18 +37,18 @@ namespace tensorflow {
 // Alias tensorflow::string to std::string.
 using std::string;
 
-static const uint8 kuint8max = ((uint8)0xFF);
-static const uint16 kuint16max = ((uint16)0xFFFF);
-static const uint32 kuint32max = ((uint32)0xFFFFFFFF);
-static const uint64 kuint64max = ((uint64)0xFFFFFFFFFFFFFFFFull);
-static const int8 kint8min = ((int8)~0x7F);
-static const int8 kint8max = ((int8)0x7F);
-static const int16 kint16min = ((int16)~0x7FFF);
-static const int16 kint16max = ((int16)0x7FFF);
-static const int32 kint32min = ((int32)~0x7FFFFFFF);
-static const int32 kint32max = ((int32)0x7FFFFFFF);
-static const int64 kint64min = ((int64)~0x7FFFFFFFFFFFFFFFll);
-static const int64 kint64max = ((int64)0x7FFFFFFFFFFFFFFFll);
+static const uint8 kuint8max = static_cast(0xFF);
+static const uint16 kuint16max = static_cast(0xFFFF);
+static const uint32 kuint32max = static_cast(0xFFFFFFFF);
+static const uint64 kuint64max = static_cast(0xFFFFFFFFFFFFFFFFull);
+static const int8 kint8min = static_cast(~0x7F);
+static const int8 kint8max = static_cast(0x7F);
+static const int16 kint16min = static_cast(~0x7FFF);
+static const int16 kint16max = static_cast(0x7FFF);
+static const int32 kint32min = static_cast(~0x7FFFFFFF);
+static const int32 kint32max = static_cast(0x7FFFFFFF);
+static const int64 kint64min = static_cast(~0x7FFFFFFFFFFFFFFFll);
+static const int64 kint64max = static_cast(0x7FFFFFFFFFFFFFFFll);
 
 // A typedef for a uint64 used as a short fingerprint.
 typedef uint64 Fprint;

From ed1d7d09aec54c8c277da957ed18d17ed6885711 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Tue, 16 Jun 2020 12:54:54 -0700
Subject: [PATCH 0317/1390] Implement Mul(Convert(Pred), operand) =>
 select(pred, operand, 0) optimization.

PiperOrigin-RevId: 316739811
Change-Id: Ica5e50c6639a9792ae1dd47eefd713021fb97533
---
 .../xla/service/algebraic_simplifier.cc       | 19 +++++++++++++++++++
 .../xla/service/hlo_creation_utils.cc         |  9 +++++++++
 .../compiler/xla/service/hlo_creation_utils.h |  5 +++++
 3 files changed, 33 insertions(+)

diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
index 98e3229b062..ce2a801fccd 100755
--- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc
+++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -2455,6 +2455,25 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) {
     return Status::OK();
   }
 
+  {
+    HloInstruction *convert_operand, *operand;
+    // Mul(Convert(Pred), operand) => select(pred, operand, 0)
+    if (Match(multiply,
+              m::MultiplyAnyOrder(
+                  m::Op(&operand),
+                  m::Convert(
+                      m::Op(&convert_operand)
+                          .WithShape(m::Shape().WithElementType(PRED)))))) {
+      HloInstruction* zero_like_multiply =
+          BroadcastZeros(computation_, multiply->shape().element_type(),
+                         multiply->shape().dimensions());
+      return ReplaceWithNewInstruction(
+          multiply, HloInstruction::CreateTernary(
+                        multiply->shape(), HloOpcode::kSelect, convert_operand,
+                        operand, zero_like_multiply));
+    }
+  }
+
   VLOG(10) << "trying transform [(A * C1) * C2 => A * (C1 * C2)]";
   HloInstruction *a, *c1, *c2;
   if (Match(multiply,
diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc
index dd174772c62..0f5267e9fbc 100644
--- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc
+++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc
@@ -539,6 +539,15 @@ HloInstruction* BroadcastZeros(HloComputation* computation,
                           /*result_shape_bounds=*/broadcast_dimensions);
 }
 
+HloInstruction* BroadcastOnes(HloComputation* computation,
+                              PrimitiveType element_type,
+                              absl::Span broadcast_dimensions) {
+  HloInstruction* one = computation->AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::One(element_type)));
+  return MakeBroadcastHlo(one, /*broadcast_dimensions=*/{},
+                          /*result_shape_bounds=*/broadcast_dimensions);
+}
+
 // Recursively creates a dummy op given a shape. Leaf nodes are broadcasted zero
 // while internal nodes are tuples.
 HloInstruction* CreateDummyOp(HloComputation::Builder* b, const Shape& shape) {
diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.h b/tensorflow/compiler/xla/service/hlo_creation_utils.h
index 3f2e3aa25a1..2ba753d3cdb 100644
--- a/tensorflow/compiler/xla/service/hlo_creation_utils.h
+++ b/tensorflow/compiler/xla/service/hlo_creation_utils.h
@@ -276,6 +276,11 @@ HloInstruction* BroadcastZeros(HloComputation* computation,
                                PrimitiveType element_type,
                                absl::Span broadcast_dimensions);
 
+// Same as above, but fill the tensor with ones.
+HloInstruction* BroadcastOnes(HloComputation* computation,
+                              PrimitiveType element_type,
+                              absl::Span broadcast_dimensions);
+
 // Creates a HLO computation that takes arguments of type `domain` and produces
 // a value of type `range`.
 StatusOr> CreateComputationWithSignature(

From 5f06c799857622c0a16d47ff985a6b8f1b5559bf Mon Sep 17 00:00:00 2001
From: Sung Jin Hwang 
Date: Tue, 16 Jun 2020 12:59:30 -0700
Subject: [PATCH 0318/1390] Limit reserve size by 2**16 in dataset batch op
 when drop_remainder is false.

Dataset batch is sometimes used to stack all the elements in the dataset. A
common pattern is to pass a very large batch size with drop_remainder == false.
The batch size should be larger than the dataset cardinality, but at the same
time, it should not be too large otherwise vector reserve in the batch op ends
up OOM.

This change limits the reserve size by 2**16 when drop_remainder is false. Then
the users may pass a large enough number like INT32_MAX or INT64_MAX to stack
all elements.

PiperOrigin-RevId: 316740746
Change-Id: I445ff66eff088363c802ec31c359e81f188d8047
---
 tensorflow/core/kernels/data/batch_dataset_op.cc | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/batch_dataset_op.cc b/tensorflow/core/kernels/data/batch_dataset_op.cc
index c915f80c2c6..cfeb63a4242 100644
--- a/tensorflow/core/kernels/data/batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/batch_dataset_op.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/kernels/data/batch_dataset_op.h"
 
+#include 
 #include 
 
 #include "tensorflow/core/framework/op_kernel.h"
@@ -49,6 +50,12 @@ class BatchDatasetOp::Dataset : public DatasetBase {
           bool parallel_copy, const DatasetBase* input, int op_version)
       : DatasetBase(DatasetContext(ctx)),
         batch_size_(batch_size),
+        // Dataset batch is sometimes used to stack all elements in the
+        // dataset. In such cases, a very large batch size (e.g., INT32_MAX)
+        // is passed with drop_remainder set to false. Avoid OOM in such case
+        // by limiting `reserve()` size by 2**16.
+        reserve_size_(drop_remainder ? batch_size
+                                     : std::min(batch_size, 1 << 16)),
         drop_remainder_(drop_remainder),
         parallel_copy_(parallel_copy),
         input_(input),
@@ -153,7 +160,7 @@ class BatchDatasetOp::Dataset : public DatasetBase {
           *end_of_sequence = true;
           return Status::OK();
         }
-        batch_elements.reserve(dataset()->batch_size_);
+        batch_elements.reserve(dataset()->reserve_size_);
         *end_of_sequence = false;
         for (int i = 0; i < dataset()->batch_size_ && !*end_of_sequence; ++i) {
           std::vector batch_element_tuple;
@@ -289,6 +296,7 @@ class BatchDatasetOp::Dataset : public DatasetBase {
   };
 
   const int64 batch_size_;
+  const int64 reserve_size_;
   const bool drop_remainder_;
   const bool parallel_copy_;
   const DatasetBase* const input_;

From 93a441910f93117238a6b5bc74a765d05081dfd3 Mon Sep 17 00:00:00 2001
From: Amit Patankar 
Date: Tue, 16 Jun 2020 13:01:22 -0700
Subject: [PATCH 0319/1390] Add a link to GitHub README for libtensorflow
 nightly binaries.

PiperOrigin-RevId: 316741133
Change-Id: Id38352fc0e52f3f1cbeef4256e4a380b8e0ed07e
---
 README.md | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index a76b1bfd0b7..54c9470b04b 100644
--- a/README.md
+++ b/README.md
@@ -61,7 +61,6 @@ commands.
 *Nightly binaries are available for testing using the
 [tf-nightly](https://pypi.python.org/pypi/tf-nightly) and
 [tf-nightly-cpu](https://pypi.python.org/pypi/tf-nightly-cpu) packages on PyPi.*
-
 #### *Try your first TensorFlow program*
 
 ```shell
@@ -114,6 +113,12 @@ Build Type               | Status
 **Android**              | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/android.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/android.html)               | [![Download](https://api.bintray.com/packages/google/tensorflow/tensorflow/images/download.svg)](https://bintray.com/google/tensorflow/tensorflow/_latestVersion)
 **Raspberry Pi 0 and 1** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/rpi01-py3.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/rpi01-py3.html)           | [Py3](https://storage.googleapis.com/tensorflow-nightly/tensorflow-1.10.0-cp34-none-linux_armv6l.whl)
 **Raspberry Pi 2 and 3** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/rpi23-py3.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/rpi23-py3.html)           | [Py3](https://storage.googleapis.com/tensorflow-nightly/tensorflow-1.10.0-cp34-none-linux_armv7l.whl)
+**Libtensorflow MacOS CPU** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-mac-cpu.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-mac-cpu.html)           | [GCS](https://storage.googleapis.com/libtensorflow-nightly)
+**Libtensorflow Linux CPU** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-linux-cpu.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-linux-cpu.html)           | [GCS](https://storage.googleapis.com/libtensorflow-nightly)
+**Libtensorflow Linux GPU** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-linux-gpu.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-linux-gpu.html)           | [GCS](https://storage.googleapis.com/libtensorflow-nightly)
+**Libtensorflow Windows CPU** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-win-cpu.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-win-cpu.html)           | [GCS](https://storage.googleapis.com/libtensorflow-nightly)
+**Libtensorflow Windows GPU** | [![Status](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-win-gpu.svg)](https://storage.googleapis.com/tensorflow-kokoro-build-badges/libtensorflow-win-gpu.html)           | [GCS](https://storage.googleapis.com/libtensorflow-nightly)
+
 
 ### Community Supported Builds
 

From 2a0ad4792645fde8563baa42f9ac476125e60fad Mon Sep 17 00:00:00 2001
From: Bruce Fontaine 
Date: Tue, 16 Jun 2020 13:07:46 -0700
Subject: [PATCH 0320/1390] Update tpu_embedding_v2.py to use the new API for
 prefetching data to host memory.

PiperOrigin-RevId: 316742491
Change-Id: I6803c798256578a284d9ef190d79bf2e35f9ce6a
---
 tensorflow/python/tpu/tpu_embedding_v2.py     |  54 +++++++-
 .../python/tpu/tpu_embedding_v2_test.py       | 126 +++++++++++++++---
 2 files changed, 152 insertions(+), 28 deletions(-)

diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index 90b43c1ebf4..f7a383c440c 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -31,6 +31,7 @@ from tensorflow.python.distribute import sharded_variable
 from tensorflow.python.distribute import tpu_strategy
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
@@ -139,6 +140,18 @@ class TPUEmbedding(tracking.AutoTrackable):
         optimizer=tf.tpu.experimental.embedding.SGD(0.1))
   ```
 
+  When creating a distributed dataset that is to be passed to the enqueue
+  operation a special input option must be specified:
+
+  ```python
+  distributed_dataset = (
+      strategy.experimental_distribute_datasets_from_function(
+          dataset_fn=...,
+          options=tf.distribute.InputOptions(
+              experimental_prefetch_to_device=False))
+  dataset_iterator = iter(distributed_dataset)
+  ```
+
   To use this API on TPU you should use a custom training loop. Below is an
   example of a training and evaluation step:
 
@@ -309,10 +322,6 @@ class TPUEmbedding(tracking.AutoTrackable):
       # We need to list of host devices for the load/retrieve operations.
       self._hosts = get_list_of_hosts(self._strategy)
 
-      # TODO(bfontain) Remove this once we have an official way of splitting
-      # prefetch between host and device.
-      self._strategy.extended._set_prefetch_on_host(True)  # pylint: disable=protected-access
-
       # We generally use the per core batch size, but will have the user pass
       # in a global batch size.
       self._batch_size = batch_size // self._strategy.num_replicas_in_sync
@@ -507,7 +516,11 @@ class TPUEmbedding(tracking.AutoTrackable):
     with strategy.scope():
       embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
 
-    distributed_dataset = strategy.experimental_distribute_dataset(...)
+    distributed_dataset = (
+        strategy.experimental_distribute_datasets_from_function(
+            dataset_fn=...,
+            options=tf.distribute.InputOptions(
+                experimental_prefetch_to_device=False))
     dataset_iterator = iter(distributed_dataset)
 
     @tf.function
@@ -594,7 +607,11 @@ class TPUEmbedding(tracking.AutoTrackable):
     with strategy.scope():
       embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
 
-    distributed_dataset = strategy.experimental_distribute_dataset(...)
+    distributed_dataset = (
+        strategy.experimental_distribute_datasets_from_function(
+            dataset_fn=...,
+            options=tf.distribute.InputOptions(
+                experimental_prefetch_to_device=False))
     dataset_iterator = iter(distributed_dataset)
 
     @tf.function
@@ -1004,6 +1021,24 @@ class TPUEmbedding(tracking.AutoTrackable):
                                                  input_tensor.op.name,
                                                  input_tensor.op.type))
 
+  def _raise_error_for_inputs_not_on_cpu(self, features):
+    """Checks all tensors in features to see are placed on the CPU."""
+
+    # expand_composites here is important, we need to check the device of each
+    # underlying tensor.
+    for path, input_tensor in nest.flatten_with_joined_string_paths(
+        features, expand_composites=True):
+      spec = tf_device.DeviceSpec.from_string(input_tensor.device)
+      if spec.device_type == "TPU":
+        raise ValueError(
+            "Received input tensor {} which is on a TPU input device {}. Input "
+            "tensors for TPU embeddings must be placed on the CPU. Please "
+            "ensure that your dataset is prefetching tensors to the host by "
+            "setting the 'experimental_prefetch_to_device' option of the "
+            "dataset distribution function. See the documentation of the "
+            "enqueue method for an example.".format(
+                path, input_tensor.device))
+
   def enqueue(self, features, weights=None, training=True, name=None):
     """Enqueues id tensors for embedding lookup.
 
@@ -1021,7 +1056,11 @@ class TPUEmbedding(tracking.AutoTrackable):
     with strategy.scope():
       embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
 
-    distributed_dataset = strategy.experimental_distribute_dataset(...)
+    distributed_dataset = (
+        strategy.experimental_distribute_datasets_from_function(
+            dataset_fn=...,
+            options=tf.distribute.InputOptions(
+                experimental_prefetch_to_device=False))
     dataset_iterator = iter(distributed_dataset)
 
     @tf.function
@@ -1091,6 +1130,7 @@ class TPUEmbedding(tracking.AutoTrackable):
       flat_weights = nest.flatten(weights)
     flat_features = nest.flatten_with_joined_string_paths(self._feature_config)
 
+    self._raise_error_for_inputs_not_on_cpu(features)
     in_tpu_context = self._raise_error_for_incorrect_control_flow_context()
     # If we are in a tpu_context, automatically apply outside compilation.
     if in_tpu_context:
diff --git a/tensorflow/python/tpu/tpu_embedding_v2_test.py b/tensorflow/python/tpu/tpu_embedding_v2_test.py
index 0c257010f6a..ebaf2791055 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2_test.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2_test.py
@@ -28,6 +28,7 @@ import numpy as np
 
 from tensorflow.python.compat import v2_compat
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
 from tensorflow.python.distribute import distribution_strategy_context
 from tensorflow.python.distribute import tpu_strategy
 from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
@@ -443,7 +444,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
   def test_pass_none_to_apply_gradients(self):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
     dataset = self._create_sparse_dataset(strategy)
-    data = next(iter(strategy.experimental_distribute_dataset(dataset)))
+    data = next(iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))))
 
     @def_function.function
     def embedding_and_set_gradients(data):
@@ -527,7 +531,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
 
     input_fn = self._create_dense_input_fn(strategy, include_weights=True)
-    dist = strategy.experimental_distribute_datasets_from_function(input_fn)
+    dist = strategy.experimental_distribute_datasets_from_function(
+        input_fn,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))
     dist_iter = iter(dist)
 
     @def_function.function
@@ -547,8 +554,14 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     sparse = self._create_sparse_dataset(strategy)
     ragged = self._create_ragged_dataset(strategy, include_weights=True)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
-    ragged_iter = iter(strategy.experimental_distribute_dataset(ragged))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
+    ragged_iter = iter(strategy.experimental_distribute_dataset(
+        ragged,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -569,8 +582,14 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     sparse = self._create_sparse_dataset(strategy, include_weights=True)
     ragged = self._create_ragged_dataset(strategy)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
-    ragged_iter = iter(strategy.experimental_distribute_dataset(ragged))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
+    ragged_iter = iter(strategy.experimental_distribute_dataset(
+        ragged,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -591,8 +610,14 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     sparse = self._create_sparse_dataset(strategy)
     ragged = self._create_ragged_dataset(strategy)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
-    ragged_iter = iter(strategy.experimental_distribute_dataset(ragged))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
+    ragged_iter = iter(strategy.experimental_distribute_dataset(
+        ragged,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -613,7 +638,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
 
     sparse = self._create_sparse_dataset(strategy)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -633,7 +661,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
 
     sparse = self._create_sparse_dataset(strategy, include_weights=True)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -654,8 +685,14 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     sparse = self._create_sparse_dataset(strategy)
     ragged = self._create_ragged_dataset(strategy)
-    sparse_iter = iter(strategy.experimental_distribute_dataset(sparse))
-    ragged_iter = iter(strategy.experimental_distribute_dataset(ragged))
+    sparse_iter = iter(strategy.experimental_distribute_dataset(
+        sparse,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
+    ragged_iter = iter(strategy.experimental_distribute_dataset(
+        ragged,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def test_fn():
@@ -678,6 +715,26 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     ragged0 = self._get_replica_numpy(ragged_activations, strategy, 0)
     self.assertAllClose(sparse0, ragged0)
 
+  def test_enqueue_cpu_tensor(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    input_fn = self._create_dense_input_fn(strategy)
+    sparse_iter = iter(strategy.experimental_distribute_datasets_from_function(
+        input_fn))
+
+    @def_function.function
+    def test_fn():
+      def get_activations():
+        return mid_level_api.dequeue()
+
+      sparse_features = next(sparse_iter)
+      mid_level_api.enqueue(sparse_features, training=False)
+      sparse_activations = strategy.run(get_activations)
+      return sparse_activations
+
+    with self.assertRaisesRegex(ValueError, 'which is on a TPU input device'):
+      test_fn()
+
   @parameterized.parameters(True, False)
   def test_enqueue_with_weights(self, ragged):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
@@ -689,7 +746,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
       dataset = self._create_sparse_dataset(strategy, include_weights=True,
                                             weight=weight)
 
-    dataset_iter = iter(strategy.experimental_distribute_dataset(dataset))
+    dataset_iter = iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def enqueue_and_get(features, weights):
@@ -727,7 +787,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
     dataset = self._create_sparse_dataset(strategy)
-    dataset_iter = iter(strategy.experimental_distribute_dataset(dataset))
+    dataset_iter = iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def enqueue_with_outside_compilation(data):
@@ -761,7 +824,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
 
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
     dataset = self._create_sparse_dataset(strategy)
-    dataset_iter = iter(strategy.experimental_distribute_dataset(dataset))
+    dataset_iter = iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     # This is one way to force the enqueue in some control flow. @tf.functions
     # aren't inlined in the calling tf.function. An alternative would be to
@@ -785,7 +851,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
   def test_enqueue_with_outside_compilation_non_direct_input(self):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
     dataset = self._create_sparse_dataset(strategy)
-    dataset_iter = iter(strategy.experimental_distribute_dataset(dataset))
+    dataset_iter = iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def enqueue_with_outside_compilation():
@@ -804,7 +873,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
   def test_enqueue_with_outside_compilation_auto_mode(self):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
     dataset = self._create_sparse_dataset(strategy)
-    dataset_iter = iter(strategy.experimental_distribute_dataset(dataset))
+    dataset_iter = iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False)))
 
     @def_function.function
     def enqueue_with_no_gradient_apply(data):
@@ -883,7 +955,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
         self._create_strategy_and_mid_level(optimizer_name))
 
     dataset = self._create_sparse_dataset(strategy)
-    dist = strategy.experimental_distribute_dataset(dataset)
+    dist = strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))
     dist_iter = iter(dist)
 
     @def_function.function
@@ -1175,7 +1250,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
 
     input_fn = self._create_dense_input_fn(strategy)
-    dist = strategy.experimental_distribute_datasets_from_function(input_fn)
+    dist = strategy.experimental_distribute_datasets_from_function(
+        input_fn,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))
     dist_iter = iter(dist)
 
     @def_function.function
@@ -1235,7 +1313,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
     def input_fn(ctx):
       del ctx
       return dataset_ops.DatasetV2.from_tensors(feature).repeat()
-    dist = strategy.experimental_distribute_datasets_from_function(input_fn)
+    dist = strategy.experimental_distribute_datasets_from_function(
+        input_fn,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))
     dist_iter = iter(dist)
 
     @def_function.function
@@ -1364,7 +1445,10 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
           optimizer=optimizer)
 
     dataset = self._create_sparse_dataset(strategy)
-    data = next(iter(strategy.experimental_distribute_dataset(dataset)))
+    data = next(iter(strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(
+            experimental_prefetch_to_device=False))))
 
     @def_function.function
     def embedding_and_set_gradients(data):

From 89df3ddcd5451aefe52c7deb684f220b0520a6b1 Mon Sep 17 00:00:00 2001
From: Katherine Wu 
Date: Tue, 16 Jun 2020 13:08:10 -0700
Subject: [PATCH 0321/1390] Raise error when calling loaded model with layers
 that are not fully serialized.

PiperOrigin-RevId: 316742578
Change-Id: Iefc40d21374388ed99f7ff40bb09436830b46cbe
---
 .../python/keras/saving/saved_model/load.py   | 51 ++++++++++++++++---
 .../saving/saved_model/saved_model_test.py    | 30 +++++++++++
 2 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/keras/saving/saved_model/load.py b/tensorflow/python/keras/saving/saved_model/load.py
index 313eea4342e..a378c1b98e7 100644
--- a/tensorflow/python/keras/saving/saved_model/load.py
+++ b/tensorflow/python/keras/saving/saved_model/load.py
@@ -690,18 +690,22 @@ def _finalize_saved_model_layers(layers):
           layer, _get_keras_attr(layer).call_and_return_conditional_losses,
           return_method=True)
       layer._init_call_fn_args()
+    else:
+      layer.call = types.MethodType(
+          _unable_to_call_layer_due_to_serialization_issue, layer)
 
   for layer in layers:
     # 2. Set model inputs and outputs.
     if isinstance(layer, RevivedNetwork):
       _set_network_attributes_from_metadata(layer)
 
-      call_fn = _get_keras_attr(layer).call_and_return_conditional_losses
-      if call_fn.input_signature is None:
-        inputs = infer_inputs_from_restored_call_function(call_fn)
-      else:
-        inputs = call_fn.input_signature[0]
-      layer._set_inputs(inputs)
+      if hasattr(_get_keras_attr(layer), 'call_and_return_conditional_losses'):
+        call_fn = _get_keras_attr(layer).call_and_return_conditional_losses
+        if call_fn.input_signature is None:
+          inputs = infer_inputs_from_restored_call_function(call_fn)
+        else:
+          inputs = call_fn.input_signature[0]
+        layer._set_inputs(inputs)  # pylint: disable=protected-access
 
     # 3. Add losses that aren't generated by the layer.call function.
     _restore_layer_unconditional_losses(layer)
@@ -713,6 +717,41 @@ def _finalize_saved_model_layers(layers):
   # pylint: enable=protected-access
 
 
+def _unable_to_call_layer_due_to_serialization_issue(
+    layer, *unused_args, **unused_kwargs):
+  """Replaces the `layer.call` if the layer was not fully serialized.
+
+  Keras Model/Layer serialization is relatively relaxed because SavedModels
+  are not always loaded back as keras models. Thus, when there is an issue
+  tracing a non-signature function, a warning is logged instead of raising an
+  error. This results in a SavedModel where the model's call function is saved,
+  but the internal layer call functions are not.
+
+  When deserialized with `tf.keras.models.load_model`, the internal layers
+  which do not have serialized call functions should raise an error when called.
+
+  Args:
+    layer: Layer without the serialized call function.
+
+  Raises:
+    ValueError
+  """
+
+  raise ValueError(
+      'Cannot call {} ({}), because the call function was not serialized to '
+      'the SavedModel (due to lack information about the inputs). Please try '
+      'one of the following methods to fix the serialization:'
+      '\n\n(1) Implement `get_config` and `from_config` in the layer/model '
+      'class, and pass the object to the `custom_objects` argument when '
+      'loading the model. For more details, see: '
+      'https://www.tensorflow.org/guide/keras/save_and_serialize'
+      '\n\n(2) Ensure that the subclassed model or layer overwrites `call` '
+      'and not `__call__`. The input shape and dtype will be automatically '
+      'recorded when the object is called, and used when saving. To manually '
+      'specify the input shape/dtype, decorate the call function with '
+      '`@tf.function(input_signature=...)`.'.format(layer.name, layer))
+
+
 def _finalize_config_layers(layers):
   """Runs the final steps of loading Keras Layers from config."""
   for layer in layers:
diff --git a/tensorflow/python/keras/saving/saved_model/saved_model_test.py b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
index 7eaa75b78e2..8d4d27e2357 100644
--- a/tensorflow/python/keras/saving/saved_model/saved_model_test.py
+++ b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
@@ -809,6 +809,36 @@ class TestModelSavingAndLoadingV2(keras_parameterized.TestCase):
     self.evaluate(variables.variables_initializer(loaded.variables))
     self.assertAllClose(model.predict(f), loaded.predict(f))
 
+  def test_load_with_partially_failed_serialization(self):
+
+    class BadCustomLayer(keras.layers.Layer):
+
+      def __call__(self, inputs):
+        return inputs
+
+    class Model(keras.models.Model):
+
+      def __init__(self):
+        super(Model, self).__init__()
+        self.layer = BadCustomLayer()
+
+      @def_function.function(
+          input_signature=[tensor_spec.TensorSpec([None, 1])])
+      def call(self, inputs):
+        return self.layer(inputs)
+
+    model = Model()
+    inp = constant_op.constant([[1.0]])
+    model(inp)
+    saved_model_dir = self._save_model_dir()
+    tf_save.save(model, saved_model_dir)
+
+    loaded = keras_load.load(saved_model_dir)
+    self.assertAllEqual([[1.0]], self.evaluate(loaded(inp)))
+    with self.assertRaisesRegexp(ValueError,
+                                 'call function was not serialized'):
+      loaded.layer(inp)
+
 
 class TestLayerCallTracing(test.TestCase, parameterized.TestCase):
 

From 1a342fb760f6a3b5ff29b0b4c47bbf07a7bb1221 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Tue, 16 Jun 2020 13:11:14 -0700
Subject: [PATCH 0322/1390] Explicitly take the set of devices in
 CollectiveAllReduce

We used to infer the devices from the inputs, but sometimes the inputs don't
have device placement. E.g. when passing into or returning from tf.function, the
device placement may be lost.

Instead of inferring from the inputs we should just be explicit about the
collective devices.

PiperOrigin-RevId: 316743112
Change-Id: I2f6995f2f4cc86864723e203deb7562363cdbc38
---
 tensorflow/python/distribute/BUILD            |  9 +-
 .../collective_all_reduce_strategy.py         | 45 +++++++--
 .../collective_all_reduce_strategy_test.py    |  3 +-
 .../python/distribute/cross_device_ops.py     | 77 +++++++++-------
 .../distribute/cross_device_ops_test.py       | 80 ++++++++++++----
 .../python/distribute/cross_device_utils.py   | 91 +++++++++----------
 .../python/distribute/mirrored_strategy.py    | 19 ++--
 .../python/distribute/strategy_common_test.py | 47 +++++++++-
 8 files changed, 252 insertions(+), 119 deletions(-)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 77ef98d1cb7..a3655699669 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -1029,14 +1029,21 @@ cuda_py_test(
         ":collective_util",
         ":combinations",
         ":cross_device_ops",
-        ":mirrored_strategy",
+        ":cross_device_utils",
+        ":device_util",
         ":multi_worker_test_base",
+        ":multi_worker_util",
+        ":reduce_util",
         ":strategy_combinations",
         ":values",
         "//tensorflow/python:array_ops",
+        "//tensorflow/python:collective_ops",
         "//tensorflow/python:constant_op",
         "//tensorflow/python:framework_ops",
+        "//tensorflow/python:kernels",
         "//tensorflow/python:math_ops",
+        "//tensorflow/python:variables",
+        "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:test",
         "@absl_py//absl/testing:parameterized",
diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy.py b/tensorflow/python/distribute/collective_all_reduce_strategy.py
index 23ed16c5cfd..68cc421c21b 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy.py
@@ -175,7 +175,7 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
     self._communication = communication
     self._initialize_strategy(self._cluster_resolver)
     self._cfer_fn_cache = weakref.WeakKeyDictionary()
-    assert isinstance(self._get_cross_device_ops(),
+    assert isinstance(self._cross_device_ops,
                       cross_device_ops_lib.CollectiveAllReduce)
 
   def _initialize_strategy(self, cluster_resolver):
@@ -217,12 +217,18 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
     self._host_input_device = numpy_dataset.SingleDevice(self._worker_device)
 
     self._collective_keys = cross_device_utils.CollectiveKeys()
-    # TODO(yuefengz): remove num_gpus_per_worker from CollectiveAllReduce.
     self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
-        num_workers=self._num_workers,
-        num_gpus_per_worker=num_gpus,
+        devices=local_devices,
+        group_size=len(local_devices),
         collective_keys=self._collective_keys,
         communication=self._communication)
+    # CrossDeviceOps for per host tensors.
+    self._host_cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
+        devices=[self._worker_device],
+        group_size=self._num_workers,
+        collective_keys=self._collective_keys,
+        communication=cross_device_ops_lib.CollectiveCommunication.RING,
+    )
     super(CollectiveAllReduceExtended, self)._initialize_single_worker(
         local_devices)
 
@@ -324,10 +330,17 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
 
     self._collective_keys = cross_device_utils.CollectiveKeys()
     self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
-        num_workers=self._num_workers,
-        num_gpus_per_worker=num_gpus,
+        devices=local_devices,
+        group_size=len(local_devices) * self._num_workers,
         collective_keys=self._collective_keys,
         communication=self._communication)
+    # CrossDeviceOps for per host tensors.
+    self._host_cross_device_ops = cross_device_ops_lib.CollectiveAllReduce(
+        devices=[self._worker_device],
+        group_size=self._num_workers,
+        collective_keys=self._collective_keys,
+        communication=cross_device_ops_lib.CollectiveCommunication.RING,
+    )
     super(CollectiveAllReduceExtended, self)._initialize_single_worker(
         local_devices)
     host_device = device_util.get_host_for_device(self._worker_device)
@@ -474,7 +487,7 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
           num_accelerators={"GPU": self._num_gpus_per_worker},
           rpc_layer=self._rpc_layer)
       self._initialize_multi_worker(cluster_resolver)
-      assert isinstance(self._get_cross_device_ops(),
+      assert isinstance(self._cross_device_ops,
                         cross_device_ops_lib.CollectiveAllReduce)
 
     if session_config:
@@ -518,6 +531,22 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
 
     return updated_config
 
+  def _get_cross_device_ops(self, value):
+    # CollectiveAllReduce works on a predefined set of devices. In most cases
+    # they should be the compute devices, but certain use cases may reduce host
+    # tensors as well (e.g. early stopping). We infer the cross_device_ops to
+    # use based on the number of devices, since inputs don't always have device
+    # annotations. The compute devices one is preferred since we can potentially
+    # leverage NCCL.
+    if isinstance(value, values.DistributedValues):
+      num_devices = len(value._values)  # pylint: disable=protected-access
+    else:
+      num_devices = 1
+    if num_devices == len(self.worker_devices):
+      return self._cross_device_ops
+    else:
+      return self._host_cross_device_ops
+
   def _reduce_to(self, reduce_op, value, destinations, experimental_hints):
     if (isinstance(value, values.Mirrored) and
         reduce_op == reduce_util.ReduceOp.MEAN):
@@ -538,7 +567,7 @@ class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended):
       # be 0.
       return cross_device_ops_lib.reduce_non_distributed_value(
           reduce_op, value, destinations, len(self.worker_devices))
-    return self._get_cross_device_ops().reduce(
+    return self._get_cross_device_ops(value).reduce(
         reduce_op,
         value,
         destinations=destinations,
diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy_test.py b/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
index 63c0db4c3b3..a9f7bc74e9e 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
@@ -116,7 +116,8 @@ class CollectiveAllReduceStrategyTestBase(
         variable_instance_key_start=10000 +
         CollectiveAllReduceStrategyTestBase.collective_key_base)
     strategy.extended._collective_keys = collective_keys
-    strategy.extended._cross_device_ops._collective_keys = (collective_keys)
+    strategy.extended._cross_device_ops._collective_keys = collective_keys
+    strategy.extended._host_cross_device_ops._collective_keys = collective_keys
 
     return strategy, target, session_config
 
diff --git a/tensorflow/python/distribute/cross_device_ops.py b/tensorflow/python/distribute/cross_device_ops.py
index b88357e0ea6..ed6b0558b46 100644
--- a/tensorflow/python/distribute/cross_device_ops.py
+++ b/tensorflow/python/distribute/cross_device_ops.py
@@ -165,7 +165,8 @@ def get_devices_from(destinations):
 
 
 def _devices_match(left, right):
-  return set(get_devices_from(left)) == set(get_devices_from(right))
+  return left is right or set(get_devices_from(left)) == set(
+      get_devices_from(right))
 
 
 def _all_devices_match(value_destination_pairs):
@@ -936,20 +937,24 @@ class CollectiveAllReduce(CrossDeviceOps):
   """
 
   def __init__(self,
-               num_workers=1,
-               num_gpus_per_worker=0,
+               devices,
+               group_size,
                collective_keys=None,
                communication=CollectiveCommunication.AUTO):
     """Initializes the object.
 
     Args:
-      num_workers: number of workers in the between-graph replicated training.
-      num_gpus_per_worker: number of GPUs per worker.
+      devices: a list of device strings to run collectives on.
+      group_size: the global group size. For between-graph replicated training
+        it's the total number of devices across all workers.
       collective_keys: an optional CollectiveKey object.
       communication: indicates which collective communication to use.
     """
-    self._num_workers = num_workers
-    self._num_gpus_per_worker = num_gpus_per_worker
+    if group_size % len(devices) > 0:
+      raise ValueError("group_size must be divisible by the number of devices.")
+
+    self._devices = tuple(device_util.canonicalize(d) for d in devices)
+    self._group_size = group_size
     self._collective_keys = (collective_keys or
                              cross_device_utils.CollectiveKeys())
     self._communication = communication
@@ -963,15 +968,15 @@ class CollectiveAllReduce(CrossDeviceOps):
     # async executor operations are still executed sequentially. In graph or
     # function building, the executors are not used.
     self._executors = []
-    for _ in range(self._num_gpus_per_worker or 1):
-      # If num_gpus_per_worker is zero, we assume there's only one device (CPU).
+    for _ in range(len(devices)):
       self._executors.append(executor.new_executor(enable_async=True))
 
     super(CollectiveAllReduce, self).__init__()
 
   @property
   def _num_between_graph_workers(self):
-    return self._num_workers
+    # Currently we only support equal number of devices on each worker.
+    return self._group_size / len(self._devices)
 
   def reduce_implementation(self, reduce_op, per_replica_value, destinations,
                             experimental_hints):
@@ -979,8 +984,7 @@ class CollectiveAllReduce(CrossDeviceOps):
                                          experimental_hints)[0]
     devices = get_devices_from(destinations)
 
-    if (isinstance(all_reduced, value_lib.Mirrored) and
-        (all_reduced._devices == devices)):  # pylint: disable=protected-access
+    if _devices_match(per_replica_value, destinations):
       return all_reduced
 
     # Convert `all_reduced` to a `Mirrored` object, as a simple and uniform
@@ -1069,14 +1073,16 @@ class CollectiveAllReduce(CrossDeviceOps):
 
     if batch_size > 1:
       logging.info(
-          "Collective batch_all_reduce: %d all-reduces, num_workers = %d, "
-          "communication_hint = %s, num_packs = %d", batch_size,
-          self._num_workers, communication, len(packs))
+          "Collective batch_all_reduce: %d all-reduces, num_devices = %d, "
+          "group_size = %d, communication_hint = %s, num_packs = %d",
+          batch_size, len(self._devices), self._group_size, communication,
+          len(packs))
     else:
       logging.log_first_n(
           logging.INFO, "Collective batch_all_reduce: %d all-reduces, "
-          "num_workers = %d, communication_hint = %s, num_packs = %d" %
-          (batch_size, self._num_workers, communication, len(packs)), 10)
+          "num_devices = %d, group_size = %d, communication_hint = %s, "
+          "num_packs = %d" % (batch_size, len(
+              self._devices), self._group_size, communication, len(packs)), 10)
 
     reduced_values = []
     for pack in packs:
@@ -1094,21 +1100,25 @@ class CollectiveAllReduce(CrossDeviceOps):
             control_inputs = None
           reduced_values.append(
               cross_device_utils.build_collective_reduce(
-                  per_replica.values, self._num_workers,
-                  self._collective_keys, "Add", "Id", communication,
-                  control_inputs, executors=self._executors))
+                  per_replica.values,
+                  self._devices,
+                  self._group_size,
+                  self._collective_keys,
+                  "Add",
+                  "Id",
+                  communication,
+                  control_inputs,
+                  executors=self._executors))
 
     mirrored = []
     # Reverse the order of reduced value to recover the order in the input.
     for value in reversed(reduced_values):
       if reduce_op == reduce_util.ReduceOp.MEAN:
-        # Assume each worker has the same number of replicas.
-        num_replicas = len(value) * self._num_workers
         for i, v in enumerate(value):
           with ops.device(v.device):
-            value[i] = v / num_replicas
-      mirrored.append(distribute_utils.regroup(value,
-                                               wrap_class=value_lib.Mirrored))
+            value[i] = v / self._group_size
+      mirrored.append(
+          distribute_utils.regroup(value, wrap_class=value_lib.Mirrored))
     return mirrored
 
   def _do_batch_all_reduce_sparse(self, reduce_op, per_replica_values):
@@ -1116,8 +1126,8 @@ class CollectiveAllReduce(CrossDeviceOps):
 
     logging.log_first_n(
         logging.INFO, "Collective batch_all_reduce for IndexedSlices: "
-        "%d all-reduces, num_workers = %d" %
-        (len(per_replica_values), self._num_workers), 10)
+        "%d all-reduces, group_size = %d" %
+        (len(per_replica_values), self._group_size), 10)
 
     # Pass self._communication to the runtime as a communication hint.
     communication_hint = self._communication.value
@@ -1133,25 +1143,24 @@ class CollectiveAllReduce(CrossDeviceOps):
       for per_replica in per_replica_values:
         gathered_values.append(
             cross_device_utils.build_collective_gather_indexed_slices(
-                per_replica.values, self._num_workers, self._collective_keys,
-                communication_hint))
+                per_replica.values, self._devices, self._group_size,
+                self._collective_keys, communication_hint))
 
     mirrored = []
     for value in gathered_values:
       if reduce_op == reduce_util.ReduceOp.MEAN:
         # Assume each worker has the same number of replicas.
-        num_replicas = len(value) * self._num_workers
         for i, v in enumerate(value):
           with ops.device(v.device):
-            value[i].values = value[i].values / num_replicas
-      mirrored.append(distribute_utils.regroup(value,
-                                               wrap_class=value_lib.Mirrored))
+            value[i].values = value[i].values / self._group_size
+      mirrored.append(
+          distribute_utils.regroup(value, wrap_class=value_lib.Mirrored))
     return mirrored
 
   def __deepcopy__(self, memo):
     # distribute_coordinator deep-copies the strategy object, so
     # CollectiveAllReduce needs to support deep copy as well.
-    return CollectiveAllReduce(self._num_workers, self._num_gpus_per_worker,
+    return CollectiveAllReduce(self._devices, self._group_size,
                                self._collective_keys, self._communication)
 
 
diff --git a/tensorflow/python/distribute/cross_device_ops_test.py b/tensorflow/python/distribute/cross_device_ops_test.py
index 9554de41a6e..4b6943e8971 100644
--- a/tensorflow/python/distribute/cross_device_ops_test.py
+++ b/tensorflow/python/distribute/cross_device_ops_test.py
@@ -26,6 +26,7 @@ import time
 from absl.testing import parameterized
 import numpy as np
 from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.distribute import cluster_resolver
 from tensorflow.python.distribute import collective_all_reduce_strategy
 from tensorflow.python.distribute import collective_util
 from tensorflow.python.distribute import combinations
@@ -34,10 +35,12 @@ from tensorflow.python.distribute import cross_device_utils
 from tensorflow.python.distribute import device_util
 from tensorflow.python.distribute import distribute_utils
 from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import multi_worker_util
 from tensorflow.python.distribute import reduce_util
 from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.distribute import values as value_lib
 from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import kernels
@@ -125,7 +128,10 @@ class CrossDeviceOpsTestBase(test.TestCase, parameterized.TestCase):
         self.evaluate(ops.convert_to_tensor(left)),
         self.evaluate(ops.convert_to_tensor(right)))
 
-  def _assert_mirrored_equal(self, left_list, right_list, sess,
+  def _assert_mirrored_equal(self,
+                             left_list,
+                             right_list,
+                             sess=None,
                              run_options=None):
     if not isinstance(left_list, list):
       left_list, right_list = [left_list], [right_list]
@@ -142,17 +148,14 @@ class CrossDeviceOpsTestBase(test.TestCase, parameterized.TestCase):
         left, right = [left], [right]
 
       for left_value, right_value in zip(left, right):
-        self.assertEqual(left_value.device, right_value.device)
+        self.assertEqual(
+            device_util.resolve(left_value.device),
+            device_util.resolve(right_value.device))
 
       # Densify IndexedSlices.
       left = [ops.convert_to_tensor(v) for v in left]
       right = [ops.convert_to_tensor(v) for v in right]
-      if context.executing_eagerly():
-        # Optional args in session run are not supported when eager execution
-        # is enabled.
-        assert run_options is None
-        left, right = sess.run((left, right))
-      else:
+      if not context.executing_eagerly():
         left, right = sess.run((left, right), options=run_options)
       for left_value, right_value in zip(left, right):
         self.assertAllEqual(left_value, right_value)
@@ -525,8 +528,8 @@ class CollectiveAllReduceTest(multi_worker_test_base.MultiWorkerTestBase,
         return strategy, devices, ""
       else:
         collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce(
-            1,
-            num_gpus,
+            devices=devices,
+            group_size=len(devices),
             collective_keys=collective_keys,
             communication=communication)
         return collective_all_reduce_ops, devices, ""
@@ -545,26 +548,28 @@ class CollectiveAllReduceTest(multi_worker_test_base.MultiWorkerTestBase,
         ]
 
       if use_strategy_object:
-        strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy(
-            communication=communication)
-        strategy.configure(
-            cluster_spec=self._cluster_spec,
+        resolver = cluster_resolver.SimpleClusterResolver(
+            cluster_spec=multi_worker_util.normalize_cluster_spec(
+                self._cluster_spec),
             task_type=task_type,
-            task_id=task_id)
+            task_id=task_id,
+            num_accelerators={"GPU": num_gpus})
+        strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy(
+            cluster_resolver=resolver, communication=communication)
         strategy.extended._collective_keys = collective_keys
         strategy.extended._cross_device_ops._collective_keys = collective_keys
         return (strategy, devices,
                 "grpc://" + self._cluster_spec[task_type][task_id])
       else:
         collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce(
-            NUM_WORKERS,
-            num_gpus,
+            devices=devices,
+            group_size=len(devices) * NUM_WORKERS,
             collective_keys=collective_keys,
             communication=communication)
         return (collective_all_reduce_ops, devices,
                 "grpc://" + self._cluster_spec[task_type][task_id])
 
-  def _assert_mirrored_equal(self, left_list, right_list, sess):
+  def _assert_mirrored_equal(self, left_list, right_list, sess=None):
     if context.executing_eagerly():
       run_options = None
     else:
@@ -895,6 +900,45 @@ class CollectiveAllReduceTest(multi_worker_test_base.MultiWorkerTestBase,
       self.assertAllEqual(reduced[1].values, [4.0, 4.0])
       t.join()
 
+  @combinations.generate(
+      combinations.combine(
+          required_gpus=2,
+          mode="eager",
+          communication=[
+              CollectiveCommunication.NCCL, CollectiveCommunication.RING
+          ]))
+  def testInputsAreFunctionArgs(self, communication):
+    # Function inputs don't have device placement.
+    hints = collective_util.Hints(bytes_per_pack=1)
+    collective, devices, _ = self._get_test_objects(
+        None,
+        None,
+        num_gpus=2,
+        communication=communication,
+        use_strategy_object=False,
+        local_mode=True)
+    devices = [device_util.canonicalize(d) for d in devices]
+
+    @def_function.function
+    def reduce_fn(v):
+      self.assertEqual(v.values[0].device, "")
+      self.assertEqual(v.values[1].device, "")
+      # We only use NCCL for batch reduce with two or more values, so we use two
+      # values here.
+      reduced = collective.batch_reduce(
+          reduce_util.ReduceOp.SUM, [(v, v), (v, v)], experimental_hints=hints)
+      self.assertEqual(reduced[0].values[0].device, devices[0])
+      self.assertEqual(reduced[0].values[1].device, devices[1])
+      self.assertEqual(reduced[1].values[0].device, devices[0])
+      self.assertEqual(reduced[1].values[1].device, devices[1])
+      # Returning Mirrored only evaluates the primary value, which causes
+      # hanging,
+      return [reduced[0].values, reduced[1].values]
+
+    v = _make_per_replica([1.0, 2.0], devices)
+    reduced = reduce_fn(v)
+    self.assertAllEqual(self.evaluate(reduced), [[3.0, 3.0], [3.0, 3.0]])
+
 
 if __name__ == "__main__":
   # Set default inter op thread pool size to one to ensure we don't exhaust the
diff --git a/tensorflow/python/distribute/cross_device_utils.py b/tensorflow/python/distribute/cross_device_utils.py
index d7be93ae2c4..9dc24b16e6a 100644
--- a/tensorflow/python/distribute/cross_device_utils.py
+++ b/tensorflow/python/distribute/cross_device_utils.py
@@ -305,19 +305,6 @@ class CollectiveKeys(object):
         self._group_key_table[key_id] = new_key
     return self._group_key_table[key_id]
 
-  def get_group_key_of_tensors(self, tensors):
-    """Returns a group key for set of tensors.
-
-    Args:
-      tensors: list of `Tensor`s in a collective group. Each tensor must be on a
-        different device.
-
-    Returns:
-      int key uniquely identifying the set of devices of these tensors.
-    """
-    devices = [t.device for t in tensors]
-    return self.get_group_key(devices)
-
   def get_op_instance_key(self):
     """Returns a new instance key for use in defining a collective op."""
     v = self._get_thread_local_object().op_instance_key
@@ -332,7 +319,8 @@ class CollectiveKeys(object):
 
 
 def build_collective_reduce(input_tensors,
-                            num_workers,
+                            devices,
+                            group_size,
                             collective_keys,
                             reduction_op='Add',
                             unary_op='Id',
@@ -347,9 +335,10 @@ def build_collective_reduce(input_tensors,
   Args:
     input_tensors: tensors within a single worker graph that are to be reduced
       together; must be one per device.
-    num_workers: total number of workers with identical independent graphs that
-      will be doing this same reduction.  The reduction will actually include
-      the corresponding tensors at all these workers.
+    devices: a list of device strings to run the collective on.
+    group_size: total number of devices globally that will be doing this same
+      reduction.  The reduction will actually include the corresponding tensors
+      at all these workers.
     collective_keys: a CollectiveKeys object.
     reduction_op: string naming the reduction op.
     unary_op: string naming the unary final op.
@@ -370,11 +359,14 @@ def build_collective_reduce(input_tensors,
         not all(e.is_async() for e in executors)):
       raise ValueError(
           'collectives requires async executors for each device in eager mode')
+  if len(input_tensors) != len(devices):
+    raise ValueError('collective requires one input tensor for each device, '
+                     'len(input_tensors) = %d, len(devices) = %d' %
+                     (len(input_tensors), len(devices)))
 
-  group_size = len(input_tensors) * num_workers
   if group_size < 2:
     return input_tensors
-  group_key = collective_keys.get_group_key_of_tensors(input_tensors)
+  group_key = collective_keys.get_group_key(devices)
   instance_key = collective_keys.get_op_instance_key()
   subdiv_offsets = [0]  # TODO(tucker): maybe support non-default subdiv spec
 
@@ -385,9 +377,9 @@ def build_collective_reduce(input_tensors,
     else:
       executor_scope = ops.NullContextmanager()
     with executor_scope, \
-         ops.device(input_tensor.device), \
+         ops.device(devices[idx]), \
          ops.control_dependencies(
-             _control_input(input_tensors, control_inputs, idx)):
+             _control_input(devices, control_inputs, idx)):
       out_tensor = collective_ops.all_reduce(input_tensor, group_size,
                                              group_key, instance_key,
                                              reduction_op, unary_op,
@@ -397,7 +389,8 @@ def build_collective_reduce(input_tensors,
 
 
 def build_collective_gather(input_tensors,
-                            num_workers,
+                            devices,
+                            group_size,
                             collective_keys,
                             communication_hint='AUTO',
                             control_inputs=None):
@@ -408,9 +401,10 @@ def build_collective_gather(input_tensors,
   Args:
     input_tensors: tensors within a single worker graph that are to be gathered
       together; must be one per device.
-    num_workers: total number of workers with identical independent graphs that
-      will be doing this same reduction.  The reduction will actually include
-      the corresponding tensors at all these workers.
+    devices: a list of device strings to run the collective on.
+    group_size: total number of devices globally that will be doing this same
+      gathering. The gathering will actually include the corresponding tensors
+      at all these workers.
     collective_keys: a CollectiveKeys object.
     communication_hint: string providing hint to runtime for choosing collective
       implementation.
@@ -423,18 +417,21 @@ def build_collective_gather(input_tensors,
   assert not context.executing_eagerly(), (
       'build_collective_gather can only be called in graph mode or inside '
       'tf.function')
+  if len(input_tensors) != len(devices):
+    raise ValueError(
+        'collective requires one input tensor for each device, %d != %d' %
+        (len(input_tensors), len(devices)))
 
-  group_size = len(input_tensors) * num_workers
   if group_size < 2:
     return input_tensors
-  group_key = collective_keys.get_group_key_of_tensors(input_tensors)
+  group_key = collective_keys.get_group_key(devices)
   instance_key = collective_keys.get_op_instance_key()
 
   out_tensors = []
   for idx, input_tensor in enumerate(input_tensors):
-    with ops.device(input_tensor.device):
+    with ops.device(devices[idx]):
       with ops.control_dependencies(
-          _control_input(input_tensors, control_inputs, idx)):
+          _control_input(devices, control_inputs, idx)):
         out_tensor = collective_ops.all_gather(input_tensor, group_size,
                                                group_key, instance_key,
                                                communication_hint)
@@ -443,7 +440,8 @@ def build_collective_gather(input_tensors,
 
 
 def build_collective_gather_indexed_slices(input_slices_list,
-                                           num_workers,
+                                           devices,
+                                           group_size,
                                            collective_keys,
                                            communication_hint='AUTO',
                                            control_inputs=None):
@@ -454,9 +452,10 @@ def build_collective_gather_indexed_slices(input_slices_list,
   Args:
     input_slices_list: a list of IndexedSlices within a single worker graph that
       are to be gathered together; must be one per device.
-    num_workers: total number of workers with identical independent graphs that
-      will be doing this same reduction.  The reduction will actually include
-      the corresponding tensors at all these workers.
+    devices: a list of device strings to run the collective on.
+    group_size: total number of devices globally that will be doing this same
+      gathering. The gathering will actually include the corresponding tensors
+      at all these workers.
     collective_keys: a CollectiveKeys object.
     communication_hint: string providing hint to runtime for choosing collective
       implementation.
@@ -474,12 +473,15 @@ def build_collective_gather_indexed_slices(input_slices_list,
   assert not context.executing_eagerly(), (
       'build_collective_gather_indexed_slices can only be called in graph mode'
       ' or inside tf.function')
+  if len(input_slices_list) != len(devices):
+    raise ValueError(
+        'collective requires one input IndexedSlice for each device, %d != %d' %
+        (len(input_slices_list), len(devices)))
 
-  group_size = len(input_slices_list) * num_workers
   if group_size < 2:
     return input_slices_list
 
-  group_key = collective_keys.get_group_key_of_tensors(input_slices_list)
+  group_key = collective_keys.get_group_key(devices)
   gather_length_key = collective_keys.get_op_instance_key()
   gather_indices_key = collective_keys.get_op_instance_key()
   gather_values_key = collective_keys.get_op_instance_key()
@@ -495,7 +497,7 @@ def build_collective_gather_indexed_slices(input_slices_list,
   out_slices_list = []
   for idx, input_slices in enumerate(input_slices_list):
     # pylint: disable = cell-var-from-loop
-    with ops.device(input_slices.device):
+    with ops.device(devices[idx]):
 
       def all_gather():
         """Use all_gather to aggregate `IndexedSlices`."""
@@ -967,14 +969,13 @@ def pack_by_size(per_replica_list, bytes_per_pack):
   return packs
 
 
-def _control_input(inputs, control_inputs, idx):
+def _control_input(devices, control_inputs, idx):
   """Returns the `idx`-th item in control_inputs to be used in ops.control_dependencies.
 
-  This is a helper function for building collective ops.  The function checks
-  that the devices of control_inputs and inputs match.
+  This is a helper function for building collective ops.
 
   Args:
-    inputs: a list of `Tensor`s
+    devices: a list of device strings the collective run on.
     control_inputs: a list or None.
     idx: the index into `inputs` and `control_inputs`.
 
@@ -984,12 +985,8 @@ def _control_input(inputs, control_inputs, idx):
   """
   if control_inputs is None:
     return []
-  if len(control_inputs) != len(inputs):
+  if len(control_inputs) != len(devices):
     raise ValueError(
-        'control_inputs must match the length of the inputs, %s != %s' %
-        (len(control_inputs), len(inputs)))
-  if control_inputs[idx].device != inputs[idx].device:
-    raise ValueError(
-        'control_inputs must match the device of the inputs, %s != %s' %
-        (control_inputs[idx].device, inputs[idx].device))
+        'control_inputs must match the length of the devices, %s != %s' %
+        (len(control_inputs), len(devices)))
   return [control_inputs[idx]]
diff --git a/tensorflow/python/distribute/mirrored_strategy.py b/tensorflow/python/distribute/mirrored_strategy.py
index ac9045d2322..36598634fac 100644
--- a/tensorflow/python/distribute/mirrored_strategy.py
+++ b/tensorflow/python/distribute/mirrored_strategy.py
@@ -579,7 +579,7 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
     if not destinations:
       # TODO(josh11b): Use current logical device instead of 0 here.
       destinations = self._devices
-    return self._get_cross_device_ops().broadcast(tensor, destinations)
+    return self._get_cross_device_ops(tensor).broadcast(tensor, destinations)
 
   def _call_for_each_replica(self, fn, args, kwargs):
     return mirrored_run.call_for_each_replica(
@@ -608,7 +608,8 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
     updated_config.isolate_session_state = True
     return updated_config
 
-  def _get_cross_device_ops(self):
+  def _get_cross_device_ops(self, value):
+    del value  # Unused.
     return self._cross_device_ops or self._inferred_cross_device_ops
 
   def _reduce_to(self, reduce_op, value, destinations, experimental_hints):
@@ -623,7 +624,7 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
       # be 0.
       return cross_device_ops_lib.reduce_non_distributed_value(
           reduce_op, value, destinations, self._num_replicas_in_sync)
-    return self._get_cross_device_ops().reduce(
+    return self._get_cross_device_ops(value).reduce(
         reduce_op,
         value,
         destinations=destinations,
@@ -631,9 +632,15 @@ class MirroredExtended(distribute_lib.StrategyExtendedV1):
 
   def _batch_reduce_to(self, reduce_op, value_destination_pairs,
                        experimental_hints):
-    return self._get_cross_device_ops().batch_reduce(reduce_op,
-                                                     value_destination_pairs,
-                                                     experimental_hints)
+    cross_device_ops = None
+    for value, _ in value_destination_pairs:
+      if cross_device_ops is None:
+        cross_device_ops = self._get_cross_device_ops(value)
+      elif cross_device_ops is not self._get_cross_device_ops(value):
+        raise ValueError("inputs to batch_reduce_to must be either all on the "
+                         "the host or all on the compute devices")
+    return cross_device_ops.batch_reduce(reduce_op, value_destination_pairs,
+                                         experimental_hints)
 
   def _update(self, var, fn, args, kwargs, group):
     # TODO(josh11b): In eager mode, use one thread per device.
diff --git a/tensorflow/python/distribute/strategy_common_test.py b/tensorflow/python/distribute/strategy_common_test.py
index 4ed5054af2d..ed52a4794ee 100644
--- a/tensorflow/python/distribute/strategy_common_test.py
+++ b/tensorflow/python/distribute/strategy_common_test.py
@@ -65,14 +65,14 @@ class StrategyReduceTest(test.TestCase, parameterized.TestCase):
     self.assertEqual(fn_graph().numpy(), 1.0 * strategy.num_replicas_in_sync)
 
 
+@combinations.generate(
+    combinations.combine(
+        strategy=[strategy_combinations.multi_worker_mirrored_two_workers],
+        mode=['eager']))
 class DistributedCollectiveAllReduceStrategyTest(
     strategy_test_lib.DistributionTestBase,
     parameterized.TestCase):
 
-  @combinations.generate(
-      combinations.combine(
-          strategy=[strategy_combinations.multi_worker_mirrored_two_workers],
-          mode=['eager']))
   def testDatasetFromFunction(self, strategy):
     def dataset_fn(input_context):
       global_batch_size = 10
@@ -95,6 +95,45 @@ class DistributedCollectiveAllReduceStrategyTest(
         sum_value.numpy(),
         expected_sum_on_workers[multi_worker_test_base.get_task_index()])
 
+  def testReduceHostTensor(self, strategy):
+    reduced = strategy.reduce(
+        reduce_util.ReduceOp.SUM, array_ops.identity(1.), axis=None)
+    self.assertEqual(reduced.numpy(), 2.)
+
+  def testReduceToHostTensor(self, strategy):
+    value = array_ops.identity(1.)
+    reduced = strategy.extended.reduce_to(reduce_util.ReduceOp.SUM, value,
+                                          value)
+    self.assertEqual(reduced.numpy(), 2.)
+
+  def testBatchReduceToHostTensor(self, strategy):
+    value = array_ops.identity(1.)
+    reduced = strategy.extended.batch_reduce_to(reduce_util.ReduceOp.SUM,
+                                                [(value, value),
+                                                 (value, value)])
+    self.assertAllEqual(reduced, [2., 2.])
+
+  def testReduceDeviceTensors(self, strategy):
+    value = strategy.run(lambda: array_ops.identity(1.))
+    reduced = strategy.reduce(reduce_util.ReduceOp.SUM, value, axis=None)
+    self.assertEqual(reduced.numpy(), 2.)
+
+  def testReduceToDeviceTensors(self, strategy):
+    value = strategy.run(lambda: array_ops.identity(1.))
+    reduced = strategy.extended.reduce_to(reduce_util.ReduceOp.SUM, value,
+                                          value)
+    self.assertEqual(reduced.numpy(), 2.)
+
+  def testBatchReduceToDeviceTensors(self, strategy):
+    value = strategy.run(lambda: array_ops.identity(1.))
+    reduced = strategy.extended.batch_reduce_to(reduce_util.ReduceOp.SUM,
+                                                [(value, value),
+                                                 (value, value)])
+    self.assertAllEqual(reduced, [2., 2.])
+
+  # TODO(crccw): add a test that mixes device and host tensors after multi
+  # worker strategy combinations can run on a fixed number of GPUs.
+
 
 if __name__ == '__main__':
   combinations.main()

From a71c78bcf91d404de37188a8a7a73016729dd2a0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 13:14:37 -0700
Subject: [PATCH 0323/1390] PR #40169: ensure model initialized on ANY
 trackable attr set

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/40169

In particular, empty tuples should not trigger this.
Copybara import of the project:

--
17b7e169135127e0e866b50577ad8b213abc1d97 by Dominic Jack :

ensure model initialized on ANY trackable attr set

--
57eccc7bc29ddb105dcaa2f6a413163461ad9987 by Dominic Jack :

added test

PiperOrigin-RevId: 316743715
Change-Id: I038a0261fbb3a0dac50c62a50c787bade10abb6a
---
 tensorflow/python/keras/engine/training.py      |  8 ++++----
 tensorflow/python/keras/engine/training_test.py | 12 ------------
 2 files changed, 4 insertions(+), 16 deletions(-)

diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index 682c0272da5..5567e1733a7 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -324,10 +324,10 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
       super(Model, self).__setattr__(name, value)
       return
 
-    if any(
-        isinstance(v, (base_layer.Layer, data_structures.TrackableDataStructure
-                      )) or trackable_layer_utils.has_weights(v)
-        for v in nest.flatten(value)):
+    if all(
+        isinstance(v, (base_layer.Layer,
+                       data_structures.TrackableDataStructure)) or
+        trackable_layer_utils.has_weights(v) for v in nest.flatten(value)):
       try:
         self._base_model_initialized
       except AttributeError:
diff --git a/tensorflow/python/keras/engine/training_test.py b/tensorflow/python/keras/engine/training_test.py
index 63943f4f720..5cf15926bfb 100644
--- a/tensorflow/python/keras/engine/training_test.py
+++ b/tensorflow/python/keras/engine/training_test.py
@@ -3383,18 +3383,6 @@ class TestTrainingWithMetrics(keras_parameterized.TestCase):
     self.assertEqual([m.name for m in outer_model.metrics],
                      ['loss', 'acc2', 'mean', 'mean1', 'mean2'])
 
-  def test_subclassed_model_with_empty_list_attr(self):
-
-    class ModelSubclass(training_module.Model):
-
-      def __init__(self):
-        self.empty_list = []
-        inputs = layers_module.Input(shape=())
-        outputs = inputs + 1
-        super(ModelSubclass, self).__init__(inputs, outputs)
-
-    ModelSubclass()  # empty_list attr assignment should not raise
-
 
 class BareUpdateLayer(layers_module.Layer):
 

From 4db7ec52010ef737300a00e669973fecea5c0603 Mon Sep 17 00:00:00 2001
From: Andy Ly 
Date: Tue, 16 Jun 2020 13:35:30 -0700
Subject: [PATCH 0324/1390] Replace const llvm::SmallVector<>& with
 llvm::ArrayRef and const std::string& with llvm::StringRef in
 TPUExtractOutsideCompilation. (NFC)

PiperOrigin-RevId: 316748196
Change-Id: Icdfcaa5a808ae69e5a6286d5bd7c6a988dbbe616
---
 .../tpu_extract_outside_compilation.cc        | 20 +++++++++----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
index 93e5cc22c30..54600faca4b 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
@@ -73,9 +73,8 @@ LogicalResult CollectAndGroupOutsideClusterOps(Block* block,
 }
 
 // Moves `cluster_ops` to associated `launch_op` body.
-void MoveOutsideClusterOpsToLaunchOp(
-    tf_device::LaunchOp launch_op,
-    const llvm::SmallVector& cluster_ops) {
+void MoveOutsideClusterOpsToLaunchOp(tf_device::LaunchOp launch_op,
+                                     llvm::ArrayRef cluster_ops) {
   MLIRContext* context = launch_op.getContext();
   Operation* terminator = launch_op.GetBody().getTerminator();
 
@@ -123,7 +122,7 @@ void PropagateParallelExecuteReturnToReplicate(
 
 // Extracts all externally provided operands of `cluster_ops`.
 llvm::SmallSetVector GetExternalOperands(
-    const llvm::SmallVector& cluster_ops) {
+    llvm::ArrayRef cluster_ops) {
   llvm::SmallSetVector external_values;
 
   for (Operation* op : cluster_ops) {
@@ -143,7 +142,7 @@ llvm::SmallSetVector GetExternalOperands(
 
 // Extracts all externally used outputs of `cluster_ops`.
 llvm::SmallVector GetExternalOutputs(
-    const llvm::SmallVector& cluster_ops) {
+    llvm::ArrayRef cluster_ops) {
   llvm::SmallSetVector external_outputs;
 
   for (Operation* op : cluster_ops) {
@@ -166,7 +165,7 @@ llvm::SmallVector GetExternalOutputs(
 // as an operand.  If there are no external_inputs, set insertion point to first
 // cluster_op.
 void SetHostComputeInsertion(
-    OpBuilder* builder, const llvm::SmallVector& cluster_ops,
+    OpBuilder* builder, llvm::ArrayRef cluster_ops,
     const llvm::SmallSetVector& external_inputs) {
   if (external_inputs.empty()) builder->setInsertionPoint(cluster_ops.front());
   for (const auto& cluster_op : cluster_ops) {
@@ -183,9 +182,9 @@ void SetHostComputeInsertion(
 // using `communication_key`.
 TF::_HostComputeMlirOp CreateHostCompute(
     OpBuilder* builder, tf_device::ClusterOp tpu_cluster,
-    const llvm::SmallVector& cluster_ops,
+    llvm::ArrayRef cluster_ops,
     const llvm::SmallSetVector& inputs, llvm::ArrayRef outputs,
-    const std::string& communication_key) {
+    llvm::StringRef communication_key) {
   llvm::SmallVector device_output_types;
   for (const auto& output : outputs)
     device_output_types.push_back(output.getType());
@@ -201,10 +200,9 @@ TF::_HostComputeMlirOp CreateHostCompute(
 
 void MoveOutsideCompiledOps(
     tf_device::ClusterOp tpu_cluster, llvm::StringRef outside_cluster_name,
-    tf_device::LaunchOp host_launch_op,
-    const llvm::SmallVector& cluster_ops,
+    tf_device::LaunchOp host_launch_op, llvm::ArrayRef cluster_ops,
     const llvm::SmallSetVector& external_inputs,
-    const llvm::SmallVector& external_outputs) {
+    llvm::ArrayRef external_outputs) {
   if (external_inputs.empty() && external_outputs.empty()) {
     MoveOutsideClusterOpsToLaunchOp(host_launch_op, cluster_ops);
     return;

From 14ee01957cb6450261f6efec77a8997ca4d8c3c5 Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Tue, 16 Jun 2020 13:37:06 -0700
Subject: [PATCH 0325/1390] [TF-numpy] Adds @np_doc/@np_doc_only to all public
 ops.

PiperOrigin-RevId: 316748465
Change-Id: I40a49431d6075e47ac05f2946c745ab1c1222214
---
 .../python/ops/numpy_ops/np_array_ops.py      | 310 ++----------------
 1 file changed, 36 insertions(+), 274 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index e97bb61613b..906e53c556d 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -39,51 +39,18 @@ from tensorflow.python.ops.numpy_ops import np_utils
 from tensorflow.python.util import nest
 
 
+@np_utils.np_doc(np.empty)
 def empty(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an empty array with the specified shape and dtype.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   return zeros(shape, dtype)
 
 
+@np_utils.np_doc(np.empty_like)
 def empty_like(a, dtype=None):
-  """Returns an empty array with the shape and possibly type of the input array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
   return zeros_like(a, dtype)
 
 
+@np_utils.np_doc(np.zeros)
 def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an ndarray with the given shape and type filled with zeros.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   dtype = (
       np_utils.result_type(dtype) if dtype else np_dtypes.default_float_type())
   if isinstance(shape, np_arrays.ndarray):
@@ -91,19 +58,8 @@ def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.zeros(shape, dtype=dtype))
 
 
-def zeros_like(a, dtype=None):
-  """Returns an array of zeros with the shape and type of the input array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
+@np_utils.np_doc(np.zeros_like)
+def zeros_like(a, dtype=None):  # pylint: disable=missing-docstring
   if isinstance(a, np_arrays.ndarray):
     a = a.data
   if dtype is None:
@@ -117,19 +73,8 @@ def zeros_like(a, dtype=None):
   return np_arrays.tensor_to_ndarray(array_ops.zeros_like(a, dtype))
 
 
+@np_utils.np_doc(np.ones)
 def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
-  """Returns an ndarray with the given shape and type filled with ones.
-
-  Args:
-    shape: A fully defined shape. Could be - NumPy array or a python scalar,
-      list or tuple of integers, - TensorFlow tensor/ndarray of integer type and
-      rank <=1.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray.
-  """
   if dtype:
     dtype = np_utils.result_type(dtype)
   if isinstance(shape, np_arrays.ndarray):
@@ -137,19 +82,8 @@ def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.ones(shape, dtype=dtype))
 
 
+@np_utils.np_doc(np.ones_like)
 def ones_like(a, dtype=None):
-  """Returns an array of ones with the shape and type of the input array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the input array. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-  """
   if isinstance(a, np_arrays.ndarray):
     a = a.data
   if dtype is None:
@@ -191,38 +125,13 @@ def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-do
       array_ops.matrix_diag(diagonal=diagonal_, num_rows=N, num_cols=M, k=k))
 
 
+@np_utils.np_doc(np.identity)
 def identity(n, dtype=float):
-  """Returns a square array with ones on the main diagonal and zeros elsewhere.
-
-  Args:
-    n: number of rows/cols.
-    dtype: Optional, defaults to float. The type of the resulting ndarray. Could
-      be a python type, a NumPy type or a TensorFlow `DType`.
-
-  Returns:
-    An ndarray of shape (n, n) and requested type.
-  """
   return eye(N=n, M=n, dtype=dtype)
 
 
+@np_utils.np_doc(np.full)
 def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
-  """Returns an array with given shape and dtype filled with `fill_value`.
-
-  Args:
-    shape: A valid shape object. Could be a native python object or an object of
-      type ndarray, numpy.ndarray or tf.TensorShape.
-    fill_value: array_like. Could be an ndarray, a Tensor or any object that can
-      be converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the `fill_value`. The type of the
-      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
-      `DType`.
-
-  Returns:
-    An ndarray.
-
-  Raises:
-    ValueError: if `fill_value` can not be broadcast to shape `shape`.
-  """
   if not isinstance(shape, np_arrays.ndarray):
     shape = asarray(np_arrays.convert_to_tensor(shape, dtype_hint=np.int32))
   shape = atleast_1d(shape).data
@@ -251,26 +160,13 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  #
 
 
 # TODO(wangpeng): investigate whether we can make `copy` default to False.
-# TODO(wangpeng): np_utils.np_doc can't handle np.array because np.array is a
-#   builtin function. Make np_utils.np_doc support builtin functions.
+# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
+@np_utils.np_doc_only(np.array)
 def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
-  """Creates an ndarray with the contents of val.
-
-  Args:
-    val: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    dtype: Optional, defaults to dtype of the `val`. The type of the resulting
-      ndarray. Could be a python type, a NumPy type or a TensorFlow `DType`.
-    copy: Determines whether to create a copy of the backing buffer. Since
-      Tensors are immutable, a copy is made only if val is placed on a different
-      device than the current one. Even if `copy` is False, a new Tensor may
-      need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
-      is an ndarray or a Tensor.
-    ndmin: The minimum rank of the returned array.
-
-  Returns:
-    An ndarray.
-  """
+  """Since Tensors are immutable, a copy is made only if val is placed on a
+  different device than the current one. Even if `copy` is False, a new Tensor
+  may need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
+  is an ndarray or a Tensor."""  # pylint:disable=g-docstring-missing-newline
   if dtype:
     dtype = np_utils.result_type(dtype)
   if isinstance(val, np_arrays.ndarray):
@@ -319,6 +215,7 @@ def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-out
   result_t = np_utils.cond(
       np_utils.greater(ndmin, ndims), true_fn, lambda: result_t)
   return np_arrays.tensor_to_ndarray(result_t)
+# pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
 
 
 @np_utils.np_doc(np.asarray)
@@ -341,6 +238,7 @@ def ascontiguousarray(a, dtype=None):
 
 
 # Numerical ranges.
+@np_utils.np_doc(np.arange)
 def arange(start, stop=None, step=1, dtype=None):
   """Returns `step`-separated values in the range [start, stop).
 
@@ -448,20 +346,8 @@ def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstrin
   return a
 
 
+@np_utils.np_doc(np.diagflat)
 def diagflat(v, k=0):
-  """Returns a 2-d array with flattened `v` as diagonal.
-
-  Args:
-    v: array_like of any rank. Gets flattened when setting as diagonal. Could be
-      an ndarray, a Tensor or any object that can be converted to a Tensor using
-      `tf.convert_to_tensor`.
-    k: Position of the diagonal. Defaults to 0, the main diagonal. Positive
-      values refer to diagonals shifted right, negative values refer to
-      diagonals shifted left.
-
-  Returns:
-    2-d ndarray.
-  """
   v = asarray(v)
   return diag(array_ops.reshape(v.data, [-1]), k)
 
@@ -471,69 +357,22 @@ def _promote_dtype(*arrays):
   return [asarray(a, dtype=dtype) for a in arrays]
 
 
+@np_utils.np_doc(np.all)
 def all(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
-  """Whether all array elements or those along an axis evaluate to true.
-
-  Casts the array to bool type if it is not already and uses `tf.reduce_all` to
-  compute the result.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Could be an int or a tuple of integers. If not specified,
-      the reduction is performed over all array indices.
-    keepdims: If true, retains reduced dimensions with length 1.
-
-  Returns:
-    An ndarray. Note that unlike NumPy this does not return a scalar bool if
-    `axis` is None.
-  """
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_all(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
+@np_utils.np_doc(np.any)
 def any(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
-  """Whether any element in the entire array or in an axis evaluates to true.
-
-  Casts the array to bool type if it is not already and uses `tf.reduce_any` to
-  compute the result.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Could be an int or a tuple of integers. If not specified,
-      the reduction is performed over all array indices.
-    keepdims: If true, retains reduced dimensions with length 1.
-
-  Returns:
-    An ndarray. Note that unlike NumPy this does not return a scalar bool if
-    `axis` is None.
-  """
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_any(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
-def compress(condition, a, axis=None):
-  """Compresses `a` by selecting values along `axis` with `condition` true.
-
-  Uses `tf.boolean_mask`.
-
-  Args:
-    condition: 1-d array of bools. If `condition` is shorter than the array axis
-      (or the flattened array if axis is None), it is padded with False.
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: Optional. Axis along which to select elements. If None, `condition` is
-      applied on flattened array.
-
-  Returns:
-    An ndarray.
-
-  Raises:
-    ValueError: if `condition` is not of rank 1.
-  """
+@np_utils.np_doc(np.compress)
+def compress(condition, a, axis=None):  # pylint: disable=redefined-outer-name,missing-function-docstring
   condition = asarray(condition, dtype=bool)
   a = asarray(a)
 
@@ -563,8 +402,8 @@ def compress(condition, a, axis=None):
       array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis))
 
 
+@np_utils.np_doc(np.copy)
 def copy(a):
-  """Returns a copy of the array."""
   return array(a, copy=True)
 
 
@@ -611,18 +450,8 @@ def cumsum(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(math_ops.cumsum(a.data, axis))
 
 
+@np_utils.np_doc(np.imag)
 def imag(a):
-  """Returns imaginary parts of all elements in `a`.
-
-  Uses `tf.imag`.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-
-  Returns:
-    An ndarray with the same shape as `a`.
-  """
   a = asarray(a)
   # TODO(srbs): np.imag returns a scalar if a is a scalar, whereas we always
   # return an ndarray.
@@ -760,6 +589,7 @@ def amin(a, axis=None, keepdims=None):
       preserve_bool=True)
 
 
+@np_utils.np_doc(np.var)
 def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: disable=missing-docstring
   if dtype:
     working_dtype = np_utils.result_type(a, dtype)
@@ -829,18 +659,8 @@ def ravel(a):  # pylint: disable=missing-docstring
 setattr(np_arrays.ndarray, 'ravel', ravel)
 
 
+@np_utils.np_doc(np.real)
 def real(val):
-  """Returns real parts of all elements in `a`.
-
-  Uses `tf.real`.
-
-  Args:
-    val: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-
-  Returns:
-    An ndarray with the same shape as `a`.
-  """
   val = asarray(val)
   # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always
   # return an ndarray.
@@ -897,7 +717,6 @@ def around(a, decimals=0):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(a).astype(dtype)
 
 
-round_ = around
 setattr(np_arrays.ndarray, '__round__', around)
 
 
@@ -933,51 +752,20 @@ def _reshape_method_wrapper(a, *newshape, **kwargs):
   return reshape(a, newshape, order=order)
 
 
+@np_utils.np_doc(np.expand_dims)
 def expand_dims(a, axis):
-  """Expand the shape of an array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: int. axis on which to expand the shape.
-
-  Returns:
-    An ndarray with the contents and dtype of `a` and shape expanded on axis.
-  """
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.expand_dims(a.data, axis=axis))
 
 
+@np_utils.np_doc(np.squeeze)
 def squeeze(a, axis=None):
-  """Removes single-element axes from the array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axis: scalar or list/tuple of ints.
-  TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution is
-    enabled. So we cannot allow axis to be array_like here. Fix.
-
-  Returns:
-    An ndarray.
-  """
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.squeeze(a, axis))
 
 
+@np_utils.np_doc(np.transpose)
 def transpose(a, axes=None):
-  """Permutes dimensions of the array.
-
-  Args:
-    a: array_like. Could be an ndarray, a Tensor or any object that can be
-      converted to a Tensor using `tf.convert_to_tensor`.
-    axes: array_like. A list of ints with length rank(a) or None specifying the
-      order of permutation. The i'th dimension of the output array corresponds
-      to axes[i]'th dimension of the `a`. If None, the axes are reversed.
-
-  Returns:
-    An ndarray.
-  """
   a = asarray(a)
   if axes is not None:
     axes = asarray(axes)
@@ -1113,37 +901,16 @@ def _setitem(arr, index, value):
         [prefix_t, array_ops.expand_dims(subarray.data, 0), postfix_t], 0)
 
 
+# TODO(wangpeng): Make a custom `setattr` that also sets docstring for the
+#   method.
 setattr(np_arrays.ndarray, 'transpose', transpose)
 setattr(np_arrays.ndarray, 'reshape', _reshape_method_wrapper)
 setattr(np_arrays.ndarray, '__setitem__', _setitem)
 
 
+@np_utils.np_doc(np.pad)
 def pad(ary, pad_width, mode, constant_values=0):
-  """Pads an array.
-
-  Args:
-    ary: array_like of rank N. Input array.
-    pad_width: {sequence, array_like, int}. Number of values padded to the edges
-      of each axis. ((before_1, after_1), ... (before_N, after_N)) unique pad
-      widths for each axis. ((before, after),) yields same before and after pad
-      for each axis. (pad,) or int is a shortcut for before = after = pad width
-      for all axes.
-    mode: string. One of the following string values: 'constant' Pads with a
-      constant value. 'reflect' Pads with the reflection of the vector mirrored
-      on the first and last values of the vector along each axis. 'symmetric'
-      Pads with the reflection of the vector mirrored along the edge of the
-      array.
-      **NOTE**: The supported list of `mode` does not match that of numpy's.
-    constant_values: scalar with same dtype as `array`. Used in 'constant' mode
-      as the pad value.  Default is 0.
-
-  Returns:
-    An ndarray padded array of rank equal to `array` with shape increased
-    according to `pad_width`.
-
-  Raises:
-    ValueError if `mode` is not supported.
-  """
+  """Only supports modes 'constant', 'reflect' and 'symmetric' currently."""
   if not (mode == 'constant' or mode == 'reflect' or mode == 'symmetric'):
     raise ValueError('Unsupported padding mode: ' + mode)
   mode = mode.upper()
@@ -1214,24 +981,19 @@ def select(condlist, choicelist, default=0):  # pylint: disable=missing-docstrin
   return output
 
 
+@np_utils.np_doc(np.shape)
 def shape(a):
-  """Return the shape of an array.
-
-  Args:
-    a: array_like. Input array.
-
-  Returns:
-    Tuple of ints.
-  """
   a = asarray(a)
   return a.shape
 
 
+@np_utils.np_doc(np.ndim)
 def ndim(a):
   a = asarray(a)
   return a.ndim
 
 
+@np_utils.np_doc(np.isscalar)
 def isscalar(a):
   return ndim(a) == 0
 

From 10d5e26a9009151639ee1c4812901c3ac96df304 Mon Sep 17 00:00:00 2001
From: Yunlu Li 
Date: Tue, 16 Jun 2020 13:38:11 -0700
Subject: [PATCH 0326/1390] Internal change

PiperOrigin-RevId: 316748654
Change-Id: Iec1a40e1523cbe2712e1dccf7a947de0e3156d46
---
 tensorflow/lite/python/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD
index c1f37c81b7f..d25e7d5ef8d 100644
--- a/tensorflow/lite/python/BUILD
+++ b/tensorflow/lite/python/BUILD
@@ -40,7 +40,6 @@ py_test(
         "no_windows",
         "noasan",  # TODO(b/137568139): enable after this is fixed.
         "nomsan",  # TODO(b/137568139): enable after this is fixed.
-        "notsan",  # TODO(b/149882556): enable after this is fixed.
     ],
     deps = [
         ":interpreter",

From 8cf2895fcce8fb47ef1b603f9784a34a72d5ae54 Mon Sep 17 00:00:00 2001
From: Suharsh Sivakumar 
Date: Tue, 16 Jun 2020 13:38:34 -0700
Subject: [PATCH 0327/1390] Add dynamic range test to op_tests A-B.

PiperOrigin-RevId: 316748721
Change-Id: I2e803777a160197f3b7a6026c5e94ce11f47ab92
---
 tensorflow/lite/testing/op_tests/add_n.py     |  9 ++++
 .../lite/testing/op_tests/arg_min_max.py      |  3 +-
 .../testing/op_tests/batch_to_space_nd.py     |  4 ++
 tensorflow/lite/testing/op_tests/binary_op.py | 54 +++++++++++++++++++
 4 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/testing/op_tests/add_n.py b/tensorflow/lite/testing/op_tests/add_n.py
index 2385bd89600..bef0a0632b9 100644
--- a/tensorflow/lite/testing/op_tests/add_n.py
+++ b/tensorflow/lite/testing/op_tests/add_n.py
@@ -32,16 +32,25 @@ def make_add_n_tests(options):
           "dtype": [tf.float32, tf.int32],
           "input_shape": [[2, 5, 3, 1]],
           "num_inputs": [2, 3, 4, 5],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32, tf.int32],
           "input_shape": [[5]],
           "num_inputs": [2, 3, 4, 5],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32, tf.int32],
           "input_shape": [[]],
           "num_inputs": [2, 3, 4, 5],
+          "dynamic_range_quantize": [False],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape": [[]],
+          "num_inputs": [2, 3, 4, 5],
+          "dynamic_range_quantize": [True],
       },
   ]
 
diff --git a/tensorflow/lite/testing/op_tests/arg_min_max.py b/tensorflow/lite/testing/op_tests/arg_min_max.py
index e693ce6f44a..ec0013225e0 100644
--- a/tensorflow/lite/testing/op_tests/arg_min_max.py
+++ b/tensorflow/lite/testing/op_tests/arg_min_max.py
@@ -34,6 +34,7 @@ def make_arg_min_max_tests(options):
       "input_shape": [[], [1, 1, 1, 3], [2, 3, 4, 5], [2, 3, 3], [5, 5], [10]],
       "output_type": [tf.int32, tf.int64],
       "is_arg_max": [True],
+      "dynamic_range_quantize": [False, True],
   }]
 
   def build_graph(parameters):
@@ -62,4 +63,4 @@ def make_arg_min_max_tests(options):
       test_parameters,
       build_graph,
       build_inputs,
-      expected_tf_failures=4)
+      expected_tf_failures=8)
diff --git a/tensorflow/lite/testing/op_tests/batch_to_space_nd.py b/tensorflow/lite/testing/op_tests/batch_to_space_nd.py
index e3f05697b0b..2180bfa5b10 100644
--- a/tensorflow/lite/testing/op_tests/batch_to_space_nd.py
+++ b/tensorflow/lite/testing/op_tests/batch_to_space_nd.py
@@ -36,6 +36,7 @@ def make_batch_to_space_nd_tests(options):
           "crops": [[[0, 0], [0, 0]], [[1, 1], [1, 1]]],
           "constant_block_shape": [True, False],
           "constant_crops": [True, False],
+          "dynamic_range_quantize": [False],
       },
       # Single batch (no-op)
       {
@@ -45,6 +46,7 @@ def make_batch_to_space_nd_tests(options):
           "crops": [[[0, 0], [0, 0]], [[1, 1], [1, 1]]],
           "constant_block_shape": [True],
           "constant_crops": [True],
+          "dynamic_range_quantize": [True, False],
       },
       # 3D use case.
       {
@@ -54,6 +56,7 @@ def make_batch_to_space_nd_tests(options):
           "crops": [[[0, 0]], [[1, 1]]],
           "constant_block_shape": [True],
           "constant_crops": [True],
+          "dynamic_range_quantize": [True, False],
       },
   ]
 
@@ -66,6 +69,7 @@ def make_batch_to_space_nd_tests(options):
         "crops": [[[0, 0], [0, 0], [0, 0]]],
         "constant_block_shape": [True, False],
         "constant_crops": [True, False],
+        "dynamic_range_quantize": [False],
     }]
 
   def build_graph(parameters):
diff --git a/tensorflow/lite/testing/op_tests/binary_op.py b/tensorflow/lite/testing/op_tests/binary_op.py
index 9d0c85e35aa..17ed2f3522d 100644
--- a/tensorflow/lite/testing/op_tests/binary_op.py
+++ b/tensorflow/lite/testing/op_tests/binary_op.py
@@ -41,6 +41,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[1, 3, 4, 3]],
           "activation": [True],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -48,6 +49,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[5]],
           "activation": [False, True],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32, tf.int32, tf.int64],
@@ -55,6 +57,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[3]],
           "activation": [True, False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32, tf.int32],
@@ -62,6 +65,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[1, 3, 4, 3]],
           "activation": [True, False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -69,6 +73,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[]],
           "activation": [False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -76,6 +81,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[1]],
           "activation": [False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -83,6 +89,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[1, 3, 4, 3]],
           "activation": [False],
           "fully_quantize": [True],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -90,6 +97,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[5]],
           "activation": [False],
           "fully_quantize": [True],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -97,6 +105,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[3]],
           "activation": [False],
           "fully_quantize": [True],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -104,6 +113,7 @@ def make_binary_op_tests(options,
           "input_shape_2": [[1, 3, 4, 3]],
           "activation": [False],
           "fully_quantize": [True],
+          "dynamic_range_quantize": [False],
       },
       {
           "dtype": [tf.float32],
@@ -111,6 +121,47 @@ def make_binary_op_tests(options,
           "input_shape_2": [[]],
           "activation": [False],
           "fully_quantize": [True],
+          "dynamic_range_quantize": [False],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape_1": [[1, 3, 4, 3]],
+          "input_shape_2": [[1, 3, 4, 3]],
+          "activation": [False],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape_1": [[5]],
+          "input_shape_2": [[5]],
+          "activation": [False],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape_1": [[1, 3, 4, 3]],
+          "input_shape_2": [[3]],
+          "activation": [False],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape_1": [[3]],
+          "input_shape_2": [[1, 3, 4, 3]],
+          "activation": [False],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
+      },
+      {
+          "dtype": [tf.float32],
+          "input_shape_1": [[]],
+          "input_shape_2": [[]],
+          "activation": [False],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
       },
   ]
 
@@ -123,6 +174,7 @@ def make_binary_op_tests(options,
             "input_shape_2": [[7]],
             "activation": [False],
             "fully_quantize": [False],
+            "dynamic_range_quantize": [False],
         },
     ]
 
@@ -204,6 +256,7 @@ def make_div_tests(options):
           "input_shape_2": [[3]],
           "activation": [False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False, True],
       },
   ]
   make_binary_op_tests(
@@ -220,6 +273,7 @@ def make_sub_tests(options):
           "input_shape_2": [[3]],
           "activation": [False],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False, True],
       },
   ]
   make_binary_op_tests(

From 23910c191f29c8ea060b0d24671fccb356ed6600 Mon Sep 17 00:00:00 2001
From: Khanh LeViet 
Date: Tue, 16 Jun 2020 13:48:29 -0700
Subject: [PATCH 0328/1390] Recommend Netron for TF Lite model visualization

PiperOrigin-RevId: 316750470
Change-Id: Id794ed7b2a8405cf5821fb106e8861d8aacef22f
---
 tensorflow/lite/g3doc/guide/faq.md | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/g3doc/guide/faq.md b/tensorflow/lite/g3doc/guide/faq.md
index 197fc2d4a8f..491e3252635 100644
--- a/tensorflow/lite/g3doc/guide/faq.md
+++ b/tensorflow/lite/g3doc/guide/faq.md
@@ -45,30 +45,37 @@ or file a [new one](https://github.com/tensorflow/tensorflow/issues).
 
 #### How do I determine the inputs/outputs for GraphDef protocol buffer?
 
-The easiest way to inspect a graph from a `.pb` file is to use the
+The easiest way to inspect a graph from a `.pb` file is to use
+[Netron](https://github.com/lutzroeder/netron), an open-source viewer for
+machine learning models.
+
+If Netron cannot open the graph, you can try the
 [summarize_graph](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md#inspecting-graphs)
 tool.
 
-If that approach yields an error, you can visualize the GraphDef with
+If the summarize_graph tool yields an error, you can visualize the GraphDef with
 [TensorBoard](https://www.tensorflow.org/guide/summaries_and_tensorboard) and
 look for the inputs and outputs in the graph. To visualize a `.pb` file, use the
 [`import_pb_to_tensorboard.py`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/import_pb_to_tensorboard.py)
 script like below:
 
-```sh
+```shell
 python import_pb_to_tensorboard.py --model_dir  --log_dir 
 ```
 
 #### How do I inspect a `.tflite` file?
 
-TensorFlow Lite models can be visualized using the
+[Netron](https://github.com/lutzroeder/netron) is the easiest way to visualize a
+TensorFlow Lite model.
+
+If Netron cannot open your TensorFlow Lite model, you can try the
 [visualize.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tools/visualize.py)
 script in our repository.
 
 *   [Clone the TensorFlow repository](https://www.tensorflow.org/install/source)
 *   Run the `visualize.py` script with bazel:
 
-```sh
+```shell
 bazel run //tensorflow/lite/tools:visualize model.tflite visualized_model.html
 ```
 

From 295ee8ab72f907a37850a26dbaa684e4f992fa2b Mon Sep 17 00:00:00 2001
From: Anjali Sridhar 
Date: Tue, 16 Jun 2020 14:07:19 -0700
Subject: [PATCH 0329/1390] Assign to all component TPUMirroredVariables when
 assigning in replica context and aggregation=NONE.

PiperOrigin-RevId: 316754219
Change-Id: I791f392b892886404cb80868368ae4a167d8b3d8
---
 .../distribute/mirrored_strategy_test.py      |  2 --
 tensorflow/python/distribute/tpu_values.py    | 30 +++++++++++++++++++
 tensorflow/python/distribute/values_test.py   | 22 ++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/distribute/mirrored_strategy_test.py b/tensorflow/python/distribute/mirrored_strategy_test.py
index 6009eece14e..950b6f2446b 100644
--- a/tensorflow/python/distribute/mirrored_strategy_test.py
+++ b/tensorflow/python/distribute/mirrored_strategy_test.py
@@ -635,8 +635,6 @@ class MirroredVariableUpdateTest(test.TestCase):
 
   def testAssignMirroredVarReplicaContextWithoutAggregationType(self,
                                                                 distribution):
-    # Test that we always have an aggregation type set on the mirrored variable
-    # if we assign to it in replica mode.
     def var_fn():
       v = variable_scope.variable(1.0, name="foo")
       return v
diff --git a/tensorflow/python/distribute/tpu_values.py b/tensorflow/python/distribute/tpu_values.py
index 3a4290a80dd..40ab058ac7c 100644
--- a/tensorflow/python/distribute/tpu_values.py
+++ b/tensorflow/python/distribute/tpu_values.py
@@ -30,6 +30,7 @@ from tensorflow.python.eager import tape
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import gen_resource_variable_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_scope
 from tensorflow.python.tpu import tpu
 
 
@@ -173,6 +174,16 @@ class TPUMirroredVariable(TPUVariableMixin, values.MirroredVariable):
   """Holds a map from replica to TPU variables whose values are kept in sync."""
 
   def assign_sub(self, value, use_locking=False, name=None, read_value=True):
+    if (enclosing_tpu_context() and
+        self.aggregation == variable_scope.VariableAggregation.NONE):
+      return _make_raw_assign_fn(
+          gen_resource_variable_ops.assign_sub_variable_op)(
+              self,
+              value=value,
+              use_locking=use_locking,
+              name=name,
+              read_value=read_value)
+
     assign_sub_fn = _make_raw_assign_fn(
         gen_resource_variable_ops.assign_sub_variable_op)
     return self._update(
@@ -183,6 +194,16 @@ class TPUMirroredVariable(TPUVariableMixin, values.MirroredVariable):
         read_value=read_value)
 
   def assign_add(self, value, use_locking=False, name=None, read_value=True):
+    if (enclosing_tpu_context() and
+        self.aggregation == variable_scope.VariableAggregation.NONE):
+      return _make_raw_assign_fn(
+          gen_resource_variable_ops.assign_add_variable_op)(
+              self,
+              value=value,
+              use_locking=use_locking,
+              name=name,
+              read_value=read_value)
+
     assign_add_fn = _make_raw_assign_fn(
         gen_resource_variable_ops.assign_add_variable_op)
     return self._update(
@@ -193,6 +214,15 @@ class TPUMirroredVariable(TPUVariableMixin, values.MirroredVariable):
         read_value=read_value)
 
   def assign(self, value, use_locking=False, name=None, read_value=True):
+    if (enclosing_tpu_context() and
+        self.aggregation == variable_scope.VariableAggregation.NONE):
+      return _make_raw_assign_fn(gen_resource_variable_ops.assign_variable_op)(
+          self,
+          value=value,
+          use_locking=use_locking,
+          name=name,
+          read_value=read_value)
+
     assign_fn = _make_raw_assign_fn(
         gen_resource_variable_ops.assign_variable_op)
     return self._update(
diff --git a/tensorflow/python/distribute/values_test.py b/tensorflow/python/distribute/values_test.py
index 180a4d3f278..0cb4d6ddd2a 100644
--- a/tensorflow/python/distribute/values_test.py
+++ b/tensorflow/python/distribute/values_test.py
@@ -915,6 +915,28 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
       sess.run(variables_lib.global_variables_initializer())
       sess.run({"complicated": mirrored})
 
+  @combinations.generate(
+      combinations.combine(
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          mode=["eager"]))
+  def testAssignValueInReplicaContextWithoutAggregation(self, distribution):
+    with distribution.scope():
+      v = variables_lib.Variable(1.0, name="foo")
+
+    @def_function.function
+    def mytest():
+      def model_fn():
+        v.assign(5.0)
+        return v.read_value()
+
+      return distribution.run(model_fn)
+
+    mytest()
+    self.assertAllEqual([5.0, 5.0], self.evaluate(v.values))
+
   @combinations.generate(
       combinations.combine(
           distribution=[

From 389405a77946410400ed410246e4cc7257802dde Mon Sep 17 00:00:00 2001
From: Suharsh Sivakumar 
Date: Tue, 16 Jun 2020 14:16:01 -0700
Subject: [PATCH 0330/1390] Add dynamic range tests to relevant op_tests
 (letter C).

PiperOrigin-RevId: 316755889
Change-Id: If087d4bb5db715f2ccefa1aec15a89e9a915c5ad
---
 tensorflow/lite/testing/op_tests/concat.py       | 13 +++++++++++--
 tensorflow/lite/testing/op_tests/conv.py         | 16 +++++++++++++++-
 .../lite/testing/op_tests/conv_activation.py     | 16 +++++++++++++++-
 .../testing/op_tests/conv_with_shared_weights.py |  1 +
 4 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/testing/op_tests/concat.py b/tensorflow/lite/testing/op_tests/concat.py
index 1cb726ceb1d..0d6936a0e48 100644
--- a/tensorflow/lite/testing/op_tests/concat.py
+++ b/tensorflow/lite/testing/op_tests/concat.py
@@ -32,13 +32,22 @@ def make_concat_tests(options):
       "num_tensors": [1, 2, 3, 4, 5, 6],
       "axis": [0, 1, 2, 3, -3, -2, -1],
       "type": [tf.float32, tf.uint8, tf.int32, tf.int64],
-      "fully_quantize": [False]
+      "fully_quantize": [False],
+      "dynamic_range_quantize": [False],
   }, {
       "base_shape": [[1, 3, 4, 3], [3, 4], [2, 3, 4, 3]],
       "num_tensors": [1, 2, 3, 4, 5, 6],
       "axis": [1, 2, 3, -3, -2, -1],
       "type": [tf.float32],
-      "fully_quantize": [True]
+      "fully_quantize": [True],
+      "dynamic_range_quantize": [False],
+  }, {
+      "base_shape": [[1, 3, 4, 3]],
+      "num_tensors": [6],
+      "axis": [1],
+      "type": [tf.float32],
+      "fully_quantize": [False],
+      "dynamic_range_quantize": [True],
   }]
 
   def get_shape(parameters, delta):
diff --git a/tensorflow/lite/testing/op_tests/conv.py b/tensorflow/lite/testing/op_tests/conv.py
index 3a12cafe478..ee0b589b5ca 100644
--- a/tensorflow/lite/testing/op_tests/conv.py
+++ b/tensorflow/lite/testing/op_tests/conv.py
@@ -39,6 +39,7 @@ def make_conv_tests(options):
           "constant_filter": [True, False],
           "channel_multiplier": [1, 2],
           "fully_quantize": [False],
+          "dynamic_range_quantize": [False],
       },
       # TODO(b/134702301): The fully_quantize param is just ignored by the MLIR
       # testing path now, resulting in duplicate tests. Either ignore these
@@ -53,7 +54,20 @@ def make_conv_tests(options):
           "constant_filter": [True],
           "channel_multiplier": [1, 2],
           "fully_quantize": [True],
-      }
+          "dynamic_range_quantize": [False],
+      },
+      {
+          "input_shape": [[1, 3, 4, 3]],
+          "filter_shape": [[1, 1]],
+          "strides": [[1, 1, 1, 1], [1, 2, 3, 1]],
+          "dilations": [[1, 1, 1, 1]],
+          "padding": ["SAME", "VALID"],
+          "data_format": ["NHWC"],
+          "constant_filter": [True],
+          "channel_multiplier": [2],
+          "fully_quantize": [False],
+          "dynamic_range_quantize": [True],
+      },
   ]
 
   def get_tensor_shapes(parameters):
diff --git a/tensorflow/lite/testing/op_tests/conv_activation.py b/tensorflow/lite/testing/op_tests/conv_activation.py
index b4cc4c6ba58..1ee1210ec9e 100644
--- a/tensorflow/lite/testing/op_tests/conv_activation.py
+++ b/tensorflow/lite/testing/op_tests/conv_activation.py
@@ -40,6 +40,7 @@ def make_conv_activation_tests(activation_op):
             "constant_filter": [True, False],
             "channel_multiplier": [1, 2],
             "fully_quantize": [False],
+            "dynamic_range_quantize": [False],
         },
         # TODO(b/134702301): The fully_quantize param is just ignored by the
         # MLIR testing path now, resulting in duplicate tests. Either ignore
@@ -54,7 +55,20 @@ def make_conv_activation_tests(activation_op):
             "constant_filter": [True],
             "channel_multiplier": [1, 2],
             "fully_quantize": [True],
-        }
+            "dynamic_range_quantize": [False],
+        },
+        {
+            "input_shape": [[1, 3, 4, 3]],
+            "filter_shape": [[1, 1], [2, 3], [3, 3]],
+            "strides": [[1, 1, 1, 1], [1, 2, 3, 1]],
+            "dilations": [[1, 1, 1, 1]],
+            "padding": ["SAME", "VALID"],
+            "data_format": ["NHWC"],
+            "constant_filter": [True],
+            "channel_multiplier": [1, 2],
+            "fully_quantize": [False],
+            "dynamic_range_quantize": [True],
+        },
     ]
 
     def get_tensor_shapes(parameters):
diff --git a/tensorflow/lite/testing/op_tests/conv_with_shared_weights.py b/tensorflow/lite/testing/op_tests/conv_with_shared_weights.py
index 6f4499af9d4..bca775ec20e 100644
--- a/tensorflow/lite/testing/op_tests/conv_with_shared_weights.py
+++ b/tensorflow/lite/testing/op_tests/conv_with_shared_weights.py
@@ -36,6 +36,7 @@ def make_conv_with_shared_weights_tests(options):
       "padding": ["SAME"],
       "data_format": ["NHWC"],
       "channel_multiplier": [1],
+      "dynamic_range_quantize": [False, True],
   }]
 
   def get_tensor_shapes(parameters):

From 8019570ea041bf274b42568c65d66b388e4746af Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 14:16:21 -0700
Subject: [PATCH 0331/1390] load nnapi 1.3 memory domain functions in shim
 layer.

PiperOrigin-RevId: 316755954
Change-Id: If4fc4a7c1001e4b47479531914d2631ee2e31fcd
---
 tensorflow/lite/nnapi/NeuralNetworksShim.h    | 325 ++++++++++++++++++
 tensorflow/lite/nnapi/NeuralNetworksTypes.h   |  72 ++++
 tensorflow/lite/nnapi/nnapi_implementation.cc |  12 +
 tensorflow/lite/nnapi/nnapi_implementation.h  | 306 ++++++++++++++++-
 4 files changed, 714 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/nnapi/NeuralNetworksShim.h b/tensorflow/lite/nnapi/NeuralNetworksShim.h
index 01b597ce36f..9a5c44dfd61 100644
--- a/tensorflow/lite/nnapi/NeuralNetworksShim.h
+++ b/tensorflow/lite/nnapi/NeuralNetworksShim.h
@@ -1237,6 +1237,331 @@ inline int ANeuralNetworksModel_setOperandExtensionData(
   EXECUTE_FUNCTION_RETURN(model, index, data, length);
 }
 
+/**
+ * Create a {@link ANeuralNetworksMemoryDesc} with no properties.
+ *
+ * This only creates the memory descriptor. Its properties should be set with
+ * calls to
+ * {@link ANeuralNetworksMemoryDesc_addInputRole},
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole}, and
+ * {@link ANeuralNetworksMemoryDesc_setDimensions}.
+ *
+ * {@link ANeuralNetworksMemoryDesc_finish} must be called once all properties
+ * have been set.
+ *
+ * {@link ANeuralNetworksMemoryDesc_free} must be called once the memory
+ * descriptor is no longer needed.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The {@link ANeuralNetworksMemoryDesc} to be created.
+ *             Set to NULL if unsuccessful.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemoryDesc_create(ANeuralNetworksMemoryDesc** desc) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_create);
+  EXECUTE_FUNCTION_RETURN(desc);
+}
+
+/**
+ * Destroy a memory descriptor.
+ *
+ * The memory descriptor need not have been finished by a call to
+ * {@link ANeuralNetworksMemoryDesc_finish}.
+ *
+ * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded usage.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor to be destroyed. Passing NULL is acceptable
+ * and results in no operation.
+ */
+inline void ANeuralNetworksMemoryDesc_free(ANeuralNetworksMemoryDesc* desc) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_free);
+  EXECUTE_FUNCTION(desc);
+}
+
+/**
+ * Specify that a memory object will be playing the role of an output to an
+ * execution created from a particular compilation.
+ *
+ * The compilation and the output index fully specify an output operand. This
+ * function may be invoked multiple times on the same memory descriptor with
+ * different output operands, and the same output operand may be specified on
+ * multiple memory descriptors. However, specifying the same output operand on
+ * the same memory descriptor object more than once will return an error.
+ *
+ * The dimensions of the corresponding model operands of all the roles specified
+ * by
+ * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be compatible with each
+ * other. Two dimensions are incompatible if both ranks are fully specified but
+ * have different values, or if there is at least one axis that is fully
+ * specified in both but has different values.
+ *
+ * At least one of {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be called on the memory
+ * descriptor before invoking {@link ANeuralNetworksMemoryDesc_finish}.
+ *
+ * Attempting to modify a memory descriptor once
+ * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+ * error.
+ *
+ * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded usage.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor to be modified.
+ * @param compilation The compilation object. It must already have been finished
+ * by calling {@link ANeuralNetworksCompilation_finish}, and must outlive the
+ * memory descriptor.
+ * @param index The index of the output argument we are referencing from the
+ *              compilation. It is an index into the outputs list passed to
+ *              {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not
+ *              the index associated with {@link
+ * ANeuralNetworksModel_addOperand}.
+ * @param frequency A floating-point value within the range (0.0, 1.0].
+ * Describes how likely the memory is to be used in the specified role. This is
+ *                  provided as a hint to optimize the case when multiple roles
+ * prefer different memory locations or data layouts.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemoryDesc_addOutputRole(
+    ANeuralNetworksMemoryDesc* desc,
+    const ANeuralNetworksCompilation* compilation, int32_t index,
+    float frequency) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_addOutputRole);
+  EXECUTE_FUNCTION_RETURN(desc, compilation, index, frequency);
+}
+
+/**
+ * Specify that a memory object will be playing the role of an input to an
+ * execution created from a particular compilation.
+ *
+ * The compilation and the input index fully specify an input operand. This
+ * function may be invoked multiple times on the same memory descriptor with
+ * different input operands, and the same input operand may be specified on
+ * multiple memory descriptors. However, specifying the same input operand on
+ * the same memory descriptor more than once will return an error.
+ *
+ * The dimensions of the corresponding model operands of all the roles specified
+ * by
+ * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be compatible with each
+ * other. Two dimensions are incompatible if both ranks are fully specified but
+ * have different values, or if there is at least one axis that is fully
+ * specified in both but has different values.
+ *
+ * At least one of {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be called on a memory
+ * descriptor before invoking {@link ANeuralNetworksMemoryDesc_finish}.
+ *
+ * Attempting to modify a memory descriptor once
+ * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+ * error.
+ *
+ * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded usage.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor to be modified.
+ * @param compilation The compilation object. It must already have been finished
+ * by calling {@link ANeuralNetworksCompilation_finish}, and must outlive the
+ * memory descriptor.
+ * @param index The index of the input argument we are referencing from the
+ * compilation. It is an index into the inputs list passed to
+ *              {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not
+ *              the index associated with {@link
+ * ANeuralNetworksModel_addOperand}.
+ * @param frequency A floating-point value within the range (0.0, 1.0].
+ * Describes how likely the memory is to be used in the specified role. This is
+ *                  provided as a hint to optimize the case when different roles
+ * prefer different memory locations or data layouts.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemoryDesc_addInputRole(
+    ANeuralNetworksMemoryDesc* desc,
+    const ANeuralNetworksCompilation* compilation, uint32_t index,
+    float frequency) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_addInputRole);
+  EXECUTE_FUNCTION_RETURN(desc, compilation, index, frequency);
+}
+
+/**
+ * Set the dimensional information of the memory descriptor.
+ *
+ * The specified dimensions must be compatible with the dimensions of the
+ * corresponding model operands of all the roles specified by
+ * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole}. Two dimensions are
+ * incompatible if both ranks are fully specified but have different values, or
+ * if there is at least one axis that is fully specified in both but has
+ * different values.
+ *
+ * Attempting to modify a memory descriptor once
+ * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+ * error.
+ *
+ * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded usage.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor to be modified.
+ * @param rank The number of dimensions. Must be 0 for scalars.
+ * @param dimensions An array of dimensions. An entry with the value 0 indicates
+ * that the corresponding axis has an unknown size.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemoryDesc_setDimensions(
+    ANeuralNetworksMemoryDesc* desc, uint32_t rank,
+    const uint32_t* dimensions) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_setDimensions);
+  EXECUTE_FUNCTION_RETURN(desc, rank, dimensions);
+}
+
+/**
+ * Indicate that we have finished modifying a memory descriptor. Required before
+ * calling
+ * {@link ANeuralNetworksMemory_createFromDesc}.
+ *
+ * This function must only be called once for a given memory descriptor.
+ *
+ * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded usage.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor to be finished.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemoryDesc_finish(ANeuralNetworksMemoryDesc* desc) {
+  LOAD_FUNCTION(ANeuralNetworksMemoryDesc_finish);
+  EXECUTE_FUNCTION_RETURN(desc);
+}
+
+/**
+ * Creates a memory object from a memory descriptor.
+ *
+ * The memory object is created with an uninitialized buffer. A memory object
+ * with an uninitialized buffer may only be used according to the roles
+ * specified by
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole}, or as the destination memory
+ * in
+ * {@link ANeuralNetworksMemory_copy}. The buffer of a memory object is
+ * initialized after the memory object is used as an output in a successful
+ * execution, or used as the destination memory in a successful {@link
+ * ANeuralNetworksMemory_copy}. A memory object with an initialized buffer may
+ * be used according to all roles specified in
+ * {@link ANeuralNetworksMemoryDesc}, or as the source or destination memory in
+ * {@link ANeuralNetworksMemory_copy}. The buffer of a memory object will return
+ * to the uninitialized state if the memory object is used as an output in a
+ * failed execution, or used as the destination memory in a failed {@link
+ * ANeuralNetworksMemory_copy}.
+ *
+ * The dimensions of the memory descriptor are deduced from the dimensions of
+ * the corresponding model operands of all the roles specified by
+ * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ * {@link ANeuralNetworksMemoryDesc_addOutputRole}, as well as the dimensions
+ * set by the call to {@link ANeuralNetworksMemoryDesc_setDimensions}, if any.
+ * The memory descriptor may have unspecified dimensions or rank. In such a
+ * case, the same memory object may be used with different shapes of outputs in
+ * different executions. When the memory is used as an input, the input shape
+ * must be the same as the output shape from the last execution using this
+ * memory object as an output, or the last
+ * {@link ANeuralNetworkMemory_copy} using this memory object as the destination
+ * memory. Creating a memory object with unspecified dimensions or rank may fail
+ * for certain sets of roles.
+ *
+ * Using the memory in roles or shapes that are not compatible with the rules
+ * specified above will return an error.
+ *
+ * When calling {@link ANeuralNetworksExecution_setInputFromMemory} or
+ * {@link ANeuralNetworksExecution_setOutputFromMemory} with the memory object,
+ * both offset and length must be set to zero and the entire memory region will
+ * be associated with the specified input or output operand.
+ *
+ * Calling {@link ANeuralNetworksModel_setOperandValueFromMemory} with the
+ * memory created from this function will return an error.
+ *
+ * {@link ANeuralNetworksMemory_free} must be called once the memory is no
+ * longer needed.
+ *
+ * Attempting to create memory from an unfinished memory descriptor will return
+ * an error.
+ *
+ * The provided {@link ANeuralNetworksMemoryDesc} need not outlive the
+ * {@link ANeuralNetworksMemory} object.
+ *
+ * Available since API level 30.
+ *
+ * @param desc The memory descriptor.
+ * @param memory The memory object to be created.
+ *               Set to NULL if unsuccessful.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful; ANEURALNETWORKS_OP_FAILED if
+ * the memory is created with unspecified dimensions or rank and it is not
+ * supported for this set of roles.
+ */
+inline int ANeuralNetworksMemory_createFromDesc(
+    const ANeuralNetworksMemoryDesc* desc, ANeuralNetworksMemory** memory) {
+  LOAD_FUNCTION(ANeuralNetworksMemory_createFromDesc);
+  EXECUTE_FUNCTION_RETURN(desc, memory);
+}
+
+/**
+ * Copies data from one memory object to another.
+ *
+ * If at most one of the src and dst is created from
+ * {@link ANeuralNetworksMemory_createFromDesc}, the src and dst must have the
+ * same logical size:
+ * - If the memory is created from {@link ANeuralNetworksMemory_createFromFd},
+ * or if it is created from {@link
+ * ANeuralNetworksMemory_createFromAHardwareBuffer} with format of
+ * AHARDWAREBUFFER_FORMAT_BLOB, the logical size equals the size of the memory.
+ * - If the memory is created from
+ *   {@link ANeuralNetworksMemory_createFromAHardwareBuffer} with a format other
+ * than AHARDWAREBUFFER_FORMAT_BLOB, the logical size equals the size when there
+ * is no padding and the data is tightly packed. This function may fail if the
+ *   AHardwareBuffer cannot be accessed.
+ * - If the memory is created from {@link ANeuralNetworksMemory_createFromDesc},
+ * the logical size equals the size indicated by the {@link OperandCode}
+ * multiplied by the number of elements. This function will fail if the number
+ * of elements is unknown.
+ *
+ * If both src and dst are created from {@link
+ * ANeuralNetworksMemory_createFromDesc}, they must have compatible dimensions.
+ * Two dimensions are incompatible if both ranks are fully specified but have
+ * different values, or if there is at least one axis that is fully specified in
+ * both but has different values. The dst may have unspecified dimensions or
+ * rank. In such a case, the dimensions of dst will get updated according to the
+ * dimensions of the src.
+ *
+ * In both cases, if the src is created from
+ * {@link ANeuralNetworksMemory_createFromDesc}, it must have been used as an
+ * output in a successful execution, or used as the destination memory in a
+ * successful
+ * {@link ANeuralNetworksMemory_copy}.
+ *
+ * The src and dst may have different data layout, in which case the data
+ * copying is performed logically with data layout transformation.
+ *
+ * Available since API level 30.
+ *
+ * @param src The source memory object.
+ * @param dst The destination memory object.
+ *
+ * @return ANEURALNETWORKS_NO_ERROR if successful.
+ */
+inline int ANeuralNetworksMemory_copy(const ANeuralNetworksMemory* src,
+                                      const ANeuralNetworksMemory* dst) {
+  LOAD_FUNCTION(ANeuralNetworksMemory_copy);
+  EXECUTE_FUNCTION_RETURN(src, dst);
+}
+
 /**/
 
 #endif  // TENSORFLOW_LITE_NNAPI_NEURALNETWORKSSHIM_H_
diff --git a/tensorflow/lite/nnapi/NeuralNetworksTypes.h b/tensorflow/lite/nnapi/NeuralNetworksTypes.h
index 6739838e4d1..3c30a0479fa 100644
--- a/tensorflow/lite/nnapi/NeuralNetworksTypes.h
+++ b/tensorflow/lite/nnapi/NeuralNetworksTypes.h
@@ -226,6 +226,50 @@ enum {
   ANEURALNETWORKS_PRIORITY_HIGH = 110,
   ANEURALNETWORKS_PRIORITY_DEFAULT = ANEURALNETWORKS_PRIORITY_MEDIUM,
 };
+/**
+ * ANeuralNetworksMemoryDesc is an opaque type that represents a memory
+ * descriptor.
+ *
+ * A memory descriptor describes the properties of a memory object, and is used
+ * by
+ * {@link ANeuralNetworksMemory_createFromDesc}.
+ *
+ * To use:
+ *   - Create a new memory descriptor by calling
+ *     {@link ANeuralNetworksMemoryDesc_create}.
+ *   - Specify all of the intended input and output roles by calling
+ *     {@link ANeuralNetworksMemoryDesc_addInputRole} and
+ *     {@link ANeuralNetworksMemoryDesc_addOutputRole}.
+ *   - Optionally, specify the memory dimensions by calling
+ *     {@link ANeuralNetworksMemoryDesc_setDimensions}.
+ *   - Complete the memory descriptor with {@link
+ * ANeuralNetworksMemoryDesc_finish}.
+ *   - Use the memory descriptor as many times as needed with
+ *     {@link ANeuralNetworksMemory_createFromDesc}.
+ *   - Destroy the memory descriptor with {@link
+ * ANeuralNetworksMemoryDesc_free}.
+ *
+ * A memory descriptor is completed by calling {@link
+ * ANeuralNetworksMemoryDesc_finish}. A memory descriptor is destroyed by
+ * calling {@link ANeuralNetworksMemoryDesc_free}.
+ *
+ * A memory descriptor must not be modified once
+ * {@link ANeuralNetworksMemoryDesc_finish}
+ * has been called on it.
+ *
+ * It is the application's responsibility to make sure that only
+ * one thread modifies a memory descriptor at a given time. It is however
+ * safe for more than one thread to use the memory descriptor once
+ * {@link ANeuralNetworksMemoryDesc_finish} has returned.
+ *
+ * It is also the application's responsibility to ensure that there are no other
+ * uses of the memory descriptor after calling {@link
+ * ANeuralNetworksMemoryDesc_free}. It is however safe to continue using a
+ * {@link ANeuralNetworksMemory} object created from the memory descriptor.
+ *
+ * Available since API level 30.
+ */
+typedef struct ANeuralNetworksMemoryDesc ANeuralNetworksMemoryDesc;
 
 /**
  * ANeuralNetworksMemory is an opaque type that represents memory.
@@ -604,4 +648,32 @@ typedef int (*ANeuralNetworksModel_setOperandExtensionData_fn)(
     ANeuralNetworksModel* model, int32_t index, const void* data,
     size_t length);
 
+typedef int (*ANeuralNetworksMemoryDesc_create_fn)(
+    ANeuralNetworksMemoryDesc** desc);
+
+typedef void (*ANeuralNetworksMemoryDesc_free_fn)(
+    ANeuralNetworksMemoryDesc* desc);
+
+typedef int (*ANeuralNetworksMemoryDesc_addInputRole_fn)(
+    ANeuralNetworksMemoryDesc* desc,
+    const ANeuralNetworksCompilation* compilation, int32_t index,
+    float frequency);
+
+typedef int (*ANeuralNetworksMemoryDesc_addOutputRole_fn)(
+    ANeuralNetworksMemoryDesc* desc,
+    const ANeuralNetworksCompilation* compilation, uint32_t index,
+    float frequency);
+
+typedef int (*ANeuralNetworksMemoryDesc_setDimensions_fn)(
+    ANeuralNetworksMemoryDesc* desc, uint32_t rank, const uint32_t* dimensions);
+
+typedef int (*ANeuralNetworksMemoryDesc_finish_fn)(
+    ANeuralNetworksMemoryDesc* desc);
+
+typedef int (*ANeuralNetworksMemory_createFromDesc_fn)(
+    const ANeuralNetworksMemoryDesc* desc, ANeuralNetworksMemory** memory);
+
+typedef int (*ANeuralNetworksMemory_copy_fn)(const ANeuralNetworksMemory* src,
+                                             const ANeuralNetworksMemory* dst);
+
 #endif  // TENSORFLOW_LITE_NNAPI_NEURALNETWORKSTYPES_H_
diff --git a/tensorflow/lite/nnapi/nnapi_implementation.cc b/tensorflow/lite/nnapi/nnapi_implementation.cc
index ad5869fec04..862c4ba2499 100644
--- a/tensorflow/lite/nnapi/nnapi_implementation.cc
+++ b/tensorflow/lite/nnapi/nnapi_implementation.cc
@@ -225,6 +225,18 @@ const NnApi LoadNnApi() {
                          ANeuralNetworksExecution_setTimeout);
   LOAD_FUNCTION_OPTIONAL(libneuralnetworks,
                          ANeuralNetworksExecution_setLoopTimeout);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks, ANeuralNetworksMemoryDesc_create);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks, ANeuralNetworksMemoryDesc_free);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks,
+                         ANeuralNetworksMemoryDesc_addInputRole);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks,
+                         ANeuralNetworksMemoryDesc_addOutputRole);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks,
+                         ANeuralNetworksMemoryDesc_setDimensions);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks, ANeuralNetworksMemoryDesc_finish);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks,
+                         ANeuralNetworksMemory_createFromDesc);
+  LOAD_FUNCTION_OPTIONAL(libneuralnetworks, ANeuralNetworksMemory_copy);
 
   return nnapi;
 }
diff --git a/tensorflow/lite/nnapi/nnapi_implementation.h b/tensorflow/lite/nnapi/nnapi_implementation.h
index abee0fbdef3..9f481cded9b 100644
--- a/tensorflow/lite/nnapi/nnapi_implementation.h
+++ b/tensorflow/lite/nnapi/nnapi_implementation.h
@@ -1225,7 +1225,311 @@ struct NnApi {
       ANeuralNetworksModel* model, int32_t index, const void* data,
       size_t length);
 
-  /**/
+  /**
+   * Create a {@link ANeuralNetworksMemoryDesc} with no properties.
+   *
+   * This only creates the memory descriptor. Its properties should be set with
+   * calls to
+   * {@link ANeuralNetworksMemoryDesc_addInputRole},
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole}, and
+   * {@link ANeuralNetworksMemoryDesc_setDimensions}.
+   *
+   * {@link ANeuralNetworksMemoryDesc_finish} must be called once all properties
+   * have been set.
+   *
+   * {@link ANeuralNetworksMemoryDesc_free} must be called once the memory
+   * descriptor is no longer needed.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The {@link ANeuralNetworksMemoryDesc} to be created.
+   *             Set to NULL if unsuccessful.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemoryDesc_create)(ANeuralNetworksMemoryDesc** desc);
+
+  /**
+   * Destroy a memory descriptor.
+   *
+   * The memory descriptor need not have been finished by a call to
+   * {@link ANeuralNetworksMemoryDesc_finish}.
+   *
+   * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded
+   * usage.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor to be destroyed. Passing NULL is
+   * acceptable and results in no operation.
+   */
+  void (*ANeuralNetworksMemoryDesc_free)(ANeuralNetworksMemoryDesc* desc);
+
+  /**
+   * Specify that a memory object will be playing the role of an input to an
+   * execution created from a particular compilation.
+   *
+   * The compilation and the input index fully specify an input operand. This
+   * function may be invoked multiple times on the same memory descriptor with
+   * different input operands, and the same input operand may be specified on
+   * multiple memory descriptors. However, specifying the same input operand on
+   * the same memory descriptor more than once will return an error.
+   *
+   * The dimensions of the corresponding model operands of all the roles
+   * specified by
+   * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be compatible with
+   * each other. Two dimensions are incompatible if both ranks are fully
+   * specified but have different values, or if there is at least one axis that
+   * is fully specified in both but has different values.
+   *
+   * At least one of {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be called on a memory
+   * descriptor before invoking {@link ANeuralNetworksMemoryDesc_finish}.
+   *
+   * Attempting to modify a memory descriptor once
+   * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+   * error.
+   *
+   * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded
+   * usage.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor to be modified.
+   * @param compilation The compilation object. It must already have been
+   * finished by calling {@link ANeuralNetworksCompilation_finish}, and must
+   * outlive the memory descriptor.
+   * @param index The index of the input argument we are referencing from the
+   * compilation. It is an index into the inputs list passed to
+   *              {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is
+   * not the index associated with {@link ANeuralNetworksModel_addOperand}.
+   * @param frequency A floating-point value within the range (0.0, 1.0].
+   * Describes how likely the memory is to be used in the specified role. This
+   * is provided as a hint to optimize the case when different roles prefer
+   * different memory locations or data layouts.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemoryDesc_addInputRole)(
+      ANeuralNetworksMemoryDesc* desc,
+      const ANeuralNetworksCompilation* compilation, int32_t index,
+      float frequency);
+
+  /**
+   * Specify that a memory object will be playing the role of an output to an
+   * execution created from a particular compilation.
+   *
+   * The compilation and the output index fully specify an output operand. This
+   * function may be invoked multiple times on the same memory descriptor with
+   * different output operands, and the same output operand may be specified on
+   * multiple memory descriptors. However, specifying the same output operand on
+   * the same memory descriptor object more than once will return an error.
+   *
+   * The dimensions of the corresponding model operands of all the roles
+   * specified by
+   * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be compatible with
+   * each other. Two dimensions are incompatible if both ranks are fully
+   * specified but have different values, or if there is at least one axis that
+   * is fully specified in both but has different values.
+   *
+   * At least one of {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole} must be called on the
+   * memory descriptor before invoking {@link ANeuralNetworksMemoryDesc_finish}.
+   *
+   * Attempting to modify a memory descriptor once
+   * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+   * error.
+   *
+   * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded
+   * usage.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor to be modified.
+   * @param compilation The compilation object. It must already have been
+   * finished by calling {@link ANeuralNetworksCompilation_finish}, and must
+   * outlive the memory descriptor.
+   * @param index The index of the output argument we are referencing from the
+   *              compilation. It is an index into the outputs list passed to
+   *              {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is
+   * not the index associated with {@link ANeuralNetworksModel_addOperand}.
+   * @param frequency A floating-point value within the range (0.0, 1.0].
+   * Describes how likely the memory is to be used in the specified role. This
+   * is provided as a hint to optimize the case when multiple roles prefer
+   * different memory locations or data layouts.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemoryDesc_addOutputRole)(
+      ANeuralNetworksMemoryDesc* desc,
+      const ANeuralNetworksCompilation* compilation, uint32_t index,
+      float frequency);
+
+  /**
+   * Set the dimensional information of the memory descriptor.
+   *
+   * The specified dimensions must be compatible with the dimensions of the
+   * corresponding model operands of all the roles specified by
+   * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole}. Two dimensions are
+   * incompatible if both ranks are fully specified but have different values,
+   * or if there is at least one axis that is fully specified in both but has
+   * different values.
+   *
+   * Attempting to modify a memory descriptor once
+   * {@link ANeuralNetworksMemoryDesc_finish} has been called will return an
+   * error.
+   *
+   * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded
+   * usage.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor to be modified.
+   * @param rank The number of dimensions. Must be 0 for scalars.
+   * @param dimensions An array of dimensions. An entry with the value 0
+   * indicates that the corresponding axis has an unknown size.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemoryDesc_setDimensions)(
+      ANeuralNetworksMemoryDesc* desc, uint32_t rank,
+      const uint32_t* dimensions);
+
+  /**
+   * Indicate that we have finished modifying a memory descriptor. Required
+   * before calling
+   * {@link ANeuralNetworksMemory_createFromDesc}.
+   *
+   * This function must only be called once for a given memory descriptor.
+   *
+   * See {@link ANeuralNetworksMemoryDesc} for information on multithreaded
+   * usage.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor to be finished.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemoryDesc_finish)(ANeuralNetworksMemoryDesc* desc);
+
+  /**
+   * Creates a memory object from a memory descriptor.
+   *
+   * The memory object is created with an uninitialized buffer. A memory object
+   * with an uninitialized buffer may only be used according to the roles
+   * specified by
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole}, or as the destination
+   * memory in
+   * {@link ANeuralNetworksMemory_copy}. The buffer of a memory object is
+   * initialized after the memory object is used as an output in a successful
+   * execution, or used as the destination memory in a successful {@link
+   * ANeuralNetworksMemory_copy}. A memory object with an initialized buffer may
+   * be used according to all roles specified in
+   * {@link ANeuralNetworksMemoryDesc}, or as the source or destination memory
+   * in
+   * {@link ANeuralNetworksMemory_copy}. The buffer of a memory object will
+   * return to the uninitialized state if the memory object is used as an output
+   * in a failed execution, or used as the destination memory in a failed {@link
+   * ANeuralNetworksMemory_copy}.
+   *
+   * The dimensions of the memory descriptor are deduced from the dimensions of
+   * the corresponding model operands of all the roles specified by
+   * {@link ANeuralNetworksMemoryDesc_addInputRole} and
+   * {@link ANeuralNetworksMemoryDesc_addOutputRole}, as well as the dimensions
+   * set by the call to {@link ANeuralNetworksMemoryDesc_setDimensions}, if any.
+   * The memory descriptor may have unspecified dimensions or rank. In such a
+   * case, the same memory object may be used with different shapes of outputs
+   * in different executions. When the memory is used as an input, the input
+   * shape must be the same as the output shape from the last execution using
+   * this memory object as an output, or the last
+   * {@link ANeuralNetworkMemory_copy} using this memory object as the
+   * destination memory. Creating a memory object with unspecified dimensions or
+   * rank may fail for certain sets of roles.
+   *
+   * Using the memory in roles or shapes that are not compatible with the rules
+   * specified above will return an error.
+   *
+   * When calling {@link ANeuralNetworksExecution_setInputFromMemory} or
+   * {@link ANeuralNetworksExecution_setOutputFromMemory} with the memory
+   * object, both offset and length must be set to zero and the entire memory
+   * region will be associated with the specified input or output operand.
+   *
+   * Calling {@link ANeuralNetworksModel_setOperandValueFromMemory} with the
+   * memory created from this function will return an error.
+   *
+   * {@link ANeuralNetworksMemory_free} must be called once the memory is no
+   * longer needed.
+   *
+   * Attempting to create memory from an unfinished memory descriptor will
+   * return an error.
+   *
+   * The provided {@link ANeuralNetworksMemoryDesc} need not outlive the
+   * {@link ANeuralNetworksMemory} object.
+   *
+   * Available since API level 30.
+   *
+   * @param desc The memory descriptor.
+   * @param memory The memory object to be created.
+   *               Set to NULL if unsuccessful.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful; ANEURALNETWORKS_OP_FAILED
+   * if the memory is created with unspecified dimensions or rank and it is not
+   * supported for this set of roles.
+   */
+  int (*ANeuralNetworksMemory_createFromDesc)(
+      const ANeuralNetworksMemoryDesc* desc, ANeuralNetworksMemory** memory);
+
+  /**
+   * Copies data from one memory object to another.
+   *
+   * If at most one of the src and dst is created from
+   * {@link ANeuralNetworksMemory_createFromDesc}, the src and dst must have the
+   * same logical size:
+   * - If the memory is created from {@link ANeuralNetworksMemory_createFromFd},
+   * or if it is created from {@link
+   * ANeuralNetworksMemory_createFromAHardwareBuffer} with format of
+   * AHARDWAREBUFFER_FORMAT_BLOB, the logical size equals the size of the
+   * memory.
+   * - If the memory is created from
+   *   {@link ANeuralNetworksMemory_createFromAHardwareBuffer} with a format
+   * other than AHARDWAREBUFFER_FORMAT_BLOB, the logical size equals the size
+   * when there is no padding and the data is tightly packed. This function may
+   * fail if the AHardwareBuffer cannot be accessed.
+   * - If the memory is created from {@link
+   * ANeuralNetworksMemory_createFromDesc}, the logical size equals the size
+   * indicated by the {@link OperandCode} multiplied by the number of elements.
+   * This function will fail if the number of elements is unknown.
+   *
+   * If both src and dst are created from {@link
+   * ANeuralNetworksMemory_createFromDesc}, they must have compatible
+   * dimensions. Two dimensions are incompatible if both ranks are fully
+   * specified but have different values, or if there is at least one axis that
+   * is fully specified in both but has different values. The dst may have
+   * unspecified dimensions or rank. In such a case, the dimensions of dst will
+   * get updated according to the dimensions of the src.
+   *
+   * In both cases, if the src is created from
+   * {@link ANeuralNetworksMemory_createFromDesc}, it must have been used as an
+   * output in a successful execution, or used as the destination memory in a
+   * successful
+   * {@link ANeuralNetworksMemory_copy}.
+   *
+   * The src and dst may have different data layout, in which case the data
+   * copying is performed logically with data layout transformation.
+   *
+   * Available since API level 30.
+   *
+   * @param src The source memory object.
+   * @param dst The destination memory object.
+   *
+   * @return ANEURALNETWORKS_NO_ERROR if successful.
+   */
+  int (*ANeuralNetworksMemory_copy)(const ANeuralNetworksMemory* src,
+                                    const ANeuralNetworksMemory* dst);
 };
 
 /**

From 4334d8b06f333bee5dc17c942ea41beb5bd9c32d Mon Sep 17 00:00:00 2001
From: Mahmoud Abuzaina 
Date: Tue, 16 Jun 2020 14:26:13 -0700
Subject: [PATCH 0332/1390] Fixing build issue on windows

---
 third_party/mkl_dnn/mkldnn_v1.BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/third_party/mkl_dnn/mkldnn_v1.BUILD b/third_party/mkl_dnn/mkldnn_v1.BUILD
index 47a7efecda3..7bdec138b99 100644
--- a/third_party/mkl_dnn/mkldnn_v1.BUILD
+++ b/third_party/mkl_dnn/mkldnn_v1.BUILD
@@ -71,6 +71,8 @@ cc_library(
         "src/cpu/**/*.cpp",
         "src/cpu/**/*.hpp",
         "src/cpu/xbyak/*.h",
+        "src/cpu/jit_utils/jitprofiling/*.c",
+        "src/cpu/jit_utils/jitprofiling/*.h",
     ]) + [
         ":dnnl_config_h",
         ":dnnl_version_h",

From aea06cde550055304a2c1b5f23d1ff1dd9a24107 Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Tue, 16 Jun 2020 14:50:48 -0700
Subject: [PATCH 0333/1390] Update tolerance for Log1p test for Windows Build.

PiperOrigin-RevId: 316762917
Change-Id: Ia6620bb7e03bd6f71db0a8090b79eef78c0cb15d
---
 tensorflow/compiler/tests/special_math_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/tests/special_math_test.py b/tensorflow/compiler/tests/special_math_test.py
index 246ab2a1641..bd105bb5e95 100644
--- a/tensorflow/compiler/tests/special_math_test.py
+++ b/tensorflow/compiler/tests/special_math_test.py
@@ -106,13 +106,13 @@ class Log1pTest(xla_test.XLATestCase, parameterized.TestCase):
     self._test_range(-40., -20., dtype, rtol, atol, is_negative=False)
     self._test_range(-40., -20., dtype, rtol, atol, is_negative=True)
 
-  @parameterized.parameters((np.float32, 1e-7, 0.),
+  @parameterized.parameters((np.float32, 2e-7, 0.),
                             (np.float64, 1e-15, 0.))
   def testGreaterThanNegativeTwentyExponent(self, dtype, rtol, atol):
     self._test_range(-20., -10., dtype, rtol, atol, is_negative=False)
     self._test_range(-20., -10., dtype, rtol, atol, is_negative=True)
 
-  @parameterized.parameters((np.float32, 1e-7, 0.),
+  @parameterized.parameters((np.float32, 2e-7, 0.),
                             (np.float64, 1e-15, 0.))
   def testGreaterThanNegativeTenExponent(self, dtype, rtol, atol):
     self._test_range(-10., -5., dtype, rtol, atol, is_negative=False)

From dd789e77833942bf7ffccb9e2e2e93e1c6dff436 Mon Sep 17 00:00:00 2001
From: Jose Baiocchi 
Date: Tue, 16 Jun 2020 15:01:16 -0700
Subject: [PATCH 0334/1390] Remove XPlane to trace viewer PID mapping from
 XPlane schema

PiperOrigin-RevId: 316764981
Change-Id: If8e23e0136493dc721fc41095fd9425b8226b060
---
 tensorflow/core/profiler/convert/BUILD        |   7 +-
 .../convert/xplane_to_trace_events.cc         | 106 ++++++++++--------
 .../convert/xplane_to_trace_events_test.cc    |  16 +--
 .../core/profiler/internal/cpu/host_tracer.cc |   1 -
 .../internal/cpu/metadata_collector.cc        |   1 -
 tensorflow/core/profiler/internal/gpu/BUILD   |   7 +-
 .../profiler/internal/gpu/device_tracer.cc    |   4 +-
 .../internal/gpu/device_tracer_test.cc        |   2 -
 tensorflow/core/profiler/utils/BUILD          |   3 +
 tensorflow/core/profiler/utils/trace_utils.h  |  15 ++-
 .../core/profiler/utils/xplane_schema.cc      |   9 --
 .../core/profiler/utils/xplane_schema.h       |  15 ---
 12 files changed, 93 insertions(+), 93 deletions(-)

diff --git a/tensorflow/core/profiler/convert/BUILD b/tensorflow/core/profiler/convert/BUILD
index 5f287a14267..abf0176bf6f 100644
--- a/tensorflow/core/profiler/convert/BUILD
+++ b/tensorflow/core/profiler/convert/BUILD
@@ -377,7 +377,9 @@ cc_library(
         "//tensorflow/core/profiler/protobuf:trace_events_proto_cc",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "//tensorflow/core/profiler/utils:tf_xplane_visitor",
+        "//tensorflow/core/profiler/utils:trace_utils",
         "//tensorflow/core/profiler/utils:xplane_schema",
+        "//tensorflow/core/profiler/utils:xplane_utils",
         "//tensorflow/core/profiler/utils:xplane_visitor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
@@ -390,17 +392,14 @@ tf_cc_test(
     srcs = ["xplane_to_trace_events_test.cc"],
     deps = [
         ":xplane_to_trace_events",
-        "//tensorflow/core:lib",
-        "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
-        "//tensorflow/core:testlib",
         "//tensorflow/core/profiler/protobuf:trace_events_proto_cc",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+        "//tensorflow/core/profiler/utils:trace_utils",
         "//tensorflow/core/profiler/utils:xplane_builder",
         "//tensorflow/core/profiler/utils:xplane_schema",
-        "//tensorflow/core/profiler/utils:xplane_utils",
     ],
 )
 
diff --git a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
index f4a0145d8f6..882f50e6080 100644
--- a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc
@@ -28,7 +28,9 @@ limitations under the License.
 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 #include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
+#include "tensorflow/core/profiler/utils/trace_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
+#include "tensorflow/core/profiler/utils/xplane_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_visitor.h"
 
 namespace tensorflow {
@@ -36,25 +38,63 @@ namespace profiler {
 
 namespace {
 
-Device BuildDeviceAndResource(const XPlaneVisitor& plane) {
-  Device device;
-  device.set_name(std::string(plane.Name()));
-  device.set_device_id(plane.Id());
+void BuildDeviceAndResources(uint32 device_id, const XPlaneVisitor& plane,
+                             Device* device) {
+  device->set_name(std::string(plane.Name()));
+  device->set_device_id(device_id);
 
-  bool sort_by_ordinal = (plane.Name() == kHostThreadsPlaneName);
+  bool sort_by_ordinal = (device_id == kHostThreadsDeviceId);
   int ordinal = 0;
   plane.ForEachLine([&](const XLineVisitor& line) {
-    Resource resource;
-    resource.set_resource_id(line.Id());
+    uint32 resource_id = line.DisplayId();
+    Resource& resource = (*device->mutable_resources())[resource_id];
+    resource.set_resource_id(resource_id);
     resource.set_name(std::string(line.DisplayName()));
     if (sort_by_ordinal) {
       // When sort_index is absent (i.e. 0), resource id will be used.
       // Therefore sort_index starts with 1.
       resource.set_sort_index(++ordinal);
     }
-    (*device.mutable_resources())[line.Id()] = resource;
   });
-  return device;
+}
+
+void ConvertXPlaneToTraceEvents(uint32 device_id, const XPlaneVisitor& xplane,
+                                Trace* trace) {
+  // Convert devices and resources.
+  BuildDeviceAndResources(device_id, xplane,
+                          &(*trace->mutable_devices())[device_id]);
+
+  // Convert events.
+  xplane.ForEachLine([device_id, trace](const XLineVisitor& xline) {
+    uint32 resource_id = xline.DisplayId();
+    xline.ForEachEvent(
+        [device_id, resource_id, trace](const XEventVisitor& xevent) {
+          int64 event_type =
+              xevent.Type().value_or(HostEventType::kUnknownHostEventType);
+          if (event_type == HostEventType::kMemoryAllocation ||
+              event_type == HostEventType::kMemoryDeallocation) {
+            return;
+          }
+          auto* event = trace->add_trace_events();
+          auto& args = *event->mutable_args();
+          event->set_device_id(device_id);
+          event->set_resource_id(resource_id);
+          if (xevent.HasDisplayName()) {
+            event->set_name(std::string(xevent.DisplayName()));
+            args["long_name"] = std::string(xevent.Name());
+          } else {
+            event->set_name(std::string(xevent.Name()));
+          }
+          event->set_timestamp_ps(xevent.TimestampPs());
+          event->set_duration_ps(xevent.DurationPs());
+
+          xevent.ForEachStat([&](const XStatVisitor& stat) {
+            if (stat.ValueCase() == XStat::VALUE_NOT_SET) return;
+            if (IsInternalStat(stat.Type())) return;
+            args[std::string(stat.Name())] = stat.ToString();
+          });
+        });
+  });
 }
 
 }  // namespace
@@ -81,44 +121,18 @@ void MaybeDropEventsForTraceViewer(Trace* trace, uint32 limit) {
 }
 
 void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace) {
-  auto* trace_devices = trace->mutable_devices();
+  const XPlane* host_plane = FindPlaneWithName(xspace, kHostThreadsPlaneName);
+  if (host_plane != nullptr) {
+    XPlaneVisitor xplane = CreateTfXPlaneVisitor(host_plane);
+    ConvertXPlaneToTraceEvents(kHostThreadsDeviceId, xplane, trace);
+  }
 
-  for (const auto& raw_plane : xspace.planes()) {
-    XPlaneVisitor xplane = CreateTfXPlaneVisitor(&raw_plane);
-    // Convert devices and resources.
-    int64 device_id = xplane.Id();
-    (*trace_devices)[device_id] = BuildDeviceAndResource(xplane);
-
-    // Convert events.
-    xplane.ForEachLine([&](const XLineVisitor& xline) {
-      int64 resource_id = xline.Id();  // Either thread id or CUDA stream id.
-      xline.ForEachEvent([&](const XEventVisitor& xevent) {
-        int64 event_type =
-            xevent.Type().value_or(HostEventType::kUnknownHostEventType);
-        if (event_type == HostEventType::kMemoryAllocation ||
-            event_type == HostEventType::kMemoryDeallocation) {
-          return;
-        }
-        auto* event = trace->add_trace_events();
-        auto& args = *event->mutable_args();
-        event->set_device_id(device_id);
-        event->set_resource_id(resource_id);
-        if (xevent.HasDisplayName()) {
-          event->set_name(std::string(xevent.DisplayName()));
-          args["long_name"] = std::string(xevent.Name());
-        } else {
-          event->set_name(std::string(xevent.Name()));
-        }
-        event->set_timestamp_ps(xevent.TimestampPs());
-        event->set_duration_ps(xevent.DurationPs());
-
-        xevent.ForEachStat([&](const XStatVisitor& stat) {
-          if (stat.ValueCase() == XStat::VALUE_NOT_SET) return;
-          if (IsInternalStat(stat.Type())) return;
-          args[std::string(stat.Name())] = stat.ToString();
-        });
-      });
-    });
+  const std::vector device_planes =
+      FindPlanesWithPrefix(xspace, kGpuPlanePrefix);
+  for (const XPlane* device_plane : device_planes) {
+    XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane);
+    uint32 device_id = kFirstDeviceId + xplane.Id();
+    ConvertXPlaneToTraceEvents(device_id, xplane, trace);
   }
 
   // Trace viewer (non-streaming) has scalability issues, we need to drop
diff --git a/tensorflow/core/profiler/convert/xplane_to_trace_events_test.cc b/tensorflow/core/profiler/convert/xplane_to_trace_events_test.cc
index b9a9fe09981..1e0d27ae68a 100644
--- a/tensorflow/core/profiler/convert/xplane_to_trace_events_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_trace_events_test.cc
@@ -18,7 +18,9 @@ limitations under the License.
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/profiler/protobuf/trace_events.pb.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+#include "tensorflow/core/profiler/utils/trace_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
+#include "tensorflow/core/profiler/utils/xplane_schema.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -26,10 +28,7 @@ namespace {
 
 void CreateXSpace(XSpace* space) {
   XPlaneBuilder host_plane(space->add_planes());
-  XPlaneBuilder device_plane(space->add_planes());
-
-  host_plane.SetName("cpu");
-  host_plane.SetId(0);
+  host_plane.SetName(kHostThreadsPlaneName);
   XLineBuilder thread1 = host_plane.GetOrCreateLine(10);
   thread1.SetName("thread1");
   XEventBuilder event1 =
@@ -47,8 +46,9 @@ void CreateXSpace(XSpace* space) {
   event2.ParseAndAddStatValue(*host_plane.GetOrCreateStatMetadata("tf_op"),
                               "Conv2D");
 
-  device_plane.SetName("gpu:0");
-  device_plane.SetId(1);
+  XPlaneBuilder device_plane(space->add_planes());
+  device_plane.SetName(GpuPlaneName(0));
+  device_plane.SetId(0);
   XLineBuilder stream1 = device_plane.GetOrCreateLine(30);
   stream1.SetName("gpu stream 1");
   XEventBuilder event3 =
@@ -67,8 +67,8 @@ TEST(ConvertXPlaneToTraceEvents, Convert) {
   ConvertXSpaceToTraceEvents(xspace, &trace);
 
   ASSERT_EQ(trace.devices_size(), 2);
-  EXPECT_EQ(trace.devices().at(0).resources_size(), 2);
-  EXPECT_EQ(trace.devices().at(1).resources_size(), 1);
+  EXPECT_EQ(trace.devices().at(kHostThreadsDeviceId).resources_size(), 2);
+  EXPECT_EQ(trace.devices().at(kFirstDeviceId).resources_size(), 1);
   EXPECT_EQ(trace.trace_events_size(), 3);
 }
 
diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer.cc b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
index 37f7baca1d3..fa21df004df 100644
--- a/tensorflow/core/profiler/internal/cpu/host_tracer.cc
+++ b/tensorflow/core/profiler/internal/cpu/host_tracer.cc
@@ -147,7 +147,6 @@ Status HostTracer::CollectData(XSpace* space) {
   }
   MakeCompleteEvents(&events_);
   XPlane* plane = FindOrAddMutablePlaneWithName(space, kHostThreadsPlaneName);
-  plane->set_id(kHostPlaneId);
   ConvertCompleteEventsToXPlane(start_timestamp_ns_, events_, plane);
   events_.clear();
   return Status::OK();
diff --git a/tensorflow/core/profiler/internal/cpu/metadata_collector.cc b/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
index 58e6385a7ec..c9a593f101b 100644
--- a/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
+++ b/tensorflow/core/profiler/internal/cpu/metadata_collector.cc
@@ -66,7 +66,6 @@ class MetadataCollector : public ProfilerInterface {
   Status CollectData(XSpace* space) override {
     if (!debug_info_.empty()) {
       XPlane* plane = FindOrAddMutablePlaneWithName(space, kMetadataPlaneName);
-      plane->set_id(kMetadataPlaneId);
       XPlaneBuilder xplane(plane);
       const XStatMetadata& hlo_proto_stat =
           *xplane.GetOrCreateStatMetadata(kHloProto);
diff --git a/tensorflow/core/profiler/internal/gpu/BUILD b/tensorflow/core/profiler/internal/gpu/BUILD
index 987bc5ea336..dd8ff1e9a27 100644
--- a/tensorflow/core/profiler/internal/gpu/BUILD
+++ b/tensorflow/core/profiler/internal/gpu/BUILD
@@ -36,17 +36,21 @@ tf_cuda_library(
     deps = [
         ":cupti_utils",
         "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/profiler/internal:annotation_stack",
         "//tensorflow/core/profiler/internal:parse_annotation",
         "//tensorflow/core/profiler/internal:profiler_factory",
         "//tensorflow/core/profiler/internal:profiler_interface",
-        "//tensorflow/core/profiler/lib:traceme",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
         "//tensorflow/core/profiler/utils:xplane_builder",
         "//tensorflow/core/profiler/utils:xplane_schema",
         "//tensorflow/core/profiler/utils:xplane_utils",
+        "@com_google_absl//absl/container:fixed_array",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
     ],
     alwayslink = 1,
@@ -145,4 +149,5 @@ tf_cuda_library(
         ":cupti_wrapper",
     ] + tf_additional_cupti_utils_cuda_deps(),
     visibility = ["//visibility:public"],
+    alwayslink = 1,
 )
diff --git a/tensorflow/core/profiler/internal/gpu/device_tracer.cc b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
index bc9952302e8..73d2a278ea4 100644
--- a/tensorflow/core/profiler/internal/gpu/device_tracer.cc
+++ b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/internal/parse_annotation.h"
 #include "tensorflow/core/profiler/internal/profiler_factory.h"
 #include "tensorflow/core/profiler/internal/profiler_interface.h"
+#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
 #include "tensorflow/core/profiler/utils/xplane_utils.h"
@@ -225,11 +226,10 @@ class CuptiTraceCollectorImpl : public CuptiTraceCollector {
     uint64 end_gpu_ns = CuptiTracer::GetTimestamp();
     XPlaneBuilder host_plane(
         FindOrAddMutablePlaneWithName(space, kCuptiDriverApiPlaneName));
-    host_plane.SetId(kCuptiDriverApiPlaneId);
     for (int device_ordinal = 0; device_ordinal < num_gpus_; ++device_ordinal) {
       std::string name = GpuPlaneName(device_ordinal);
       XPlaneBuilder device_plane(FindOrAddMutablePlaneWithName(space, name));
-      device_plane.SetId(kGpuPlaneBaseId + device_ordinal);
+      device_plane.SetId(device_ordinal);
       per_device_collector_[device_ordinal].Flush(start_gpu_ns_, end_gpu_ns,
                                                   &device_plane, &host_plane);
       per_device_collector_[device_ordinal].GetDeviceCapabilities(
diff --git a/tensorflow/core/profiler/internal/gpu/device_tracer_test.cc b/tensorflow/core/profiler/internal/gpu/device_tracer_test.cc
index 6fc19e776e1..973167ff51b 100644
--- a/tensorflow/core/profiler/internal/gpu/device_tracer_test.cc
+++ b/tensorflow/core/profiler/internal/gpu/device_tracer_test.cc
@@ -263,12 +263,10 @@ TEST_F(DeviceTracerTest, TraceToXSpace) {
   // At least one gpu plane and one host plane for launching events.
   const XPlane* host_plane = FindPlaneWithName(space, kCuptiDriverApiPlaneName);
   ASSERT_NE(host_plane, nullptr);
-  EXPECT_EQ(host_plane->id(), kCuptiDriverApiPlaneId);
 
   const XPlane* device_plane =
       FindPlaneWithName(space, strings::StrCat(kGpuPlanePrefix, 0));
   ASSERT_NE(device_plane, nullptr);  // Check if device plane is serialized.
-  EXPECT_EQ(device_plane->id(), kGpuPlaneBaseId);
   // one for MemcpyH2D, one for MemcpyD2H, two for Matmul (one from Eigen, one
   // from cudnn).
   EXPECT_EQ(device_plane->event_metadata_size(), 4);
diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index 6942f3ea306..ece58802661 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -145,6 +145,9 @@ cc_library(
 cc_library(
     name = "trace_utils",
     hdrs = ["trace_utils.h"],
+    deps = [
+        "//tensorflow/core:lib",
+    ],
 )
 
 cc_library(
diff --git a/tensorflow/core/profiler/utils/trace_utils.h b/tensorflow/core/profiler/utils/trace_utils.h
index 024330faa79..72e6f23feb9 100644
--- a/tensorflow/core/profiler/utils/trace_utils.h
+++ b/tensorflow/core/profiler/utils/trace_utils.h
@@ -16,11 +16,20 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_TRACE_UTILS_H_
 #define TENSORFLOW_CORE_PROFILER_UTILS_TRACE_UTILS_H_
 
+#include "tensorflow/core/platform/types.h"
+
 namespace tensorflow {
 namespace profiler {
 
-// The thread id used for step information in GPU trace viewer.
-// First derived stream/thread id.
+// Constants used as trace_viewer PID (device_id in trace_events.proto).
+// PID 0 is unused.
+// Support up to 500 accelerator devices.
+constexpr uint32 kFirstDeviceId = 1;
+constexpr uint32 kLastDeviceId = 500;
+// Host threads are shown as a single fake device.
+constexpr uint32 kHostThreadsDeviceId = kLastDeviceId + 1;
+
+// Constants used as trace_viewer TID (resource_id in trace_events.proto).
 constexpr int kThreadIdDerivedMin = 0xdeadbeef;
 constexpr int kThreadIdStepInfo = kThreadIdDerivedMin;
 constexpr int kThreadIdKernelLaunch = kThreadIdDerivedMin + 1;
@@ -29,8 +38,6 @@ constexpr int kThreadIdTfOp = kThreadIdDerivedMin + 3;
 constexpr int kThreadIdHloModule = kThreadIdDerivedMin + 4;
 constexpr int kThreadIdHloOp = kThreadIdDerivedMin + 5;
 constexpr int kThreadIdOverhead = kThreadIdDerivedMin + 6;
-
-// Last derived stream/thread id.
 constexpr int kThreadIdDerivedMax = kThreadIdOverhead;
 
 static inline bool IsDerivedThreadId(int thread_id) {
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index 2c79df7980f..197dab75d3b 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -39,15 +39,6 @@ const absl::string_view kXlaModuleLineName = "XLA Modules";
 const absl::string_view kXlaOpLineName = "XLA Ops";
 const absl::string_view kKernelLaunchLineName = "Launch Stats";
 
-const int32 kHostPlaneId = 49;
-const int32 kGpuPlaneBaseId = 0;
-const int32 kCuptiDriverApiPlaneId = 50;
-const int32 kMetadataPlaneId = 99;
-const int32 kTFStreamzPlaneId = 98;
-
-const int32 kThreadGroupMinPlaneId = kCuptiDriverApiPlaneId + 1;
-const int32 kThreadGroupMaxPlaneId = kTFStreamzPlaneId - 1;
-
 namespace {
 
 constexpr int kNumHostEventTypes =
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index a045e20d8de..8b999dc6f9f 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -45,21 +45,6 @@ ABSL_CONST_INIT extern const absl::string_view kXlaModuleLineName;
 ABSL_CONST_INIT extern const absl::string_view kXlaOpLineName;
 ABSL_CONST_INIT extern const absl::string_view kKernelLaunchLineName;
 
-// Id of XPlane that contains TraceMe events.
-ABSL_CONST_INIT extern const int32 kHostPlaneId;
-// Ids prefix of XPlane that contains GPU events.
-ABSL_CONST_INIT extern const int32 kGpuPlaneBaseId;
-// Id of XPlane that contains CUPTI driver API generated events which happens
-// on CPU host threads, e.g. Kernel launch.
-ABSL_CONST_INIT extern const int32 kCuptiDriverApiPlaneId;
-// Id of XPlane that contains profile metadata such as XLA debug info.
-ABSL_CONST_INIT extern const int32 kMetadataPlaneId;
-// Id of XPlane that contains kpi related metrics.
-ABSL_CONST_INIT extern const int32 kTFStreamzPlaneId;
-
-ABSL_CONST_INIT extern const int32 kThreadGroupMinPlaneId;
-ABSL_CONST_INIT extern const int32 kThreadGroupMaxPlaneId;
-
 // Interesting event types (i.e., TraceMe names).
 enum HostEventType {
   kFirstHostEventType = 0,

From d4c2030375551acc9ad36356e090dcc1af8baa51 Mon Sep 17 00:00:00 2001
From: David Majnemer 
Date: Tue, 16 Jun 2020 15:01:21 -0700
Subject: [PATCH 0335/1390] [XLA] Make IdenticalSlowPath more strict

- Map did not check |dimensions|. This has not resulted in miscompliation
  because none of the backends/frontends support/generate map's with dimensions.
- DynamicSlice did not check |dynamic_slice_sizes|. This has not resulted in
  miscompliation because it is not possible for two DynamicSlice operations
  with different |dynamic_slice_sizes| to have the same shape.
- HloCollective did not check |constrain_layout|.

PiperOrigin-RevId: 316765004
Change-Id: Id2629fef71446842eeae18901142a502a634b010
---
 tensorflow/compiler/xla/service/hlo_instructions.cc | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc
index bcc00d806da..2a53841fd34 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.cc
+++ b/tensorflow/compiler/xla/service/hlo_instructions.cc
@@ -550,6 +550,7 @@ bool HloCollectiveInstruction::IdenticalSlowPath(
   const auto& casted_other =
       static_cast(other);
   return HloChannelInstruction::IdenticalSlowPath(other, eq_computations) &&
+         constrain_layout() == casted_other.constrain_layout() &&
          absl::c_equal(replica_groups(), casted_other.replica_groups(),
                        [](const ReplicaGroup& a, const ReplicaGroup& b) {
                          return absl::c_equal(a.replica_ids(), b.replica_ids());
@@ -1101,7 +1102,9 @@ bool HloMapInstruction::IdenticalSlowPath(
     const HloInstruction& other,
     const std::function&
         eq_computations) const {
-  return eq_computations(to_apply(), other.to_apply());
+  const auto& casted_other = static_cast(other);
+  return eq_computations(to_apply(), casted_other.to_apply()) &&
+         dimensions() == casted_other.dimensions();
 }
 
 std::unique_ptr HloMapInstruction::CloneWithNewOperandsImpl(
@@ -2515,7 +2518,8 @@ bool HloDynamicSliceInstruction::IdenticalSlowPath(
     const HloInstruction& other,
     const std::function&
         eq_computations) const {
-  return true;
+  const auto& casted_other = static_cast(other);
+  return dynamic_slice_sizes() == casted_other.dynamic_slice_sizes();
 }
 
 std::unique_ptr

From ef68eff5378aa52ec52064a0b630b79876e90c5a Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 15:06:59 -0700
Subject: [PATCH 0336/1390] Some fixes for private mem broadcast convolution.

PiperOrigin-RevId: 316766133
Change-Id: Icf2273be31c299c49664b9b9fe17df55a0c53693
---
 .../delegates/gpu/cl/kernels/conv_powervr.cc  | 87 ++++++++++---------
 .../delegates/gpu/cl/kernels/conv_powervr.h   | 33 +++++++
 2 files changed, 78 insertions(+), 42 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
index 6ab22bf545d..363a0157420 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
@@ -196,6 +196,9 @@ absl::Status ConvPowerVR::Compile(const CreationContext& creation_context) {
       creation_context.device->IsPowerVR()) {
     options.push_back(CompilerOptions::POWERVR_FP16);
   }
+  if (conv_params_.IsPrivateMemBroadcast()) {
+    options.push_back(CompilerOptions::CL_2_0);
+  }
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", options, *creation_context.context,
       *creation_context.device, &kernel_);
@@ -311,37 +314,10 @@ std::string GenerateConv(
   const int local_mem_size =
       conv_params.block_size.z * 4 * conv_params.src_depth_loop_size;
 
-  const bool use_simd_broadcast =
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST ||
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST ||
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST ||
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST ||
-      conv_params.weights_upload_type ==
-          ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST;
+  const bool use_simd_broadcast = conv_params.IsPrivateMemBroadcast();
+  const int simd_size = conv_params.GetSimdSize();
 
-  int simd_size = 1;
-  if (conv_params.weights_upload_type ==
-      ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST) {
-    simd_size = 8;
-  } else if (conv_params.weights_upload_type ==
-             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST) {
-    simd_size = 16;
-  } else if (conv_params.weights_upload_type ==
-             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST) {
-    simd_size = 32;
-  } else if (conv_params.weights_upload_type ==
-             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST) {
-    simd_size = 64;
-  } else if (conv_params.weights_upload_type ==
-             ConvPowerVR::WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST) {
-    simd_size = 128;
-  }
-
-  bool late_oob_check = need_local_mem || use_simd_broadcast;
+  const bool late_oob_check = need_local_mem || use_simd_broadcast;
 
   const std::string weights_space =
       conv_params.weights_upload_type ==
@@ -355,6 +331,12 @@ std::string GenerateConv(
   const std::string weights_global_ptr =
       weights_space + " " + weights_data_type + "*";
 
+  if (use_simd_broadcast) {
+    if (device.cl_version() == OpenCLVersion::CL_2_0) {
+      c += "#pragma OPENCL EXTENSION cl_khr_subgroups : enable\n";
+    }
+  }
+
   const int3 work_group_size = conv_params.work_group_size;
   const int3 block_size = conv_params.block_size;
   if (conv_params.fixed_work_group_size) {
@@ -364,7 +346,7 @@ std::string GenerateConv(
          std::to_string(work_group_size.z) + ")))\n";
   }
   if (use_simd_broadcast && device.IsIntel()) {
-    c += "__attribute__((intel_reqd_work_group_size(" +
+    c += "__attribute__((intel_reqd_sub_group_size(" +
          std::to_string(simd_size) + ")))\n";
   }
   c += "__kernel void main_function(\n";
@@ -408,6 +390,9 @@ std::string GenerateConv(
            std::to_string(work_group_size.x) + " + get_local_id(0);\n";
     }
   }
+  if (use_simd_broadcast) {
+    c += "  int simd_id = get_sub_group_local_id();\n";
+  }
   for (int z = 0; z < block_size.z; ++z) {
     for (int y = 0; y < block_size.y; ++y) {
       for (int x = 0; x < block_size.x; ++x) {
@@ -555,17 +540,36 @@ std::string GenerateConv(
           for (int y = 0; y < block_size.y; ++y) {
             for (int x = 0; x < block_size.x; ++x) {
               std::string id = std::to_string(y) + std::to_string(x);
-              std::string w_val = "weights_cache[" +
-                                  std::to_string(z * 4 + ch + shared_offset) +
-                                  "]";
               if (use_simd_broadcast) {
                 int simd_id = (z * 4 + ch + shared_offset) / simd_size;
                 int thread_id = (z * 4 + ch + shared_offset) % simd_size;
-                w_val = "sub_group_broadcast(simd_w" + std::to_string(simd_id) +
-                        ", " + std::to_string(thread_id) + "u)";
+                std::string w_val_x = "sub_group_broadcast(simd_w" +
+                                      std::to_string(simd_id) + ".x, " +
+                                      std::to_string(thread_id) + "u)";
+                std::string w_val_y = "sub_group_broadcast(simd_w" +
+                                      std::to_string(simd_id) + ".y, " +
+                                      std::to_string(thread_id) + "u)";
+                std::string w_val_z = "sub_group_broadcast(simd_w" +
+                                      std::to_string(simd_id) + ".z, " +
+                                      std::to_string(thread_id) + "u)";
+                std::string w_val_w = "sub_group_broadcast(simd_w" +
+                                      std::to_string(simd_id) + ".w, " +
+                                      std::to_string(thread_id) + "u)";
+                c += "    r" + std::to_string(z) + id + ".x += " + w_val_x +
+                     " * src" + id + "." + channels[ch] + ";\n";
+                c += "    r" + std::to_string(z) + id + ".y += " + w_val_y +
+                     " * src" + id + "." + channels[ch] + ";\n";
+                c += "    r" + std::to_string(z) + id + ".z += " + w_val_z +
+                     " * src" + id + "." + channels[ch] + ";\n";
+                c += "    r" + std::to_string(z) + id + ".w += " + w_val_w +
+                     " * src" + id + "." + channels[ch] + ";\n";
+              } else {
+                std::string w_val = "weights_cache[" +
+                                    std::to_string(z * 4 + ch + shared_offset) +
+                                    "]";
+                c += "    r" + std::to_string(z) + id + " += " + w_val +
+                     " * src" + id + "." + channels[ch] + ";\n";
               }
-              c += "    r" + std::to_string(z) + id + " += " + w_val +
-                   " * src" + id + "." + channels[ch] + ";\n";
             }
           }
         }
@@ -608,16 +612,15 @@ std::string GenerateConv(
     int parts = local_mem_size / simd_size;
     int reminder = local_mem_size % simd_size;
     for (int i = 0; i < parts; ++i) {
-      c += "    FLT4 simd_w" + std::to_string(i) +
-           " = filters_loc[get_sub_group_local_id() + " +
+      c += "    FLT4 simd_w" + std::to_string(i) + " = filters_loc[simd_id + " +
            std::to_string(i * simd_size) + "];\n";
     }
     if (reminder) {
       c += "    FLT4 simd_w" + std::to_string(parts) + ";\n";
       c += "    if (simd_id < " + std::to_string(reminder) + ") {\n";
       c += "      simd_w" + std::to_string(parts) +
-           " = filters_loc[get_sub_group_local_id() + " +
-           std::to_string(parts * simd_size) + "];\n";
+           " = filters_loc[simd_id + " + std::to_string(parts * simd_size) +
+           "];\n";
       c += "    }\n";
     }
   } else {  // GLOBAL_MEM/CONSTANT_MEM
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
index a729098bded..cf182404e23 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.h
@@ -90,6 +90,39 @@ class ConvPowerVR : public GPUOperation {
     WeightsUploadType weights_upload_type;
     bool x_kernel_is_1;
     bool y_kernel_is_1;
+
+    bool IsPrivateMemBroadcast() const {
+      return weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST ||
+             weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST ||
+             weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST ||
+             weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST ||
+             weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST;
+    }
+
+    int GetSimdSize() const {
+      if (weights_upload_type ==
+          WeightsUploadType::PRIVATE_MEM_SIMD8_BROADCAST) {
+        return 8;
+      } else if (weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST) {
+        return 16;
+      } else if (weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD32_BROADCAST) {
+        return 32;
+      } else if (weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD64_BROADCAST) {
+        return 64;
+      } else if (weights_upload_type ==
+                 WeightsUploadType::PRIVATE_MEM_SIMD128_BROADCAST) {
+        return 128;
+      }
+      return 1;
+    }
   };
 
   ConvPowerVR(const OperationDef& definition,

From 0ea67f633b6272c809244925c9aa05f16adf2401 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Tue, 16 Jun 2020 15:17:40 -0700
Subject: [PATCH 0337/1390] Try to deduce job, replica and task from
 config.list_logical_devices() again

We used to assume that job is localhost, which may not be true in a multi worker environment.

list_logical_devices() isn't always safe since it may return remote devices as well, but we don't have list_local_devices() yet, and we're already using list_logical_devices() in MirorredStrategy (when user doesn't explicit pass devices), so this shouldn't make things worse.

An alternative is to re-consider whether we need device_util.canonicalize at all, but there are multiple usage of it. This requires further effort.

PiperOrigin-RevId: 316768135
Change-Id: Idb2653497e0c4e746a61c44b4d452ab1d052d014
---
 tensorflow/python/distribute/BUILD            |  8 ++++
 tensorflow/python/distribute/device_util.py   | 13 ++++-
 .../python/distribute/device_util_test.py     | 48 +++++++++++++++----
 3 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index a3655699669..e39631d634f 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -121,9 +121,17 @@ cuda_py_test(
     name = "device_util_test",
     srcs = ["device_util_test.py"],
     deps = [
+        ":combinations",
         ":device_util",
+        ":multi_worker_test_base",
+        ":multi_worker_util",
+        "//tensorflow/core:protos_all_py",
         "//tensorflow/python:client_testlib",
+        "//tensorflow/python:extra_py_tests_deps",
         "//tensorflow/python:framework_ops",
+        "//tensorflow/python:training_server_lib",
+        "//tensorflow/python/eager:context",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
diff --git a/tensorflow/python/distribute/device_util.py b/tensorflow/python/distribute/device_util.py
index 7f32ed39aed..0aee34e33ae 100644
--- a/tensorflow/python/distribute/device_util.py
+++ b/tensorflow/python/distribute/device_util.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.eager import context
+from tensorflow.python.framework import config
 from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import ops
 
@@ -55,8 +56,16 @@ def canonicalize(d, default=None):
   result = tf_device.DeviceSpec(
       replica=0, task=0, device_type="CPU", device_index=0)
   if ops.executing_eagerly_outside_functions():
-    # The default job is localhost if eager execution is enabled
-    result = result.replace(job="localhost")
+    # Try to deduce job, replica and task in case it's in a multi worker setup.
+    # TODO(b/151452748): Using list_logical_devices is not always safe since it
+    # may return remote devices as well, but we're already doing this elsewhere.
+    host_cpu = tf_device.DeviceSpec.from_string(
+        config.list_logical_devices("CPU")[0].name)
+    if host_cpu.job:
+      result = result.make_merged_spec(host_cpu)
+    else:
+      # The default job is localhost if eager execution is enabled
+      result = result.replace(job="localhost")
   if default:
     # Overrides any defaults with values from the default device if given.
     result = result.make_merged_spec(
diff --git a/tensorflow/python/distribute/device_util_test.py b/tensorflow/python/distribute/device_util_test.py
index 2f0d7ed3b31..df53fe0288a 100644
--- a/tensorflow/python/distribute/device_util_test.py
+++ b/tensorflow/python/distribute/device_util_test.py
@@ -18,16 +18,27 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import multi_worker_test_base
 from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
 
 
-class DeviceUtilTest(test.TestCase):
+class DeviceUtilTest(test.TestCase, parameterized.TestCase):
 
-  @test_util.run_deprecated_v1
+  def setUp(self):
+    super(DeviceUtilTest, self).setUp()
+    context._reset_context()  # pylint: disable=protected-access
+
+  @combinations.generate(
+      combinations.combine(mode="graph")
+  )
   def testCurrentDeviceWithGlobalGraph(self):
     with ops.device("/cpu:0"):
       self.assertEqual(device_util.current(), "/device:CPU:0")
@@ -51,11 +62,16 @@ class DeviceUtilTest(test.TestCase):
         self.assertEqual(device_util.current(),
                          "/job:localhost/replica:0/task:0/device:CPU:0")
 
-  @test_util.run_deprecated_v1
-  def testCanonicalizeWithoutDefaultDevice(self):
-    self.assertEqual(
-        device_util.canonicalize("/cpu:0"),
-        "/replica:0/task:0/device:CPU:0")
+  @combinations.generate(combinations.combine(mode=["graph", "eager"]))
+  def testCanonicalizeWithoutDefaultDevice(self, mode):
+    if mode == "graph":
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/replica:0/task:0/device:CPU:0")
+    else:
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/job:localhost/replica:0/task:0/device:CPU:0")
     self.assertEqual(
         device_util.canonicalize("/job:worker/cpu:0"),
         "/job:worker/replica:0/task:0/device:CPU:0")
@@ -63,6 +79,22 @@ class DeviceUtilTest(test.TestCase):
         device_util.canonicalize("/job:worker/task:1/cpu:0"),
         "/job:worker/replica:0/task:1/device:CPU:0")
 
+  @combinations.generate(combinations.combine(mode=["eager"]))
+  def testCanonicalizeWithoutDefaultDeviceCollectiveEnabled(self):
+    cluster_spec = server_lib.ClusterSpec(
+        multi_worker_test_base.create_cluster_spec(
+            has_chief=False, num_workers=1, num_ps=0, has_eval=False))
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_spec.as_cluster_def(),
+        job_name="worker",
+        task_index=0,
+        protocol="grpc",
+        port=0)
+    context.context().enable_collective_ops(server_def)
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+
   def testCanonicalizeWithDefaultDevice(self):
     self.assertEqual(
         device_util.canonicalize("/job:worker/task:1/cpu:0", default="/gpu:0"),

From 31e83f2f68b42025939113f5778a528d4ebd0417 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Tue, 16 Jun 2020 15:27:59 -0700
Subject: [PATCH 0338/1390] Fix cross_device_ops_test with multi GPU

The test doesn't limit the number of GPUs when creating
CollectiveAllReduceStrategy, which results in mismatch in the number of devices.

PiperOrigin-RevId: 316769983
Change-Id: I50c11107a99348162b37615f209fd6fa6ee659d2
---
 .../python/distribute/collective_all_reduce_strategy.py    | 7 +++++--
 tensorflow/python/distribute/cross_device_ops_test.py      | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy.py b/tensorflow/python/distribute/collective_all_reduce_strategy.py
index 68cc421c21b..f73474c4af4 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy.py
@@ -112,9 +112,12 @@ class CollectiveAllReduceStrategy(distribute_lib.Strategy):
         "num_replicas_per_worker").set(self.extended._num_gpus_per_worker)
 
   @classmethod
-  def _from_local_devices(cls, devices):
+  def _from_local_devices(
+      cls,
+      devices,
+      communication=cross_device_ops_lib.CollectiveCommunication.AUTO):
     """A convenience method to create an object with a list of devices."""
-    obj = cls()
+    obj = cls(communication)
     obj.extended._initialize_local(TFConfigClusterResolver(), devices=devices)  # pylint: disable=protected-access
     return obj
 
diff --git a/tensorflow/python/distribute/cross_device_ops_test.py b/tensorflow/python/distribute/cross_device_ops_test.py
index 4b6943e8971..d54e16c2748 100644
--- a/tensorflow/python/distribute/cross_device_ops_test.py
+++ b/tensorflow/python/distribute/cross_device_ops_test.py
@@ -521,10 +521,13 @@ class CollectiveAllReduceTest(multi_worker_test_base.MultiWorkerTestBase,
         devices = ["/device:CPU:0"]
 
       if use_strategy_object:
-        strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy(
-            communication=communication)
+        strategy = (
+            collective_all_reduce_strategy.CollectiveAllReduceStrategy
+            ._from_local_devices(devices, communication=communication))  # pylint: disable=protected-access
         strategy.extended._collective_keys = collective_keys
         strategy.extended._cross_device_ops._collective_keys = collective_keys
+        strategy.extended._host_cross_device_ops._collective_keys = (
+            collective_keys)
         return strategy, devices, ""
       else:
         collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce(

From ade14486097e5ab81839f42c537db626d0bf5e88 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher 
Date: Thu, 7 May 2020 18:17:52 +0200
Subject: [PATCH 0339/1390] Test ConvertSquaredDifference in dynamic shape mode

Additionally the runtime error handling is improved:
- Do not check runtime error details in TF-TRT op converter test
- Report an error if invalid dimension is found
---
 .../tf2tensorrt/convert/convert_nodes_test.cc | 131 ++++++++----------
 .../tf2tensorrt/utils/trt_engine_utils.cc     |   8 ++
 2 files changed, 66 insertions(+), 73 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
index aeadbdbf012..f5e61dc0061 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
@@ -1803,7 +1803,9 @@ class ParameterizedOpConverterTestBase
     const int batch_size = input_data_[0].tensor.shape().dim_size(0);
     Status stat =
         OpConverterTest::BuildAndRun(input_data_, &output_data, batch_size);
-    ASSERT_EQ(expected_runtime_status, stat);
+    ASSERT_EQ(expected_runtime_status.ok(), stat.ok())
+        << "expected status: " << expected_runtime_status
+        << ", actual status: " << stat;
     if (expected_runtime_status.ok() && stat.ok()) {
       for (int i = 0; i < n_output; i++) {
         // Check the shape of the actual output tensors
@@ -6440,87 +6442,70 @@ NodeDef GetSquaredDifferenceNodeDef(DataType dtype) {
   return squared_diff.operation.node()->def();
 }
 
-template 
-void TestConvertSquaredDifference(OpConverterTest* test) {
-  typedef typename EnumToDataType::Type CType;
-
-  struct TestParams {
-    std::vector dims_x;
-    std::vector dims_y;
-    std::vector value_x;
-    std::vector value_y;
-    std::vector expected_output_dims;
-    std::vector expected_output;
-  };
-
-  const std::vector common_input = InitTestVector(6);
-  std::vector params = {
-      {
-          /*dims_x=*/{1, 2, 3},
-          /*dims_y=*/{1, 2, 3},
-          /*value_x=*/common_input,
-          /*value_y=*/CastTestVector({0, -1, 3, 0, 10, -7}),
-          /*expected_output_dims=*/{1, 2, 3},
-          /*expected_output=*/CastTestVector({0, 4, 1, 9, 36, 144}),
-      },
-      {
-          /*dims_x=*/{1, 2, 3},
-          /*dims_y=*/{1, 1, 3},
-          /*value_x=*/common_input,
-          /*value_y=*/CastTestVector({0, 1, 2}),
-          /*expected_output_dims=*/{1, 2, 3},
-          /*expected_output=*/CastTestVector({0, 0, 0, 9, 9, 9}),
-      },
-  };
-
-  for (int i = 0; i < params.size(); ++i) {
-    test->Reset();
-
-    NodeDef node_def = GetSquaredDifferenceNodeDef(dtype);
-    test->AddTestTensor("x", params[i].dims_x, 1, TfDataTypeToTrt(dtype));
-    test->AddTestTensor("y", params[i].dims_y, 1, TfDataTypeToTrt(dtype));
-    test->RunValidationAndConversion(node_def);
-
-    TRT_TensorOrWeights output;
-    TF_EXPECT_OK(test->GetTensorOrWeights("my_squared_diff", &output));
-    EXPECT_TRUE(output.is_tensor());
-    ExpectTrtDimsEqualsArray(params[i].expected_output_dims,
-                             output.tensor()->getDimensions());
-
-    DataVec input_data{{"x", test->AsTensor(params[i].value_x)},
-                       {"y", test->AsTensor(params[i].value_y)}};
-    DataVec output_data{
-        {"my_squared_diff",
-         test->ConstructTensor(params[i].expected_output.size())}};
-    TF_EXPECT_OK(test->BuildAndRun(input_data, &output_data));
-    EXPECT_THAT(GetSpanForData(output_data[0]),
-                ElementsAreArray(params[i].expected_output));
-  }
-}
-
-TEST_F(OpConverterTest, ConvertSquaredDifference) {
+TEST_P(OpConverterTest2, ConvertSquaredDifference) {
   {
     // Input is a weight, should fail.
     Reset();
-    NodeDef node_def = GetSquaredDifferenceNodeDef(DT_FLOAT);
+    NodeDef node_def = GetSquaredDifferenceNodeDef(tf_type);
     AddTestWeights("x", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
-    AddTestTensor("y", {1, 2, 3});
+    AddTestTensor("y", {1, 1, 2, 3});
     RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
                                "The input \"x\" for SquaredDifference must be "
                                "a tensor, at my_squared_diff");
   }
-  {
-    // Shapes are not broadcastable, should fail.
-    Reset();
-    NodeDef node_def = GetSquaredDifferenceNodeDef(DT_FLOAT);
-    AddTestTensor("x", {2, 3});
-    AddTestTensor("y", {7, 5});
-    RunValidationAndConversion(node_def, error::INVALID_ARGUMENT,
-                               "Infeasible broadcast scheme");
-  }
 
-  TestConvertSquaredDifference(this);
-  TestConvertSquaredDifference(this);
+  struct TestParams {
+    std::vector dims_x;
+    std::vector dims_y;
+    std::vector value_x;
+    std::vector value_y;
+    std::vector expected_output_dims;
+    std::vector expected_output;
+    Status status;
+    Status runtime_status;
+  };
+
+  const std::vector common_input = InitTestVector(6);
+  std::vector params = {
+      {/*dims_x=*/{1, 2, 3},
+       /*dims_y=*/{1, 7, 5},
+       /*value_x=*/common_input,
+       /*value_y=*/std::vector(7 * 5, 0),
+       /*expected_output_dims=*/{1, 1, 2, 3},
+       /*expected_output=*/common_input,
+       trt_mode == TrtTestMode::kDynamicShape
+           ? Status::OK()
+           : errors::InvalidArgument("Infeasible broadcast scheme"),
+       errors::Internal(
+           "Binding index out of range. This can happen if profile is not set, "
+           "or the network is invalid for the current profile.")},
+      {
+          /*dims_x=*/{1, 1, 2, 3},
+          /*dims_y=*/{1, 1, 2, 3},
+          /*value_x=*/common_input,
+          /*value_y=*/{0, -1, 3, 0, 10, -7},
+          /*expected_output_dims=*/{1, 1, 2, 3},
+          /*expected_output=*/{0, 4, 1, 9, 36, 144},
+      },
+      {
+          /*dims_x=*/{1, 1, 2, 3},
+          /*dims_y=*/{1, 1, 1, 3},
+          /*value_x=*/common_input,
+          /*value_y=*/{0, 1, 2},
+          /*expected_output_dims=*/{1, 1, 2, 3},
+          /*expected_output=*/{0, 0, 0, 9, 9, 9},
+      },
+  };
+
+  for (auto p : params) {
+    Reset();
+    NodeDef node_def = GetSquaredDifferenceNodeDef(tf_type);
+    AddTestTensor("x", p.dims_x, p.value_x);
+    AddTestTensor("y", p.dims_y, p.value_y);
+    TestOpConverter("my_squared_diff", node_def, p.expected_output_dims,
+                    p.status, p.runtime_status,
+                    ElementsAreArray(p.expected_output));
+  }
 }
 
 #if IS_TRT_VERSION_GE(6, 0, 0, 0)
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
index 213c1732e59..ed997b267b1 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
@@ -46,6 +46,14 @@ Status GetTrtBindingShape(const nvinfer1::ICudaEngine* cuda_engine,
     // Get dims from context instead of engine in explicit batch mode because
     // the engine might have dynamic shapes.
     dims = execution_context->getBindingDimensions(binding_index);
+    if (dims.nbDims == -1) {
+      // Invalid dimensions. There can be multiple reasons for this. If we have
+      // incompatible input shapes (network invalid for the current profile)
+      // that can trigger this error.
+      return errors::Internal(
+          "Binding index out of range. This can happen if profile is not set, "
+          "or the network is invalid for the current profile.");
+    }
 #else
     return errors::Internal(
         "Explicit batch mode is only supported with TensorRT 6 and above.");

From 6b126156d962547f4b4f66cf19aabfe14d75dd88 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Tue, 16 Jun 2020 15:51:18 -0700
Subject: [PATCH 0340/1390] Fix clang-tidy warnings in `logging.{cc,h}`.

PiperOrigin-RevId: 316773944
Change-Id: Ie35223f385d0e9038e0a0ea223825b5ea78f4769
---
 tensorflow/core/platform/default/logging.cc |  8 +++---
 tensorflow/core/platform/default/logging.h  | 30 +++++++++------------
 2 files changed, 17 insertions(+), 21 deletions(-)

diff --git a/tensorflow/core/platform/default/logging.cc b/tensorflow/core/platform/default/logging.cc
index 9bf809131a9..6d2af607748 100644
--- a/tensorflow/core/platform/default/logging.cc
+++ b/tensorflow/core/platform/default/logging.cc
@@ -303,7 +303,7 @@ void MakeCheckOpValueString(std::ostream* os, const char& v) {
   if (v >= 32 && v <= 126) {
     (*os) << "'" << v << "'";
   } else {
-    (*os) << "char value " << static_cast(v);
+    (*os) << "char value " << static_cast(v);
   }
 }
 
@@ -312,7 +312,7 @@ void MakeCheckOpValueString(std::ostream* os, const signed char& v) {
   if (v >= 32 && v <= 126) {
     (*os) << "'" << v << "'";
   } else {
-    (*os) << "signed char value " << static_cast(v);
+    (*os) << "signed char value " << static_cast(v);
   }
 }
 
@@ -321,13 +321,13 @@ void MakeCheckOpValueString(std::ostream* os, const unsigned char& v) {
   if (v >= 32 && v <= 126) {
     (*os) << "'" << v << "'";
   } else {
-    (*os) << "unsigned char value " << static_cast(v);
+    (*os) << "unsigned char value " << static_cast(v);
   }
 }
 
 #if LANG_CXX11
 template <>
-void MakeCheckOpValueString(std::ostream* os, const std::nullptr_t& p) {
+void MakeCheckOpValueString(std::ostream* os, const std::nullptr_t& v) {
   (*os) << "nullptr";
 }
 #endif
diff --git a/tensorflow/core/platform/default/logging.h b/tensorflow/core/platform/default/logging.h
index 8b171a15f7c..f60deb43683 100644
--- a/tensorflow/core/platform/default/logging.h
+++ b/tensorflow/core/platform/default/logging.h
@@ -265,16 +265,12 @@ inline const T& GetReferenceableValue(const T& t) {
 inline char GetReferenceableValue(char t) { return t; }
 inline unsigned char GetReferenceableValue(unsigned char t) { return t; }
 inline signed char GetReferenceableValue(signed char t) { return t; }
-inline short GetReferenceableValue(short t) { return t; }
-inline unsigned short GetReferenceableValue(unsigned short t) { return t; }
+inline int16 GetReferenceableValue(int16 t) { return t; }
+inline uint16 GetReferenceableValue(uint16 t) { return t; }
 inline int GetReferenceableValue(int t) { return t; }
 inline unsigned int GetReferenceableValue(unsigned int t) { return t; }
-inline long GetReferenceableValue(long t) { return t; }
-inline unsigned long GetReferenceableValue(unsigned long t) { return t; }
-inline long long GetReferenceableValue(long long t) { return t; }
-inline unsigned long long GetReferenceableValue(unsigned long long t) {
-  return t;
-}
+inline int64 GetReferenceableValue(int64 t) { return t; }
+inline uint64 GetReferenceableValue(uint64 t) { return t; }
 
 // This formats a value for a failing CHECK_XX statement.  Ordinarily,
 // it uses the definition for operator<<, with a few special cases below.
@@ -295,16 +291,16 @@ void MakeCheckOpValueString(std::ostream* os, const unsigned char& v);
 #if LANG_CXX11
 // We need an explicit specialization for std::nullptr_t.
 template <>
-void MakeCheckOpValueString(std::ostream* os, const std::nullptr_t& p);
+void MakeCheckOpValueString(std::ostream* os, const std::nullptr_t& v);
 #endif
 
 // A container for a string pointer which can be evaluated to a bool -
 // true iff the pointer is non-NULL.
 struct CheckOpString {
-  CheckOpString(string* str) : str_(str) {}
+  explicit CheckOpString(string* str) : str_(str) {}
   // No destructor: if str_ is non-NULL, we're about to LOG(FATAL),
   // so there's no point in cleaning up str_.
-  operator bool() const { return TF_PREDICT_FALSE(str_ != NULL); }
+  explicit operator bool() const { return TF_PREDICT_FALSE(str_ != nullptr); }
   string* str_;
 };
 
@@ -393,12 +389,12 @@ TF_DEFINE_CHECK_OP_IMPL(Check_GT, >)
 
 // In optimized mode, use CheckOpString to hint to compiler that
 // the while condition is unlikely.
-#define CHECK_OP_LOG(name, op, val1, val2)                            \
-  while (::tensorflow::internal::CheckOpString _result =              \
-             ::tensorflow::internal::name##Impl(                      \
-                 ::tensorflow::internal::GetReferenceableValue(val1), \
-                 ::tensorflow::internal::GetReferenceableValue(val2), \
-                 #val1 " " #op " " #val2))                            \
+#define CHECK_OP_LOG(name, op, val1, val2)                     \
+  while (::tensorflow::internal::CheckOpString _result{        \
+      ::tensorflow::internal::name##Impl(                      \
+          ::tensorflow::internal::GetReferenceableValue(val1), \
+          ::tensorflow::internal::GetReferenceableValue(val2), \
+          #val1 " " #op " " #val2)})                           \
   ::tensorflow::internal::LogMessageFatal(__FILE__, __LINE__) << *(_result.str_)
 
 #define CHECK_OP(name, op, val1, val2) CHECK_OP_LOG(name, op, val1, val2)

From a6945b9b0f0c43b6e6557c3c6a441dad5f98a866 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 16:02:34 -0700
Subject: [PATCH 0341/1390] SpaceToDepth converted to new style.

PiperOrigin-RevId: 316775897
Change-Id: I2715cadbf112dffc93b2b45570d2220444af31a4
---
 .../gpu/cl/kernels/space_to_depth.cc          | 76 ++++++++++---------
 1 file changed, 41 insertions(+), 35 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/space_to_depth.cc b/tensorflow/lite/delegates/gpu/cl/kernels/space_to_depth.cc
index 34227f6b887..439b7d0fc15 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/space_to_depth.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/space_to_depth.cc
@@ -27,28 +27,33 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetSpaceToDepthCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+std::string GetSpaceToDepthCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+  args->AddInt("block_size");
+
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,\n";
-  c += "    int4 dst_size,\n";
-  c += "    int src_channels,\n";
-  c += "    int block_size) {\n";
-  c += "  int X = get_global_id(0);\n";
+  c += "$0) {\n";
+  if (op_def.IsBatchSupported()) {
+    c += "  int linear_id = get_global_id(0);\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = get_global_id(0);\n";
+  }
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
   c += "  FLT tmp[4];\n";
   c += "  tmp[0] = (FLT)(0.0f);\n";
   c += "  tmp[1] = (FLT)(0.0f);\n";
@@ -56,19 +61,17 @@ std::string GetSpaceToDepthCode(
   c += "  tmp[3] = (FLT)(0.0f);\n";
   c += "  for (int i = 0; i < 4; ++i) {\n";
   c += "    int dst_c = 4 * Z + i;\n";
-  c += "    int block_id = dst_c / src_channels;\n";
-  c += "    int src_x = X * block_size + block_id % block_size;\n";
-  c += "    int src_y = Y * block_size + block_id / block_size;\n";
-  c += "    int src_c = dst_c % src_channels;\n";
+  c += "    int block_id = dst_c / args.src_tensor.Channels();\n";
+  c += "    int src_x = X * args.block_size + block_id % args.block_size;\n";
+  c += "    int src_y = Y * args.block_size + block_id / args.block_size;\n";
+  c += "    int src_c = dst_c % args.src_tensor.Channels();\n";
   c += "    int src_z = src_c / 4;\n";
-  c += "    FLT4 t = " + src_tensor.ReadWHS("src_x", "src_y", "src_z") + ";\n";
+  c += "    FLT4 t =  args.src_tensor.Read(src_x, src_y, src_z);\n";
   c += "    FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
   c += "    tmp[i] = t_ar[src_c % 4];\n";
   c += "  }\n";
   c += "  FLT4 result = (FLT4)(tmp[0], tmp[1], tmp[2], tmp[3]);\n";
-  const LinkingContext context{"result", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("result", "X", "Y", "Z");
+  c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
   c += "}\n";
   return c;
 }
@@ -92,21 +95,24 @@ SpaceToDepth& SpaceToDepth::operator=(SpaceToDepth&& operation) {
 }
 
 absl::Status SpaceToDepth::Compile(const CreationContext& creation_context) {
-  const auto code = GetSpaceToDepthCode(definition_, linked_operations_);
+  std::string code = GetSpaceToDepthCode(definition_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status SpaceToDepth::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->Channels()));
-  return kernel_.SetBytesAuto(attr_.block_size);
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(args_.SetInt("block_size", attr_.block_size));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 SpaceToDepth::GetGridSize() const {

From fad7b3a33b476ed2b4d7b6a901b8dc5ab02f38b0 Mon Sep 17 00:00:00 2001
From: Berkin Ilbeyi 
Date: Tue, 16 Jun 2020 16:12:25 -0700
Subject: [PATCH 0342/1390] [XLA] More tweaks and fixes to memory space
 assignment:

- When sorting the BufferIntervals, now we also correctly include aliased
  HloValues into the benefit computation.
- To avoid double counting, we now skip over while and conditional HLOs when
  using cost analysis since we already look inside the called computations.
- When trying to find a preferred alternate memory offset, we now find the
  latest use that can be contiguously allocated (determined by the max overlap
  ratio heuristic) instead of trying to find an allocation that is as
  long-living as possible. This should improve fragmentation slightly.

PiperOrigin-RevId: 316777648
Change-Id: If9d40a79283b644db975ebe62b1bb6c545fea89d
---
 .../xla/service/memory_space_assignment.cc    | 113 +++++++++++++-----
 .../xla/service/memory_space_assignment.h     |  35 ++++--
 .../service/memory_space_assignment_test.cc   |  16 +--
 3 files changed, 109 insertions(+), 55 deletions(-)

diff --git a/tensorflow/compiler/xla/service/memory_space_assignment.cc b/tensorflow/compiler/xla/service/memory_space_assignment.cc
index 388a2e18f38..ea1438380a6 100644
--- a/tensorflow/compiler/xla/service/memory_space_assignment.cc
+++ b/tensorflow/compiler/xla/service/memory_space_assignment.cc
@@ -31,6 +31,22 @@ const int kWhileExecutionCount = 5;
 
 }  // namespace
 
+/*static*/ StatusOr>
+MemorySpaceAssignmentCostAnalysis::Create(
+    const HloCostAnalysis& cost_analysis,
+    float async_copy_bandwidth_bytes_per_second,
+    float alternate_mem_bandwidth_bytes_per_second, const HloModule& module) {
+  TF_ASSIGN_OR_RETURN(auto alias_analysis, HloAliasAnalysis::Run(&module));
+  TF_ASSIGN_OR_RETURN(auto hlo_live_range,
+                      HloLiveRange::Run(module.schedule(), *alias_analysis,
+                                        module.entry_computation()));
+  auto call_graph = CallGraph::Build(&module);
+  return absl::WrapUnique(new MemorySpaceAssignmentCostAnalysis(
+      cost_analysis, async_copy_bandwidth_bytes_per_second,
+      alternate_mem_bandwidth_bytes_per_second, std::move(alias_analysis),
+      std::move(hlo_live_range), std::move(call_graph)));
+}
+
 float MemorySpaceAssignmentCostAnalysis::GetAlternateMemoryBenefit(
     const HloInstruction& instruction, float elapsed_time_due_to_alternate_mem,
     MemorySpaceAssignmentCostAnalysis::Cache* cache) const {
@@ -74,19 +90,32 @@ float MemorySpaceAssignmentCostAnalysis::GetMemoryBoundedness(
                                        /*operand_in_alternate_mem=*/{},
                                        /*output_in_alternate_mem=*/true),
       cache);
-  for (const HloUse& use : interval.buffer->uses()) {
-    float use_alternate_mem_benefit = GetAlternateMemoryBenefit(
-        *use.instruction,
-        GetInstructionElapsedDueToMemory(*use.instruction, use.operand_number),
-        cache);
-    // If the benefit is positive (memory bound), add it to this buffer's
-    // benefit. If the benefit is negative (compute bound), calculate the
-    // maximum.
-    if (alternate_mem_benefit > 0 && use_alternate_mem_benefit > 0) {
-      alternate_mem_benefit += use_alternate_mem_benefit;
-    } else {
-      alternate_mem_benefit =
-          std::max(alternate_mem_benefit, use_alternate_mem_benefit);
+  for (const HloBuffer* buffer : alias_analysis_->ComputeBuffersAt(
+           interval.buffer->defining_position().instruction,
+           interval.buffer->defining_position().index)) {
+    for (const HloValue* value : buffer->values()) {
+      for (const HloUse& use : value->uses()) {
+        // We look inside the called computations of while and conditional, so
+        // don't use the benefit of while and conditional directly.
+        if (use.instruction->opcode() == HloOpcode::kWhile ||
+            use.instruction->opcode() == HloOpcode::kConditional) {
+          continue;
+        }
+        float use_alternate_mem_benefit =
+            GetAlternateMemoryBenefit(*use.instruction,
+                                      GetInstructionElapsedDueToMemory(
+                                          *use.instruction, use.operand_number),
+                                      cache);
+        // If the benefit is positive (memory bound), add it to this buffer's
+        // benefit. If the benefit is negative (compute bound), calculate the
+        // maximum.
+        if (alternate_mem_benefit > 0 && use_alternate_mem_benefit > 0) {
+          alternate_mem_benefit += use_alternate_mem_benefit;
+        } else {
+          alternate_mem_benefit =
+              std::max(alternate_mem_benefit, use_alternate_mem_benefit);
+        }
+      }
     }
   }
 
@@ -95,17 +124,9 @@ float MemorySpaceAssignmentCostAnalysis::GetMemoryBoundedness(
   float alternate_mem_slowdown =
       GetInstructionElapsedDueToMemorySlowdown(interval.size);
 
-  // Scale the slowdown based on the time of this buffer. We would want earlier
-  // buffers have lower slowdown values, because they are less likely to overlap
-  // with other HLOs.
-  // TODO(yuemmawang): We may want a piecewise function, where a lower slowdown
-  // for early HLOs, and full slowdown for mid-to-late HLOs.
-  // TODO(yuemmawang): Further in a smarter way, we want buffers overlapped with
-  // more HLOs have higher slowdown, and vice versa.
-  float scale = interval.start * 1.0 / GetScheduleEndTime();
-  alternate_mem_slowdown *= scale;
-
-  return alternate_mem_benefit - alternate_mem_slowdown;
+  // Divide by the size of the buffer to prioritize smaller buffers that will
+  // give the largest alternate memory benefit.
+  return (alternate_mem_benefit - alternate_mem_slowdown) / interval.size;
 }
 
 int MemorySpaceAssignmentCostAnalysis::CalculateWhileLoopNestLevel(
@@ -113,7 +134,7 @@ int MemorySpaceAssignmentCostAnalysis::CalculateWhileLoopNestLevel(
   int nest_level = 0;
   const HloComputation* computation = instruction->parent();
   while (!computation->IsEntryComputation()) {
-    auto node = call_graph_.GetNode(computation);
+    auto node = call_graph_->GetNode(computation);
     auto callsites = node.caller_callsites();
     CHECK_EQ(callsites.size(), 1) << "The module is not flattened!";
     auto callsite = callsites[0];
@@ -195,7 +216,7 @@ float MemorySpaceAssignmentCostAnalysis::GetAsyncCopyElapsed(
 }
 
 int64 MemorySpaceAssignmentCostAnalysis::GetScheduleEndTime() const {
-  return hlo_live_range_.schedule_end_time();
+  return hlo_live_range_->schedule_end_time();
 }
 
 bool InstructionCountPrefetchIntervalPicker::CanAllocateInAlternateMemoryNoCopy(
@@ -253,6 +274,13 @@ CostAnalysisPrefetchIntervalPicker::CostAnalysisPrefetchIntervalPicker(
   std::vector instructions_elapsed_time(instruction_schedule_->size(),
                                                0.0);
   for (const auto& instruction_and_logical_time : *instruction_schedule_) {
+    // To avoid double counting, don't include the elapsed time of while and
+    // conditional HLOs.
+    const HloInstruction* instruction = instruction_and_logical_time.first;
+    if (instruction->opcode() == HloOpcode::kWhile ||
+        instruction->opcode() == HloOpcode::kConditional) {
+      continue;
+    }
     float elapsed_time = cost_analysis_.cost_analysis().optimal_seconds(
         *instruction_and_logical_time.first);
     int64 logical_time = instruction_and_logical_time.second;
@@ -1937,17 +1965,38 @@ AlternateMemoryBestFitHeap::FindBestChunkCandidate(
     BufferInterval* alternate_mem_interval) const {
   int64 end_time = request.end_time;
   if (!preferred_offset) {
+    // First find the earliest use that is the same or later than the end time.
+    const auto& uses = request.allocation_value->uses();
+    auto use_it = uses.begin();
+    for (; use_it->time < end_time; ++use_it) {
+    }
+    CHECK(use_it != uses.end());
+    int64 earliest_use = use_it->time;
+
+    // Then find the latest use that can be allocated contiguously without
+    // copies.
+    const Shape& shape = request.allocation_value->defining_position().shape();
+    for (;
+         (use_it + 1) != uses.end() &&
+         options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy(
+             shape, use_it->time, (use_it + 1)->time);
+         ++use_it) {
+    }
+    CHECK(use_it != uses.end());
+    int64 latest_contiguous_use = use_it->time;
+
     // Find a chunk that's as long living as possible iterating in reverse over
     // the use times.
-    for (auto use_it = request.allocation_value->uses().rbegin();
-         use_it != request.allocation_value->uses().rend() &&
-         use_it->time >= end_time;
-         ++use_it) {
+    for (; use_it >= uses.begin() && use_it->time >= end_time; --use_it) {
       alternate_mem_interval->end = use_it->time;
       ChunkCandidate chunk_candidate =
           FindChunkCandidate(*alternate_mem_interval);
       if (chunk_candidate.heap_size <= available_heap_size()) {
         alternate_mem_interval->end = end_time;
+        VLOG(3) << "FindBestChunkCandidate earliest use = " << earliest_use
+                << ", latest contiguous use = " << latest_contiguous_use
+                << ", use with available mem = " << use_it->time
+                << ", offset = " << chunk_candidate.chunk.offset;
         return chunk_candidate;
       }
     }
@@ -2005,8 +2054,8 @@ MemorySpaceAssignment::CalculateAsyncCopyStats() const {
 MemorySpaceAssignment::GetMemoryBoundednessBufferIntervalCompare(
     const MemorySpaceAssignmentCostAnalysis& cost_analysis,
     MemorySpaceAssignmentCostAnalysis::Cache* cache) {
-  return [cost_analysis, cache](const BufferInterval& x,
-                                const BufferInterval& y) {
+  return [&cost_analysis, cache](const BufferInterval& x,
+                                 const BufferInterval& y) {
     float x_memory_boundedness = cost_analysis.GetMemoryBoundedness(x, cache);
     float y_memory_boundedness = cost_analysis.GetMemoryBoundedness(y, cache);
     if (x_memory_boundedness != y_memory_boundedness) {
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment.h b/tensorflow/compiler/xla/service/memory_space_assignment.h
index f9e5738d17e..5e34f755fe9 100644
--- a/tensorflow/compiler/xla/service/memory_space_assignment.h
+++ b/tensorflow/compiler/xla/service/memory_space_assignment.h
@@ -84,18 +84,10 @@ class MemorySpaceAssignmentCostAnalysis {
     absl::flat_hash_map while_nest_multiplier;
   };
 
-  MemorySpaceAssignmentCostAnalysis(
+  static StatusOr> Create(
       const HloCostAnalysis& cost_analysis,
       float async_copy_bandwidth_bytes_per_second,
-      float alternate_mem_bandwidth_bytes_per_second,
-      const HloLiveRange& hlo_live_range, const CallGraph& call_graph)
-      : cost_analysis_(cost_analysis),
-        async_copy_bandwidth_bytes_per_second_(
-            async_copy_bandwidth_bytes_per_second),
-        alternate_mem_bandwidth_bytes_per_second_(
-            alternate_mem_bandwidth_bytes_per_second),
-        hlo_live_range_(hlo_live_range),
-        call_graph_(call_graph) {}
+      float alternate_mem_bandwidth_bytes_per_second, const HloModule& module);
 
   const HloCostAnalysis& cost_analysis() const { return cost_analysis_; }
 
@@ -153,14 +145,31 @@ class MemorySpaceAssignmentCostAnalysis {
   // 0 means it is not in a while loop.
   int CalculateWhileLoopNestLevel(const HloInstruction* instruction) const;
 
-  const HloLiveRange& hlo_live_range() const { return hlo_live_range_; }
+  const HloLiveRange& hlo_live_range() const { return *hlo_live_range_; }
 
  private:
+  MemorySpaceAssignmentCostAnalysis(
+      const HloCostAnalysis& cost_analysis,
+      float async_copy_bandwidth_bytes_per_second,
+      float alternate_mem_bandwidth_bytes_per_second,
+      std::unique_ptr alias_analysis,
+      std::unique_ptr hlo_live_range,
+      std::unique_ptr call_graph)
+      : cost_analysis_(cost_analysis),
+        async_copy_bandwidth_bytes_per_second_(
+            async_copy_bandwidth_bytes_per_second),
+        alternate_mem_bandwidth_bytes_per_second_(
+            alternate_mem_bandwidth_bytes_per_second),
+        alias_analysis_(std::move(alias_analysis)),
+        hlo_live_range_(std::move(hlo_live_range)),
+        call_graph_(std::move(call_graph)) {}
+
   const HloCostAnalysis& cost_analysis_;
   float async_copy_bandwidth_bytes_per_second_;
   float alternate_mem_bandwidth_bytes_per_second_;
-  const HloLiveRange& hlo_live_range_;
-  const CallGraph& call_graph_;
+  std::unique_ptr alias_analysis_;
+  std::unique_ptr hlo_live_range_;
+  std::unique_ptr call_graph_;
 };
 
 // Abstract base class that memory space assignment uses to pick prefetch
diff --git a/tensorflow/compiler/xla/service/memory_space_assignment_test.cc b/tensorflow/compiler/xla/service/memory_space_assignment_test.cc
index 032a3f53479..398f07d4a40 100644
--- a/tensorflow/compiler/xla/service/memory_space_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/memory_space_assignment_test.cc
@@ -53,22 +53,18 @@ class MemorySpaceAssignmentTest : public HloTestBase,
       TF_CHECK_OK(computation->Accept(&hlo_cost_analysis));
     }
     auto alias_analysis = HloAliasAnalysis::Run(module).ValueOrDie();
-    std::unique_ptr hlo_live_range =
-        HloLiveRange::Run(module->schedule(), *alias_analysis,
-                          module->entry_computation())
-            .ValueOrDie();
-    std::unique_ptr call_graph = CallGraph::Build(module);
-    MemorySpaceAssignmentCostAnalysis cost_analysis(
-        hlo_cost_analysis, kAsyncCopyBandwidth, kAlternateMemBandwidth,
-        *hlo_live_range, *call_graph);
+    auto cost_analysis = MemorySpaceAssignmentCostAnalysis::Create(
+                             hlo_cost_analysis, kAsyncCopyBandwidth,
+                             kAlternateMemBandwidth, *module)
+                             .ValueOrDie();
     CostAnalysisPrefetchIntervalPicker prefetch_interval_picker(
         CostAnalysisPrefetchIntervalPicker(
-            cost_analysis, /*min_async_copy_to_overlap_ratio=*/0.8,
+            *cost_analysis, /*min_async_copy_to_overlap_ratio=*/0.8,
             /*max_async_copy_to_overlap_ratio=*/10.0));
     return AssignMemorySpace(
         module, /*max_outstanding_async_copies=*/-1,
         MemorySpaceAssignment::GetMemoryBoundednessBufferIntervalCompare(
-            cost_analysis, &cache_),
+            *cost_analysis, &cache_),
         &prefetch_interval_picker);
   }
 

From 89a1d3f4e9dcea93bff6cdcd469e5e10c26376bc Mon Sep 17 00:00:00 2001
From: Jaesung Chung 
Date: Tue, 16 Jun 2020 16:24:49 -0700
Subject: [PATCH 0343/1390] Implement a pass for lifting variables.

This pass creates GlobalTensorOp for each variable from function
arguments and converts the function arguments to the corresponding saved
model arguments.

This change fixes the Kokoro builds by avoiding the common_runtime
dependency.

PiperOrigin-RevId: 316779884
Change-Id: I07c83bf12486748e4350717d94928a75bad92342
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |  69 +++++++
 .../tests/tf_saved_model_lift_variables.mlir  |  61 ++++++
 ..._model_lift_variables_invalid_session.mlir |  33 ++++
 .../tensorflow/transforms/lift_variables.cc   | 183 ++++++++++++++++++
 .../tensorflow/transforms/lift_variables.h    |  33 ++++
 .../transforms/lift_variables_pass.h          |  57 ++++++
 .../transforms/lift_variables_test_pass.h     | 146 ++++++++++++++
 .../lift_variables_test_pass_registration.cc  |  32 +++
 8 files changed, 614 insertions(+)
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables.mlir
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables_invalid_session.mlir
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.cc
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass_registration.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 9e5688cd230..904ccb7e820 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -397,6 +397,73 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "lift_variables_lib",
+    srcs = [
+        "transforms/lift_variables.cc",
+    ],
+    hdrs = [
+        "transforms/lift_variables.h",
+    ],
+    deps = [
+        ":convert_tensor",
+        ":tensorflow",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:threadpool_options",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "lift_variables_pass",
+    hdrs = [
+        "transforms/lift_variables_pass.h",
+    ],
+    deps = [
+        ":lift_variables_lib",
+        ":tensorflow",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "lift_variables_test_pass",
+    hdrs = [
+        "transforms/lift_variables_test_pass.h",
+    ],
+    deps = [
+        ":lift_variables_lib",
+        ":tensorflow",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:threadpool_options",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "tensorflow_passes",
     srcs = [
@@ -520,9 +587,11 @@ cc_library(
 cc_library(
     name = "tensorflow_test_passes",
     srcs = [
+        "transforms/lift_variables_test_pass_registration.cc",
         "transforms/lower_tf_pass.cc",
     ],
     deps = [
+        ":lift_variables_test_pass",
         ":lower_tf_lib",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables.mlir
new file mode 100644
index 00000000000..0c04a0d738c
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables.mlir
@@ -0,0 +1,61 @@
+// RUN: tf-opt -verify-diagnostics -tf-saved-model-lift-variables-test -split-input-file %s | FileCheck %s --dump-input=fail
+
+module attributes {tf_saved_model.semantics} {
+
+  // Test case: Freezing VarHandleOp ops.
+
+  func @serving_default(%arg0: tensor>> {tf.resource_name = "dense/kernel"}, %arg1: tensor>> {tf.resource_name = "dense/bias"}) -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})
+  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor>>
+    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor>>) -> tensor<100x50xf32>
+    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor>>
+    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor>>) -> tensor<50xf32>
+    %4 = "tf.Add"(%1, %3) {device = ""} : (tensor<100x50xf32>, tensor<50xf32>) -> tensor<100x50xf32>
+    return %4 : tensor<100x50xf32>
+  }
+  // CHECK: "tf_saved_model.global_tensor"()
+  // CHECK:    sym_name = "dense/kernel"
+  // CHECK: "tf_saved_model.global_tensor"()
+  // CHECK:    sym_name = "dense/bias"
+  // CHECK:  func @serving_default(
+  // CHECK:    %arg0: tensor>> {tf_saved_model.bound_input = @"dense/kernel"},
+  // CHECK:    %arg1: tensor>> {tf_saved_model.bound_input = @"dense/bias"})
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // Test case: Freezing shared VarHandleOp ops.
+
+  func @f(%arg0: tensor>> {tf.resource_name = "dense/kernel"}, %arg1: tensor>> {tf.resource_name = "dense/bias"}) -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})
+  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["f"]} {
+    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor>>
+    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor>>) -> tensor<100x50xf32>
+    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor>>
+    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor>>) -> tensor<50xf32>
+    %4 = "tf.Add"(%1, %3) {device = ""} : (tensor<100x50xf32>, tensor<50xf32>) -> tensor<100x50xf32>
+    return %4 : tensor<100x50xf32>
+  }
+
+  func @f2(%arg0: tensor>> {tf.resource_name = "dense/kernel"}, %arg1: tensor>> {tf.resource_name = "dense/bias"}) -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})
+  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["f2"]} {
+    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor>>
+    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor>>) -> tensor<100x50xf32>
+    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor>>
+    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor>>) -> tensor<50xf32>
+    %4 = "tf.Add"(%1, %3) {device = ""} : (tensor<100x50xf32>, tensor<50xf32>) -> tensor<100x50xf32>
+    return %4 : tensor<100x50xf32>
+  }
+  // CHECK: "tf_saved_model.global_tensor"()
+  // CHECK:    sym_name = "dense/kernel"
+  // CHECK: "tf_saved_model.global_tensor"()
+  // CHECK:    sym_name = "dense/bias"
+  // CHECK:  func @f(
+  // CHECK:    %arg0: tensor>> {tf_saved_model.bound_input = @"dense/kernel"},
+  // CHECK:    %arg1: tensor>> {tf_saved_model.bound_input = @"dense/bias"})
+
+  // CHECK:  func @f2(
+  // CHECK:    %arg0: tensor>> {tf_saved_model.bound_input = @"dense/kernel"},
+  // CHECK:    %arg1: tensor>> {tf_saved_model.bound_input = @"dense/bias"})
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables_invalid_session.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables_invalid_session.mlir
new file mode 100644
index 00000000000..17244d8481a
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_lift_variables_invalid_session.mlir
@@ -0,0 +1,33 @@
+// RUN: tf-opt -verify-diagnostics -tf-saved-model-lift-variables-invalid-session-test -split-input-file %s | FileCheck %s --dump-input=fail
+
+// Test case: Invalid session.
+// expected-error @+1 {{'module' op no session provided}}
+module attributes {tf_saved_model.semantics} {
+
+  func @serving_default(%arg0: tensor>> {tf.resource_name = "dense/kernel"}, %arg1: tensor>> {tf.resource_name = "dense/bias"}) -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})
+  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor>>
+    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor>>) -> tensor<100x50xf32>
+    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor>>
+    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor>>) -> tensor<50xf32>
+    %4 = "tf.Add"(%1, %3) {device = ""} : (tensor<100x50xf32>, tensor<50xf32>) -> tensor<100x50xf32>
+    return %4 : tensor<100x50xf32>
+  }
+}
+
+// -----
+
+// Test case: No errors on no resource arguments.
+module attributes {tf_saved_model.semantics} {
+
+	// CHECK-LABEL: @serving_default
+  func @serving_default() -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})
+  attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "dense_2/Add:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+    %0 = "tf.VarHandleOp"() {_class = ["loc:@dense/kernel"], allowed_devices = [], container = "", device = "", shared_name = "dense/kernel"} : () -> tensor>>
+    %1 = "tf.ReadVariableOp"(%0) {device = ""} : (tensor>>) -> tensor<100x50xf32>
+    %2 = "tf.VarHandleOp"() {_class = ["loc:@dense/bias"], allowed_devices = [], container = "", device = "", shared_name = "dense/bias"} : () -> tensor>>
+    %3 = "tf.ReadVariableOp"(%2) {device = ""} : (tensor>>) -> tensor<50xf32>
+    %4 = "tf.Add"(%1, %3) {device = ""} : (tensor<100x50xf32>, tensor<50xf32>) -> tensor<100x50xf32>
+    return %4 : tensor<100x50xf32>
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.cc b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.cc
new file mode 100644
index 00000000000..6686b340be9
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.cc
@@ -0,0 +1,183 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h"
+
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+#include 
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Module.h"  // from @llvm-project
+#include "mlir/IR/UseDefLists.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/resource_var.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/threadpool_options.h"
+#include "tensorflow/core/public/session.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+using llvm::SmallSet;
+using ::tensorflow::Device;
+using ::tensorflow::DeviceMgr;
+using ::tensorflow::mutex_lock;
+using ::tensorflow::ResourceHandle;
+using ::tensorflow::Session;
+using ::tensorflow::Status;
+using ::tensorflow::StatusOr;
+using ::tensorflow::Tensor;
+using ::tensorflow::Var;
+
+namespace {
+
+constexpr char kResourceNameArgAttr[] = "tf.resource_name";
+constexpr char kSavedModelArgAttr[] = "tf_saved_model.bound_input";
+
+LogicalResult LiftVariablesFromSession(
+    ModuleOp module, Session* session,
+    const SmallSet& resource_names) {
+  OpBuilder builder(module.getBodyRegion());
+  MLIRContext* context = module.getContext();
+
+  if (!session) return module.emitOpError() << "no session provided";
+
+  // Read all resource variables from the session.
+  std::vector variable_names;
+  variable_names.reserve(resource_names.size());
+  for (StringRef name : resource_names) variable_names.push_back(name.str());
+
+  std::vector resource_tensors;
+  Status status = session->Run(
+      /*inputs=*/{}, variable_names,
+      /*target_node_names=*/{}, &resource_tensors);
+  if (!status.ok()) {
+    return module.emitOpError()
+           << "failed to run the provided session: " << status.error_message();
+  }
+
+  const DeviceMgr* device_manager;
+  if (!(session->LocalDeviceManager(&device_manager).ok())) {
+    return module.emitOpError() << "failed to get local device manager";
+  }
+
+  // Read all underlying tensors of the variables from the session.
+  std::vector tensors;
+  tensors.reserve(resource_tensors.size());
+  for (const Tensor& resource_tensor : resource_tensors) {
+    if (resource_tensor.dtype() != tensorflow::DT_RESOURCE) {
+      tensors.push_back(resource_tensor);
+      continue;
+    }
+
+    const ResourceHandle& resource_handle =
+        resource_tensor.scalar()();
+
+    Device* device;
+    if (!(device_manager->LookupDevice(resource_handle.device(), &device)
+              .ok())) {
+      return module.emitOpError() << "failed to look up device";
+    }
+
+    tensorflow::Var* var_ptr;
+    if (!(device->resource_manager()
+              ->Lookup(resource_handle.container(), resource_handle.name(),
+                       &var_ptr)
+              .ok())) {
+      return module.emitOpError() << "failed to look up resource value";
+    }
+    tensorflow::core::RefCountPtr var(var_ptr);
+
+    // The variable tensor is already loaded into corresponding device's
+    // resource manager when we load the saved model using LoadSavedModel().
+    // Here we just read its value.
+    mutex_lock ml(*var->mu());
+    tensors.push_back(*var->tensor());
+  }
+
+  for (const auto iter : llvm::zip(resource_names, tensors)) {
+    const StringRef name = std::get<0>(iter);
+    const Tensor& tensor = std::get<1>(iter);
+
+    // Create tensor attribute for this variable.
+    StatusOr tensor_attr_or = ConvertTensor(tensor, &builder);
+    if (!tensor_attr_or.ok()) {
+      return module.emitOpError()
+             << "failed to convert tensor (name: " << name.str() << ")";
+    }
+    ElementsAttr tensor_attr = tensor_attr_or.ValueOrDie();
+
+    builder.create(
+        NameLoc::get(builder.getIdentifier(name.str()), context),
+        builder.getStringAttr(name), tensor_attr,
+        TypeAttr::get(tensor_attr.getType()), builder.getUnitAttr());
+  }
+
+  return success();
+}
+
+}  // namespace
+
+LogicalResult LiftVariables(ModuleOp module, Session* session) {
+  MLIRContext* context = module.getContext();
+  mlir::Builder builder(context);
+  Identifier resource_name_id = builder.getIdentifier(kResourceNameArgAttr);
+
+  SmallSet resource_names;
+
+  for (FuncOp func : module.getOps()) {
+    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
+      auto resource_arg =
+          func.getArgAttrOfType(i, kResourceNameArgAttr);
+      if (!resource_arg) continue;
+
+      StringRef resource_name = resource_arg.getValue();
+      auto flat_symbol_ref_attr =
+          FlatSymbolRefAttr::get(resource_name, context);
+
+      // Add the corresponding `tf_saved_model.bound_input` attribute.
+      func.setArgAttr(i, kSavedModelArgAttr, flat_symbol_ref_attr);
+
+      resource_names.insert(flat_symbol_ref_attr.getValue());
+
+      // Remove the existing `tf.resource_name` attribute.
+      func.removeArgAttr(i, resource_name_id);
+    }
+  }
+
+  if (resource_names.empty()) return success();
+
+  return LiftVariablesFromSession(module, session, resource_names);
+}
+
+}  // namespace tf_saved_model
+}  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h
new file mode 100644
index 00000000000..12dc787fbcf
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h
@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_H_
+#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_H_
+
+#include "mlir/IR/Module.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/core/public/session.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+// Creates GlobalTensorOp for each variable from function arguments and converts
+// them to the corresponding saved model arguments.
+LogicalResult LiftVariables(ModuleOp module, ::tensorflow::Session* session);
+
+}  // namespace tf_saved_model
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_H_
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
new file mode 100644
index 00000000000..0eaee959c77
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_pass.h
@@ -0,0 +1,57 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_PASS_H_
+
+#include "mlir/IR/Module.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h"
+#include "tensorflow/core/public/session.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+// This pass takes care of finding all variables from the function arguments and
+// converting them to the corresponding global tensors, that will be located out
+// of function. Also it converts resource arguments from function types to the
+// corresponding saved model arguments accordingly.
+class LiftVariablesPass
+    : public PassWrapper> {
+ public:
+  explicit LiftVariablesPass(::tensorflow::Session* session)
+      : session_(session) {}
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    if (failed(LiftVariables(module, session_))) signalPassFailure();
+  }
+
+ private:
+  ::tensorflow::Session* session_;
+};
+
+// Creates as pass that creates GlobalTensorOp for each variable from function
+// arguments and converts the function arguments to the corresponding saved
+// model arguments.
+std::unique_ptr> CreateLiftVariablesPass(
+    ::tensorflow::Session* session);
+
+}  // namespace tf_saved_model
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_PASS_H_
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
new file mode 100644
index 00000000000..faecdf04368
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
@@ -0,0 +1,146 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_TEST_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_TEST_PASS_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/threadpool_options.h"
+#include "tensorflow/core/public/session.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+using ::tensorflow::DeviceMgr;
+using ::tensorflow::Session;
+using ::tensorflow::Status;
+using ::tensorflow::Tensor;
+
+// FakeSession is for testing only.
+class FakeSession : public tensorflow::Session {
+ public:
+  FakeSession() {}
+  ~FakeSession() override = default;
+
+  Status Create(const tensorflow::GraphDef& graph) override {
+    return tensorflow::errors::Unimplemented("not available");
+  }
+  Status Extend(const tensorflow::GraphDef& graph) override {
+    return tensorflow::errors::Unimplemented("not available");
+  }
+
+  Status Close() override {
+    return tensorflow::errors::Unimplemented("not available");
+  }
+
+  Status ListDevices(
+      std::vector* response) override {
+    return tensorflow::errors::Unimplemented("not available");
+  }
+
+  Status LocalDeviceManager(
+      const tensorflow::DeviceMgr** deviceMgrPtr) override {
+    // This method returns a null device manager without making an error.
+    // Users of this method will be notified since it will have a fake data.
+    *deviceMgrPtr = nullptr;
+    return Status::OK();
+  }
+
+  Status Run(const std::vector>& inputs,
+             const std::vector& output_names,
+             const std::vector& target_nodes,
+             std::vector* outputs) override {
+    tensorflow::RunMetadata run_metadata;
+    return Run(tensorflow::RunOptions(), inputs, output_names, target_nodes,
+               outputs, &run_metadata);
+  }
+
+  Status Run(const tensorflow::RunOptions& run_options,
+             const std::vector>& inputs,
+             const std::vector& output_names,
+             const std::vector& target_nodes,
+             std::vector* outputs,
+             tensorflow::RunMetadata* run_metadata) override {
+    return Run(run_options, inputs, output_names, target_nodes, outputs,
+               run_metadata, tensorflow::thread::ThreadPoolOptions());
+  }
+
+  Status Run(const tensorflow::RunOptions& run_options,
+             const std::vector>& inputs,
+             const std::vector& output_names,
+             const std::vector& target_nodes,
+             std::vector* outputs,
+             tensorflow::RunMetadata* run_metadata,
+             const tensorflow::thread::ThreadPoolOptions& thread_pool_options)
+      override {
+    for (const std::string& output_name : output_names) {
+      Tensor output;
+      if (output_name == "dense/bias") {
+        outputs->push_back(
+            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({50})));
+      } else if (output_name == "dense/kernel") {
+        outputs->push_back(
+            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({100, 50})));
+      } else {
+        // Create a scalar float tensor.
+        outputs->push_back(
+            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({})));
+      }
+    }
+    return Status::OK();
+  }
+};
+
+// This pass is only available in the tf-opt binary for testing.
+class LiftVariablesTestPass
+    : public PassWrapper> {
+ public:
+  LiftVariablesTestPass() { session_ = new FakeSession(); }
+
+  ~LiftVariablesTestPass() override { delete session_; }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    if (failed(LiftVariables(module, session_))) signalPassFailure();
+  }
+
+ private:
+  Session* session_;
+};
+
+// This pass is only available in the tf-opt binary for testing.
+class LiftVariablesInvalidSessionTestPass
+    : public PassWrapper> {
+ public:
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    // Pass an invalid session argument, which is a nullptr.
+    if (failed(LiftVariables(module, /*session=*/nullptr))) signalPassFailure();
+  }
+};
+
+}  // namespace tf_saved_model
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_LIFT_VARIABLES_TEST_PASS_H_
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass_registration.cc b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass_registration.cc
new file mode 100644
index 00000000000..19c367c6d46
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass_registration.cc
@@ -0,0 +1,32 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h"
+
+namespace mlir {
+namespace tf_saved_model {
+
+static PassRegistration lift_variables_test_pass(
+    "tf-saved-model-lift-variables-test",
+    "Lift variables and save them as global tensors");
+
+static PassRegistration
+    lift_variables_invalid_session_test_pass(
+        "tf-saved-model-lift-variables-invalid-session-test",
+        "Lift variables and save them as global tensors with an invalid "
+        "session");
+
+}  // namespace tf_saved_model
+}  // namespace mlir

From e3a423a9493af054d87f7ce45feb228a3ccfe6e6 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Tue, 16 Jun 2020 16:34:38 -0700
Subject: [PATCH 0344/1390] Added FuzzedDataProvider to split fuzzer data

Switched manual data splicing to FuzzedDataProvider @mihaimaruseac
---
 tensorflow/security/fuzzing/status_fuzz.cc | 26 +++++++++-------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/tensorflow/security/fuzzing/status_fuzz.cc b/tensorflow/security/fuzzing/status_fuzz.cc
index 7b161645148..8b5949009b1 100644
--- a/tensorflow/security/fuzzing/status_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_fuzz.cc
@@ -17,6 +17,8 @@ limitations under the License.
 
 #include "tensorflow/core/platform/status.h"
 
+#include 
+
 // This is a fuzzer for `tensorflow::Status`. Since `Status` is used almost
 // everywhere, we need to ensure that the common functionality is safe. We don't
 // expect many crashes from this fuzzer since we only create a status and then
@@ -26,9 +28,7 @@ limitations under the License.
 
 namespace {
 
-tensorflow::error::Code BuildRandomErrorCode(uint8_t a, uint8_t b, uint8_t c,
-                                             uint8_t d) {
-  int code = (a << 24) | (b << 16) | (c << 8) | d;
+tensorflow::error::Code BuildRandomErrorCode(uint32_t code){
 
   // We cannot build a `Status` with error_code of 0 and a message, so force
   // error code to be non-zero.
@@ -39,22 +39,16 @@ tensorflow::error::Code BuildRandomErrorCode(uint8_t a, uint8_t b, uint8_t c,
   return static_cast(code);
 }
 
-std::string GetRandomErrorString(const uint8_t *data, size_t size) {
-  const char *p = reinterpret_cast(data);
-  return std::string(p, size);
-}
-
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
-  // TODO(mihaimaruseac): Use `FuzzedDataProvider` and then make these `const`
   tensorflow::error::Code error_code;
   std::string error_message;
-  if (size < 4) {
-    error_code = BuildRandomErrorCode(0, 0, 0, 0);
-    error_message = GetRandomErrorString(data, size);
-  } else {
-    error_code = BuildRandomErrorCode(data[0], data[1], data[2], data[3]);
-    error_message = GetRandomErrorString(data + 4, size - 4);
-  }
+
+  FuzzedDataProvider fuzzed_data(data, size);
+
+  uint32_t code = fuzzed_data.ConsumeIntegral();
+  error_code = BuildRandomErrorCode(code);
+
+  error_message = fuzzed_data.ConsumeRemainingBytesAsString();
 
   tensorflow::Status s = tensorflow::Status(error_code, error_message);
   const std::string actual_message = s.ToString();

From dac169cd7f2f618abc6511efcae0bb87d033ec7e Mon Sep 17 00:00:00 2001
From: YoungSeok Yoon 
Date: Tue, 16 Jun 2020 16:34:47 -0700
Subject: [PATCH 0345/1390] Include TensorFlow LICENSE file to TFLite aars

PiperOrigin-RevId: 316781613
Change-Id: I6323349321f5b009b1511e0cec12dbaca6a8c770
---
 tensorflow/lite/java/aar_with_jni.bzl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/java/aar_with_jni.bzl b/tensorflow/lite/java/aar_with_jni.bzl
index 71da735703d..34706c19c54 100644
--- a/tensorflow/lite/java/aar_with_jni.bzl
+++ b/tensorflow/lite/java/aar_with_jni.bzl
@@ -52,7 +52,11 @@ EOF
         ],
     )
 
-    srcs = [android_library + ".aar", name + "_dummy_app_for_so_unsigned.apk"]
+    srcs = [
+        android_library + ".aar",
+        name + "_dummy_app_for_so_unsigned.apk",
+        "//:LICENSE",
+    ]
 
     cmd = """
 cp $(location {0}.aar) $(location :{1}.aar)
@@ -62,6 +66,8 @@ cd $$(mktemp -d)
 unzip $$origdir/$(location :{1}_dummy_app_for_so_unsigned.apk) "lib/*"
 cp -r lib jni
 zip -r $$origdir/$(location :{1}.aar) jni/*/*.so
+cp $$origdir/$(location //:LICENSE) ./
+zip $$origdir/$(location :{1}.aar) LICENSE
 """.format(android_library, name)
 
     if headers:

From 897e3c0ecad3b45f5e96615173e7511619eebc93 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 16:39:31 -0700
Subject: [PATCH 0346/1390] Softmax1x1 converted to new style.

PiperOrigin-RevId: 316782370
Change-Id: I1f7761c0520d72876f352c9f156341b349b90cbe
---
 .../delegates/gpu/cl/kernels/softmax1x1.cc    | 87 ++++++++++---------
 1 file changed, 45 insertions(+), 42 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/softmax1x1.cc b/tensorflow/lite/delegates/gpu/cl/kernels/softmax1x1.cc
index 192bee771d6..fcfe4a1810c 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/softmax1x1.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/softmax1x1.cc
@@ -25,47 +25,45 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetSoftmaxKernelCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor("src_data",
-                                 WHSBPoint{"tensor_size.x", "tensor_size.y",
-                                           "tensor_size.z", "tensor_size.w"},
-                                 op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor("dst_data",
-                                 WHSBPoint{"tensor_size.x", "tensor_size.y",
-                                           "tensor_size.z", "tensor_size.w"},
-                                 op_def.dst_tensors[0]);
+std::string GetSoftmaxKernelCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+  args->AddFloat("mask_x");
+  args->AddFloat("mask_y");
+  args->AddFloat("mask_z");
+  args->AddFloat("mask_w");
+  args->AddInt("slices_x32");
 
-  const std::string batch_id = op_def.IsBatchSupported() ? "batch_id" : "";
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 tensor_size,\n";
-  c += "    int2 size,\n";
-  c += "    float4 mask\n";
-  c += ") {\n";
+  c += "$0) {\n";
   if (op_def.IsBatchSupported()) {
     c += "  int batch_id = get_global_id(1);\n";
-    c += "  if (batch_id >= tensor_size.w) return;\n";
+    c += "  if (batch_id >= args.dst_tensor.Batch()) return;\n";
+    c += "  args.dst_tensor.SetBatchRef(batch_id);\n";
+    c += "  args.src_tensor.SetBatchRef(batch_id);\n";
   }
+  c += "  float4 mask = (float4)(args.mask_x, args.mask_y, args.mask_z, "
+       "args.mask_w);\n";
   c += "  int offset = 0;\n";
   c += "  float sum = 0.0f;\n";
   c += "  int s = 0;\n";
   c += "  int tid = get_local_id(0);\n";
   c += "  do {\n";
   c += "    int z = offset + tid;\n";
-  c += "    if (z < size.x) {\n";
-  c += "      float4 mask_temp = z == size.x - 1 ? mask : (float4)(1.0f);\n";
-  c += "      float4 src = " +
-       src_tensor.ReadAsFloatWHSB("0", "0", "z", batch_id) + ";\n";
+  c += "    if (z < args.dst_tensor.Slices()) {\n";
+  c += "      float4 mask_temp = z == args.dst_tensor.Slices() - 1 ? mask : "
+       "(float4)(1.0f);\n";
+  c += "      float4 src = args.src_tensor.Read(0, 0, z);\n";
   c += "      sum += dot(mask_temp, exp(src));\n";
   c += "      offset += 32;\n";
   c += "    }\n";
   c += "    s++;\n";
-  c += "  } while (s < size.y);\n";
+  c += "  } while (s < args.slices_x32);\n";
   c += "\n";
   c += "  __local float4 tmp[8];\n";
   c += "  __local float* tmpx1 = (__local float*)tmp;\n";
@@ -89,16 +87,14 @@ std::string GetSoftmaxKernelCode(
   c += "  s = 0;\n";
   c += "  do {\n";
   c += "    int z = offset + tid;\n";
-  c += "    if (z < size.x) {\n";
-  c += "      FLT4 res = TO_FLT4(exp(" +
-       src_tensor.ReadAsFloatWHSB("0", "0", "z", batch_id) + ")*sum);\n";
-  const LinkingContext context{"res", "0", "0", "z"};
-  c += PostProcess(linked_operations, context);
-  c += "    " + dst_tensor.WriteWHSB("res", "0", "0", "z", batch_id);
+  c += "    if (z < args.dst_tensor.Slices()) {\n";
+  c += "      FLT4 res = TO_FLT4(exp(args.src_tensor.Read(0, 0, "
+       "z))*sum);\n";
+  c += "      args.dst_tensor.Write(res, 0, 0, z);\n";
   c += "      offset += 32;\n";
   c += "    }\n";
   c += "    s++;\n";
-  c += "  } while (s < size.y);\n";
+  c += "  } while (s < args.slices_x32);\n";
   c += "}\n";
   return c;
 }
@@ -116,23 +112,30 @@ Softmax1x1& Softmax1x1::operator=(Softmax1x1&& kernel) {
 }
 
 absl::Status Softmax1x1::Compile(const CreationContext& creation_context) {
-  const auto code = GetSoftmaxKernelCode(definition_, linked_operations_);
+  std::string code = GetSoftmaxKernelCode(definition_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Softmax1x1::AddToQueue(CLCommandQueue* queue) {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  const int depth = src_[0]->Slices();
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(int2(depth, DivideRoundUp(depth, 32))));
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  float4 mask = GetMaskForLastPlane(src_[0]->Channels());
+  RETURN_IF_ERROR(args_.SetFloat("mask_x", mask.x));
+  RETURN_IF_ERROR(args_.SetFloat("mask_y", mask.y));
+  RETURN_IF_ERROR(args_.SetFloat("mask_z", mask.z));
+  RETURN_IF_ERROR(args_.SetFloat("mask_w", mask.w));
   RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(GetMaskForLastPlane(src_[0]->Channels())));
-
+      args_.SetInt("slices_x32", DivideRoundUp(src_[0]->Slices(), 32)));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  RETURN_IF_ERROR(args_.Bind(kernel_.kernel()));
   return queue->DispatchImplicit(kernel_, {32, dst_[0]->Batch(), 1},
                                  {32, 1, 1});
 }

From 267f956246750357f8eff2b88b2e6c8741b46775 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 16:39:49 -0700
Subject: [PATCH 0347/1390] load libcupti.so shared object if libcupti.so.10.1
 is not installed. This makes it possible for tensorflow to load CUPTI 10.2
 via libcupti.so symlink in the event that libcupti.so.10.1 is missing.

PiperOrigin-RevId: 316782424
Change-Id: I930912f1d3a8fa80e8b91c41214c374650f08847
---
 .../stream_executor/platform/default/dso_loader.cc     | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensorflow/stream_executor/platform/default/dso_loader.cc b/tensorflow/stream_executor/platform/default/dso_loader.cc
index fb7d88aaedb..01af4114536 100644
--- a/tensorflow/stream_executor/platform/default/dso_loader.cc
+++ b/tensorflow/stream_executor/platform/default/dso_loader.cc
@@ -101,13 +101,11 @@ port::StatusOr GetCurandDsoHandle() {
 }
 
 port::StatusOr GetCuptiDsoHandle() {
-#if defined(ANDROID_TEGRA)
-  // On Android devices the CUDA version number is not added to the library
-  // name.
+  // Load specific version of CUPTI this is built.
+  auto status_or_handle = GetDsoHandle("cupti", GetCudaVersion());
+  if (status_or_handle.ok()) return status_or_handle;
+  // Load whatever libcupti.so user specified.
   return GetDsoHandle("cupti", "");
-#else
-  return GetDsoHandle("cupti", GetCudaVersion());
-#endif
 }
 
 port::StatusOr GetCudnnDsoHandle() {

From f8657c62c60dffe01e27f8d47028b533c0837d2c Mon Sep 17 00:00:00 2001
From: Allen Lavoie 
Date: Tue, 16 Jun 2020 16:48:19 -0700
Subject: [PATCH 0348/1390] Parallel device: avoid deadlocks when the
 EagerContext's default executor is async

Creates one sync executor per thread.

Requires fixing a tangential use-after-free where the context assumed all of the thread-local executors were still allocated at shutdown.

PiperOrigin-RevId: 316783819
Change-Id: I62e7a91dcccb847d4e1c2a5f08e30c2877556618
---
 tensorflow/c/eager/c_api_experimental_test.cc | 29 +++++++++++++++++
 .../parallel_device/parallel_device_lib.cc    | 18 +++++++++++
 .../parallel_device/parallel_device_test.cc   |  6 +---
 .../core/common_runtime/eager/context.cc      | 32 ++++++++++++++++++-
 .../core/common_runtime/eager/context.h       |  2 ++
 .../common_runtime/eager/eager_executor.cc    | 11 +++++++
 .../common_runtime/eager/eager_executor.h     | 10 ++++++
 .../parallel_device/parallel_device_test.py   | 24 +++++++++++++-
 8 files changed, 125 insertions(+), 7 deletions(-)

diff --git a/tensorflow/c/eager/c_api_experimental_test.cc b/tensorflow/c/eager/c_api_experimental_test.cc
index 0c058398299..a4d31417073 100644
--- a/tensorflow/c/eager/c_api_experimental_test.cc
+++ b/tensorflow/c/eager/c_api_experimental_test.cc
@@ -212,6 +212,35 @@ TEST(CAPI, CancellationManager) {
   TFE_DeleteCancellationManager(c_mgr);
 }
 
+TEST(CAPI, ExecutorContextDestructionOrder) {
+  TF_Status* status = TF_NewStatus();
+
+  {
+    TFE_ContextOptions* opts = TFE_NewContextOptions();
+    TFE_Context* ctx = TFE_NewContext(opts, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TFE_DeleteContextOptions(opts);
+    TFE_Executor* executor = TFE_NewExecutor(/*is_async=*/false);
+    TFE_ContextSetExecutorForThread(ctx, executor);
+
+    TFE_DeleteContext(ctx);
+    TFE_DeleteExecutor(executor);
+  }
+
+  {
+    TFE_ContextOptions* opts = TFE_NewContextOptions();
+    TFE_Context* ctx = TFE_NewContext(opts, status);
+    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
+    TFE_DeleteContextOptions(opts);
+    TFE_Executor* executor = TFE_NewExecutor(/*is_async=*/false);
+    TFE_ContextSetExecutorForThread(ctx, executor);
+
+    TFE_DeleteExecutor(executor);
+    TFE_DeleteContext(ctx);
+  }
+  TF_DeleteStatus(status);
+}
+
 TEST(CAPI, Function_ident_CPU) {
   // First create a simple identity function.
   TF_Graph* function_graph = TF_NewGraph();
diff --git a/tensorflow/c/eager/parallel_device/parallel_device_lib.cc b/tensorflow/c/eager/parallel_device/parallel_device_lib.cc
index 98cd4812610..d0149b29c08 100644
--- a/tensorflow/c/eager/parallel_device/parallel_device_lib.cc
+++ b/tensorflow/c/eager/parallel_device/parallel_device_lib.cc
@@ -37,6 +37,15 @@ class StatusDeleter {
 
 using StatusPtr = std::unique_ptr;
 
+class ExecutorDeleter {
+ public:
+  void operator()(TFE_Executor* to_delete) const {
+    TFE_DeleteExecutor(to_delete);
+  }
+};
+
+using ExecutorPtr = std::unique_ptr;
+
 }  // namespace
 
 // Allows a single op at a time to be launched without blocking.
@@ -51,6 +60,13 @@ class DeviceThread {
   explicit DeviceThread(const std::string& device)
       : status_(TF_NewStatus()),
         device_(device),
+        // If the context's default exector is set to async, re-using that in
+        // each thread would cause collectives to deadlock. For consistency we
+        // create a new sync executor for every thread.
+        //
+        // TODO(allenl): We should have an async API that works with the
+        // parallel device.
+        executor_(TFE_NewExecutor(/*is_async=*/false)),
         op_(nullptr),
         thread_(tensorflow::Env::Default()->StartThread(
             tensorflow::ThreadOptions(), "parallel_device_execute",
@@ -105,6 +121,7 @@ class DeviceThread {
   StatusPtr status_ TF_GUARDED_BY(execution_mutex_);
 
   const std::string device_;
+  ExecutorPtr executor_ TF_GUARDED_BY(execution_mutex_);
   mutable OpPtr op_ TF_GUARDED_BY(execution_mutex_);
   std::unique_ptr thread_;
 };
@@ -186,6 +203,7 @@ void DeviceThread::Execute(TFE_Context* context, const char* operation_name,
                            std::vector* outputs,
                            TF_Status* status) const {
   if (op_ == nullptr) {
+    TFE_ContextSetExecutorForThread(context, executor_.get());
     op_.reset(TFE_NewOp(context, operation_name, status));
     if (TF_GetCode(status) != TF_OK) return;
     TFE_OpSetDevice(op_.get(), device_.c_str(), status);
diff --git a/tensorflow/c/eager/parallel_device/parallel_device_test.cc b/tensorflow/c/eager/parallel_device/parallel_device_test.cc
index e5412dbba61..2fa183d50f6 100644
--- a/tensorflow/c/eager/parallel_device/parallel_device_test.cc
+++ b/tensorflow/c/eager/parallel_device/parallel_device_test.cc
@@ -412,6 +412,7 @@ void TestCollective(bool async) {
       TF_NewStatus(), TF_DeleteStatus);
   std::unique_ptr opts(
       TFE_NewContextOptions(), TFE_DeleteContextOptions);
+  TFE_ContextOptionsSetAsync(opts.get(), async);
   std::unique_ptr config(
       TF_CreateConfig(
           /*xla*/ false,
@@ -423,9 +424,6 @@ void TestCollective(bool async) {
   std::unique_ptr context(
       TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
   ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
-  std::unique_ptr executor(
-      TFE_NewExecutor(async), TFE_DeleteExecutor);
-  TFE_ContextSetExecutorForThread(context.get(), executor.get());
 
   const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
   std::array underlying_devices{
@@ -455,8 +453,6 @@ void TestCollective(bool async) {
   ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
   ExpectScalarEq(result_components[0].get(), 3.);
   ExpectScalarEq(result_components[1].get(), 3.);
-  // Destroying the context's default executor first isn't safe.
-  context.reset();
 }
 
 TEST(PARALLEL_DEVICE, TestCollectiveSync) { TestCollective(/*async=*/false); }
diff --git a/tensorflow/core/common_runtime/eager/context.cc b/tensorflow/core/common_runtime/eager/context.cc
index 5d8cb3da6bc..970c2bcbb89 100644
--- a/tensorflow/core/common_runtime/eager/context.cc
+++ b/tensorflow/core/common_runtime/eager/context.cc
@@ -341,7 +341,28 @@ void EagerContext::SetExecutorForThread(EagerExecutor* executor) {
   if (executor == &default_executor_) {
     thread_local_executor_.erase(std::this_thread::get_id());
   } else {
-    thread_local_executor_[std::this_thread::get_id()] = executor;
+    auto thread_id = std::this_thread::get_id();
+    thread_local_executor_[thread_id] = executor;
+    auto& executors_with_cleanups = has_cleanup_[thread_id];
+    if (executors_with_cleanups.find(executor) ==
+        executors_with_cleanups.end()) {
+      executors_with_cleanups.insert(executor);
+      // If the executor is deleted before this context, we need to remove it
+      // from the map to avoid attempting to sync it in our destructor.
+      std::function cleanup([this, thread_id, executor]() {
+        {
+          tensorflow::mutex_lock l(executor_map_mu_);
+          auto existing = thread_local_executor_.find(thread_id);
+          if (existing != thread_local_executor_.end() &&
+              existing->second == executor) {
+            thread_local_executor_.erase(thread_id);
+          }
+          has_cleanup_[thread_id].erase(executor);
+        }
+      });
+      executor->AddCleanup(reinterpret_cast(this),
+                           std::move(cleanup));
+    }
   }
 }
 
@@ -525,6 +546,15 @@ EagerContext::~EagerContext() {
   custom_devices_.clear();
 
   ClearCachesAndThreadExecutors();
+  std::unordered_map executors_copy;
+  {
+    mutex_lock l(executor_map_mu_);
+    executors_copy = thread_local_executor_;
+  }
+  for (const auto& entry : executors_copy) {
+    // Let the executor know that its cleanup closure is no longer valid.
+    entry.second->RemoveCleanups(reinterpret_cast(this));
+  }
   for (auto& entry : registered_functions_) {
     while (!entry.second->Unref()) {
       // remove all references.
diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h
index fa57afecbaf..cb6d09f8f1d 100644
--- a/tensorflow/core/common_runtime/eager/context.h
+++ b/tensorflow/core/common_runtime/eager/context.h
@@ -639,6 +639,8 @@ class EagerContext : public AbstractContextInterface, public core::RefCounted {
   // Not owned.
   std::unordered_map thread_local_executor_
       TF_GUARDED_BY(executor_map_mu_);
+  std::unordered_map>
+      has_cleanup_ TF_GUARDED_BY(executor_map_mu_);
 
   const bool log_memory_;
 
diff --git a/tensorflow/core/common_runtime/eager/eager_executor.cc b/tensorflow/core/common_runtime/eager/eager_executor.cc
index ddfdabf9472..7fe321edffd 100644
--- a/tensorflow/core/common_runtime/eager/eager_executor.cc
+++ b/tensorflow/core/common_runtime/eager/eager_executor.cc
@@ -46,6 +46,11 @@ EagerExecutor::~EagerExecutor() {
   tensorflow::mutex_lock l(node_queue_mutex_);
   state_ = ExecutorState::kShutDown;
   nodes_pending_.notify_all();
+  for (const auto& cleanups_for_key : cleanups_) {
+    for (const std::function& cleanup : cleanups_for_key.second) {
+      cleanup();
+    }
+  }
 }
 
 Status EagerExecutor::ShutDown() {
@@ -413,4 +418,10 @@ Status EagerExecutor::MoveToUnfinished(core::RefCountPtr item,
   return Status::OK();
 }
 
+void EagerExecutor::AddCleanup(intptr_t key, std::function callback) {
+  cleanups_[key].push_back(callback);
+}
+
+void EagerExecutor::RemoveCleanups(intptr_t key) { cleanups_.erase(key); }
+
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/eager_executor.h b/tensorflow/core/common_runtime/eager/eager_executor.h
index aa8864c7ad6..34847abc26a 100644
--- a/tensorflow/core/common_runtime/eager/eager_executor.h
+++ b/tensorflow/core/common_runtime/eager/eager_executor.h
@@ -153,6 +153,13 @@ class EagerExecutor {
 
   bool ok() const TF_NO_THREAD_SAFETY_ANALYSIS { return ok_; }
 
+  // On destruction, runs `callback`. Used by the EagerContext for clearing
+  // thread-local executors.
+  void AddCleanup(intptr_t key, std::function callback);
+  // If `key` (e.g. a context) is destroyed before the executor, the associated
+  // callbacks are no longer safe to run.
+  void RemoveCleanups(intptr_t key);
+
  private:
   // Possible states for this executor.
   // Executor starts in kActive state. When Shutdown() is called, Executor
@@ -250,6 +257,9 @@ class EagerExecutor {
   const eager::EagerClient* last_eager_client_;
 
   const bool enable_async_wait_for_remote_function_;
+
+  // Callbacks to run on destruction.
+  std::unordered_map>> cleanups_;
 };
 
 inline bool EagerExecutor::Async() const { return thread_ != nullptr; }
diff --git a/tensorflow/python/distribute/parallel_device/parallel_device_test.py b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
index 9dbf258f70f..8fc3dcb5816 100644
--- a/tensorflow/python/distribute/parallel_device/parallel_device_test.py
+++ b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
@@ -23,6 +23,7 @@ import threading
 from tensorflow.python.distribute.parallel_device import parallel_device
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
+from tensorflow.python.framework import config
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.module import module
@@ -136,7 +137,7 @@ class ParallelDeviceTests(_VirtualDeviceTestCase):
     self.assertIn(self.device.components[0], outputs[0].backing_device)
     self.assertIn(self.device.components[1], outputs[1].backing_device)
 
-  def test_collective_reduce_async(self):
+  def test_collective_reduce_async_scope(self):
     # Note that ops on the parallel device currently don't execute
     # asynchronously. The test is just that we don't get deadlocks.
     with context.async_scope(), ops.device(self.device.name):
@@ -149,6 +150,27 @@ class ParallelDeviceTests(_VirtualDeviceTestCase):
     self.assertIn(self.device.components[0], outputs[0].backing_device)
     self.assertIn(self.device.components[1], outputs[1].backing_device)
 
+  def test_collective_reduce_async_context(self):
+    previous = config.get_synchronous_execution()
+    try:
+      context._reset_context()
+      config.set_synchronous_execution(False)
+      self.setUp()
+      # Note that ops on the parallel device currently don't execute
+      # asynchronously. The test is just that we don't get deadlocks.
+      with ops.device(self.device.name):
+        x = self.device.pack(
+            [constant_op.constant(-1.5),
+             constant_op.constant(3.5)])
+        reduced = _collective_sum(x, num_replicas=2)
+        outputs = self.device.unpack(reduced)
+      self.assertAllClose([2., 2.], outputs)
+      self.assertIn(self.device.components[0], outputs[0].backing_device)
+      self.assertIn(self.device.components[1], outputs[1].backing_device)
+    finally:
+      context._reset_context()
+      config.set_synchronous_execution(previous)
+
   def test_checkpointing(self):
     prefix = os.path.join(self.get_temp_dir(), "ckpt")
     with self.device.scope():

From 950cffcd8deb881dcbfdf92f22c37eaa36f61e04 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Tue, 16 Jun 2020 16:51:59 -0700
Subject: [PATCH 0349/1390] Enable builds of TF to link against `libc++`.

This should now enable more fuzzers and a nicer/stabler OSSFuzz integration.

PiperOrigin-RevId: 316784432
Change-Id: Iaef6c288221c3f7214d7806aa6913f0370a63544
---
 .bazelrc | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/.bazelrc b/.bazelrc
index 5ea8048d5d9..e21a1a32917 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -30,6 +30,7 @@
 #     short_logs:       Only log errors during build, skip warnings.
 #     monolithic:       Build all TF C++ code into a single shared object.
 #     dynamic_kernels:  Try to link all kernels dynamically (experimental).
+#     libc++:           Link against libc++ instead of stdlibc++
 #
 #
 # TF version options;
@@ -79,6 +80,14 @@
 #     elinux_armhf:    Embedded Linux options for armhf (ARMv7) CPU support.
 
 
+# Allow builds using libc++ as a linker library
+# This is mostly for OSSFuzz, so we also pass in the flags from environment to clean build file
+build:libc++ --action_env=CC
+build:libc++ --action_env=CXX
+build:libc++ --action_env=CXXFLAGS=-stdlib=libc++
+build:libc++ --action_env=PATH
+build:libc++ --define force_libcpp=enabled
+build:libc++ --linkopt -fuse-ld=lld
 
 # Android configs. Bazel needs to have --cpu and --fat_apk_cpu both set to the
 # target CPU to build transient dependencies correctly. See

From c49404dc26d9daa82a012253ca55f95eab6891eb Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Tue, 16 Jun 2020 17:00:16 -0700
Subject: [PATCH 0350/1390] TPU library refactor.

PiperOrigin-RevId: 316785838
Change-Id: I408dd6be75ed6a6ccfdbfad704a1a695906f2aa9
---
 .../core/tpu/kernels/tpu_compile_op_common.cc | 22 +++++++++----------
 .../core/tpu/kernels/tpu_compile_op_common.h  | 11 +++++-----
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
index ae090913dc7..c8faba1d975 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
@@ -118,7 +118,7 @@ Status SetPerCoreArgShapes(
 
 }  // namespace
 
-Status TPUCompileOpKernelCommon::AssignReturnValueToCore(
+Status TpuCompileOpKernelCommon::AssignReturnValueToCore(
     std::vector* retval_core_mapping) {
   std::vector per_core_retval_counts(metadata_.num_cores_per_replica(), 0);
   for (int i = 0; i < metadata_.retvals_size(); ++i) {
@@ -149,7 +149,7 @@ Status TPUCompileOpKernelCommon::AssignReturnValueToCore(
   return Status::OK();
 }
 
-Status TPUCompileOpKernelCommon::BuildComputationArgumentDescriptions(
+Status TpuCompileOpKernelCommon::BuildComputationArgumentDescriptions(
     const std::vector& arg_shapes,
     const OpInputList& guaranteed_constants, const XlaCompiler& compiler,
     std::vector* args,
@@ -207,7 +207,7 @@ Status TPUCompileOpKernelCommon::BuildComputationArgumentDescriptions(
   return Status::OK();
 }
 
-Status TPUCompileOpKernelCommon::GetShardingInfo(
+Status TpuCompileOpKernelCommon::GetShardingInfo(
     absl::Span arg_shapes,
     const XlaCompiler::ShapeRepresentationFn shape_representation_fn,
     std::vector* arg_core_mapping,
@@ -230,7 +230,7 @@ Status TPUCompileOpKernelCommon::GetShardingInfo(
   return Status::OK();
 }
 
-Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
+Status TpuCompileOpKernelCommon::CompileTFFunctionToHlo(
     const FunctionLibraryDefinition& flib_def, int graph_def_version,
     const XlaCompiler::ShapeRepresentationFn shape_representation_fn,
     const std::vector& arg_shapes,
@@ -260,7 +260,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
   std::vector retval_core_mapping(
       metadata_.retvals_size());
   TF_RETURN_IF_ERROR(
-      TPUCompileOpKernelCommon::AssignReturnValueToCore(&retval_core_mapping));
+      TpuCompileOpKernelCommon::AssignReturnValueToCore(&retval_core_mapping));
 
   LOG(INFO) << "Instantiating function:" << function.name();
   FunctionLibraryRuntime::Handle handle;
@@ -341,7 +341,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
                                 args, compilation_result);
 }
 
-/* static */ void TPUCompileOpKernelCommon::ExitCountdown(
+/* static */ void TpuCompileOpKernelCommon::ExitCountdown(
     OpKernelContext* ctx, std::shared_ptr> done) {
   const int kSleepSeconds = 300;
   LOG(INFO) << "TpuCompileOp was cancelled. Sleeping for " << kSleepSeconds
@@ -355,7 +355,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
   LogAndExit(42);
 }
 
-/* static */ Status TPUCompileOpKernelCommon::GetDynamicShapes(
+/* static */ Status TpuCompileOpKernelCommon::GetDynamicShapes(
     OpKernelContext* ctx, std::vector* shapes) {
   OpInputList dynamic_shapes;
   TF_RETURN_IF_ERROR(ctx->input_list("dynamic_shapes", &dynamic_shapes));
@@ -368,7 +368,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
   return Status::OK();
 }
 
-/* static */ Status TPUCompileOpKernelCommon::ComputeArgumentShapes(
+/* static */ Status TpuCompileOpKernelCommon::ComputeArgumentShapes(
     const tpu::TPUCompileMetadataProto& metadata,
     const std::vector& dynamic_shapes,
     std::vector* arg_shapes) {
@@ -409,7 +409,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
 
 // Function arguments and return values lose their device assignments, so we
 // must recreate them.
-/* static */ Status TPUCompileOpKernelCommon::AssignDevicesToArgsAndRetvals(
+/* static */ Status TpuCompileOpKernelCommon::AssignDevicesToArgsAndRetvals(
     absl::Span arg_core_mapping,
     absl::Span retval_core_mapping, Graph* graph) {
   auto assign = [&](Node* node, const xla::OpSharding& sharding) -> Status {
@@ -444,7 +444,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
 
 // Performs shape inference on the body of `graph`. Shapes for arguments
 // are taken from `metadata` and `arg_shapes`.
-/* static */ Status TPUCompileOpKernelCommon::RunShapeInferenceOnComputation(
+/* static */ Status TpuCompileOpKernelCommon::RunShapeInferenceOnComputation(
     const tpu::TPUCompileMetadataProto& metadata,
     const std::vector& arg_shapes, Graph* graph,
     FunctionLibraryRuntime* flr, GraphShapeInfo* shape_info) {
@@ -476,7 +476,7 @@ Status TPUCompileOpKernelCommon::CompileTFFunctionToHlo(
       shape_info);
 }
 
-Status TPUCompileOpKernelCommon::OptimizeGraph(
+Status TpuCompileOpKernelCommon::OptimizeGraph(
     const tpu::TPUCompileMetadataProto& metadata,
     const std::vector& arg_shapes,
     std::unique_ptr* graph, FunctionLibraryRuntime* flr,
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.h b/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
index 74ae8729f8b..2c8d90643ef 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
@@ -33,9 +33,9 @@ namespace tensorflow {
 namespace tpu {
 
 // Abstract base class for TpuCompileOpKernel implementation.
-class TPUCompileOpKernelCommon {
+class TpuCompileOpKernelCommon {
  public:
-  TPUCompileOpKernelCommon(const std::string& mlir_module,
+  TpuCompileOpKernelCommon(const std::string& mlir_module,
                            const tpu::TPUCompileMetadataProto metadata,
                            int num_computations)
       : metadata_(metadata),
@@ -43,7 +43,7 @@ class TPUCompileOpKernelCommon {
         mlir_module_(mlir_module),
         num_computations_(num_computations) {}
 
-  TPUCompileOpKernelCommon(const NameAttrList& function,
+  TpuCompileOpKernelCommon(const NameAttrList& function,
                            const tpu::TPUCompileMetadataProto metadata,
                            int num_computations)
       : metadata_(metadata),
@@ -51,7 +51,7 @@ class TPUCompileOpKernelCommon {
         function_(function),
         num_computations_(num_computations) {}
 
-  virtual ~TPUCompileOpKernelCommon() = default;
+  virtual ~TpuCompileOpKernelCommon() = default;
 
   virtual void Compute(OpKernelContext* ctx) = 0;
 
@@ -153,8 +153,7 @@ class TPUCompileOpKernelCommon {
   int num_computations_;
 
  private:
-  TPUCompileOpKernelCommon(const TPUCompileOpKernelCommon&) = delete;
-  TPUCompileOpKernelCommon& operator=(const TPUCompileOpKernelCommon&) = delete;
+  DISALLOW_COPY_AND_ASSIGN(TpuCompileOpKernelCommon);
 };
 
 }  // namespace tpu

From 6ad3d3c05cb0dff9d6253a6089acaa7ac6c57604 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Tue, 16 Jun 2020 17:00:56 -0700
Subject: [PATCH 0351/1390] Make device_util_test as PY3 only

PiperOrigin-RevId: 316785955
Change-Id: Ibef166f3b3d095f9c7aaceb68c1c15952d91d250
---
 tensorflow/python/distribute/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index e39631d634f..96559a9a740 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -120,6 +120,8 @@ py_library(
 cuda_py_test(
     name = "device_util_test",
     srcs = ["device_util_test.py"],
+    python_version = "PY3",
+    tags = ["no_oss_py2"],
     deps = [
         ":combinations",
         ":device_util",

From de0d8ddc270f33b899d56b04eed622c2a0906006 Mon Sep 17 00:00:00 2001
From: Yunlu Li 
Date: Tue, 16 Jun 2020 17:01:18 -0700
Subject: [PATCH 0352/1390] Set the type of sparsity metadata accordingly to
 further reduce the model size.

PiperOrigin-RevId: 316786000
Change-Id: I8748ba453802ea86d586170bc2450ed970ed316c
---
 .../compiler/mlir/lite/flatbuffer_export.cc   | 65 ++++++++++++++++---
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
index df84b028f63..a260670015a 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
@@ -1406,22 +1406,67 @@ BufferOffset Translator::BuildSparsityParameters(
       for (int j = 0; j < segments.size(); j++) {
         vector_segments[j] = segments[j].dyn_cast().getInt();
       }
-      auto array_segments =
-          tflite::CreateInt32Vector(builder_,
-                                    builder_.CreateVector(vector_segments))
-              .Union();
+      tflite::SparseIndexVector segments_type;
+      BufferOffset array_segments;
+      // The segment array is sorted.
+      // TODO(b/147449640): Clean this up with util functions.
+      int max_of_segments = vector_segments[segments.size() - 1];
+      if (max_of_segments <= UINT8_MAX) {
+        segments_type = tflite::SparseIndexVector_Uint8Vector;
+        std::vector uint8_vector(vector_segments.begin(),
+                                          vector_segments.end());
+        array_segments = tflite::CreateUint8Vector(
+                             builder_, builder_.CreateVector(uint8_vector))
+                             .Union();
+      } else if (max_of_segments <= UINT16_MAX) {
+        segments_type = tflite::SparseIndexVector_Uint16Vector;
+        std::vector uint16_vector(vector_segments.begin(),
+                                            vector_segments.end());
+        array_segments = tflite::CreateUint16Vector(
+                             builder_, builder_.CreateVector(uint16_vector))
+                             .Union();
+      } else {
+        segments_type = tflite::SparseIndexVector_Int32Vector;
+        array_segments = tflite::CreateInt32Vector(
+                             builder_, builder_.CreateVector(vector_segments))
+                             .Union();
+      }
+
       auto indices = dim_metadata.indices();
       std::vector vector_indices(indices.size(), 0);
+      int max_of_indices = 0;
       for (int j = 0; j < indices.size(); j++) {
         vector_indices[j] = indices[j].dyn_cast().getInt();
+        if (vector_indices[j] > max_of_indices) {
+          max_of_indices = vector_indices[j];
+        }
       }
-      auto array_indices = tflite::CreateInt32Vector(
-                               builder_, builder_.CreateVector(vector_indices))
-                               .Union();
+      tflite::SparseIndexVector indices_type;
+      BufferOffset array_indices;
+      if (max_of_indices <= UINT8_MAX) {
+        indices_type = tflite::SparseIndexVector_Uint8Vector;
+        std::vector uint8_vector(vector_indices.begin(),
+                                          vector_indices.end());
+        array_indices = tflite::CreateUint8Vector(
+                            builder_, builder_.CreateVector(uint8_vector))
+                            .Union();
+      } else if (max_of_indices <= UINT16_MAX) {
+        indices_type = tflite::SparseIndexVector_Uint16Vector;
+        std::vector uint16_vector(vector_indices.begin(),
+                                            vector_indices.end());
+        array_indices = tflite::CreateUint16Vector(
+                            builder_, builder_.CreateVector(uint16_vector))
+                            .Union();
+      } else {
+        indices_type = tflite::SparseIndexVector_Int32Vector;
+        array_indices = tflite::CreateInt32Vector(
+                            builder_, builder_.CreateVector(vector_indices))
+                            .Union();
+      }
+
       fb_dim_metadata[i] = tflite::CreateDimensionMetadata(
-          builder_, tflite::DimensionType_SPARSE_CSR, 0,
-          tflite::SparseIndexVector_Int32Vector, array_segments,
-          tflite::SparseIndexVector_Int32Vector, array_indices);
+          builder_, tflite::DimensionType_SPARSE_CSR, 0, segments_type,
+          array_segments, indices_type, array_indices);
     }
   }
 

From 6ffd6820ff844917e5f3d673695158396806d1be Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Wed, 17 Jun 2020 07:13:24 +0700
Subject: [PATCH 0353/1390] Add dependecies

---
 tensorflow/c/env.cc | 7 ++++---
 tensorflow/c/env.h  | 2 ++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc
index 3d490d95e66..43879a18359 100644
--- a/tensorflow/c/env.cc
+++ b/tensorflow/c/env.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/c/c_api_internal.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/types.h"
 
 struct TF_StringStream {
@@ -148,9 +149,9 @@ TF_StringStream* TF_GetLocalTempDirectories() {
 
 void TF_GetTempFileName(const char* extension, std::string* name,
                         TF_Status* status) {
-  *name = ::tensorflow::Env::Default()->GetTempFilename(extension);
-  if (*name.length() == 0) {
-    TF_SetStatus(status, TF_INTERNAL, "Can not create temp file name");
+  *name = ::tensorflow::io::GetTempFilename(extension);
+  if (name->length() == 0) {
+    TF_SetStatus(status, TF_INTERNAL, "Can not get temp file name");
   } else {
     TF_SetStatus(status, TF_OK, "");
   }
diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h
index b50d0fdec03..273a3b5e142 100644
--- a/tensorflow/c/env.h
+++ b/tensorflow/c/env.h
@@ -20,6 +20,8 @@ limitations under the License.
 #include 
 #include 
 
+#include 
+
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/tf_file_statistics.h"
 

From 01cfc8a8a3d6176b1f028886087de4eaaa64ce2f Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 17:07:23 -0700
Subject: [PATCH 0354/1390] Resize converted to new style.

PiperOrigin-RevId: 316787130
Change-Id: I67db63fa6eaec2bccc87031f2e202da65a2ce439
---
 .../lite/delegates/gpu/cl/kernels/resize.cc   | 287 +++++++++---------
 1 file changed, 147 insertions(+), 140 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/resize.cc b/tensorflow/lite/delegates/gpu/cl/kernels/resize.cc
index 5d578fe6e09..6aa2d1d2570 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/resize.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/resize.cc
@@ -25,168 +25,166 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetResizeCode(
-    const OperationDef& op_def, SamplingType sampling_type,
-    bool half_pixel_centers,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+std::string GetResizeCode(const OperationDef& op_def,
+                          SamplingType sampling_type, bool half_pixel_centers,
+                          Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    src_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    dst_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
+  args->AddInt("border_x");
+  args->AddInt("border_y");
+  args->AddFloat("scale_factor_x");
+  args->AddFloat("scale_factor_y");
 
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,         \n";
-  c += "    int4 dst_size,         \n";
-  c += "    int2 border,           \n";
-  c += "    float2 scale_factor    \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
   if (op_def.IsBatchSupported()) {
     c += "  int linear_id = get_global_id(0);\n";
-    c += "  int X = linear_id / dst_size.w;\n";
-    c += "  int B = linear_id % dst_size.w;\n";
-    c += "  if (get_global_id(0) >= dst_size.x || Y >= dst_size.y || Z >= "
-         "dst_size.z) return;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  if (linear_id >= args.dst_tensor.Width() || Y >= "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Slices()) return;\n";
   } else {
     c += "  int X = get_global_id(0);\n";
-    c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) "
-         "return;\n";
+    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() "
+         "|| Z >= args.dst_tensor.Slices()) return;\n";
   }
   if (sampling_type == SamplingType::NEAREST) {
-    c += "  int2 coord = (int2)(X * scale_factor.x, Y * scale_factor.y);\n";
+    c += "  int2 coord = (int2)(X * args.scale_factor_x, Y * "
+         "args.scale_factor_y);\n";
     if (op_def.IsBatchSupported()) {
-      c += "  coord.x = coord.x * src_size.w + B;\n";
-      c += "  X = X * src_size.w + B;\n";
+      c += "  coord.x = coord.x * args.src_tensor.Batch() + B;\n";
+      c += "  X = X * args.src_tensor.Batch() + B;\n";
     }
-    c += "  FLT4 r0 = " + src_tensor.ReadWHS("coord.x", "coord.y", "Z") + ";\n";
+    c += "  FLT4 r0 = args.src_tensor.Read(coord.x, coord.y, Z);\n";
   } else {
     if (half_pixel_centers) {
-      c += "  float2 f_coords = ((float2)(X, Y) + 0.5f) * scale_factor - "
+      c += "  float2 f_coords = ((float2)(X, Y) + 0.5f) * "
+           "(float2)(args.scale_factor_x, args.scale_factor_y) - "
            "0.5f;\n";
     } else {
-      c += "  float2 f_coords = (float2)(X, Y) * scale_factor;\n";
+      c += "  float2 f_coords = (float2)(X, Y) * (float2)(args.scale_factor_x, "
+           "args.scale_factor_y);\n";
     }
     c += "  float2 f_coords_floor = floor(f_coords);\n";
     c += "  int2 coords_floor = (int2)(f_coords_floor.x, f_coords_floor.y);\n";
     c += "  int4 st;\n";
     c += "  st.xy = max(coords_floor, (int2)(0, 0));\n";
-    c += "  st.zw = min(coords_floor + (int2)(1, 1), border);\n";
+    c += "  st.zw = min(coords_floor + (int2)(1, 1), (int2)(args.border_x, "
+         "args.border_y));\n";
     c += "  float2 t = f_coords - f_coords_floor;\n";
     if (op_def.IsBatchSupported()) {
-      c += "  st.x = st.x * src_size.w + B;\n";
-      c += "  st.z = st.z * src_size.w + B;\n";
-      c += "  X = X * src_size.w + B;\n";
+      c += "  st.x = st.x * args.src_tensor.Batch() + B;\n";
+      c += "  st.z = st.z * args.src_tensor.Batch() + B;\n";
+      c += "  X = X * args.src_tensor.Batch() + B;\n";
     }
-    c += "  float4 src0 = " + src_tensor.ReadAsFloatWHS("st.x", "st.y", "Z") +
-         ";\n";
-    c += "  float4 src1 = " + src_tensor.ReadAsFloatWHS("st.z", "st.y", "Z") +
-         ";\n";
-    c += "  float4 src2 = " + src_tensor.ReadAsFloatWHS("st.x", "st.w", "Z") +
-         ";\n";
-    c += "  float4 src3 = " + src_tensor.ReadAsFloatWHS("st.z", "st.w", "Z") +
-         ";\n";
+    c += "  float4 src0 = args.src_tensor.Read(st.x, st.y, Z);\n";
+    c += "  float4 src1 = args.src_tensor.Read(st.z, st.y, Z);\n";
+    c += "  float4 src2 = args.src_tensor.Read(st.x, st.w, Z);\n";
+    c += "  float4 src3 = args.src_tensor.Read(st.z, st.w, Z);\n";
     c += "  FLT4 r0 = TO_FLT4(mix(mix(src0, src1, t.x), mix(src2, src3, t.x), "
          "t.y));\n";
   }
-  const LinkingContext context{"r0", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("r0", "X", "Y", "Z");
+  c += "  args.dst_tensor.Write(r0, X, Y, Z);\n";
   c += "}\n";
   return c;
 }
 
-std::string GetResize3DCode(
-    const OperationDef& op_def, SamplingType sampling_type,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHDSPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
+std::string GetResize3DCode(const OperationDef& op_def,
+                            SamplingType sampling_type, Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    src_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    dst_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
+  args->AddInt("border_x");
+  args->AddInt("border_y");
+  args->AddInt("border_z");
+  args->AddFloat("scale_factor_x");
+  args->AddFloat("scale_factor_y");
+  args->AddFloat("scale_factor_z");
 
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,         \n";
-  c += "    int4 dst_size,         \n";
-  if (op_def.IsBatchSupported()) {
-    c += "    int batch_size,      \n";
-  }
-  c += "    int4 border,           \n";
-  c += "    float4 scale_factor    \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int Y = get_global_id(1);\n";
   c += "  int linear_id_z = get_global_id(2);\n";
-  c += "  int S = linear_id_z % dst_size.w;\n";
-  c += "  int Z = linear_id_z / dst_size.w;\n";
+  c += "  int S = linear_id_z % args.dst_tensor.Slices();\n";
+  c += "  int Z = linear_id_z / args.dst_tensor.Slices();\n";
   if (op_def.IsBatchSupported()) {
     c += "  int linear_id = get_global_id(0);\n";
-    c += "  int X = linear_id / batch_size;\n";
-    c += "  int B = linear_id % batch_size;\n";
-    c += "  if (linear_id >= dst_size.x || Y >= dst_size.y || Z >= "
-         "dst_size.z) return;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  if (linear_id >= args.dst_tensor.Width() || Y >= "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Depth()) return;\n";
   } else {
     c += "  int X = get_global_id(0);\n";
-    c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) "
-         "return;\n";
+    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() "
+         "|| Z >= args.dst_tensor.Depth()) return;\n";
   }
   if (sampling_type == SamplingType::NEAREST) {
-    c += "  int4 coord = (int4)(X * scale_factor.x, Y * scale_factor.y, Z * "
-         "scale_factor.z, 0);\n";
+    c += "  int4 coord = (int4)(X * args.scale_factor_x, Y * "
+         "args.scale_factor_y, Z * "
+         "args.scale_factor_z, 0);\n";
     if (op_def.IsBatchSupported()) {
-      c += "  coord.x = coord.x * batch_size + B;\n";
-      c += "  X = X * batch_size + B;\n";
+      c += "  coord.x = coord.x * args.src_tensor.Batch() + B;\n";
+      c += "  X = X * args.src_tensor.Batch() + B;\n";
     }
-    c += "  FLT4 r0 = " +
-         src_tensor.ReadWHDS("coord.x", "coord.y", "coord.z", "S") + ";\n";
+    c += "  FLT4 r0 = args.src_tensor.Read(coord.x, coord.y, coord.z, S);\n";
   } else {
-    c += "  float4 f_coords = (float4)(X, Y, Z, 0) * scale_factor;\n";
+    c += "  float4 f_coords;\n";
+    c += "  f_coords.x = (float)(X) * args.scale_factor_x;\n";
+    c += "  f_coords.y = (float)(Y) * args.scale_factor_y;\n";
+    c += "  f_coords.z = (float)(Z) * args.scale_factor_z;\n";
     c += "  int4 start = (int4)(f_coords.x, f_coords.y, f_coords.z, 0);\n";
-    c += "  int4 end = min(start + (int4)(1, 1, 1, 0), border);\n";
+    c += "  int4 end;\n";
+    c += "  end.x = min(start.x + 1, args.border_x);\n";
+    c += "  end.y = min(start.y + 1, args.border_y);\n";
+    c += "  end.z = min(start.z + 1, args.border_z);\n";
     c += "  float4 t = f_coords - (float4)(start.x, start.y, start.z, 0.0f);\n";
     if (op_def.IsBatchSupported()) {
-      c += "  start.x = start.x * batch_size + B;\n";
-      c += "  end.x = end.x * batch_size + B;\n";
-      c += "  X = X * batch_size + B;\n";
+      c += "  start.x = start.x * args.src_tensor.Batch() + B;\n";
+      c += "  end.x = end.x * args.src_tensor.Batch() + B;\n";
+      c += "  X = X * args.src_tensor.Batch() + B;\n";
     }
-    c += "  float4 src0 = " +
-         src_tensor.ReadAsFloatWHDS("start.x", "start.y", "start.z", "S") +
-         ";\n";
-    c += "  float4 src1 = " +
-         src_tensor.ReadAsFloatWHDS("end.x", "start.y", "start.z", "S") + ";\n";
-    c += "  float4 src2 = " +
-         src_tensor.ReadAsFloatWHDS("start.x", "end.y", "start.z", "S") + ";\n";
-    c += "  float4 src3 = " +
-         src_tensor.ReadAsFloatWHDS("end.x", "end.y", "start.z", "S") + ";\n";
-    c += "  float4 src4 = " +
-         src_tensor.ReadAsFloatWHDS("start.x", "start.y", "end.z", "S") + ";\n";
-    c += "  float4 src5 = " +
-         src_tensor.ReadAsFloatWHDS("end.x", "start.y", "end.z", "S") + ";\n";
-    c += "  float4 src6 = " +
-         src_tensor.ReadAsFloatWHDS("start.x", "end.y", "end.z", "S") + ";\n";
-    c += "  float4 src7 = " +
-         src_tensor.ReadAsFloatWHDS("end.x", "end.y", "end.z", "S") + ";\n";
+    c += "  float4 src0 = args.src_tensor.Read(start.x, start.y, "
+         "start.z, S);\n";
+    c += "  float4 src1 = args.src_tensor.Read(end.x, start.y, start.z, "
+         "S);\n";
+    c += "  float4 src2 = args.src_tensor.Read(start.x, end.y, start.z, "
+         "S);\n";
+    c += "  float4 src3 = args.src_tensor.Read(end.x, end.y, start.z, "
+         "S);\n";
+    c += "  float4 src4 = args.src_tensor.Read(start.x, start.y, end.z, "
+         "S);\n";
+    c += "  float4 src5 = args.src_tensor.Read(end.x, start.y, end.z, "
+         "S);\n";
+    c += "  float4 src6 = args.src_tensor.Read(start.x, end.y, end.z, "
+         "S);\n";
+    c += "  float4 src7 = args.src_tensor.Read(end.x, end.y, end.z, "
+         "S);\n";
     c +=
         "  float4 t0 = mix(mix(src0, src1, t.x), mix(src2, src3, t.x), t.y);\n";
     c +=
         "  float4 t1 = mix(mix(src4, src5, t.x), mix(src6, src7, t.x), t.y);\n";
     c += "  FLT4 r0 = TO_FLT4(mix(t0, t1, t.z));\n";
   }
-  const LinkingContext context{"r0", "X", "Y", "S"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHDS("r0", "X", "Y", "Z", "S");
+  c += "  args.dst_tensor.Write(r0, X, Y, Z, S);\n";
   c += "}\n";
   return c;
 }
@@ -210,27 +208,32 @@ Resize& Resize::operator=(Resize&& operation) {
 }
 
 absl::Status Resize::Compile(const CreationContext& creation_context) {
-  const auto code = GetResizeCode(definition_, attr_.type,
-                                  attr_.half_pixel_centers, linked_operations_);
+  std::string code =
+      GetResizeCode(definition_, attr_.type, attr_.half_pixel_centers, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Resize::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(src_[0]->Width() - 1, src_[0]->Height() - 1)));
-  float2 scale_factor =
-      float2(CalculateResizeScale(src_[0]->Width(), dst_[0]->Width(), attr_),
-             CalculateResizeScale(src_[0]->Height(), dst_[0]->Height(), attr_));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(scale_factor));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(args_.SetInt("border_x", src_[0]->Width() - 1));
+  RETURN_IF_ERROR(args_.SetInt("border_y", src_[0]->Height() - 1));
+  RETURN_IF_ERROR(args_.SetFloat(
+      "scale_factor_x",
+      CalculateResizeScale(src_[0]->Width(), dst_[0]->Width(), attr_)));
+  RETURN_IF_ERROR(args_.SetFloat(
+      "scale_factor_y",
+      CalculateResizeScale(src_[0]->Height(), dst_[0]->Height(), attr_)));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Resize::GetGridSize() const {
@@ -272,31 +275,35 @@ Resize3D& Resize3D::operator=(Resize3D&& operation) {
 }
 
 absl::Status Resize3D::Compile(const CreationContext& creation_context) {
-  const auto code =
-      GetResize3DCode(definition_, attr_.type, linked_operations_);
+  std::string code = GetResize3DCode(definition_, attr_.type, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Resize3D::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHDS()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHDS()));
-  if (definition_.IsBatchSupported()) {
-    RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->Batch()));
-  }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(int4(
-      src_[0]->Width() - 1, src_[0]->Height() - 1, src_[0]->Depth() - 1, 0)));
-  float4 scale_factor = float4(
-      CalculateResizeScale(src_[0]->Width(), dst_[0]->Width(), attr_),
-      CalculateResizeScale(src_[0]->Height(), dst_[0]->Height(), attr_),
-      CalculateResizeScale(src_[0]->Depth(), dst_[0]->Depth(), attr_), 1.0f);
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(scale_factor));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(args_.SetInt("border_x", src_[0]->Width() - 1));
+  RETURN_IF_ERROR(args_.SetInt("border_y", src_[0]->Height() - 1));
+  RETURN_IF_ERROR(args_.SetInt("border_z", src_[0]->Depth() - 1));
+  RETURN_IF_ERROR(args_.SetFloat(
+      "scale_factor_x",
+      CalculateResizeScale(src_[0]->Width(), dst_[0]->Width(), attr_)));
+  RETURN_IF_ERROR(args_.SetFloat(
+      "scale_factor_y",
+      CalculateResizeScale(src_[0]->Height(), dst_[0]->Height(), attr_)));
+  RETURN_IF_ERROR(args_.SetFloat(
+      "scale_factor_z",
+      CalculateResizeScale(src_[0]->Depth(), dst_[0]->Depth(), attr_)));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Resize3D::GetGridSize() const {

From 7873e14cf9900b4c01d3b6f06e36f843dbb7f05b Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Tue, 16 Jun 2020 17:16:15 -0700
Subject: [PATCH 0355/1390] [tf.data] Add note about performance cost of
 `Dataset.unbatch`.

PiperOrigin-RevId: 316788559
Change-Id: I87d0d8a2be0a6d6751baa838ea8acc1eb9ee8d9a
---
 tensorflow/python/data/ops/dataset_ops.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
index 013447dfd8d..586b82e9ca6 100644
--- a/tensorflow/python/data/ops/dataset_ops.py
+++ b/tensorflow/python/data/ops/dataset_ops.py
@@ -2099,6 +2099,10 @@ name=None))
     >>> list(dataset.as_numpy_iterator())
     [1, 2, 3, 1, 2, 1, 2, 3, 4]
 
+    Note: `unbatch` requires a data copy to slice up the batched tensor into
+    smaller, unbatched tensors. When optimizing performance, try to avoid
+    unnecessary usage of `unbatch`.
+
     Returns:
       A `Dataset`.
     """

From cdb6e80b21997c2b24336eb524134198fa6754d8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 17:25:15 -0700
Subject: [PATCH 0356/1390] Qualify uses of std::string

PiperOrigin-RevId: 316789814
Change-Id: Ice83a74e70122008e090af3b818b9920abf7f5bc
---
 .../toco/tensorflow_graph_matching/cluster.h  | 14 +++----
 .../cluster_utils.cc                          |  4 +-
 .../resolve_cluster.cc                        | 13 +++---
 .../resolve_cluster.h                         |  4 +-
 .../tensorflow_graph_matching/resolve_svdf.cc | 40 +++++++++----------
 .../tensorflow_graph_matching/resolve_svdf.h  |  6 +--
 .../resolve_svdf_test.cc                      | 24 +++++------
 7 files changed, 53 insertions(+), 52 deletions(-)

diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/cluster.h b/tensorflow/lite/toco/tensorflow_graph_matching/cluster.h
index af268ddd370..7dc79f17a6b 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/cluster.h
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/cluster.h
@@ -47,7 +47,7 @@ class Cluster {
   // 2- All the nodes in GraphDef which belong to this cluster.
   void SetGraphDefInfo(const tensorflow::GraphDef* graph_def);
 
-  const string& GetName() const { return name_; }
+  const std::string& GetName() const { return name_; }
 
   const std::vector>& GetNewNodes() const {
     return new_nodes_;
@@ -55,18 +55,18 @@ class Cluster {
 
   const std::vector& GetNodes() { return nodes_; }
 
-  void SetName(const string& name) { name_ = name; }
+  void SetName(const std::string& name) { name_ = name; }
 
-  void SetDevice(const string& device) { device_ = device; }
+  void SetDevice(const std::string& device) { device_ = device; }
 
   // Find the input(s) and output(s) of this Cluster.
   bool FindClusterInputsAndOutputs();
 
  protected:
-  string name_;
-  string device_;
-  std::vector inputs_;
-  std::vector outputs_;
+  std::string name_;
+  std::string device_;
+  std::vector inputs_;
+  std::vector outputs_;
 
   // Used to hold the pointers to nodes which are in this cluster. These nodes
   // are pointing to the nodes in graph_def_.
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/cluster_utils.cc b/tensorflow/lite/toco/tensorflow_graph_matching/cluster_utils.cc
index 8a010ef8208..fb12e62d9af 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/cluster_utils.cc
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/cluster_utils.cc
@@ -16,8 +16,8 @@ limitations under the License.
 #include "tensorflow/lite/toco/toco_types.h"
 namespace toco {
 
-bool StrContains(const string& x, const string& search_pattern) {
-  return x.find(search_pattern) != string::npos;
+bool StrContains(const std::string& x, const std::string& search_pattern) {
+  return x.find(search_pattern) != std::string::npos;
 }
 
 void Transpose2DTensor(const float* tensor, int row, int col,
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.cc b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.cc
index 7a187512078..f2645e89511 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.cc
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.cc
@@ -33,7 +33,8 @@ using tensorflow::GraphDef;
 using tensorflow::NodeDef;
 
 void AddNodeToGraph(const NodeDef& node,
-                    const std::vector& cluster_names, GraphDef* graph) {
+                    const std::vector& cluster_names,
+                    GraphDef* graph) {
   NodeDef* new_node = graph->add_node();
   new_node->set_op(node.op());
   new_node->set_name(node.name());
@@ -41,9 +42,9 @@ void AddNodeToGraph(const NodeDef& node,
   // If the inputs are coming from a node which belongs to another cluster, then
   // those inputs are renamed to the source cluster name. Otherwise the original
   // input name is used.
-  for (const string& node_input : node.input()) {
+  for (const std::string& node_input : node.input()) {
     bool input_from_cluster = false;
-    for (const string& cluster_name : cluster_names) {
+    for (const std::string& cluster_name : cluster_names) {
       if (StrContains(node_input, cluster_name) &&
           !StrContains(node.name(), cluster_name)) {
         new_node->add_input(cluster_name);
@@ -62,7 +63,7 @@ void AddNodeToGraph(const NodeDef& node,
 
 bool FindCluster(const ClusterFactoryInterface& cluster_factory,
                  const GraphDef& graph_def,
-                 std::unordered_map* is_node_in_cluster,
+                 std::unordered_map* is_node_in_cluster,
                  std::vector>* clusters) {
   for (const NodeDef& node : graph_def.node()) {
     // If the node is not assigned to any cluster, then we check if it belong to
@@ -90,12 +91,12 @@ std::unique_ptr MaybeResolveClusters(
   std::unique_ptr pruned_graph(new GraphDef);
   // The structure to keep track of which cluster each node is assigned to, and
   // to initialize them to all un-assigned,
-  std::unordered_map is_node_in_cluster;
+  std::unordered_map is_node_in_cluster;
   for (const NodeDef& node : graph_def.node()) {
     is_node_in_cluster[node.name()] = false;
   }
 
-  std::vector cluster_names;
+  std::vector cluster_names;
   std::vector> all_clusters;
   // Find the clusters for all available cluster factories.
   for (const ClusterFactoryInterface* cluster_factory : cluster_factories) {
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.h b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.h
index d7afcced7b7..5215114fccf 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.h
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_cluster.h
@@ -40,7 +40,7 @@ std::unique_ptr MaybeResolveClusters(
 // belongs to another cluster, then those inputs are renamed to the source
 // cluster name.
 void AddNodeToGraph(const tensorflow::NodeDef& node,
-                    const std::vector& cluster_names,
+                    const std::vector& cluster_names,
                     tensorflow::GraphDef* graph);
 
 // Given a graph and a cluster class, it finds all the nodes which belong to a
@@ -49,7 +49,7 @@ void AddNodeToGraph(const tensorflow::NodeDef& node,
 // they belong to the generated clusters.
 bool FindCluster(const ClusterFactoryInterface& cluster_factory,
                  const tensorflow::GraphDef& graph_def,
-                 std::unordered_map* is_node_in_cluster,
+                 std::unordered_map* is_node_in_cluster,
                  std::vector>* clusters);
 
 // Receives a graph and generates another graph by replacing the cluster of
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.cc b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.cc
index 2f9f9a8c9b0..7d83a9dbfed 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.cc
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.cc
@@ -47,11 +47,11 @@ namespace {
 // Since these nodes are connected to a Concatenate node, it makes sure the
 // axis value input of the Concatenate operator is 0.
 void FilterPartitionedConstNodes(
-    const string& const_pattern,
+    const std::string& const_pattern,
     const std::vector& cluster_nodes,
     std::vector* const_node_parts) {
   for (const NodeDef* node : cluster_nodes) {
-    string node_name_to_upper = node->name();
+    std::string node_name_to_upper = node->name();
     std::transform(node_name_to_upper.begin(), node_name_to_upper.end(),
                    node_name_to_upper.begin(), ::toupper);
     if (StrContains(node->name(), const_pattern) && node->op() == "Const") {
@@ -97,7 +97,7 @@ int SvdfCluster::InferFilterRank() {
 }
 
 void SvdfCluster::CreateNodes() {
-  for (const string& const_pattern : const_node_patterns_) {
+  for (const std::string& const_pattern : const_node_patterns_) {
     CreateConstNode(const_pattern);
   }
   std::unique_ptr svdf_node(new NodeDef);
@@ -110,14 +110,14 @@ void SvdfCluster::CreateNodes() {
 
   // Add the rest of the inputs to Svdf cell: weights and bias.
   CHECK(new_nodes_.size() == 3 || new_nodes_.size() == 2);
-  string* weights_feature_input = svdf_node->add_input();
-  string* weights_time_input = svdf_node->add_input();
-  string* bias_input;
+  std::string* weights_feature_input = svdf_node->add_input();
+  std::string* weights_time_input = svdf_node->add_input();
+  std::string* bias_input;
   if (new_nodes_.size() == 3) {
     bias_input = svdf_node->add_input();
   }
   for (const std::unique_ptr& node : new_nodes_) {
-    const string node_name = node->name();
+    const std::string node_name = node->name();
     if (StrContains(node_name, "SVDF_weights_feature")) {
       *weights_feature_input = node_name;
     } else if (StrContains(node_name, "SVDF_weights_time")) {
@@ -136,7 +136,7 @@ void SvdfCluster::CreateNodes() {
   CHECK_GT(rank, 0);
 
   // Add Svdf activation and rank.
-  string activation_function =
+  std::string activation_function =
       StrContains(outputs_[0], "Relu") ? "Relu" : "None";
   (*svdf_node->mutable_attr())["ActivationFunction"].set_s(activation_function);
   (*svdf_node->mutable_attr())["Rank"].set_i(rank);
@@ -145,7 +145,7 @@ void SvdfCluster::CreateNodes() {
   new_nodes_.push_back(std::move(svdf_node));
 }
 
-void SvdfCluster::CreateConstNode(const string& const_pattern) {
+void SvdfCluster::CreateConstNode(const std::string& const_pattern) {
   // Find the nodes with pattern like: "const_pattern"/part_xxx of type Const.
   std::vector const_node_parts;
   FilterPartitionedConstNodes(const_pattern, nodes_, &const_node_parts);
@@ -236,15 +236,15 @@ void SvdfCluster::MaybeMergeConstNodes(
 
     // Set the tensor attributes.
     allocated_tensor->set_tensor_content(
-        string(reinterpret_cast(transposed_tensor.get()),
-               allocated_content_flat_size));
+        std::string(reinterpret_cast(transposed_tensor.get()),
+                    allocated_content_flat_size));
   } else {
     tensor_shape_dim0->set_size(dim0_size);
 
     // Set the tensor attributes.
     allocated_tensor->set_tensor_content(
-        string(reinterpret_cast(allocated_content.get()),
-               allocated_content_flat_size));
+        std::string(reinterpret_cast(allocated_content.get()),
+                    allocated_content_flat_size));
   }
 }
 
@@ -252,21 +252,21 @@ void SvdfCluster::MaybeMergeConstNodes(
 
 std::unique_ptr SvdfClusterFactory::CreateCluster(
     const NodeDef& node, const GraphDef& graph_def) const {
-  std::vector node_patterns = {"SVDF_weights_feature",
-                                       "SVDF_weights_time", "SVDF_bias"};
+  std::vector node_patterns = {"SVDF_weights_feature",
+                                            "SVDF_weights_time", "SVDF_bias"};
 
-  string node_name_to_upper = node.name();
+  std::string node_name_to_upper = node.name();
   std::transform(node_name_to_upper.begin(), node_name_to_upper.end(),
                  node_name_to_upper.begin(), ::toupper);
   std::unique_ptr cluster = nullptr;
-  if (node_name_to_upper.find("SVDF", 0) != string::npos) {
+  if (node_name_to_upper.find("SVDF", 0) != std::string::npos) {
     size_t weights_pos = node.name().find(node_patterns[0]);
-    if (weights_pos != string::npos) {
+    if (weights_pos != std::string::npos) {
       // Assuming the node name has a pattern like:
       // "SOMESTRING1/CELLNAME/SEARCH_PATTERN/SOMESTRING2", we use
       // CELLNAME as the cluster name.
       size_t cell_pos = node.name().rfind("/", weights_pos - 2) + 1;
-      string cell_name =
+      std::string cell_name =
           node.name().substr(cell_pos, weights_pos - cell_pos - 1);
       cluster = std::unique_ptr(new SvdfCluster);
       cluster->SetName(cell_name);
@@ -274,7 +274,7 @@ std::unique_ptr SvdfClusterFactory::CreateCluster(
       cluster->SetGraphDefInfo(&graph_def);
       CHECK(cluster->FindClusterInputsAndOutputs());
 
-      for (const string& const_pattern : node_patterns) {
+      for (const std::string& const_pattern : node_patterns) {
         cluster->AddConstNodePattern(const_pattern);
       }
     }
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.h b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.h
index 649cadfa066..b5843016299 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.h
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf.h
@@ -36,7 +36,7 @@ class SvdfCluster : public Cluster {
 
   // A helper function to set the pattern of Const nodes which CreateNodes()
   // should handle specially.
-  void AddConstNodePattern(const string& const_pattern) {
+  void AddConstNodePattern(const std::string& const_pattern) {
     const_node_patterns_.push_back(const_pattern);
   }
 
@@ -46,7 +46,7 @@ class SvdfCluster : public Cluster {
   // The main function which is used to create Const nodes for this cluster.
   // These Const nodes are the inputs to the composite op generated for this
   // cluster.
-  void CreateConstNode(const string& const_pattern);
+  void CreateConstNode(const std::string& const_pattern);
 
   // Receives a vector of Const nodes, merge them (if necessary) and returns
   // only one Const node holding all the arrays contents. It transposes it if
@@ -61,7 +61,7 @@ class SvdfCluster : public Cluster {
   // shape to [num_units, rank, batch] shape. The 2nd shape element is rank.
   int InferFilterRank();
 
-  std::vector const_node_patterns_;
+  std::vector const_node_patterns_;
 };
 
 class SvdfClusterFactory : public ClusterFactoryInterface {
diff --git a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc
index f66b59ccce6..9828b0050b6 100644
--- a/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc
+++ b/tensorflow/lite/toco/tensorflow_graph_matching/resolve_svdf_test.cc
@@ -77,8 +77,8 @@ class ResolveSvdfTest : public ::testing::Test {
   ~ResolveSvdfTest() override {}
 
  protected:
-  void AddNewNode(const string& name, const string& op,
-                  const std::vector& inputs) {
+  void AddNewNode(const std::string& name, const std::string& op,
+                  const std::vector& inputs) {
     NodeDef* node = graph_.add_node();
     node->set_name(name);
     node->set_op(op);
@@ -89,8 +89,8 @@ class ResolveSvdfTest : public ::testing::Test {
     }
   }
 
-  void AddNewNode(const string& name, const string& op,
-                  const std::vector& inputs,
+  void AddNewNode(const std::string& name, const std::string& op,
+                  const std::vector& inputs,
                   const std::vector& values) {
     NodeDef* node = graph_.add_node();
     node->set_name(name);
@@ -109,12 +109,12 @@ class ResolveSvdfTest : public ::testing::Test {
     tensor_shape_dim0->set_size(values.size());
     allocated_tensor->set_allocated_tensor_shape(allocated_tensor_shape);
     allocated_tensor->set_tensor_content(
-        string(reinterpret_cast(values.data()),
-               values.size() * sizeof(float)));
+        std::string(reinterpret_cast(values.data()),
+                    values.size() * sizeof(float)));
     (*node->mutable_attr())["value"].set_allocated_tensor(allocated_tensor);
   }
 
-  void AddShapeNode(const string& name, const std::vector& values) {
+  void AddShapeNode(const std::string& name, const std::vector& values) {
     NodeDef* node = graph_.add_node();
     node->set_name(name);
     node->set_op("Const");
@@ -128,8 +128,8 @@ class ResolveSvdfTest : public ::testing::Test {
     tensor_shape_dim0->set_size(values.size());
     allocated_tensor->set_allocated_tensor_shape(allocated_tensor_shape);
     allocated_tensor->set_tensor_content(
-        string(reinterpret_cast(values.data()),
-               values.size() * sizeof(int)));
+        std::string(reinterpret_cast(values.data()),
+                    values.size() * sizeof(int)));
     (*node->mutable_attr())["value"].set_allocated_tensor(allocated_tensor);
   }
 
@@ -157,12 +157,12 @@ TEST_F(ResolveSvdfTest, TestTranspose2DTensor) {
 }
 
 TEST_F(ResolveSvdfTest, TestResolveSvdfFlow) {
-  std::unordered_map is_node_in_cluster;
+  std::unordered_map is_node_in_cluster;
   for (const NodeDef& node : graph_.node()) {
     is_node_in_cluster[node.name()] = false;
   }
 
-  std::vector cluster_names;
+  std::vector cluster_names;
   CHECK(FindCluster(svdf_cluster_factory_, graph_, &is_node_in_cluster,
                     &clusters_));
 
@@ -174,7 +174,7 @@ TEST_F(ResolveSvdfTest, TestResolveSvdfFlow) {
   EXPECT_THAT(cluster_names,
               testing::UnorderedElementsAreArray({"Svdf1", "Svdf2"}));
 
-  std::vector new_node_names;
+  std::vector new_node_names;
   std::vector content_array(3);
   for (const std::unique_ptr& cluster : clusters_) {
     // After CreateNodes in each cluster we have three nodes: Svdf,

From 81cff8fc905be6d5e0489e454f9d8362e391dcad Mon Sep 17 00:00:00 2001
From: Tiezhen WANG 
Date: Tue, 16 Jun 2020 17:30:30 -0700
Subject: [PATCH 0357/1390] TFLM: Rename FinishTensorAllocation to
 FinishModelAllocation in the comments.

The name of this function has been changed.

PiperOrigin-RevId: 316790532
Change-Id: I31a77d3c2e7fe5af3138c99317ccd075ecbdf4ff
---
 tensorflow/lite/micro/micro_allocator.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/micro/micro_allocator.h b/tensorflow/lite/micro/micro_allocator.h
index 09b6567ac30..ab3f2a44d18 100644
--- a/tensorflow/lite/micro/micro_allocator.h
+++ b/tensorflow/lite/micro/micro_allocator.h
@@ -40,7 +40,7 @@ TfLiteStatus InitializeTfLiteTensorFromFlatbuffer(
 
 // A handle tracking scratch buffer allocation. This handle is created by
 // `RequestScratchBufferInArena`. `data` field is populated in
-// `FinishTensorAllocation` after static memory planning.
+// `FinishModelAllocation` after static memory planning.
 // TODO(b/150257460) As a future optimization, this struct could be replaced by
 // a union, since once `data` is populated, `bytes` and `node_idx` is not
 // needed.
@@ -126,7 +126,7 @@ class MicroAllocator {
 
   // Register a scratch buffer of size `bytes` for Node with `node_id`.
   // This method only allocates a BufferHandle holding information for memory
-  // planning. The buffer ptr is ready after `FinishTensorAllocation` and can
+  // planning. The buffer ptr is ready after `FinishModelAllocation` and can
   // be retrieved by `GetScratchBuffer` method using the returned buffer_idx.
   // Note that there should be no tail allocation between two consecutive
   // `RequestScratchBufferInArena` calls.
@@ -136,7 +136,7 @@ class MicroAllocator {
   void* GetScratchBuffer(int buffer_idx) const;
 
   // Returns the arena usage in bytes, only available after
-  // `FinishTensorAllocation`. Otherwise, it will return 0.
+  // `FinishModelAllocation`. Otherwise, it will return 0.
   size_t used_bytes() const;
 
  protected:

From 451352fd4031c9ef7be9ffae49ab59132107111e Mon Sep 17 00:00:00 2001
From: Brian Zhao 
Date: Tue, 16 Jun 2020 17:38:15 -0700
Subject: [PATCH 0358/1390] Adding convenience functions for calling ops needed
 for implementing the SavedModel C API. This change starts by adding functions
 to create and destroy resource variables.

PiperOrigin-RevId: 316791867
Change-Id: Ieba37cabcc0200e48fc0b64c980c82c2ee476e1c
---
 .../c/experimental/saved_model/core/ops/BUILD | 94 +++++++++++++++++++
 .../core/ops/owned_eager_context.h            | 54 +++++++++++
 .../saved_model/core/ops/owned_eager_op.h     | 42 +++++++++
 .../saved_model/core/ops/owned_tensor.h       | 42 +++++++++
 .../core/ops/owned_tensor_handle.h            | 54 +++++++++++
 .../saved_model/core/ops/variable_ops.cc      | 74 +++++++++++++++
 .../saved_model/core/ops/variable_ops.h       | 46 +++++++++
 .../saved_model/core/ops/variable_ops_test.cc | 77 +++++++++++++++
 8 files changed, 483 insertions(+)
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/BUILD
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/owned_tensor.h
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
 create mode 100644 tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc

diff --git a/tensorflow/c/experimental/saved_model/core/ops/BUILD b/tensorflow/c/experimental/saved_model/core/ops/BUILD
new file mode 100644
index 00000000000..b42e93c3716
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/BUILD
@@ -0,0 +1,94 @@
+# This package contains written convenience helpers for Eager Operations
+# used by SavedModel. Once we autogenerate C++ Eager Op wrappers, we can remove these.
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+
+package(
+    default_visibility = [
+        # Restricting visibility for now
+        "//tensorflow/c/experimental/saved_model/core:__subpackages__",
+        "//tensorflow/c/experimental/saved_model/internal:__subpackages__",
+    ],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+cc_library(
+    name = "owned_eager_op",
+    hdrs = [
+        "owned_eager_op.h",
+    ],
+    deps = [
+        "//tensorflow/c/eager:operation_interface",
+    ],
+)
+
+cc_library(
+    name = "owned_tensor_handle",
+    hdrs = [
+        "owned_tensor_handle.h",
+    ],
+    deps = [
+        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/core/common_runtime/eager:tensor_handle",
+    ],
+)
+
+cc_library(
+    name = "owned_eager_context",
+    hdrs = ["owned_eager_context.h"],
+    deps = [
+        "//tensorflow/c/eager:context_interface",
+        "//tensorflow/core/common_runtime/eager:context",
+    ],
+)
+
+cc_library(
+    name = "owned_tensor",
+    hdrs = ["owned_tensor.h"],
+    deps = [
+        "//tensorflow/c:tensor_interface",
+    ],
+)
+
+cc_library(
+    name = "variable_ops",
+    srcs = [
+        "variable_ops.cc",
+    ],
+    hdrs = [
+        "variable_ops.h",
+    ],
+    deps = [
+        ":owned_eager_op",
+        ":owned_tensor_handle",
+        "//tensorflow/c/eager:context_interface",
+        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+tf_cc_test(
+    name = "variable_ops_test",
+    srcs = [
+        "variable_ops_test.cc",
+    ],
+    deps = [
+        ":owned_eager_context",
+        ":owned_tensor",
+        ":owned_tensor_handle",
+        ":variable_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/common_runtime:core_cpu_lib",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/common_runtime/eager:core",
+    ],
+)
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
new file mode 100644
index 00000000000..300059cd069
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
@@ -0,0 +1,54 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_EAGER_CONTEXT_H_
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_EAGER_CONTEXT_H_
+
+#include 
+
+#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+
+namespace tensorflow {
+namespace internal {
+
+struct AbstractContextInterfaceDeleter {
+  void operator()(AbstractContextInterface* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+struct EagerContextDeleter {
+  void operator()(EagerContext* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+}  // namespace internal
+
+using AbstractContextPtr =
+    std::unique_ptr;
+
+using EagerContextPtr =
+    std::unique_ptr;
+
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_EAGER_CONTEXT_H_
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
new file mode 100644
index 00000000000..c6b21578820
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
@@ -0,0 +1,42 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_EAGER_OP_H_
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_EAGER_OP_H_
+
+#include 
+
+#include "tensorflow/c/eager/operation_interface.h"
+
+namespace tensorflow {
+namespace internal {
+
+struct AbstractOperationInterfaceDeleter {
+  void operator()(AbstractOperationInterface* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+}  // namespace internal
+
+using AbstractOpPtr =
+    std::unique_ptr;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_EAGER_OP_H_
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_tensor.h b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor.h
new file mode 100644
index 00000000000..335d9e46c7a
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor.h
@@ -0,0 +1,42 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_TENSOR_H_
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_TENSOR_H_
+
+#include 
+
+#include "tensorflow/c/tensor_interface.h"
+
+namespace tensorflow {
+namespace internal {
+
+struct AbstractTensorInterfaceDeleter {
+  void operator()(AbstractTensorInterface* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+}  // namespace internal
+
+using AbstractTensorPtr =
+    std::unique_ptr;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_OWNED_TENSOR_H_
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
new file mode 100644
index 00000000000..e98d6554afb
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
@@ -0,0 +1,54 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_TENSOR_HANDLE_H_
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_TENSOR_HANDLE_H_
+
+#include 
+
+#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+
+namespace tensorflow {
+namespace internal {
+
+struct TensorHandleDeleter {
+  void operator()(TensorHandle* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+struct AbstractTensorHandleDeleter {
+  void operator()(AbstractTensorHandleInterface* p) const {
+    if (p != nullptr) {
+      p->Release();
+    }
+  }
+};
+
+}  // namespace internal
+
+using TensorHandlePtr =
+    std::unique_ptr;
+
+using AbstractTensorHandlePtr =
+    std::unique_ptr;
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OWNED_TENSOR_HANDLE_H_
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
new file mode 100644
index 00000000000..94548e553ad
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
@@ -0,0 +1,74 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/c/experimental/saved_model/core/ops/variable_ops.h"
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h"
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace internal {
+
+static const char kNoSharingResourceID[] =
+    "cd2c89b7-88b7-44c8-ad83-06c2a9158347";
+
+Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
+                                           DataType dtype, TensorShape shape,
+                                           AbstractTensorHandlePtr* handle) {
+  AbstractOpPtr varhandle_op = AbstractOpPtr(ctx->CreateOperation());
+
+  TF_RETURN_IF_ERROR(varhandle_op->Reset("VarHandleOp", nullptr));
+  TF_RETURN_IF_ERROR(varhandle_op->SetAttrType("dtype", dtype));
+
+  // Note that if shape is unknown rank, shape.dim_sizes() will be empty, and
+  // shape.dims() will be -1.
+  gtl::InlinedVector dim_sizes = shape.dim_sizes();
+  TF_RETURN_IF_ERROR(varhandle_op->SetAttrShape(
+      "shape", reinterpret_cast(dim_sizes.data()),
+      shape.dims()));
+  TF_RETURN_IF_ERROR(varhandle_op->SetAttrString("container", "", 0));
+  TF_RETURN_IF_ERROR(varhandle_op->SetAttrString(
+      "shared_name", kNoSharingResourceID, strlen(kNoSharingResourceID)));
+
+  AbstractTensorHandleInterface* var_handle = nullptr;
+  int num_retvals = 1;
+  TF_RETURN_IF_ERROR(varhandle_op->Execute(
+      absl::MakeSpan(&var_handle, num_retvals), &num_retvals));
+  handle->reset(var_handle);
+  return Status();
+}
+
+Status DestroyResource(AbstractContextInterface* ctx,
+                       AbstractTensorHandleInterface* handle) {
+  AbstractOpPtr destroy_op = AbstractOpPtr(ctx->CreateOperation());
+  TF_RETURN_IF_ERROR(destroy_op->Reset("DestroyResourceOp", nullptr));
+  TF_RETURN_IF_ERROR(destroy_op->SetAttrBool("ignore_lookup_error", true));
+  TF_RETURN_IF_ERROR(destroy_op->AddInput(handle));
+
+  int num_retvals = 0;
+  TF_RETURN_IF_ERROR(destroy_op->Execute({}, &num_retvals));
+  return Status();
+}
+
+}  // namespace internal
+}  // namespace tensorflow
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
new file mode 100644
index 00000000000..1c4d757af8c
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
@@ -0,0 +1,46 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H
+
+#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace internal {
+
+// Executes a VarHandleOp using `ctx`, and fills `handle` with the DT_RESOURCE
+// TensorHandle associated with the variable. This is equivalent to creating an
+// unitialized TF2 tf.Variable.
+// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L1867-L1872
+Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
+                                           DataType dtype, TensorShape shape,
+                                           AbstractTensorHandlePtr* handle);
+
+// Executes DestroyResourceOp on `handle`, using `ctx`. This is equivalent to
+// the cleanup that occurs in a tf.Variable's EagerResourceDeleter:
+// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L289-L290
+Status DestroyResource(AbstractContextInterface* ctx,
+                       AbstractTensorHandleInterface* handle);
+
+}  // namespace internal
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc b/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
new file mode 100644
index 00000000000..7a9486f8ebd
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
@@ -0,0 +1,77 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/c/experimental/saved_model/core/ops/variable_ops.h"
+
+#include 
+
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h"
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor.h"
+#include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+class VariableOpsTest : public ::testing::Test {
+ public:
+  VariableOpsTest()
+      : device_mgr_(std::make_unique(DeviceFactory::NewDevice(
+            "CPU", {}, "/job:localhost/replica:0/task:0"))),
+        ctx_(new EagerContext(
+            SessionOptions(),
+            tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+            tensorflow::ContextMirroringPolicy::MIRRORING_NONE,
+            /* async= */ false,
+            /* lazy_copy_function_remote_inputs= */ false, device_mgr_.get(),
+            /* device_mgr_owned= */ false, /* rendezvous= */ nullptr,
+            /* custom_kernel_creator= */ nullptr,
+            /* cluster_flr= */ nullptr)) {}
+
+  EagerContext* context() { return ctx_.get(); }
+
+ private:
+  std::unique_ptr device_mgr_;
+  EagerContextPtr ctx_;
+};
+
+// Sanity check for variable creation
+TEST_F(VariableOpsTest, CreateVariableSuccessful) {
+  // Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
+  AbstractTensorHandlePtr handle;
+  TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
+      context(), DT_FLOAT, {}, &handle));
+  // The created TensorHandle should be a DT_Resource
+  EXPECT_EQ(handle->DataType(), DT_RESOURCE);
+}
+
+// Sanity check for variable destruction
+TEST_F(VariableOpsTest, DestroyVariableSuccessful) {
+  // Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
+  AbstractTensorHandlePtr handle;
+  TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
+      context(), DT_FLOAT, {}, &handle));
+
+  // Destroy the variable
+  TF_EXPECT_OK(internal::DestroyResource(context(), handle.get()));
+}
+
+}  // namespace
+}  // namespace tensorflow

From afb849aa36145b37c70f88211890f1b968dfb6da Mon Sep 17 00:00:00 2001
From: Mark Daoust 
Date: Tue, 16 Jun 2020 17:39:29 -0700
Subject: [PATCH 0359/1390] Fix list formatting on TextVectorization.

Paren-numbered lists are not recognized on tensorflow.org

PiperOrigin-RevId: 316792053
Change-Id: I72daa1bc6da92211904b77e4c5e057efc0435732
---
 .../layers/preprocessing/text_vectorization.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/keras/layers/preprocessing/text_vectorization.py b/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
index bff7969477c..97e3ac4a63c 100644
--- a/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
+++ b/tensorflow/python/keras/layers/preprocessing/text_vectorization.py
@@ -89,23 +89,25 @@ class TextVectorization(CombinerPreprocessingLayer):
   to create the vocabulary.
 
   The processing of each sample contains the following steps:
-    1) standardize each sample (usually lowercasing + punctuation stripping)
-    2) split each sample into substrings (usually words)
-    3) recombine substrings into tokens (usually ngrams)
-    4) index tokens (associate a unique int value with each token)
-    5) transform each sample using this index, either into a vector of ints or
+
+    1. standardize each sample (usually lowercasing + punctuation stripping)
+    2. split each sample into substrings (usually words)
+    3. recombine substrings into tokens (usually ngrams)
+    4. index tokens (associate a unique int value with each token)
+    5. transform each sample using this index, either into a vector of ints or
        a dense float vector.
 
   Some notes on passing Callables to customize splitting and normalization for
   this layer:
-    1) Any callable can be passed to this Layer, but if you want to serialize
+
+    1. Any callable can be passed to this Layer, but if you want to serialize
        this object you should only pass functions that are registered Keras
        serializables (see `tf.keras.utils.register_keras_serializable` for more
        details).
-    2) When using a custom callable for `standardize`, the data received
+    2. When using a custom callable for `standardize`, the data received
        by the callable will be exactly as passed to this layer. The callable
        should return a tensor of the same shape as the input.
-    3) When using a custom callable for `split`, the data received by the
+    3. When using a custom callable for `split`, the data received by the
        callable will have the 1st dimension squeezed out - instead of
        `[["string to split"], ["another string to split"]]`, the Callable will
        see `["string to split", "another string to split"]`. The callable should

From 2bab28aecc7f44f2c1a2207362bdfacc93059ca0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 17:49:15 -0700
Subject: [PATCH 0360/1390] Explicitly cast printed integer types to the
 primitive type so as to not implicitly rely on whether int64 is long or long
 long.

PiperOrigin-RevId: 316793291
Change-Id: I6eb74c5d7987fd7ece10e8eef26397e6dd58cf33
---
 .../common_runtime/direct_session_test.cc     |  2 +-
 tensorflow/core/debug/bfc_dump_reader.cc      | 57 +++++++++++--------
 .../core/kernels/data/experimental/io_ops.cc  |  3 +-
 .../experimental/map_and_batch_dataset_op.cc  |  8 ++-
 .../data/experimental/snapshot_util.cc        | 19 +++++--
 .../profiler/internal/gpu/device_tracer.cc    |  2 +-
 6 files changed, 56 insertions(+), 35 deletions(-)

diff --git a/tensorflow/core/common_runtime/direct_session_test.cc b/tensorflow/core/common_runtime/direct_session_test.cc
index 968713f8acd..eab508662e6 100644
--- a/tensorflow/core/common_runtime/direct_session_test.cc
+++ b/tensorflow/core/common_runtime/direct_session_test.cc
@@ -2822,7 +2822,7 @@ class StatefulOutputRequiredOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     // The op counts the number of outputs required in the current subgraph,
     // and emits that number on each of its required outputs.
-    Tensor count_outputs_required_t(0LL);
+    Tensor count_outputs_required_t(int64{0});
     int64& count_outputs_required = count_outputs_required_t.scalar()();
     for (int i = 0; i < num_outputs(); ++i) {
       if (ctx->output_required(i)) ++count_outputs_required;
diff --git a/tensorflow/core/debug/bfc_dump_reader.cc b/tensorflow/core/debug/bfc_dump_reader.cc
index f28eb6c428f..722315574e2 100644
--- a/tensorflow/core/debug/bfc_dump_reader.cc
+++ b/tensorflow/core/debug/bfc_dump_reader.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include 
 #include 
 
 #include "tensorflow/core/platform/env.h"
@@ -70,17 +71,18 @@ void PrintChunk(const MemChunk& mc, const uint64 ac_offset, bool freed_at,
   // A size class corresponding approximately to log base 100.
   int size_class = floor(0.5 * log10(static_cast(mc.size())));
   *cumulative_bytes += mc.size();
-  printf("  %c %d %p bin=%d bytes=%llu %3.1f%%", mc.in_use() ? 'U' : 'F',
+  printf("  %c %d %p bin=%d bytes=%" PRIu64 " %3.1f%%", mc.in_use() ? 'U' : 'F',
          size_class, reinterpret_cast(mc.address()), mc.bin(),
-         mc.size(),
+         static_cast(mc.size()),
          100 * (*cumulative_bytes / static_cast(total_bytes)));
   if (freed_at) {
-    printf(" freed_at=%llu", mc.freed_at_count());
+    printf(" freed_at=%" PRIu64, static_cast(mc.freed_at_count()));
   }
   if (ac_offset > 0) {
-    printf(" age=%llu", ac_offset - mc.action_count());
+    printf(" age=%" PRIu64,
+           static_cast(ac_offset - mc.action_count()));
   } else {
-    printf(" ac=%llu", mc.action_count());
+    printf(" ac=%" PRIu64, static_cast(mc.action_count()));
   }
   // step_ids are random, so save space by showing only low 16 bits.
   printf(" step=%x op=%s\n", static_cast(0xFFFF & mc.step_id()),
@@ -90,18 +92,24 @@ void PrintChunk(const MemChunk& mc, const uint64 ac_offset, bool freed_at,
 void PrintSummary(const MemoryDump& md) {
   printf("MemoryMap for allocator %s\n", md.allocator_name().c_str());
   for (auto& it : md.bin_summary()) {
-    printf(
-        "   Bin %2d total bytes=%10lld \tin use=%10lld \ttotal_chunks=%6lld "
-        "\tin_use=%6lld\n",
-        it.bin(), it.total_bytes_in_bin(), it.total_bytes_in_use(),
-        it.total_chunks_in_bin(), it.total_chunks_in_use());
+    printf("   Bin %2d total bytes=%10" PRId64 " \tin use=%10" PRId64
+           " \ttotal_chunks=%6" PRId64
+           " "
+           "\tin_use=%6" PRId64 "\n",
+           it.bin(), static_cast(it.total_bytes_in_bin()),
+           static_cast(it.total_bytes_in_use()),
+           static_cast(it.total_chunks_in_bin()),
+           static_cast(it.total_chunks_in_use()));
   }
-  printf(
-      "Total num_allocs: %lld, bytes_in_use: %lld, peak_bytes_in_use: %lld,\n"
-      "largest_alloc_size: %lld, fragmentation: %f\n",
-      md.stats().num_allocs(), md.stats().bytes_in_use(),
-      md.stats().peak_bytes_in_use(), md.stats().largest_alloc_size(),
-      md.stats().fragmentation_metric());
+  printf("Total num_allocs: %" PRId64 ", bytes_in_use: %" PRId64
+         ", peak_bytes_in_use: %" PRId64
+         ",\n"
+         "largest_alloc_size: %" PRId64 ", fragmentation: %f\n",
+         static_cast(md.stats().num_allocs()),
+         static_cast(md.stats().bytes_in_use()),
+         static_cast(md.stats().peak_bytes_in_use()),
+         static_cast(md.stats().largest_alloc_size()),
+         md.stats().fragmentation_metric());
 }
 
 void PrintSortedChunks(
@@ -125,10 +133,10 @@ void PrintSortedChunks(
   for (int i = 0; i < chunks.size(); ++i) {
     const MemChunk* c = chunks[i];
     if (by_addr && i > 0 && last_end != c->address()) {
-      printf("  empty range from %p to %p  (%lld)\n",
+      printf("  empty range from %p to %p  (%" PRId64 ")\n",
              reinterpret_cast(last_end),
              reinterpret_cast(c->address()),
-             (c->address() - last_end));
+             static_cast(c->address() - last_end));
     }
     PrintChunk(*c, max_action_count, freed_at, total_bytes, &cumulative_bytes);
     last_end = c->address() + c->size();
@@ -182,8 +190,8 @@ void PrintChunksByOpName(const MemoryDump& md, const string& op_name,
       total_bytes += it.size();
     }
   }
-  printf("\t%d matching Chunks of total size %llu bytes:\n",
-         filtered.chunk_size(), total_bytes);
+  printf("\t%d matching Chunks of total size %" PRIu64 " bytes:\n",
+         filtered.chunk_size(), static_cast(total_bytes));
   PrintSortedChunks(
       filtered,
       [](const MemChunk* a, const MemChunk* b) {
@@ -205,10 +213,13 @@ void PrintSizeHistory(const MemoryDump& md, bool by_age) {
   }
   for (auto& it : md.snap_shot()) {
     if (by_age) {
-      printf("\tage=%llu, size=%lld\n", max_action_count - it.action_count(),
-             it.size());
+      printf("\tage=%" PRIu64 ", size=%" PRId64 "\n",
+             static_cast(max_action_count - it.action_count()),
+             static_cast(it.size()));
     } else {
-      printf("\tac=%llu, size=%lld\n", it.action_count(), it.size());
+      printf("\tac=%" PRIu64 ", size=%" PRId64 "\n",
+             static_cast(it.action_count()),
+             static_cast(it.size()));
     }
   }
 }
diff --git a/tensorflow/core/kernels/data/experimental/io_ops.cc b/tensorflow/core/kernels/data/experimental/io_ops.cc
index d41c604bb32..112a58a7e9a 100644
--- a/tensorflow/core/kernels/data/experimental/io_ops.cc
+++ b/tensorflow/core/kernels/data/experimental/io_ops.cc
@@ -159,7 +159,8 @@ Status SaveDatasetOp::WriteMetadataFile(Env* env, const std::string& path,
                                         uint64 num_elements, bool finalized) {
   SnapshotMetadataRecord metadata;
   metadata.set_creation_timestamp(EnvTime::NowMicros());
-  metadata.set_run_id(strings::Printf("%llu", run_id));
+  metadata.set_run_id(
+      strings::Printf("%llu", static_cast(run_id)));
   metadata.set_version(kFileFormatVersion);
   for (const auto& output_dtype : output_dtypes) {
     metadata.add_dtype(output_dtype);
diff --git a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
index 09783161091..0cf85a58985 100644
--- a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
@@ -302,9 +302,11 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
       }
       auto result = dataset()->traceme_metadata_;
       result.push_back(std::make_pair(
-          "max_batch_results", strings::Printf("%lld", max_batch_results)));
-      result.push_back(
-          std::make_pair("parallelism", strings::Printf("%lld", parallelism)));
+          "max_batch_results",
+          strings::Printf("%lld", static_cast(max_batch_results))));
+      result.push_back(std::make_pair(
+          "parallelism",
+          strings::Printf("%lld", static_cast(parallelism))));
       return result;
     }
 
diff --git a/tensorflow/core/kernels/data/experimental/snapshot_util.cc b/tensorflow/core/kernels/data/experimental/snapshot_util.cc
index 00269ab534a..3b051d7d572 100644
--- a/tensorflow/core/kernels/data/experimental/snapshot_util.cc
+++ b/tensorflow/core/kernels/data/experimental/snapshot_util.cc
@@ -50,11 +50,14 @@ namespace snapshot_util {
     CustomReader::kSnappyReaderOutputBufferSizeBytes;
 
 std::string HashDirectory(const std::string& path, uint64 hash) {
-  return io::JoinPath(path, strings::Printf("%llu", hash));
+  return io::JoinPath(
+      path, strings::Printf("%llu", static_cast(hash)));
 }
 
 std::string RunDirectory(const std::string& hash_directory, uint64 run_id) {
-  return RunDirectory(hash_directory, strings::Printf("%llu", run_id));
+  return RunDirectory(
+      hash_directory,
+      strings::Printf("%llu", static_cast(run_id)));
 }
 
 std::string RunDirectory(const std::string& hash_directory,
@@ -63,13 +66,17 @@ std::string RunDirectory(const std::string& hash_directory,
 }
 
 std::string ShardDirectory(const std::string& run_directory, int64 shard_id) {
-  return io::JoinPath(run_directory, strings::Printf("%08llu%s", shard_id,
-                                                     kShardDirectorySuffix));
+  return io::JoinPath(
+      run_directory,
+      strings::Printf("%08llu%s", static_cast(shard_id),
+                      kShardDirectorySuffix));
 }
 std::string GetCheckpointFileName(const std::string& shard_directory,
                                   uint64 checkpoint_id) {
-  return io::JoinPath(shard_directory,
-                      strings::Printf("%08llu.snapshot", checkpoint_id));
+  return io::JoinPath(
+      shard_directory,
+      strings::Printf("%08llu.snapshot",
+                      static_cast(checkpoint_id)));
 }
 
 Status Writer::Create(Env* env, const std::string& filename,
diff --git a/tensorflow/core/profiler/internal/gpu/device_tracer.cc b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
index 73d2a278ea4..9c3e2d67bf0 100644
--- a/tensorflow/core/profiler/internal/gpu/device_tracer.cc
+++ b/tensorflow/core/profiler/internal/gpu/device_tracer.cc
@@ -486,7 +486,7 @@ class CuptiTraceCollectorImpl : public CuptiTraceCollector {
         // Times 2 because HBM is DDR memory; it gets two data bits per each
         // data lane.
         auto memory_bandwidth =
-            2ULL * (*mem_clock_khz) * 1000 * (*mem_bus_width_bits) / 8;
+            uint64{2} * (*mem_clock_khz) * 1000 * (*mem_bus_width_bits) / 8;
         device_plane->AddStatValue(
             *device_plane->GetOrCreateStatMetadata(
                 GetStatTypeStr(StatType::kDevCapMemoryBandwidth)),

From ec27f797f1711adbd255e492c617d1c5ba0cb273 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 18:03:25 -0700
Subject: [PATCH 0361/1390] fix kokoro build.

PiperOrigin-RevId: 316795091
Change-Id: I5c8b121d79a0d21f0fa89ba8999f930e0d305322
---
 tensorflow/core/tpu/kernels/tpu_compile_op_common.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.h b/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
index 2c8d90643ef..5223732430a 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.h
@@ -153,7 +153,7 @@ class TpuCompileOpKernelCommon {
   int num_computations_;
 
  private:
-  DISALLOW_COPY_AND_ASSIGN(TpuCompileOpKernelCommon);
+  TF_DISALLOW_COPY_AND_ASSIGN(TpuCompileOpKernelCommon);
 };
 
 }  // namespace tpu

From f3d3480c811480c2b966df1086828e47131eb783 Mon Sep 17 00:00:00 2001
From: Koan-Sin Tan 
Date: Wed, 17 Jun 2020 09:25:28 +0800
Subject: [PATCH 0362/1390] [tflite] make label_image build on linux and macOS

label_image doesn't build on Linux and macOS platforms

```
bazel build --config opt //tensorflow/lite/examples/label_image:label_image
```
shows something like
```
ERROR: /home/freedom/work/tensorflow/tensorflow/lite/examples/label_image/BUILD:15:1: undeclared inclusion(s) in rule '//tensorflow/lite/examples/label_image:label_image':
this rule is missing dependency declarations for the following files included by 'tensorflow/lite/examples/label_image/label_image.cc':
  'external/com_google_absl/absl/strings/string_view.h'
  'external/com_google_absl/absl/base/internal/throw_delegate.h'

```

Add `"@com_google_absl//absl/strings"` to deps
---
 tensorflow/lite/examples/label_image/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/lite/examples/label_image/BUILD b/tensorflow/lite/examples/label_image/BUILD
index 01296b0b2a0..633f767c5e9 100644
--- a/tensorflow/lite/examples/label_image/BUILD
+++ b/tensorflow/lite/examples/label_image/BUILD
@@ -38,6 +38,7 @@ cc_binary(
         "//tensorflow/lite/profiling:profiler",
         "//tensorflow/lite/tools/evaluation:utils",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
     ] + select({
         "//tensorflow:android": [
             "//tensorflow/lite/delegates/gpu:delegate",

From 06ff30f7ea35098cb68a231a9eb7ff3ff4be4e1e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 18:22:14 -0700
Subject: [PATCH 0363/1390] [Profiler] Enhance the memory profiler to handle
 very large number of snapshots.

PiperOrigin-RevId: 316797391
Change-Id: Ia3c3a678241498bb480f7efdeb535d032eeb634a
---
 .../convert/xplane_to_memory_profile.cc       | 34 ++++++++++++++-----
 .../convert/xplane_to_memory_profile.h        |  7 ++--
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc b/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
index a0353d371d6..d039ca8da32 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
@@ -31,7 +31,6 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/protobuf/memory_profile.pb.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 #include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
@@ -443,9 +442,23 @@ void ProcessActiveAllocations(int64 peak_bytes_profile_step_id,
           << memory_profile->active_allocations_size();
 }
 
+void SampleSnapshots(
+    int64 max_num_snapshots,
+    protobuf::RepeatedPtrField* snapshots) {
+  if (snapshots->size() <= max_num_snapshots) return;
+  absl::c_partial_sort(
+      *snapshots, snapshots->begin() + max_num_snapshots,
+      [](const MemoryProfileSnapshot& a, const MemoryProfileSnapshot& b) {
+        return a.aggregation_stats().free_memory_bytes() <
+               b.aggregation_stats().free_memory_bytes();
+      });
+  snapshots->erase(snapshots->begin() + max_num_snapshots, snapshots->end());
+}
+
 // Post-process the memory profile to correctly update proto fields, and break
 // down peak memory usage for each allocator.
-void ProcessMemoryProfileProto(MemoryProfile* memory_profile) {
+void ProcessMemoryProfileProto(int64 max_num_snapshots,
+                               MemoryProfile* memory_profile) {
   memory_profile->set_num_hosts(1);
   // Add sorted memory ids within memory profile data to the selection list.
   for (const auto& id_and_allocator_profile :
@@ -460,12 +473,13 @@ void ProcessMemoryProfileProto(MemoryProfile* memory_profile) {
        *memory_profile->mutable_memory_profile_per_allocator()) {
     PerAllocatorMemoryProfile* allocator_memory_profile =
         &id_and_allocator_profile.second;
+    protobuf::RepeatedPtrField* snapshots =
+        allocator_memory_profile->mutable_memory_profile_snapshots();
     // Sort the memory_profile_snapshots by time_offset_ps (ascending) in proto.
-    absl::c_sort(
-        *allocator_memory_profile->mutable_memory_profile_snapshots(),
-        [](const MemoryProfileSnapshot& a, const MemoryProfileSnapshot& b) {
-          return a.time_offset_ps() < b.time_offset_ps();
-        });
+    absl::c_sort(*snapshots, [](const MemoryProfileSnapshot& a,
+                                const MemoryProfileSnapshot& b) {
+      return a.time_offset_ps() < b.time_offset_ps();
+    });
 
     UpdateStepId(memory_profile->step_count(), allocator_memory_profile);
     UpdateDeallocation(allocator_memory_profile);
@@ -476,14 +490,16 @@ void ProcessMemoryProfileProto(MemoryProfile* memory_profile) {
     int64 peak_step_id =
         GetPeakMemoryStep(peak_bytes_profile, allocator_memory_profile);
     ProcessActiveAllocations(peak_step_id, allocator_memory_profile);
+    SampleSnapshots(max_num_snapshots, snapshots);
   }
 }
 
 }  // namespace
 
-MemoryProfile ConvertXPlaneToMemoryProfile(const XPlane& host_plane) {
+MemoryProfile ConvertXPlaneToMemoryProfile(const XPlane& host_plane,
+                                           int64 max_num_snapshots) {
   MemoryProfile memory_profile = GenerateMemoryProfile(&host_plane);
-  ProcessMemoryProfileProto(&memory_profile);
+  ProcessMemoryProfileProto(max_num_snapshots, &memory_profile);
   return memory_profile;
 }
 
diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile.h b/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
index bd8a6e8df08..873ac800aa5 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
@@ -16,14 +16,17 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_CONVERT_XPLANE_TO_MEMORY_PROFILE_H_
 #define TENSORFLOW_CORE_PROFILER_CONVERT_XPLANE_TO_MEMORY_PROFILE_H_
 
+#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/protobuf/memory_profile.pb.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
 
 namespace tensorflow {
 namespace profiler {
 
-// Process the host threads XPlane and generate MemoryProfile result.
-MemoryProfile ConvertXPlaneToMemoryProfile(const XPlane& host_plane);
+// Process the host threads XPlane and generate MemoryProfile result; at most
+// max_num_snapshots will be displayed on the UI.
+MemoryProfile ConvertXPlaneToMemoryProfile(const XPlane& host_plane,
+                                           int64 max_num_snapshots = 1000);
 
 }  // namespace profiler
 }  // namespace tensorflow

From 8d8212ea8129347603e60b053f147d6bd0d1f83a Mon Sep 17 00:00:00 2001
From: tg-at-google 
Date: Tue, 16 Jun 2020 21:31:09 -0400
Subject: [PATCH 0364/1390] Update quantization_config.cc

---
 .../mlir/lite/quantization/quantization_config.cc         | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
index cdff93502f2..08cc29dd647 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_config.cc
@@ -48,9 +48,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names,
   std::vector node_mins;
   if (!min_values.empty()) {
     std::vector node_mins_str = absl::StrSplit(min_values, ',');
-    for (const std::string&node_min : node_mins_str.size()) {
+    for (int i = 0, e = node_mins_str.size(); i < e; i++) {
       double value;
-      if (!absl::SimpleAtod(node_min, &value)) {
+      if (!absl::SimpleAtod(node_mins_str[i], &value)) {	
         return true;
       }
       node_mins.push_back(value);
@@ -60,9 +60,9 @@ bool ParseInputNodeQuantSpecs(absl::string_view node_names,
   std::vector node_maxs;
   if (!max_values.empty()) {
     std::vector node_maxs_str = absl::StrSplit(max_values, ',');
-    for (const std::string&node_max : node_maxs_str.size()) {
+    for (int i = 0, e = node_maxs_str.size(); i < e; i++) {
       double value;
-      if (!absl::SimpleAtod(node_max, &value)) {
+      if (!absl::SimpleAtod(node_maxs_str[i], &value)) {
         llvm::errs() << "Unexpected mins: " << node_maxs_str[i] << "\n";
         return true;
       }

From 9338f4da0d648cb73339c202fafbbc9376bb3fcb Mon Sep 17 00:00:00 2001
From: tg-at-google 
Date: Tue, 16 Jun 2020 22:19:40 -0400
Subject: [PATCH 0365/1390] Update sqlite_query_connection.cc

---
 .../kernels/data/experimental/sql/sqlite_query_connection.cc  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc b/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
index ada94be15bf..9a7eb125f95 100644
--- a/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
+++ b/tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc
@@ -68,8 +68,8 @@ Status SqliteQueryConnection::GetNext(IteratorContext* ctx,
 
 Status SqliteQueryConnection::PrepareQuery() {
   TF_RETURN_IF_ERROR(db_->Prepare(query_, &stmt_));
-  size_t column_count = stmt_.ColumnCount();
-  if (column_count != output_types_.size()) {
+  int column_count = stmt_.ColumnCount();
+  if (column_count != static_cast(output_types_.size())) {
     stmt_ = SqliteStatement();
     return errors::InvalidArgument(tensorflow::strings::Printf(
         "The number of columns in query (%d) must match the number of "

From 1444b6fbbd880aee31426bdbe3b2b7b66cdeac6e Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Tue, 16 Jun 2020 19:39:06 -0700
Subject: [PATCH 0366/1390] [XLA] Remove TPU-specific aliasing #ifdef which is
 no longer TPU-specific

PiperOrigin-RevId: 316806281
Change-Id: I6acab4ed606cc123052b539a4b2b160fffe21b5b
---
 tensorflow/compiler/xla/tests/buffer_donation_test.cc | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/buffer_donation_test.cc b/tensorflow/compiler/xla/tests/buffer_donation_test.cc
index 9cfeefadead..18cde722a64 100644
--- a/tensorflow/compiler/xla/tests/buffer_donation_test.cc
+++ b/tensorflow/compiler/xla/tests/buffer_donation_test.cc
@@ -109,9 +109,7 @@ class BufferDonationTest : public HloTestBase {
               << "             size = " << result_root_buffer.size();
 
     // Check for expected aliasing between input and output buffers.
-    // The following aliasing pattern is only ever generated by the TPU backend
-    // at the moment.
-#if defined(XLA_TEST_BACKEND_TPU)
+#ifndef XLA_TEST_BACKEND_INTERPRETER
     for (int i = 0; i < ShapeUtil::TupleElementCount(argument_literal.shape());
          ++i) {
       const ShapeIndex index({i});

From 62beb0fc6651cd9628ce1f34f377a54ef97acfab Mon Sep 17 00:00:00 2001
From: Brian Zhao 
Date: Tue, 16 Jun 2020 19:48:23 -0700
Subject: [PATCH 0367/1390] Adding convenience functions for assigning and
 reading resource variables.

PiperOrigin-RevId: 316807134
Change-Id: Ia37fdf8299064fe1a471eb0b34568e85164808ce
---
 .../saved_model/core/ops/variable_ops.cc      | 30 +++++++++++++++++++
 .../saved_model/core/ops/variable_ops.h       | 16 ++++++++++
 .../saved_model/core/ops/variable_ops_test.cc | 30 +++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
index 94548e553ad..a3b3ace7be9 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
@@ -58,6 +58,36 @@ Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
   return Status();
 }
 
+Status AssignVariable(AbstractContextInterface* ctx,
+                      AbstractTensorHandleInterface* variable_handle,
+                      DataType dtype, AbstractTensorHandleInterface* value) {
+  AbstractOpPtr assign_op(ctx->CreateOperation());
+  TF_RETURN_IF_ERROR(assign_op->Reset("AssignVariableOp", nullptr));
+  TF_RETURN_IF_ERROR(assign_op->SetAttrType("dtype", dtype));
+  TF_RETURN_IF_ERROR(assign_op->AddInput(variable_handle));
+  TF_RETURN_IF_ERROR(assign_op->AddInput(value));
+
+  int num_retvals = 0;
+  TF_RETURN_IF_ERROR(assign_op->Execute({}, &num_retvals));
+  return Status();
+}
+
+Status ReadVariable(AbstractContextInterface* ctx,
+                    AbstractTensorHandleInterface* variable_handle,
+                    DataType dtype, AbstractTensorHandlePtr* output) {
+  AbstractOpPtr read_op = AbstractOpPtr(ctx->CreateOperation());
+  TF_RETURN_IF_ERROR(read_op->Reset("ReadVariableOp", nullptr));
+  TF_RETURN_IF_ERROR(read_op->SetAttrType("dtype", dtype));
+  TF_RETURN_IF_ERROR(read_op->AddInput(variable_handle));
+
+  AbstractTensorHandleInterface* value = nullptr;
+  int num_retvals = 1;
+  TF_RETURN_IF_ERROR(
+      read_op->Execute(absl::MakeSpan(&value, num_retvals), &num_retvals));
+  output->reset(value);
+  return Status();
+}
+
 Status DestroyResource(AbstractContextInterface* ctx,
                        AbstractTensorHandleInterface* handle) {
   AbstractOpPtr destroy_op = AbstractOpPtr(ctx->CreateOperation());
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
index 1c4d757af8c..8a410328b9e 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
@@ -34,6 +34,22 @@ Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
                                            DataType dtype, TensorShape shape,
                                            AbstractTensorHandlePtr* handle);
 
+// Executes an AssignVariableOp using `ctx`, assigning the variable associated
+// with `variable_handle` with `value`. `dtype` must be the datatype of the
+// underlying variable for `variable_handle`. Note that it is illegal to assign
+// a variable to a Tensor with a different dtype than what the variable was
+// created with.
+Status AssignVariable(AbstractContextInterface* ctx,
+                      AbstractTensorHandleInterface* variable_handle,
+                      DataType dtype, AbstractTensorHandleInterface* value);
+
+// Executes a ReadVariableOp using `ctx`. This reads the underlying variable
+// value of `variable_handle` and copies the value to `output`. `dtype` must be
+// the dtype of the variable associated with `variable_handle`.
+Status ReadVariable(AbstractContextInterface* ctx,
+                    AbstractTensorHandleInterface* variable_handle,
+                    DataType dtype, AbstractTensorHandlePtr* output);
+
 // Executes DestroyResourceOp on `handle`, using `ctx`. This is equivalent to
 // the cleanup that occurs in a tf.Variable's EagerResourceDeleter:
 // https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L289-L290
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc b/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
index 7a9486f8ebd..3c57ed4d38a 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
@@ -30,6 +30,13 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+AbstractTensorHandlePtr CreateScalarTensorHandle(EagerContext* context,
+                                                 float value) {
+  AbstractTensorPtr tensor(context->CreateFloatScalar(value));
+  AbstractTensorHandlePtr handle(context->CreateLocalHandle(tensor.get()));
+  return handle;
+}
+
 class VariableOpsTest : public ::testing::Test {
  public:
   VariableOpsTest()
@@ -73,5 +80,28 @@ TEST_F(VariableOpsTest, DestroyVariableSuccessful) {
   TF_EXPECT_OK(internal::DestroyResource(context(), handle.get()));
 }
 
+// Sanity check for handle assignment and reading
+TEST_F(VariableOpsTest, AssignVariableAndReadSuccessful) {
+  // Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
+  AbstractTensorHandlePtr variable;
+  TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
+      context(), DT_FLOAT, {}, &variable));
+
+  // Create a Scalar float TensorHandle with value 42, and assign it to
+  // the variable.
+  AbstractTensorHandlePtr my_value = CreateScalarTensorHandle(context(), 42.0);
+  TF_EXPECT_OK(internal::AssignVariable(context(), variable.get(), DT_FLOAT,
+                                        my_value.get()));
+
+  // Read back the value from the variable, and check that it is 42.
+  AbstractTensorHandlePtr read_value_handle;
+  TF_EXPECT_OK(internal::ReadVariable(context(), variable.get(), DT_FLOAT,
+                                      &read_value_handle));
+  Status status;
+  AbstractTensorPtr read_value(read_value_handle->Resolve(&status));
+  TF_EXPECT_OK(status);
+  EXPECT_FLOAT_EQ(42.0, *static_cast(read_value->Data()));
+}
+
 }  // namespace
 }  // namespace tensorflow

From 7c5ddb830f17b15708664fc4a34fe3b24f142628 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 16 Jun 2020 19:54:11 -0700
Subject: [PATCH 0368/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/4799fb63b551

PiperOrigin-RevId: 316807673
Change-Id: I246445df966079989305be15318b9b9b4cc4db3a
---
 tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
index 9e75c1b9ac5..72f3a4dfac7 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/Function.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "tensorflow/compiler/xla/service/cpu/cpu_options.h"
 #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h"
@@ -255,7 +256,8 @@ Status DotOpEmitter::EmitLinalgMatmul() {
         mlir::edsc::ScopedContext scope(*builder, function.getLoc());
         mlir::Value a = function.getArgument(0), b = function.getArgument(1),
                     c = function.getArgument(2);
-        mlir::edsc::intrinsics::linalg_matmul(b, c, a);
+        mlir::edsc::intrinsics::linalg_matmul(mlir::TypeRange{},
+                                              mlir::ValueRange{b, c, a});
         mlir::edsc::intrinsics::std_ret();
       });
 }

From d68cccb57eaa23d9af2e1e065373f084a808098b Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 16 Jun 2020 19:54:13 -0700
Subject: [PATCH 0369/1390] Removed fuse add to conv transformation. This
 transformation incorrect for border elements, when used Zero clamping.

PiperOrigin-RevId: 316807675
Change-Id: Iffbdcd3f5a46ecbdc8c4ad1c4f8f4393410f7be3
---
 .../transformations/fuse_add_to_conv.cc       | 114 ------------------
 .../common/transformations/fuse_add_to_conv.h |  23 ----
 .../transformations/fuse_add_to_conv_test.cc  | 114 ------------------
 .../general_transformations.cc                |   4 +-
 4 files changed, 1 insertion(+), 254 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.cc b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.cc
index b279e49e40c..adee86e4a64 100644
--- a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.cc
+++ b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.cc
@@ -92,66 +92,12 @@ class MergeConvolutionWithAdd : public SequenceTransformation {
   }
 };
 
-class MergeAddWithConvolution : public SequenceTransformation {
- public:
-  int ExpectedSequenceLength() const final { return 2; }
-
-  TransformResult ApplyToNodesSequence(const std::vector& sequence,
-                                       GraphFloat32* graph) final {
-    auto& conv_node = *sequence[1];
-    auto& add_node = *sequence[0];
-    if (add_node.operation.type != ToString(OperationType::ADD)) {
-      return {TransformStatus::SKIPPED, ""};
-    }
-    AddAttributes add_attr =
-        absl::any_cast(add_node.operation.attributes);
-    if (!absl::holds_alternative>(
-            add_attr.param) &&
-        !absl::holds_alternative(add_attr.param)) {
-      return {TransformStatus::DECLINED,
-              "This fuse applicable only for broadcast or scalar addition."};
-    }
-
-    if (conv_node.operation.type == ToString(OperationType::CONVOLUTION_2D)) {
-      Convolution2DAttributes* conv_attr =
-          absl::any_cast(
-              &conv_node.operation.attributes);
-      FuseAddWithConvolution2D(add_attr, conv_attr);
-    } else if (conv_node.operation.type ==
-               ToString(OperationType::DEPTHWISE_CONVOLUTION)) {
-      DepthwiseConvolution2DAttributes* conv_attr =
-          absl::any_cast(
-              &conv_node.operation.attributes);
-      FuseAddWithDepthwiseConvolution2D(add_attr, conv_attr);
-    } else if (conv_node.operation.type ==
-               ToString(OperationType::FULLY_CONNECTED)) {
-      FullyConnectedAttributes* conv_attr =
-          absl::any_cast(
-              &conv_node.operation.attributes);
-      FuseAddWithFullyConnected(add_attr, conv_attr);
-    } else {
-      return {TransformStatus::SKIPPED, ""};
-    }
-
-    absl::Status status = RemovePrecedingNode(graph, &add_node, &conv_node);
-    if (!status.ok()) {
-      return {TransformStatus::INVALID,
-              "Unable to remove add node after convolution: " +
-                  std::string(status.message())};
-    }
-    return {TransformStatus::APPLIED, ""};
-  }
-};
 }  // namespace
 
 std::unique_ptr NewMergeConvolutionWithAdd() {
   return absl::make_unique();
 }
 
-std::unique_ptr NewMergeAddWithConvolution() {
-  return absl::make_unique();
-}
-
 void FuseConvolution2DWithAdd(const AddAttributes& add_attr,
                               Convolution2DAttributes* attr) {
   FuseBiasWithAddAttributes(add_attr, attr->weights.shape.o, &attr->bias);
@@ -173,65 +119,5 @@ void FuseFullyConnectedWithAdd(const AddAttributes& add_attr,
   FuseBiasWithAddAttributes(add_attr, attr->weights.shape.o, &attr->bias);
 }
 
-void FuseAddWithConvolution2D(const AddAttributes& add_attr,
-                              Convolution2DAttributes* attr) {
-  auto add = absl::get_if>(&add_attr.param);
-  auto add_scalar = absl::get_if(&add_attr.param);
-  if (attr->bias.data.empty()) {
-    attr->bias = MakeZeroTensor(
-        Linear(attr->weights.shape.o));
-  }
-  for (int d = 0; d < attr->weights.shape.o; ++d) {
-    for (int s = 0; s < attr->weights.shape.i; ++s) {
-      const float add_value = add ? add->data[s] : *add_scalar;
-      for (int k_y = 0; k_y < attr->weights.shape.h; ++k_y) {
-        for (int k_x = 0; k_x < attr->weights.shape.w; ++k_x) {
-          const int index = attr->weights.shape.LinearIndex({{d, k_y, k_x, s}});
-          attr->bias.data[d] += attr->weights.data[index] * add_value;
-        }
-      }
-    }
-  }
-}
-
-void FuseAddWithDepthwiseConvolution2D(const AddAttributes& add_attr,
-                                       DepthwiseConvolution2DAttributes* attr) {
-  auto add = absl::get_if>(&add_attr.param);
-  auto add_scalar = absl::get_if(&add_attr.param);
-  if (attr->bias.data.empty()) {
-    attr->bias = MakeZeroTensor(
-        Linear(attr->weights.shape.o * attr->weights.shape.i));
-  }
-  for (int s = 0; s < attr->weights.shape.i; ++s) {
-    const float add_value = add ? add->data[s] : *add_scalar;
-    for (int g = 0; g < attr->weights.shape.o; ++g) {
-      const int d = s * attr->weights.shape.o + g;
-      for (int k_y = 0; k_y < attr->weights.shape.h; ++k_y) {
-        for (int k_x = 0; k_x < attr->weights.shape.w; ++k_x) {
-          const int index = attr->weights.shape.LinearIndex({{g, k_y, k_x, s}});
-          attr->bias.data[d] += attr->weights.data[index] * add_value;
-        }
-      }
-    }
-  }
-}
-
-void FuseAddWithFullyConnected(const AddAttributes& add_attr,
-                               FullyConnectedAttributes* attr) {
-  auto add = absl::get_if>(&add_attr.param);
-  auto add_scalar = absl::get_if(&add_attr.param);
-  if (attr->bias.data.empty()) {
-    attr->bias = MakeZeroTensor(
-        Linear(attr->weights.shape.o));
-  }
-  for (int d = 0; d < attr->weights.shape.o; ++d) {
-    for (int s = 0; s < attr->weights.shape.i; ++s) {
-      const float add_value = add ? add->data[s] : *add_scalar;
-      const int index = attr->weights.shape.LinearIndex({{d, 0, 0, s}});
-      attr->bias.data[d] += attr->weights.data[index] * add_value;
-    }
-  }
-}
-
 }  // namespace gpu
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.h b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.h
index 49871a815da..85014ec177e 100644
--- a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.h
+++ b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv.h
@@ -30,11 +30,6 @@ namespace gpu {
 // convolution.
 std::unique_ptr NewMergeConvolutionWithAdd();
 
-// Fuse Add Scalar or Add Broadcast before Convolution(Convolution2D,
-// DepthWise, FullyConnected) into biases of
-// convolution.
-std::unique_ptr NewMergeAddWithConvolution();
-
 // Modify Convolution2DAttributes so that after making convolution with
 // modified attributes we will have the same result as convolution
 // with old attributes and following add operation.
@@ -59,24 +54,6 @@ void FuseConvolutionTransposedWithAdd(const AddAttributes& add_attr,
 void FuseFullyConnectedWithAdd(const AddAttributes& add_attr,
                                FullyConnectedAttributes* attr);
 
-// Modify Convolution2DAttributes so that after making convolution with
-// modified attributes we will have the same result as add operation and
-// convolution with old attributes
-void FuseAddWithConvolution2D(const AddAttributes& add_attr,
-                              Convolution2DAttributes* attr);
-
-// Modify DepthwiseConvolution2DAttributes so that after making depth wise
-// convolution with modified attributes we will have the same result as add
-// operation and depth wise convolution with old attributes
-void FuseAddWithDepthwiseConvolution2D(const AddAttributes& add_attr,
-                                       DepthwiseConvolution2DAttributes* attr);
-
-// Modify FullyConnectedAttributes so that after making fully connected
-// with modified attributes we will have the same result as add operation and
-// fully connected with old attributes
-void FuseAddWithFullyConnected(const AddAttributes& add_attr,
-                               FullyConnectedAttributes* attr);
-
 }  // namespace gpu
 }  // namespace tflite
 
diff --git a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv_test.cc b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv_test.cc
index 431d8167f81..53dba56ffb8 100644
--- a/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv_test.cc
+++ b/tensorflow/lite/delegates/gpu/common/transformations/fuse_add_to_conv_test.cc
@@ -78,57 +78,6 @@ TEST(MergeConvolutionWithAddTest, Smoke) {
             graph.nodes()[0]->operation.type);
 }
 
-TEST(MergeAddWithConvolutionTest, Smoke) {
-  GraphFloat32 graph;
-  auto input = graph.NewValue();
-  input->tensor.shape = BHWC(1, 4, 4, 8);
-
-  Convolution2DAttributes conv_attr;
-  conv_attr.padding.prepended = HW(0, 0);
-  conv_attr.padding.appended = HW(0, 0);
-  conv_attr.strides = HW(1, 1);
-  conv_attr.dilations = HW(1, 1);
-  conv_attr.weights.shape = OHWI(16, 3, 2, 8);
-  conv_attr.weights.data.resize(conv_attr.weights.shape.DimensionsProduct());
-  conv_attr.bias.shape = Linear(16);
-  conv_attr.bias.data.resize(16);
-
-  Tensor add_tensor;
-  add_tensor.shape = Linear(8);
-  add_tensor.data.resize(8);
-  AddAttributes add_attr;
-  add_attr.param = add_tensor;
-
-  auto conv_node = graph.NewNode();
-  conv_node->operation.type = ToString(OperationType::CONVOLUTION_2D);
-  conv_node->operation.attributes = conv_attr;
-  auto add_node = graph.NewNode();
-  add_node->operation.type = ToString(OperationType::ADD);
-  add_node->operation.attributes = add_attr;
-
-  ASSERT_TRUE(graph.AddConsumer(add_node->id, input->id).ok());
-
-  Value* output;
-  ASSERT_TRUE(AddOutput(&graph, conv_node, &output).ok());
-  output->tensor.shape = BHWC(1, 4, 4, 16);
-
-  Value* link1;
-  ASSERT_TRUE(ConnectTwoNodes(&graph, add_node, conv_node, &link1).ok());
-  link1->tensor.shape = BHWC(1, 4, 4, 16);
-
-  ASSERT_EQ(2, graph.nodes().size());
-  ASSERT_EQ(3, graph.values().size());
-
-  auto transformation = NewMergeAddWithConvolution();
-  ModelTransformer transformer(&graph, nullptr);
-  transformer.Apply("merge_add_with_convolution", transformation.get());
-
-  EXPECT_EQ(1, graph.nodes().size());
-  EXPECT_EQ(2, graph.values().size());
-  EXPECT_EQ(ToString(OperationType::CONVOLUTION_2D),
-            graph.nodes()[0]->operation.type);
-}
-
 TEST(FuseAddAfterConvolution2DTest, Smoke) {
   Convolution2DAttributes attr;
   attr.weights.shape = OHWI(2, 1, 2, 2);
@@ -213,69 +162,6 @@ TEST(FuseAddAfterFullyConnectedTest, Smoke) {
   EXPECT_THAT(attr.bias.data, Pointwise(FloatNear(1e-6), {1.4f, 1.9f}));
 }
 
-TEST(FuseAddBeforeConvolution2DTest, Smoke) {
-  Convolution2DAttributes attr;
-  attr.weights.shape = OHWI(2, 1, 2, 2);
-  attr.weights.data = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f};
-  attr.bias.shape = Linear(2);
-  attr.bias.data = {1.1f, 1.2f};
-
-  Tensor add_tensor;
-  add_tensor.shape = Linear(2);
-  add_tensor.data = {2.0f, 0.5f};
-  AddAttributes add_attr;
-  add_attr.param = add_tensor;
-
-  FuseAddWithConvolution2D(add_attr, &attr);
-
-  EXPECT_THAT(attr.weights.data,
-              Pointwise(FloatNear(1e-6),
-                        {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f}));
-  EXPECT_THAT(attr.bias.data, Pointwise(FloatNear(1e-6), {2.2f, 4.3f}));
-}
-
-TEST(FuseAddBeforeDepthwiseConvolution2DTest, Smoke) {
-  DepthwiseConvolution2DAttributes attr;
-  attr.weights.shape = OHWI(2, 1, 2, 2);
-  attr.weights.data = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f};
-  attr.bias.shape = Linear(4);
-  attr.bias.data = {1.1f, 1.2f, 1.3f, 1.4f};
-
-  Tensor add_tensor;
-  add_tensor.shape = Linear(4);
-  add_tensor.data = {0.3f, 0.7f, 0.5f, 0.1f};
-  AddAttributes add_attr;
-  add_attr.param = add_tensor;
-
-  FuseAddWithDepthwiseConvolution2D(add_attr, &attr);
-
-  EXPECT_THAT(attr.weights.data,
-              Pointwise(FloatNear(1e-6),
-                        {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f}));
-  EXPECT_THAT(attr.bias.data,
-              Pointwise(FloatNear(1e-6), {1.22f, 1.56f, 1.72f, 2.38f}));
-}
-
-TEST(FuseAddBeforeFullyConnectedTest, Smoke) {
-  FullyConnectedAttributes attr;
-  attr.weights.shape = OHWI(2, 1, 1, 2);
-  attr.weights.data = {0.1f, 0.2f, 0.3f, 0.4f};
-  attr.bias.shape = Linear(2);
-  attr.bias.data = {1.1f, 1.2f};
-
-  Tensor add_tensor;
-  add_tensor.shape = Linear(2);
-  add_tensor.data = {0.5f, 2.0f};
-  AddAttributes add_attr;
-  add_attr.param = add_tensor;
-
-  FuseAddWithFullyConnected(add_attr, &attr);
-
-  EXPECT_THAT(attr.weights.data,
-              Pointwise(FloatNear(1e-6), {0.1f, 0.2f, 0.3f, 0.4f}));
-  EXPECT_THAT(attr.bias.data, Pointwise(FloatNear(1e-6), {1.55f, 2.15f}));
-}
-
 }  // namespace
 }  // namespace gpu
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/gpu/common/transformations/general_transformations.cc b/tensorflow/lite/delegates/gpu/common/transformations/general_transformations.cc
index 354fbcd040b..f9ae7f41f8f 100644
--- a/tensorflow/lite/delegates/gpu/common/transformations/general_transformations.cc
+++ b/tensorflow/lite/delegates/gpu/common/transformations/general_transformations.cc
@@ -54,9 +54,7 @@ bool ApplyGeneralTransformations(ModelTransformer* transformer) {
          transformer->Apply("merge_convolution_with_add",
                             NewMergeConvolutionWithAdd().get()) &&
          transformer->Apply("merge_mul_with_convolution",
-                            NewMergeMulWithConvolution().get()) &&
-         transformer->Apply("merge_add_with_convolution",
-                            NewMergeAddWithConvolution().get());
+                            NewMergeMulWithConvolution().get());
 }
 
 }  // namespace gpu

From 2b9a7d5e40aa85e84f7d6196d3e7527675d33038 Mon Sep 17 00:00:00 2001
From: peng 
Date: Wed, 17 Jun 2020 11:21:12 +0800
Subject: [PATCH 0370/1390] skip if no cuda

---
 tensorflow/core/grappler/optimizers/remapper_test.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/remapper_test.cc b/tensorflow/core/grappler/optimizers/remapper_test.cc
index 831aa51ee72..a947210d8a6 100644
--- a/tensorflow/core/grappler/optimizers/remapper_test.cc
+++ b/tensorflow/core/grappler/optimizers/remapper_test.cc
@@ -452,6 +452,9 @@ TEST_F(RemapperTest, FuseMatMulWithBias) {
 TEST_F(RemapperTest, FuseConv2DWithBiasAndActivationOnGPU) {
   using ::tensorflow::ops::Placeholder;
 
+#if !(GOOGLE_CUDA)
+  GTEST_SKIP() << "No CUDA, skip FuseConv2DWithBiasAndActivation on GPU";
+#endif  // !GOOGLE_CUDA
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
 
   auto input_shape = Placeholder::Shape({8, 32, 32, 3});

From 9bf28d66871931c16ba8269873918a914bc3efe7 Mon Sep 17 00:00:00 2001
From: peng 
Date: Wed, 17 Jun 2020 11:24:34 +0800
Subject: [PATCH 0371/1390] delete empty lines

---
 tensorflow/core/grappler/optimizers/remapper.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc
index b738a23de23..a6ddad8dfc8 100644
--- a/tensorflow/core/grappler/optimizers/remapper.cc
+++ b/tensorflow/core/grappler/optimizers/remapper.cc
@@ -1667,7 +1667,6 @@ bool RequiresInferredShapes(const RemapperContext& ctx, int node_index) {
     return false;
   };
 
-
 #ifdef INTEL_MKL
   return is_batch_norm_candidate() || is_batch_norm_fusion_candidate() ||
          IsConv2DWithAdd(ctx, node_index);
@@ -1817,7 +1816,6 @@ Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
 #endif  // !INTEL_MKL
 
-
     // Remap FusedBatchNorm++ into the _FusedBatchNormEx.
     FusedBatchNormEx fused_batch_norm_ex;
     if (allow_non_differentiable_rewrites &&

From dd417ca4477794c1dc0bc3df5b818b420417b4bf Mon Sep 17 00:00:00 2001
From: peng 
Date: Wed, 17 Jun 2020 11:28:25 +0800
Subject: [PATCH 0372/1390] delete empty lines

---
 tensorflow/core/grappler/optimizers/remapper.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc
index a6ddad8dfc8..b1158813345 100644
--- a/tensorflow/core/grappler/optimizers/remapper.cc
+++ b/tensorflow/core/grappler/optimizers/remapper.cc
@@ -1783,7 +1783,6 @@ Status Remapper::Optimize(Cluster* cluster, const GrapplerItem& item,
 // Remove this once TF-MKL supports _FusedConv2D with these operations.
 #ifndef INTEL_MKL
     // Remap Conv2D+Squeeze+BiasAdd into the _FusedConv2D+Squeeze.
-
     ContractionWithSqueezeAndBiasAdd contract_with_squeeze_and_bias;
     if (allow_non_differentiable_rewrites &&
         FindConv2DWithSqueezeAndBias(ctx, i, &contract_with_squeeze_and_bias)) {

From 1890578d7a576cd18e2044c330b1286efaf34eee Mon Sep 17 00:00:00 2001
From: Steenu Johnson 
Date: Wed, 17 Jun 2020 09:18:43 +0530
Subject: [PATCH 0373/1390] Fixing build errors.

api_test
api_compatibility_test

Signed-off-by: Steenu Johnson 
---
 tensorflow/core/api_def/base_api/api_def_CSVDatasetV2.pbtxt   | 4 ++++
 tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt       | 4 ++++
 .../golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt | 2 +-
 tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt       | 4 ++++
 4 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_CSVDatasetV2.pbtxt

diff --git a/tensorflow/core/api_def/base_api/api_def_CSVDatasetV2.pbtxt b/tensorflow/core/api_def/base_api/api_def_CSVDatasetV2.pbtxt
new file mode 100644
index 00000000000..65b6c660947
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_CSVDatasetV2.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "CSVDatasetV2"
+  visibility: HIDDEN
+}
\ No newline at end of file
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
index a5fe83e713e..96e1122f1e5 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
@@ -632,6 +632,10 @@ tf_module {
     name: "CSVDataset"
     argspec: "args=[\'filenames\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\', \'record_defaults\', \'output_shapes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
+  member_method {
+    name: "CSVDatasetV2"
+    argspec: "args=[\'filenames\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\', \'record_defaults\', \'exclude_cols\', \'output_shapes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
   member_method {
     name: "CTCBeamSearchDecoder"
     argspec: "args=[\'inputs\', \'sequence_length\', \'beam_width\', \'top_paths\', \'merge_repeated\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
index 41865c9700f..21d53a27e74 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.data.experimental.-csv-dataset.pbtxt
@@ -12,7 +12,7 @@ tf_class {
   }
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'filenames\', \'record_defaults\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \',\', \'True\', \'\', \'None\'], "
+    argspec: "args=[\'self\', \'filenames\', \'record_defaults\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\', \'exclude_cols\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \',\', \'True\', \'\', \'None\', \'None\'], "
   }
   member_method {
     name: "apply"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
index a5fe83e713e..96e1122f1e5 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
@@ -632,6 +632,10 @@ tf_module {
     name: "CSVDataset"
     argspec: "args=[\'filenames\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\', \'record_defaults\', \'output_shapes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
+  member_method {
+    name: "CSVDatasetV2"
+    argspec: "args=[\'filenames\', \'compression_type\', \'buffer_size\', \'header\', \'field_delim\', \'use_quote_delim\', \'na_value\', \'select_cols\', \'record_defaults\', \'exclude_cols\', \'output_shapes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
   member_method {
     name: "CTCBeamSearchDecoder"
     argspec: "args=[\'inputs\', \'sequence_length\', \'beam_width\', \'top_paths\', \'merge_repeated\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "

From 806e85998ab1cfbb8a969e67f1272f31604ea5f0 Mon Sep 17 00:00:00 2001
From: Mehdi Amini 
Date: Wed, 17 Jun 2020 00:14:26 -0700
Subject: [PATCH 0374/1390] Bump open source llvm revision to
 4799fb63b5513f655ca8e85416ec8fe35df49bae

PiperOrigin-RevId: 316836386
Change-Id: I26fa7c031f5057f1db5e3b26db09a02ff79f4081
---
 tensorflow/workspace.bzl                  |  4 ++--
 third_party/llvm/llvm.autogenerated.BUILD | 22 ----------------------
 2 files changed, 2 insertions(+), 24 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index acba1598d1f..cb1ea721fb0 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "9b72b47ed63351ee5ceff4c44ccd9a71dc7dad27"
-    LLVM_SHA256 = "03ce1e00901936e7259c6ee465773b7f231ca1724925460b71909868f5a61e11"
+    LLVM_COMMIT = "4799fb63b5513f655ca8e85416ec8fe35df49bae"
+    LLVM_SHA256 = "f401a61bd7f5b05bd8a3ffdfb1f32e9379cae2c8e988f3ae6772b588ad97c84a"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
diff --git a/third_party/llvm/llvm.autogenerated.BUILD b/third_party/llvm/llvm.autogenerated.BUILD
index d05fa841420..50b71a3686f 100644
--- a/third_party/llvm/llvm.autogenerated.BUILD
+++ b/third_party/llvm/llvm.autogenerated.BUILD
@@ -2753,27 +2753,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "MLPolicies",
-    srcs = glob([
-        "lib/Analysis/ML/*.c",
-        "lib/Analysis/ML/*.cpp",
-        "lib/Analysis/ML/*.inc",
-        "lib/Analysis/ML/*.h",
-    ]),
-    hdrs = glob([
-        "include/llvm/Analysis/ML/*.h",
-        "include/llvm/Analysis/ML/*.def",
-        "include/llvm/Analysis/ML/*.inc",
-    ]),
-    copts = llvm_copts,
-    deps = [
-        ":Core",
-        ":Support",
-        ":config",
-    ],
-)
-
 cc_library(
     name = "MSP430AsmParser",
     srcs = glob([
@@ -3256,7 +3235,6 @@ cc_library(
         ":IPO",
         ":InstCombine",
         ":Instrumentation",
-        ":MLPolicies",
         ":Scalar",
         ":Support",
         ":Target",

From af926984871a130eec2816815cfc98a362d4f5b6 Mon Sep 17 00:00:00 2001
From: Tres Popp 
Date: Wed, 17 Jun 2020 00:44:26 -0700
Subject: [PATCH 0375/1390] [TF:XLA] Update TF:XLA tests for
 matrix_triangular_solve to test V1 and V2.

TF:V1 raises an error on non-square coefficient matrices
TF:V2 allows non-square coefficient matrices.
PiperOrigin-RevId: 316839892
Change-Id: I34c2567ba3579c8f0fd4bc6da57abe14bc6471b2
---
 .../tests/matrix_triangular_solve_op_test.py         | 12 +++++-------
 .../tf2xla/kernels/matrix_triangular_solve_op.cc     |  8 ++++++++
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py b/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py
index b07b254c600..0202c582ef3 100644
--- a/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py
+++ b/tensorflow/compiler/tests/matrix_triangular_solve_op_test.py
@@ -135,18 +135,16 @@ class MatrixTriangularSolveOpTest(xla_test.XLATestCase):
     self._VerifyTriangularSolve(
         a.astype(np.float32), b.astype(np.float32), True, False, 1e-4)
 
-  @test_util.run_deprecated_v1
-  def testNonSquareCoefficientMatrixV1(self):
+  def testNonSquareCoefficientMatrix(self):
     rng = np.random.RandomState(0)
     for dtype in self.float_types:
       a = rng.randn(3, 4).astype(dtype)
       b = rng.randn(4, 4).astype(dtype)
-      with self.assertRaises(ValueError):
-        linalg_ops.matrix_triangular_solve(a, b)
-      with self.assertRaises(ValueError):
-        linalg_ops.matrix_triangular_solve(a, b)
+      with self.test_scope():
+        with self.assertRaises((ValueError, errors.InvalidArgumentError)):
+          linalg_ops.matrix_triangular_solve(a, b)
 
-  @test_util.run_v2_only
+  @test_util.run_v2_only  # Different error types
   def testWrongDimensionsV2(self):
     randn = np.random.RandomState(0).randn
     for dtype in self.float_types:
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
index 5a719484e05..8d222d947c9 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
@@ -50,6 +50,14 @@ class MatrixTriangularSolveOp : public XlaOpKernel {
       return;
     }
 
+    auto lhs_size = lhs_shape.dims();
+    OP_REQUIRES(
+        ctx,
+        lhs_shape.dim_size(lhs_size - 1) == lhs_shape.dim_size(lhs_size - 2),
+        errors::InvalidArgument("The coefficient matrix must be square in "
+                                "the inner-most two dimensions: ",
+                                lhs_shape.DebugString()));
+
     xla::XlaOp a = ctx->Input(0);
     xla::XlaOp b = ctx->Input(1);
     std::tie(a, b) = Broadcast(a, lhs_shape, b, rhs_shape, bcast);

From 02aad07710422ccb12a1d63093cb34c37e80e538 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Wed, 17 Jun 2020 00:55:02 -0700
Subject: [PATCH 0376/1390] Remove left-over declaration for buffer assignment
 pass that has since been removed.

PiperOrigin-RevId: 316841221
Change-Id: I9c38becf498cfaaa7edb9490948890d52e49e737
---
 .../compiler/mlir/xla/transforms/passes.h      | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/transforms/passes.h b/tensorflow/compiler/mlir/xla/transforms/passes.h
index f0c2d9b7372..a2af8124786 100644
--- a/tensorflow/compiler/mlir/xla/transforms/passes.h
+++ b/tensorflow/compiler/mlir/xla/transforms/passes.h
@@ -115,24 +115,6 @@ std::unique_ptr createLhloCopyRemovalPass();
 std::unique_ptr> createLegalizeLhloToParallelLoopsPass();
 
 }  // namespace xla_lhlo
-
-namespace xla {
-
-/// Moves alloc nodes (and their associated dealloc nodes - if any) into the
-/// right positions. If there is no associated dealloc node for a given alloc
-/// node, this pass will automatically insert a proper dealloc node in the right
-/// place. The intended use case of this pass is to store SSA values into
-/// buffers using load/store operations. For this purpose, you need to know
-/// proper positions to place the required allocs and deallocs.
-/// 1) Note that the function signatures and all types for which buffers should
-/// be allocated need to be converted in advance.
-/// 2) All required alloc nodes have the be inserted in advance.
-/// 3) Note that the current implementation does not support loops.
-/// Refer to the class mlir::xla::BufferAssignmentLegalizer for more
-/// information.
-std::unique_ptr> createBufferAssignmentPass();
-
-}  // namespace xla
 }  // namespace mlir
 
 #endif  // TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_PASSES_H_

From de907d8746c0f8909d50e8d699ff09f6248408a1 Mon Sep 17 00:00:00 2001
From: Tres Popp 
Date: Wed, 17 Jun 2020 01:20:35 -0700
Subject: [PATCH 0377/1390] Enable XLA:CPU fast math for min/max by default to
 be similar to TF's behavior.

Another big change here is changing the use of this flag to use the value in the HloModule and not the global environment variable which was bad temporary behavior.

PiperOrigin-RevId: 316844057
Change-Id: I995715ccc9009e9845fca77060b835fdc50fb4d2
---
 tensorflow/compiler/xla/debug_options_flags.cc       |  2 +-
 .../compiler/xla/service/cpu/elemental_ir_emitter.h  |  4 ++++
 .../compiler/xla/service/cpu/llvm_ir_runtime.cc      |  4 +++-
 .../xla/service/cpu/vector_support_library.cc        |  5 +++--
 .../xla/service/cpu/vector_support_library.h         |  8 +++++---
 .../compiler/xla/service/elemental_ir_emitter.cc     |  4 ++--
 .../compiler/xla/service/elemental_ir_emitter.h      |  2 ++
 .../compiler/xla/service/gpu/elemental_ir_emitter.h  |  4 ++++
 tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc | 12 ++++--------
 tensorflow/compiler/xla/service/llvm_ir/llvm_util.h  |  4 ++--
 tensorflow/compiler/xla/service/llvm_ir/math_ops.cc  |  8 ++++++--
 tensorflow/python/ops/nn_test.py                     |  2 ++
 12 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/tensorflow/compiler/xla/debug_options_flags.cc b/tensorflow/compiler/xla/debug_options_flags.cc
index 81655101701..8ca6e2b294c 100644
--- a/tensorflow/compiler/xla/debug_options_flags.cc
+++ b/tensorflow/compiler/xla/debug_options_flags.cc
@@ -64,7 +64,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_cpu_fast_math_honor_division(true);
 
   // By default, copy TF's Eigen style min_max behavior with nans.
-  opts.set_xla_cpu_enable_fast_min_max(false);
+  opts.set_xla_cpu_enable_fast_min_max(true);
 
   opts.set_xla_gpu_enable_fast_min_max(true);
 
diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h
index 5c9f6677ab3..4c3167e16d9 100644
--- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h
@@ -50,6 +50,10 @@ class CpuElementalIrEmitter : public ElementalIrEmitter {
     return ir_emitter_->EmitThreadLocalCall(callee, parameters, name);
   }
 
+  bool fast_min_max() override {
+    return hlo_module_config_.debug_options().xla_cpu_enable_fast_min_max();
+  }
+
   IrEmitter* ir_emitter_;
 };
 
diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc
index f62769cc615..8d9229c1223 100644
--- a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc
+++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc
@@ -318,7 +318,9 @@ llvm::Value* GenerateVF32Log(llvm::IRBuilder<>* b, llvm::Value* input,
   llvm::Value* is_pos_inf_mask = vsl.FCmpEQMask(input, pos_inf);
 
   // Cut off denormalized stuff.
-  llvm::Value* tmp0 = vsl.Max(min_norm_pos, input);
+  // Always allow fast max because we are checking for the nan above.
+  llvm::Value* tmp0 =
+      vsl.Max(min_norm_pos, input, /*enable_fast_min_max=*/true);
 
   // VectorSupportLibrary (intentionally) can't juggle more than one type at a
   // time so drop down to IRBuilder for this bit.
diff --git a/tensorflow/compiler/xla/service/cpu/vector_support_library.cc b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc
index b15ad1e162d..0d2eab9fd42 100644
--- a/tensorflow/compiler/xla/service/cpu/vector_support_library.cc
+++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.cc
@@ -80,10 +80,11 @@ llvm::Value* VectorSupportLibrary::Sub(llvm::Value* lhs, llvm::Value* rhs) {
   return b()->CreateFSub(lhs, rhs);
 }
 
-llvm::Value* VectorSupportLibrary::Max(llvm::Value* lhs, llvm::Value* rhs) {
+llvm::Value* VectorSupportLibrary::Max(llvm::Value* lhs, llvm::Value* rhs,
+                                       bool enable_fast_min_max) {
   AssertCorrectTypes({lhs, rhs});
   if (scalar_type_->isFloatingPointTy()) {
-    return llvm_ir::EmitFloatMax(lhs, rhs, b_);
+    return llvm_ir::EmitFloatMax(lhs, rhs, b_, enable_fast_min_max);
   } else {
     LOG(FATAL) << "Max for integers is unimplemented";
   }
diff --git a/tensorflow/compiler/xla/service/cpu/vector_support_library.h b/tensorflow/compiler/xla/service/cpu/vector_support_library.h
index cbbc4d7bf34..f1a0b0a4406 100644
--- a/tensorflow/compiler/xla/service/cpu/vector_support_library.h
+++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.h
@@ -78,9 +78,11 @@ class VectorSupportLibrary {
   llvm::Value* Sub(llvm::Value* lhs, const llvm::APFloat& rhs) {
     return Sub(lhs, GetConstantFloat(lhs->getType(), rhs));
   }
-  llvm::Value* Max(llvm::Value* lhs, llvm::Value* rhs);
-  llvm::Value* Max(const llvm::APFloat& lhs, llvm::Value* rhs) {
-    return Max(GetConstantFloat(rhs->getType(), lhs), rhs);
+  llvm::Value* Max(llvm::Value* lhs, llvm::Value* rhs,
+                   bool enable_fast_min_max);
+  llvm::Value* Max(const llvm::APFloat& lhs, llvm::Value* rhs,
+                   bool enable_fast_min_max) {
+    return Max(GetConstantFloat(rhs->getType(), lhs), rhs, enable_fast_min_max);
   }
   llvm::Value* Div(llvm::Value* lhs, llvm::Value* rhs);
 
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
index e4097b0c06f..4b6c30cadc4 100644
--- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc
@@ -1313,12 +1313,12 @@ StatusOr ElementalIrEmitter::EmitComplexBinaryOp(
 
 llvm::Value* ElementalIrEmitter::EmitFloatMax(llvm::Value* lhs_value,
                                               llvm::Value* rhs_value) {
-  return llvm_ir::EmitFloatMax(lhs_value, rhs_value, b_);
+  return llvm_ir::EmitFloatMax(lhs_value, rhs_value, b_, fast_min_max());
 }
 
 llvm::Value* ElementalIrEmitter::EmitFloatMin(llvm::Value* lhs_value,
                                               llvm::Value* rhs_value) {
-  return llvm_ir::EmitFloatMin(lhs_value, rhs_value, b_);
+  return llvm_ir::EmitFloatMin(lhs_value, rhs_value, b_, fast_min_max());
 }
 
 StatusOr ElementalIrEmitter::EmitLog(PrimitiveType prim_type,
diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/elemental_ir_emitter.h
index e39d2dd99ec..365e3f56b85 100644
--- a/tensorflow/compiler/xla/service/elemental_ir_emitter.h
+++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.h
@@ -245,6 +245,8 @@ class ElementalIrEmitter : public IrBuilderMixin {
       std::vector initial_value_generators,
       const llvm_ir::IrArray::Index& index);
 
+  virtual bool fast_min_max() = 0;
+
   llvm::IRBuilder<>* const b_;
 
   llvm::Module* module_;
diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h
index a3056b1ddad..766a4c84df5 100644
--- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h
+++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h
@@ -96,6 +96,10 @@ class GpuElementalIrEmitter : public ElementalIrEmitter {
 
   llvm::Value* EmitThreadId() override;
 
+  bool fast_min_max() override {
+    return hlo_module_config_.debug_options().xla_gpu_enable_fast_min_max();
+  }
+
  private:
   // Emits IR for op, which must have opcode kPower.
   StatusOr EmitPowerOp(const HloInstruction* op,
diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
index e4ca08f972b..b01ae2efe43 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
@@ -91,10 +91,8 @@ llvm::CallInst* EmitCallToIntrinsic(
 }
 
 llvm::Value* EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value,
-                          llvm::IRBuilder<>* b) {
-  // TODO(tpopp): Pass this information down from the HLO's ModuleConfig.
-  if (b->getFastMathFlags().noNaNs() ||
-      GetDebugOptionsFromFlags().xla_cpu_enable_fast_min_max()) {
+                          llvm::IRBuilder<>* b, bool enable_fast_min_max) {
+  if (b->getFastMathFlags().noNaNs() || enable_fast_min_max) {
     auto cmp = b->CreateFCmpUGE(lhs_value, rhs_value);
     return b->CreateSelect(cmp, lhs_value, rhs_value);
   } else {
@@ -106,10 +104,8 @@ llvm::Value* EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value,
 }
 
 llvm::Value* EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value,
-                          llvm::IRBuilder<>* b) {
-  // TODO(tpopp): Pass this information down from the HLO's ModuleConfig.
-  if (b->getFastMathFlags().noNaNs() ||
-      GetDebugOptionsFromFlags().xla_cpu_enable_fast_min_max()) {
+                          llvm::IRBuilder<>* b, bool enable_fast_min_max) {
+  if (b->getFastMathFlags().noNaNs() || enable_fast_min_max) {
     auto cmp = b->CreateFCmpULE(lhs_value, rhs_value);
     return b->CreateSelect(cmp, lhs_value, rhs_value);
   } else {
diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h
index 691898011ed..642965b6470 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h
@@ -108,12 +108,12 @@ llvm::CallInst* EmitCallToIntrinsic(
 // Emit float max. Emit maxnum intrinsic is fast math is disabled, or
 // fcmp+select otherwise
 llvm::Value* EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value,
-                          llvm::IRBuilder<>* b);
+                          llvm::IRBuilder<>* b, bool enable_fast_min_max);
 
 // Emit float min. Emit minnum intrinsic is fast math is disabled, or
 // fcmp+select otherwise
 llvm::Value* EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value,
-                          llvm::IRBuilder<>* b);
+                          llvm::IRBuilder<>* b, bool enable_fast_min_max);
 
 // Convenience methods for emitting a GEP instruction that indexes into a buffer
 // (1-dimensional array), equivalent to array[index]. The type is automatically
diff --git a/tensorflow/compiler/xla/service/llvm_ir/math_ops.cc b/tensorflow/compiler/xla/service/llvm_ir/math_ops.cc
index 333a2e8f612..0604cb848d2 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/math_ops.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/math_ops.cc
@@ -31,9 +31,13 @@ llvm::Value* EmitFastTanh(llvm::IRBuilder<>* b, llvm::Value* input) {
       b->CreateFCmpOLT(abs_x, llvm::ConstantFP::get(type, kCanUseApprox));
 
   // Clamp the input to [-9, 9].
+  //
+  // To simplify the code base until it's an issue, don't have a slow min/max in
+  // this approximation.
   llvm::Value* input_clamped = llvm_ir::EmitFloatMin(
-      llvm_ir::EmitFloatMax(input, llvm::ConstantFP::get(type, -9.0), b),
-      llvm::ConstantFP::get(type, 9.0), b);
+      llvm_ir::EmitFloatMax(input, llvm::ConstantFP::get(type, -9.0), b,
+                            /*enable_fast_min_max=*/true),
+      llvm::ConstantFP::get(type, 9.0), b, /*enable_fast_min_max=*/true);
 
   static constexpr std::array numerator_coeffs{
       -2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f,
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index 6ecd1e015d2..bfe11b63eea 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -1004,6 +1004,8 @@ class ReluTest(test_lib.TestCase):
     z = self.evaluate(nn_ops.relu(constant_op.constant(x)))
     self.assertAllEqual(y, z)
 
+  @test_util.disable_xla(
+      "This test relies on undefined behavior that XLA does not replicate")
   @test_util.run_deprecated_v1
   def testNaNs(self):
     # Test that relu(nan) = nan for various sizes.

From b0199879f09e4bd0068e627e1a43384dcff34983 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Wed, 17 Jun 2020 01:26:51 -0700
Subject: [PATCH 0378/1390] Add lowering for TanhOp that uses an approximation
 instead of lowering to intrinsics.

The same approximation is used by the XLA compiler.

PiperOrigin-RevId: 316844625
Change-Id: I1a909bd063509491a6a58ae0acc3bfb919cb34d5
---
 tensorflow/compiler/mlir/xla/BUILD            |  19 ++
 .../tests/legalize_tanh_to_approximation.mlir | 134 ++++++++++++++
 .../legalize_tanh_to_approximation.cc         | 167 ++++++++++++++++++
 .../compiler/mlir/xla/transforms/passes.h     |   7 +
 .../compiler/mlir/xla/transforms/rewriters.h  |   8 +
 5 files changed, 335 insertions(+)
 create mode 100644 tensorflow/compiler/mlir/xla/tests/legalize_tanh_to_approximation.mlir
 create mode 100644 tensorflow/compiler/mlir/xla/transforms/legalize_tanh_to_approximation.cc

diff --git a/tensorflow/compiler/mlir/xla/BUILD b/tensorflow/compiler/mlir/xla/BUILD
index 43458aab2d3..d089f80d571 100644
--- a/tensorflow/compiler/mlir/xla/BUILD
+++ b/tensorflow/compiler/mlir/xla/BUILD
@@ -515,6 +515,24 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "xla_legalize_tanh_to_approximation",
+    srcs = ["transforms/legalize_tanh_to_approximation.cc"],
+    hdrs = [
+        "transforms/passes.h",
+        "transforms/rewriters.h",
+    ],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:StandardOps",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+    alwayslink = 1,
+)
+
 gentbl(
     name = "xla_lower_complex_inc_gen",
     tbl_outs = [
@@ -946,6 +964,7 @@ cc_library(
         ":xla_hlo_fusion",
         ":xla_hlo_to_lhlo_with_xla",
         ":xla_legalize_control_flow",
+        ":xla_legalize_tanh_to_approximation",
         ":xla_legalize_tf",
         ":xla_legalize_tf_with_tf2xla",
         ":xla_legalize_to_linalg",
diff --git a/tensorflow/compiler/mlir/xla/tests/legalize_tanh_to_approximation.mlir b/tensorflow/compiler/mlir/xla/tests/legalize_tanh_to_approximation.mlir
new file mode 100644
index 00000000000..a8286c9b5a9
--- /dev/null
+++ b/tensorflow/compiler/mlir/xla/tests/legalize_tanh_to_approximation.mlir
@@ -0,0 +1,134 @@
+// RUN: xla-opt -xla-legalize-tanh-to-approximation -split-input-file %s | FileCheck %s
+
+func @tanh_f64(%arg0 : f64) -> f64 {
+  %res = tanh %arg0 : f64
+  return %res : f64
+}
+
+// CHECK-LABEL: @tanh_f64
+// CHECK: tanh
+
+// -----
+
+func @tanh_f32(%arg0 : f32) -> f32 {
+  %res = tanh %arg0 : f32
+  return %res : f32
+}
+
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+// CHECK:       module {
+
+// CHECK-LABEL:   func @tanh_f32(
+// CHECK-SAME:                   %[[VAL_0:.*]]: f32) -> f32 {
+// CHECK:           %[[VAL_1:.*]] = constant 2.000000e+01 : f32
+// CHECK:           %[[VAL_2:.*]] = constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_3:.*]] = constant 4.000000e-04 : f32
+// CHECK:           %[[VAL_4:.*]] = constant 9.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = constant -2.76076837E-16 : f32
+// CHECK:           %[[VAL_6:.*]] = constant 2.00018794E-13 : f32
+// CHECK:           %[[VAL_7:.*]] = constant -8.60467184E-11 : f32
+// CHECK:           %[[VAL_8:.*]] = constant 5.12229725E-8 : f32
+// CHECK:           %[[VAL_9:.*]] = constant 1.48572235E-5 : f32
+// CHECK:           %[[VAL_10:.*]] = constant 6.37261954E-4 : f32
+// CHECK:           %[[VAL_11:.*]] = constant 0.00489352457 : f32
+// CHECK:           %[[VAL_12:.*]] = constant 1.19825836E-6 : f32
+// CHECK:           %[[VAL_13:.*]] = constant 1.18534706E-4 : f32
+// CHECK:           %[[VAL_14:.*]] = constant 0.00226843474 : f32
+// CHECK:           %[[VAL_15:.*]] = constant 0.00489352504 : f32
+// CHECK:           %[[VAL_16:.*]] = absf %[[VAL_0]] : f32
+// CHECK:           %[[VAL_17:.*]] = copysign %[[VAL_2]], %[[VAL_0]] : f32
+// CHECK:           %[[VAL_18:.*]] = cmpf "ult", %[[VAL_16]], %[[VAL_1]] : f32
+// CHECK:           %[[VAL_19:.*]] = cmpf "olt", %[[VAL_16]], %[[VAL_3]] : f32
+// CHECK:           %[[VAL_20:.*]] = cmpf "ule", %[[VAL_16]], %[[VAL_4]] : f32
+// CHECK:           %[[VAL_21:.*]] = copysign %[[VAL_4]], %[[VAL_0]] : f32
+// CHECK:           %[[VAL_22:.*]] = select %[[VAL_20]], %[[VAL_0]], %[[VAL_21]] : f32
+// CHECK:           %[[VAL_23:.*]] = mulf %[[VAL_22]], %[[VAL_22]] : f32
+// CHECK:           %[[VAL_24:.*]] = mulf %[[VAL_23]], %[[VAL_5]] : f32
+// CHECK:           %[[VAL_25:.*]] = addf %[[VAL_24]], %[[VAL_6]] : f32
+// CHECK:           %[[VAL_26:.*]] = mulf %[[VAL_23]], %[[VAL_25]] : f32
+// CHECK:           %[[VAL_27:.*]] = addf %[[VAL_26]], %[[VAL_7]] : f32
+// CHECK:           %[[VAL_28:.*]] = mulf %[[VAL_23]], %[[VAL_27]] : f32
+// CHECK:           %[[VAL_29:.*]] = addf %[[VAL_28]], %[[VAL_8]] : f32
+// CHECK:           %[[VAL_30:.*]] = mulf %[[VAL_23]], %[[VAL_29]] : f32
+// CHECK:           %[[VAL_31:.*]] = addf %[[VAL_30]], %[[VAL_9]] : f32
+// CHECK:           %[[VAL_32:.*]] = mulf %[[VAL_23]], %[[VAL_31]] : f32
+// CHECK:           %[[VAL_33:.*]] = addf %[[VAL_32]], %[[VAL_10]] : f32
+// CHECK:           %[[VAL_34:.*]] = mulf %[[VAL_23]], %[[VAL_33]] : f32
+// CHECK:           %[[VAL_35:.*]] = addf %[[VAL_34]], %[[VAL_11]] : f32
+// CHECK:           %[[VAL_36:.*]] = mulf %[[VAL_22]], %[[VAL_35]] : f32
+// CHECK:           %[[VAL_37:.*]] = mulf %[[VAL_23]], %[[VAL_12]] : f32
+// CHECK:           %[[VAL_38:.*]] = addf %[[VAL_37]], %[[VAL_13]] : f32
+// CHECK:           %[[VAL_39:.*]] = mulf %[[VAL_23]], %[[VAL_38]] : f32
+// CHECK:           %[[VAL_40:.*]] = addf %[[VAL_39]], %[[VAL_14]] : f32
+// CHECK:           %[[VAL_41:.*]] = mulf %[[VAL_23]], %[[VAL_40]] : f32
+// CHECK:           %[[VAL_42:.*]] = addf %[[VAL_41]], %[[VAL_15]] : f32
+// CHECK:           %[[VAL_43:.*]] = divf %[[VAL_36]], %[[VAL_42]] : f32
+// CHECK:           %[[VAL_44:.*]] = select %[[VAL_19]], %[[VAL_0]], %[[VAL_43]] : f32
+// CHECK:           %[[VAL_45:.*]] = select %[[VAL_18]], %[[VAL_44]], %[[VAL_17]] : f32
+// CHECK:           return %[[VAL_45]] : f32
+// CHECK:         }
+// CHECK:       }
+
+// -----
+
+func @tanh_f16(%arg0 : f16) -> f16 {
+  %res = tanh %arg0 : f16
+  return %res : f16
+}
+
+// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
+// CHECK:       module {
+
+// CHECK-LABEL:   func @tanh_f16(
+// CHECK-SAME:                   %[[VAL_0:.*]]: f16) -> f16 {
+// CHECK:           %[[VAL_1:.*]] = constant 2.000000e+01 : f32
+// CHECK:           %[[VAL_2:.*]] = constant 1.000000e+00 : f32
+// CHECK:           %[[VAL_3:.*]] = constant 4.000000e-04 : f32
+// CHECK:           %[[VAL_4:.*]] = constant 9.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = constant -2.76076837E-16 : f32
+// CHECK:           %[[VAL_6:.*]] = constant 2.00018794E-13 : f32
+// CHECK:           %[[VAL_7:.*]] = constant -8.60467184E-11 : f32
+// CHECK:           %[[VAL_8:.*]] = constant 5.12229725E-8 : f32
+// CHECK:           %[[VAL_9:.*]] = constant 1.48572235E-5 : f32
+// CHECK:           %[[VAL_10:.*]] = constant 6.37261954E-4 : f32
+// CHECK:           %[[VAL_11:.*]] = constant 0.00489352457 : f32
+// CHECK:           %[[VAL_12:.*]] = constant 1.19825836E-6 : f32
+// CHECK:           %[[VAL_13:.*]] = constant 1.18534706E-4 : f32
+// CHECK:           %[[VAL_14:.*]] = constant 0.00226843474 : f32
+// CHECK:           %[[VAL_15:.*]] = constant 0.00489352504 : f32
+// CHECK:           %[[VAL_16:.*]] = fpext %[[VAL_0]] : f16 to f32
+// CHECK:           %[[VAL_17:.*]] = absf %[[VAL_16]] : f32
+// CHECK:           %[[VAL_18:.*]] = copysign %[[VAL_2]], %[[VAL_16]] : f32
+// CHECK:           %[[VAL_19:.*]] = cmpf "ult", %[[VAL_17]], %[[VAL_1]] : f32
+// CHECK:           %[[VAL_20:.*]] = cmpf "olt", %[[VAL_17]], %[[VAL_3]] : f32
+// CHECK:           %[[VAL_21:.*]] = cmpf "ule", %[[VAL_17]], %[[VAL_4]] : f32
+// CHECK:           %[[VAL_22:.*]] = copysign %[[VAL_4]], %[[VAL_16]] : f32
+// CHECK:           %[[VAL_23:.*]] = select %[[VAL_21]], %[[VAL_16]], %[[VAL_22]] : f32
+// CHECK:           %[[VAL_24:.*]] = mulf %[[VAL_23]], %[[VAL_23]] : f32
+// CHECK:           %[[VAL_25:.*]] = mulf %[[VAL_24]], %[[VAL_5]] : f32
+// CHECK:           %[[VAL_26:.*]] = addf %[[VAL_25]], %[[VAL_6]] : f32
+// CHECK:           %[[VAL_27:.*]] = mulf %[[VAL_24]], %[[VAL_26]] : f32
+// CHECK:           %[[VAL_28:.*]] = addf %[[VAL_27]], %[[VAL_7]] : f32
+// CHECK:           %[[VAL_29:.*]] = mulf %[[VAL_24]], %[[VAL_28]] : f32
+// CHECK:           %[[VAL_30:.*]] = addf %[[VAL_29]], %[[VAL_8]] : f32
+// CHECK:           %[[VAL_31:.*]] = mulf %[[VAL_24]], %[[VAL_30]] : f32
+// CHECK:           %[[VAL_32:.*]] = addf %[[VAL_31]], %[[VAL_9]] : f32
+// CHECK:           %[[VAL_33:.*]] = mulf %[[VAL_24]], %[[VAL_32]] : f32
+// CHECK:           %[[VAL_34:.*]] = addf %[[VAL_33]], %[[VAL_10]] : f32
+// CHECK:           %[[VAL_35:.*]] = mulf %[[VAL_24]], %[[VAL_34]] : f32
+// CHECK:           %[[VAL_36:.*]] = addf %[[VAL_35]], %[[VAL_11]] : f32
+// CHECK:           %[[VAL_37:.*]] = mulf %[[VAL_23]], %[[VAL_36]] : f32
+// CHECK:           %[[VAL_38:.*]] = mulf %[[VAL_24]], %[[VAL_12]] : f32
+// CHECK:           %[[VAL_39:.*]] = addf %[[VAL_38]], %[[VAL_13]] : f32
+// CHECK:           %[[VAL_40:.*]] = mulf %[[VAL_24]], %[[VAL_39]] : f32
+// CHECK:           %[[VAL_41:.*]] = addf %[[VAL_40]], %[[VAL_14]] : f32
+// CHECK:           %[[VAL_42:.*]] = mulf %[[VAL_24]], %[[VAL_41]] : f32
+// CHECK:           %[[VAL_43:.*]] = addf %[[VAL_42]], %[[VAL_15]] : f32
+// CHECK:           %[[VAL_44:.*]] = divf %[[VAL_37]], %[[VAL_43]] : f32
+// CHECK:           %[[VAL_45:.*]] = select %[[VAL_20]], %[[VAL_16]], %[[VAL_44]] : f32
+// CHECK:           %[[VAL_46:.*]] = select %[[VAL_19]], %[[VAL_45]], %[[VAL_18]] : f32
+// CHECK:           %[[VAL_47:.*]] = fptrunc %[[VAL_46]] : f32 to f16
+// CHECK:           return %[[VAL_47]] : f16
+// CHECK:         }
+// CHECK:       }
+
diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tanh_to_approximation.cc b/tensorflow/compiler/mlir/xla/transforms/legalize_tanh_to_approximation.cc
new file mode 100644
index 00000000000..9696db377da
--- /dev/null
+++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tanh_to_approximation.cc
@@ -0,0 +1,167 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+// This file implements logic for lowering the tanh standard ops to an
+// approximation.
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
+#include "mlir/IR/Function.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/xla/transforms/passes.h"
+#include "tensorflow/compiler/mlir/xla/transforms/rewriters.h"
+
+namespace mlir {
+namespace xla {
+namespace {
+
+/// Emits the fast tanh approximation that is also used by XLA.
+static Value EmitTanhApproximation(Value input, Value abs_value, Location loc,
+                                   PatternRewriter &rewriter) {
+  // For small values of x, we can approximate tanh(x)=x. For extremely small
+  // values of x (|x| < 1e-37), the other approximation would evaluate
+  // tanh(x) = 0.
+  constexpr float kCanUseApprox = 0.0004;
+  Value can_use_approx =
+      rewriter.create(loc, rewriter.getF32FloatAttr(kCanUseApprox));
+  Value return_input = rewriter.create(loc, CmpFPredicate::OLT,
+                                               abs_value, can_use_approx);
+
+  // Clamp the input to [-9, 9].
+  Value plus_nine =
+      rewriter.create(loc, rewriter.getF32FloatAttr(9.0));
+  Value smaller_than_nine =
+      rewriter.create(loc, CmpFPredicate::ULE, abs_value, plus_nine);
+  Value input_clamped = rewriter.create(
+      loc, smaller_than_nine, input,
+      rewriter.create(loc, plus_nine, input));
+
+  static constexpr std::array numerator_coeffs{
+      -2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f,
+      5.12229709037114e-08f,  1.48572235717979e-05f, 6.37261928875436e-04f,
+      4.89352455891786e-03f};
+
+  static constexpr std::array denominator_coeffs{
+      1.19825839466702e-06f, 1.18534705686654e-04f, 2.26843463243900e-03f,
+      4.89352518554385e-03f};
+
+  Value input_squared =
+      rewriter.create(loc, input_clamped, input_clamped);
+  Value numerator = rewriter.create(
+      loc, rewriter.getF32FloatAttr(numerator_coeffs[0]));
+  for (int i = 1; i < numerator_coeffs.size(); i++) {
+    numerator = rewriter.create(
+        loc, rewriter.create(loc, input_squared, numerator),
+        rewriter.create(
+            loc, rewriter.getF32FloatAttr(numerator_coeffs[i])));
+  }
+
+  numerator = rewriter.create(loc, input_clamped, numerator);
+
+  Value denominator = rewriter.create(
+      loc, rewriter.getF32FloatAttr(denominator_coeffs[0]));
+  for (int i = 1; i < denominator_coeffs.size(); i++) {
+    denominator = rewriter.create(
+        loc, rewriter.create(loc, input_squared, denominator),
+        rewriter.create(
+            loc, rewriter.getF32FloatAttr(denominator_coeffs[i])));
+  }
+
+  Value approx = rewriter.create(loc, numerator, denominator);
+
+  return rewriter.create(loc, return_input, input, approx);
+}
+
+class ApproximateTanhLowering : public OpRewritePattern {
+ public:
+  explicit ApproximateTanhLowering(MLIRContext *ctx)
+      : OpRewritePattern(ctx, 100) {}
+
+  LogicalResult matchAndRewrite(TanhOp tanhOp,
+                                PatternRewriter &rewriter) const override {
+    Type operand_type = tanhOp.getType();
+
+    if (operand_type.isF64()) {
+      // Similar to XLA, do not rewrite f64 as precision might matter.
+      return failure();
+    }
+
+    Location loc = tanhOp.getLoc();
+    Value input = tanhOp.operand();
+    if (operand_type.isF16()) {
+      input = rewriter.create(loc, input, rewriter.getF32Type());
+    }
+
+    // If we still do not have f32, fail.
+    if (!input.getType().isF32()) {
+      return failure();
+    }
+
+    // For |operand| > 20.0, we just return -1/1.
+    constexpr double kMaxValue = 20.0;
+    Value max_value =
+        rewriter.create(loc, rewriter.getF32FloatAttr(kMaxValue));
+    Value abs_value = rewriter.create(loc, input);
+
+    Value one = rewriter.create(loc, rewriter.getF32FloatAttr(1.0));
+    Value one_with_sign = rewriter.create(loc, one, input);
+
+    Value smaller_than_twenty =
+        rewriter.create(loc, CmpFPredicate::ULT, abs_value, max_value);
+
+    // Otherwise, we use the approximation.
+    Value approx = EmitTanhApproximation(input, abs_value, loc, rewriter);
+
+    Value result = rewriter.create(loc, smaller_than_twenty, approx,
+                                             one_with_sign);
+
+    // Truncate back if needed.
+    if (operand_type.isF16()) {
+      result = rewriter.create(loc, result, rewriter.getF16Type());
+    }
+
+    rewriter.replaceOp(tanhOp, {result});
+    return success();
+  }
+};
+
+struct LegalizeTanhToApproximation
+    : public PassWrapper {
+  /// Perform the lowering of standard dialect operations to approximations.
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    PopulateTanhToApproximationPatterns(&getContext(), &patterns);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr>
+createLegalizeTanhToApproximationPass() {
+  return std::make_unique();
+}
+
+void PopulateTanhToApproximationPatterns(mlir::MLIRContext *context,
+                                         OwningRewritePatternList *patterns) {
+  patterns->insert(context);
+}
+
+static PassRegistration legalize_pass(
+    "xla-legalize-tanh-to-approximation",
+    "Legalize tanh from standard dialect to an approximation");
+
+}  // namespace xla
+}  // namespace mlir
diff --git a/tensorflow/compiler/mlir/xla/transforms/passes.h b/tensorflow/compiler/mlir/xla/transforms/passes.h
index a2af8124786..3db0bc3b474 100644
--- a/tensorflow/compiler/mlir/xla/transforms/passes.h
+++ b/tensorflow/compiler/mlir/xla/transforms/passes.h
@@ -115,6 +115,13 @@ std::unique_ptr createLhloCopyRemovalPass();
 std::unique_ptr> createLegalizeLhloToParallelLoopsPass();
 
 }  // namespace xla_lhlo
+
+namespace xla {
+
+/// Lowers the standard TanhOp to an approximation that does not use intrinsics.
+std::unique_ptr> createLegalizeTanhToApproximationPass();
+
+}  // namespace xla
 }  // namespace mlir
 
 #endif  // TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_PASSES_H_
diff --git a/tensorflow/compiler/mlir/xla/transforms/rewriters.h b/tensorflow/compiler/mlir/xla/transforms/rewriters.h
index 59347198fe4..7303b87be75 100644
--- a/tensorflow/compiler/mlir/xla/transforms/rewriters.h
+++ b/tensorflow/compiler/mlir/xla/transforms/rewriters.h
@@ -91,6 +91,14 @@ void PopulateLegalizeChloToHloPatterns(MLIRContext *context,
 
 }  // namespace xla_chlo
 
+namespace xla {
+
+// Populates a pattern that translates the standard TanhOp to an approximation
+// that does not use intrinsics.
+void PopulateTanhToApproximationPatterns(MLIRContext *context,
+                                         OwningRewritePatternList *patterns);
+
+}  // namespace xla
 }  // namespace mlir
 
 #endif  // TENSORFLOW_COMPILER_MLIR_XLA_TRANSFORMS_REWRITERS_H_

From d5435d42c5856477cafc70103c71e4039861d990 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 02:01:56 -0700
Subject: [PATCH 0379/1390] compat: Update forward compatibility horizon to
 2020-06-17

PiperOrigin-RevId: 316848134
Change-Id: Ied10fd4074e6cc4e0a737cdf4e6ed6ab48632c94
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 314acfdd38f..8a42b3dfdd3 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 16)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 17)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 25e5d3cb1208ac9368a35ca97756a8a89110e2cf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 02:02:02 -0700
Subject: [PATCH 0380/1390] Update GraphDef version to 435.

PiperOrigin-RevId: 316848144
Change-Id: I871d54f7b4f40d8ed9ab1d947302934bb03b463a
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 8e3c66edfc2..3e4e3888d87 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 434  // Updated: 2020/6/16
+#define TF_GRAPH_DEF_VERSION 435  // Updated: 2020/6/17
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From eaf10e9a279e270d8133a41cfc9690de10423a1d Mon Sep 17 00:00:00 2001
From: Lutz Roeder 
Date: Sun, 14 Jun 2020 23:02:59 -0700
Subject: [PATCH 0381/1390] Fix Keras documentation

---
 tensorflow/python/keras/layers/core.py     | 8 ++++++--
 tensorflow/python/keras/layers/wrappers.py | 9 +++++----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py
index abfb025db30..e64a1c27bcf 100644
--- a/tensorflow/python/keras/layers/core.py
+++ b/tensorflow/python/keras/layers/core.py
@@ -825,10 +825,14 @@ class Lambda(Layer):
       returned as output mask regardless of what the input is.
     arguments: Optional dictionary of keyword arguments to be passed to the
       function.
-  Input shape: Arbitrary. Use the keyword argument input_shape (tuple of
+
+  Input shape:
+    Arbitrary. Use the keyword argument input_shape (tuple of
     integers, does not include the samples axis) when using this layer as the
     first layer in a model.
-  Output shape: Specified by `output_shape` argument
+
+  Output shape:
+    Specified by `output_shape` argument
   """
 
   @trackable.no_automatic_dependency_tracking
diff --git a/tensorflow/python/keras/layers/wrappers.py b/tensorflow/python/keras/layers/wrappers.py
index 8fe3b3b20bb..23fef467cfe 100644
--- a/tensorflow/python/keras/layers/wrappers.py
+++ b/tensorflow/python/keras/layers/wrappers.py
@@ -341,10 +341,11 @@ class Bidirectional(Wrapper):
       combined. One of {'sum', 'mul', 'concat', 'ave', None}. If None, the
       outputs will not be combined, they will be returned as a list. Default
       value is 'concat'.
-    backward_layer: Optional `keras.layers.RNN`, or keras.layers.Layer` instance
-      to be used to handle backwards input processing. If `backward_layer` is
-      not provided, the layer instance passed as the `layer` argument will be
-      used to generate the backward layer automatically.
+    backward_layer: Optional `keras.layers.RNN`, or `keras.layers.Layer`
+      instance to be used to handle backwards input processing.
+      If `backward_layer` is not provided, the layer instance passed as the
+      `layer` argument will be used to generate the backward layer
+      automatically.
       Note that the provided `backward_layer` layer should have properties
       matching those of the `layer` argument, in particular it should have the
       same values for `stateful`, `return_states`, `return_sequence`, etc.

From 2aa2332364e4b27bb4fa433f0c6bd1709fab7f34 Mon Sep 17 00:00:00 2001
From: Nishidha Panpaliya 
Date: Wed, 17 Jun 2020 09:13:57 +0000
Subject: [PATCH 0382/1390] Fixed build error

---
 tensorflow/compiler/xla/service/cpu/BUILD       |  2 +-
 tensorflow/compiler/xla/service/cpu/tests/BUILD | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD
index 7be4d3e724a..8ea3672d6a8 100644
--- a/tensorflow/compiler/xla/service/cpu/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/BUILD
@@ -1080,7 +1080,7 @@ tf_cc_test(
     deps = [
         ":cpu_compiler",
         ":cpu_transfer_manager",
-	":test_header_helper",
+        ":test_header_helper",
         "//tensorflow/compiler/xla:test",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD
index 9036c5c9024..510e7b18fba 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD
@@ -42,7 +42,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
         "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
-	"//tensorflow/compiler/xla/service/cpu:test_header_helper",
+        "//tensorflow/compiler/xla/service/cpu:test_header_helper",
         "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
@@ -138,7 +138,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
         "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
         "//tensorflow/compiler/xla/tests:test_utils",
-	"//tensorflow/compiler/xla/service/cpu:test_header_helper",
+        "//tensorflow/compiler/xla/service/cpu:test_header_helper",
         "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
@@ -218,7 +218,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service:hlo_parser",
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
         "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
-	"//tensorflow/compiler/xla/service/cpu:test_header_helper",
+        "//tensorflow/compiler/xla/service/cpu:test_header_helper",
         "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
@@ -232,7 +232,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
         "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
-	"//tensorflow/compiler/xla/service/cpu:test_header_helper",
+        "//tensorflow/compiler/xla/service/cpu:test_header_helper",
         "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
@@ -246,7 +246,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
         "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
-	"//tensorflow/compiler/xla/service/cpu:test_header_helper",
+        "//tensorflow/compiler/xla/service/cpu:test_header_helper",
         "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",

From 41970246764299cd059fee0cbb08b5ba904b1520 Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Wed, 17 Jun 2020 02:06:23 -0700
Subject: [PATCH 0383/1390] Publish TpuCompileOp.h header.

PiperOrigin-RevId: 316848726
Change-Id: I36ef7ae0f8357cf144d057851cd3267309632fe5
---
 tensorflow/core/tpu/kernels/BUILD            |  6 ++
 tensorflow/core/tpu/kernels/tpu_compile_op.h | 75 ++++++++++++++++++++
 2 files changed, 81 insertions(+)
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compile_op.h

diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index e7be7d2b062..318d60b22df 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -365,3 +365,9 @@ tf_proto_library_cc(
     srcs = ["tpu_compilation_cache.proto"],
     cc_api_version = 2,
 )
+
+cc_library(
+    name = "tpu_compile_op_hdrs",
+    hdrs = ["tpu_compile_op.h"],
+    deps = ["//tensorflow/core:framework"],
+)
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op.h b/tensorflow/core/tpu/kernels/tpu_compile_op.h
new file mode 100644
index 00000000000..8a1963dde5c
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op.h
@@ -0,0 +1,75 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_COMPILE_OP_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_TPU_COMPILE_OP_H_
+
+#include 
+
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+namespace tpu {
+// Forward declaration.
+class TpuCompileOpKernelImpl;
+}  // namespace tpu
+
+// The TPUCompile operator compiles a Tensorflow function into a
+// TPU executable to be run by TPUExecute.
+//
+class TpuCompileOp : public OpKernel {
+ public:
+  explicit TpuCompileOp(OpKernelConstruction* ctx);
+  ~TpuCompileOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override;
+
+ private:
+  std::unique_ptr impl_;
+
+  DISALLOW_COPY_AND_ASSIGN(TpuCompileOp);
+};
+
+// The TPUCompile operator compiles a MLIR module into a
+// TPU executable to be run by TPUExecute.
+//
+class TpuCompileMlirOp : public OpKernel {
+ public:
+  explicit TpuCompileMlirOp(OpKernelConstruction* ctx);
+  ~TpuCompileMlirOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override;
+
+ private:
+  std::unique_ptr impl_;
+
+  DISALLOW_COPY_AND_ASSIGN(TpuCompileMlirOp);
+};
+
+class TpuCompileSucceededAssertOp : public OpKernel {
+ public:
+  explicit TpuCompileSucceededAssertOp(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+  ~TpuCompileSucceededAssertOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override;
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(TpuCompileSucceededAssertOp);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_COMPILE_OP_H_

From 70a617943c2f54bae01ffe695a56ba8e14ab4b65 Mon Sep 17 00:00:00 2001
From: Tiezhen WANG 
Date: Wed, 17 Jun 2020 02:14:20 -0700
Subject: [PATCH 0384/1390] TFL: Selective registration for C++ target.

Usage:
- Create a tflite_custom_cc_library rule in the BUILD file with the targeted model.
- Call tflite::CreateOpResolver to get the slimmed op resolver.
PiperOrigin-RevId: 316849510
Change-Id: I3e7d75da6a9f2876b3fbefe1962e5ae09ebadb33
---
 tensorflow/lite/build_def.bzl                 | 46 +++++++++++++++++++
 tensorflow/lite/java/src/main/native/BUILD    | 13 ++++++
 .../lite/java/src/main/native/op_resolver.h   | 26 +++++++++++
 .../java/src/main/native/selected_ops_jni.cc  | 36 +++++++++++++++
 4 files changed, 121 insertions(+)
 create mode 100644 tensorflow/lite/java/src/main/native/op_resolver.h
 create mode 100644 tensorflow/lite/java/src/main/native/selected_ops_jni.cc

diff --git a/tensorflow/lite/build_def.bzl b/tensorflow/lite/build_def.bzl
index 285824a613f..5e487395355 100644
--- a/tensorflow/lite/build_def.bzl
+++ b/tensorflow/lite/build_def.bzl
@@ -732,3 +732,49 @@ def tflite_experimental_runtime_linkopts(if_eager = [], if_non_eager = [], if_no
         ] + if_non_eager,
         if_none = [] + if_none,
     )
+
+def tflite_custom_cc_library(name, models = [], srcs = [], deps = [], visibility = ["//visibility:private"]):
+    """Generates a tflite cc library, stripping off unused operators.
+
+    This library includes the TfLite runtime as well as all operators needed for the given models.
+    Op resolver can be retrieved using tflite::CreateOpResolver method.
+
+    Args:
+        name: Str, name of the target.
+        models: List of models. This TFLite build will only include
+            operators used in these models. If the list is empty, all builtin
+            operators are included.
+        srcs: List of files implementing custom operators if any.
+        deps: Additional dependencies to build all the custom operators.
+        visibility: Visibility setting for the generated target. Default to private.
+    """
+    real_srcs = []
+    real_srcs.extend(srcs)
+    real_deps = []
+    real_deps.extend(deps)
+
+    if models:
+        gen_selected_ops(
+            name = "%s_registration" % name,
+            model = models[0],
+        )
+        real_srcs.append(":%s_registration" % name)
+        real_deps.append("//tensorflow/lite/java/src/main/native:selected_ops_jni")
+    else:
+        # Support all operators if `models` not specified.
+        real_deps.append("//tensorflow/lite/java/src/main/native")
+
+    native.cc_library(
+        name = name,
+        srcs = real_srcs,
+        copts = tflite_copts(),
+        linkopts = [
+            "-lm",
+            "-ldl",
+        ],
+        deps = depset([
+            "//tensorflow/lite:framework",
+            "//tensorflow/lite/kernels:builtin_ops",
+        ] + real_deps),
+        visibility = visibility,
+    )
diff --git a/tensorflow/lite/java/src/main/native/BUILD b/tensorflow/lite/java/src/main/native/BUILD
index 0d3535b29af..fdbbc9dc72c 100644
--- a/tensorflow/lite/java/src/main/native/BUILD
+++ b/tensorflow/lite/java/src/main/native/BUILD
@@ -45,14 +45,27 @@ cc_library(
     srcs = [
         "builtin_ops_jni.cc",
     ],
+    hdrs = ["op_resolver.h"],
     copts = tflite_copts(),
     deps = [
         ":native_framework_only",
+        "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels:builtin_ops",
     ],
     alwayslink = 1,
 )
 
+# TODO(b/153652701): Generate this target to give CreateOpResolver a custom namespace.
+cc_library(
+    name = "selected_ops_jni",
+    srcs = ["selected_ops_jni.cc"],
+    hdrs = ["op_resolver.h"],
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:framework",
+    ],
+)
+
 exports_files(
     [
         "exported_symbols.lds",
diff --git a/tensorflow/lite/java/src/main/native/op_resolver.h b/tensorflow/lite/java/src/main/native/op_resolver.h
new file mode 100644
index 00000000000..ba9c1bfb487
--- /dev/null
+++ b/tensorflow/lite/java/src/main/native/op_resolver.h
@@ -0,0 +1,26 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_JAVA_SRC_MAIN_NATIVE_OP_RESOLVER_H_
+#define TENSORFLOW_LITE_JAVA_SRC_MAIN_NATIVE_OP_RESOLVER_H_
+
+#include "tensorflow/lite/op_resolver.h"
+
+namespace tflite {
+
+std::unique_ptr CreateOpResolver();
+
+}
+
+#endif  // TENSORFLOW_LITE_JAVA_SRC_MAIN_NATIVE_OP_RESOLVER_H_
diff --git a/tensorflow/lite/java/src/main/native/selected_ops_jni.cc b/tensorflow/lite/java/src/main/native/selected_ops_jni.cc
new file mode 100644
index 00000000000..d8eb233f90a
--- /dev/null
+++ b/tensorflow/lite/java/src/main/native/selected_ops_jni.cc
@@ -0,0 +1,36 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/lite/java/src/main/native/op_resolver.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+// This method is generated by `gen_selected_ops`.
+// TODO(b/153652701): Instead of relying on a global method, make
+// `gen_selected_ops` generating a header file with custom namespace.
+void RegisterSelectedOps(::tflite::MutableOpResolver* resolver);
+
+namespace tflite {
+// This interface is the unified entry point for creating op resolver
+// regardless if selective registration is being used. C++ client will call
+// this method directly and Java client will call this method indirectly via
+// JNI code in interpreter_jni.cc.
+std::unique_ptr CreateOpResolver() {
+  std::unique_ptr resolver =
+      std::unique_ptr(new MutableOpResolver());
+  RegisterSelectedOps(resolver.get());
+  return std::move(resolver);
+}
+
+}  // namespace tflite

From 17cacd3a2124616876096c94699931f6ffb64e2c Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Wed, 17 Jun 2020 02:32:08 -0700
Subject: [PATCH 0385/1390] Add support for unrolling to the build rules.

PiperOrigin-RevId: 316851326
Change-Id: Icef09025b154f6d88c94884a7970f93022bcd160
---
 .../core/kernels/cubin_headers/build_defs.bzl    | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/cubin_headers/build_defs.bzl b/tensorflow/core/kernels/cubin_headers/build_defs.bzl
index f0f4a944e74..5880cbe8add 100644
--- a/tensorflow/core/kernels/cubin_headers/build_defs.bzl
+++ b/tensorflow/core/kernels/cubin_headers/build_defs.bzl
@@ -15,9 +15,11 @@ def _gen_kernel_image_hdr_impl(ctx):
 
     name = ctx.attr.name
     tile_sizes = ctx.attr.tile_size.replace("x", ",")
-    same_shape = []
+    cmd_args = []
     if ctx.attr.same_shape:
-        same_shape.append("--same_shape=%s" % ctx.attr.same_shape)
+        cmd_args.append("--same_shape=%s" % ctx.attr.same_shape)
+    if ctx.attr.unroll_factors:
+        cmd_args.append("--unroll_factors=%s" % ctx.attr.unroll_factors)
 
     cubins = []
     images = []
@@ -30,7 +32,7 @@ def _gen_kernel_image_hdr_impl(ctx):
             inputs = [ctx.file.mlir_op],
             outputs = [cubin],
             executable = ctx.executable._tool,
-            arguments = same_shape + [
+            arguments = cmd_args + [
                 "--tile_sizes=%s" % tile_sizes,
                 "--arch=%s" % arch.split("_")[1],
                 "--input=%s" % ctx.file.mlir_op.path,
@@ -74,6 +76,7 @@ _gen_kernel_image_hdr_rule = rule(
         "mlir_op": attr.label(mandatory = True, allow_single_file = True),
         "tile_size": attr.string(mandatory = True),
         "same_shape": attr.string(),
+        "unroll_factors": attr.string(),
         "out": attr.output(mandatory = True),
         "symbol": attr.string(mandatory = True),
         "gpu_archs": attr.string_list(mandatory = True),
@@ -88,7 +91,7 @@ _gen_kernel_image_hdr_rule = rule(
     },
 )
 
-def _gen_kernel_image_hdr(name, mlir_op, tile_size, tags = [], same_shape = None):
+def _gen_kernel_image_hdr(name, mlir_op, tile_size, tags = [], same_shape = None, unroll_factors = None):
     """Generates a C header with fatbin data from a Tensorflow op."""
     if cuda_gpu_architectures():
         _gen_kernel_image_hdr_rule(
@@ -96,6 +99,7 @@ def _gen_kernel_image_hdr(name, mlir_op, tile_size, tags = [], same_shape = None
             mlir_op = mlir_op,
             tile_size = tile_size,
             same_shape = same_shape,
+            unroll_factors = unroll_factors,
             out = "%s.h" % name,
             symbol = "k%s" % name.replace("_", " ").title().replace(" ", ""),
             gpu_archs = cuda_gpu_architectures(),
@@ -131,13 +135,14 @@ def _gen_mlir_op(name, type):
         out = "{name}_{type}.mlir".format(name = name, type = type),
     )
 
-def gen_kernel_library(name, types, tile_size, tags = [], same_shape = None):
+def gen_kernel_library(name, types, tile_size, tags = [], same_shape = None, unroll_factors = None):
     """ Generate a library with kernels for a specific tensorflow op.
 
     Args:
       name: The name of the tensorflow op.
       types: The types ("f16", "f32", "f64") for which a kernel should be generated.
       tile_size: The tiling specification, e.g. "16x16".
+      unroll_factors: The unrolling specification, e.g. "4,4"
       tags: The tags which should be added to the library.
       same_shape: The information about which shapes are the same, e.g. "0,1".
     """
@@ -154,6 +159,7 @@ def gen_kernel_library(name, types, tile_size, tags = [], same_shape = None):
                 tile_size = tile_size,
                 tags = tags,
                 same_shape = same_shape,
+                unroll_factors = unroll_factors,
             )
 
     native.cc_library(

From 80e9a8d07680ffd1b138ab95306c1827c4e39865 Mon Sep 17 00:00:00 2001
From: Taehee Jeong 
Date: Wed, 17 Jun 2020 02:54:52 -0700
Subject: [PATCH 0386/1390] Add tflite_ios_static_framework rule for hiding
 symbols

Exposed C++ symbols might cause collisions with other libraries.

PiperOrigin-RevId: 316853905
Change-Id: Ib7fd3f14fa40bc5c3e2ece0bf244cfb6e0623770
---
 tensorflow/lite/experimental/ios/BUILD.apple  |  24 ++--
 .../ios/hide_symbols_with_whitelist.sh        | 135 ++++++++++++++++++
 tensorflow/lite/experimental/ios/ios.bzl      |  58 ++++++++
 .../ios/whitelist_TensorFlowLiteC.txt         |   1 +
 .../ios/whitelist_TensorFlowLiteCCoreML.txt   |   2 +
 .../ios/whitelist_TensorFlowLiteCMetal.txt    |   2 +
 6 files changed, 212 insertions(+), 10 deletions(-)
 create mode 100755 tensorflow/lite/experimental/ios/hide_symbols_with_whitelist.sh
 create mode 100644 tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteC.txt
 create mode 100644 tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCCoreML.txt
 create mode 100644 tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCMetal.txt

diff --git a/tensorflow/lite/experimental/ios/BUILD.apple b/tensorflow/lite/experimental/ios/BUILD.apple
index aa41b9e2d62..1a85b604f9b 100644
--- a/tensorflow/lite/experimental/ios/BUILD.apple
+++ b/tensorflow/lite/experimental/ios/BUILD.apple
@@ -1,7 +1,7 @@
 # TensorFlow Lite for iOS
 
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
-load("//tensorflow/lite/experimental/ios:ios.bzl", "TFL_MINIMUM_OS_VERSION")
+load("//tensorflow/lite/experimental/ios:ios.bzl", "TFL_MINIMUM_OS_VERSION", "tflite_ios_static_framework")
 load("@build_bazel_rules_apple//apple:ios.bzl", "ios_static_framework")
 
 package(
@@ -11,8 +11,15 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+sh_binary(
+    name = "hide_symbols_with_whitelist",
+    srcs = [
+        "hide_symbols_with_whitelist.sh",
+    ],
+)
+
 # bazel build -c opt --config=ios_fat //tensorflow/lite/experimental/ios:TensorFlowLiteC_framework
-ios_static_framework(
+tflite_ios_static_framework(
     name = "TensorFlowLiteC_framework",
     hdrs = [
         "//tensorflow/lite/c:c_api.h",
@@ -20,6 +27,7 @@ ios_static_framework(
     ],
     bundle_name = "TensorFlowLiteC",
     minimum_os_version = TFL_MINIMUM_OS_VERSION,
+    whitelist_symbols_file = ":whitelist_TensorFlowLiteC.txt",
     deps = [
         ":tensorflow_lite_c",
     ],
@@ -60,16 +68,14 @@ genrule(
 # TensorFlowLiteC framework above in a composable way.
 #
 # bazel build -c opt --config=ios_fat //tensorflow/lite/experimental/ios:TensorFlowLiteCCoreMl_framework
-ios_static_framework(
+tflite_ios_static_framework(
     name = "TensorFlowLiteCCoreML_framework",
     hdrs = [
         ":coreml_delegate.h",
     ],
-    avoid_deps = [
-        ":tensorflow_lite_c",
-    ],
     bundle_name = "TensorFlowLiteCCoreML",
     minimum_os_version = TFL_MINIMUM_OS_VERSION,
+    whitelist_symbols_file = ":whitelist_TensorFlowLiteCCoreML.txt",
     deps = [
         "//tensorflow/lite/experimental/delegates/coreml:coreml_delegate",
     ],
@@ -81,16 +87,14 @@ ios_static_framework(
 # TensorFlowLiteC framework above in a composable way.
 #
 # bazel build -c opt --config=ios_fat //tensorflow/lite/experimental/ios:TensorFlowLiteCMetal_framework
-ios_static_framework(
+tflite_ios_static_framework(
     name = "TensorFlowLiteCMetal_framework",
     hdrs = [
         "//tensorflow/lite/delegates/gpu:metal_delegate.h",
     ],
-    avoid_deps = [
-        ":tensorflow_lite_c",
-    ],
     bundle_name = "TensorFlowLiteCMetal",
     minimum_os_version = TFL_MINIMUM_OS_VERSION,
+    whitelist_symbols_file = ":whitelist_TensorFlowLiteCMetal.txt",
     deps = [
         "//tensorflow/lite/delegates/gpu:metal_delegate",
     ],
diff --git a/tensorflow/lite/experimental/ios/hide_symbols_with_whitelist.sh b/tensorflow/lite/experimental/ios/hide_symbols_with_whitelist.sh
new file mode 100755
index 00000000000..2fa0fc53c33
--- /dev/null
+++ b/tensorflow/lite/experimental/ios/hide_symbols_with_whitelist.sh
@@ -0,0 +1,135 @@
+#!/bin/bash
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+#
+# A script to merge Mach-O object files into a single object file and hide
+# their internal symbols. Only whitelisted symbols will be visible in the
+# symbol table after this script.
+
+# To run this script, you must set several variables:
+#   INPUT_FRAMEWORK: a zip file containing the iOS static framework.
+#   BUNDLE_NAME: the pod/bundle name of the iOS static framework.
+#   WHITELIST_FILE_PATH: contains the whitelisted symbols.
+#   OUTPUT: the output zip file.
+
+# Halt on any error or any unknown variable.
+set -ue
+
+LD_DEBUGGABLE_FLAGS="-x"
+# Uncomment the below to get debuggable output. This can only be done for one
+# library at a time.
+# LD_DEBUGGABLE_FLAGS="-d"
+
+# Exits if C++ symbols are found in the whitelist list.
+if grep -q "^__Z" "${WHITELIST_FILE_PATH}"
+then
+  echo "ERROR: Failed in symbol hiding. This rule does not permit hiding of" \
+       "C++ symbols due to possible serious problems mixing symbol hiding," \
+       "shared libraries and the C++ runtime." \
+       "More info can be found in go/ios-symbols-hiding." \
+       "Please recheck the whitelist list and remove C++ symbols:"
+  echo "$(grep "^__Z" "${WHITELIST_FILE_PATH}")"
+  exit 1 # terminate and indicate error
+fi
+# Unzips the framework zip file into a temp workspace.
+framework=$(mktemp -t framework -d)
+unzip "${INPUT_FRAMEWORK}" -d "${framework}"/
+
+# Executable file in the framework.
+executable_file="${BUNDLE_NAME}.framework/${BUNDLE_NAME}"
+
+# Extracts architectures from the framework binary.
+archs_str=$(xcrun lipo -info "${framework}/${executable_file}" |
+sed -En -e 's/^(Non-|Architectures in the )fat file: .+( is architecture| are): (.*)$/\3/p')
+
+IFS=' ' read -r -a archs <<< "${archs_str}"
+
+merge_cmd=(xcrun lipo)
+
+# Merges object files and hide symbols for each architecture.
+for arch in "${archs[@]}"
+do
+    archdir=$(mktemp -t "${arch}" -d)
+    arch_file="${archdir}/${arch}"
+
+    # Handles the binary differently if they are fat or thin.
+    if [[ "${#archs[@]}" -gt 1 ]]; then
+       xcrun lipo "${framework}/${executable_file}" -thin "${arch}" -output "${arch_file}"
+    else
+       mv "${framework}/${executable_file}" "${arch_file}"
+    fi
+    if [[ "$arch" == "armv7" ]]; then
+      # Check that there are no thread local variables in the input, as they get broken.
+      # See b/124533863.
+      thread_locals=$(xcrun nm -m -g "${arch_file}" | awk '/__DATA,__thread_vars/ { print $5 }' | c++filt)
+      if [[ -n "${thread_locals}" ]]; then
+         echo
+         echo "WARNING: This symbol hiding script breaks thread local variables on 32-bit arm, you had:"
+         echo "${thread_locals}"
+         echo
+         echo "Your build will crash if these variables are actually used at runtime."
+         echo
+      fi
+    fi
+    xcrun ar -x "${arch_file}"
+    mv *.o "${archdir}"/
+
+    objects_file_list=$(mktemp)
+    # Hides the symbols except the whitelisted ones.
+    find "${archdir}" -name "*.o" >> "${objects_file_list}"
+
+    # Checks whether bitcode is enabled in the framework.
+    all_objects_have_bitcode=true
+    for object_file in $(cat "$objects_file_list"); do
+      if otool -arch "${arch}" -l "${object_file}" | grep -q __LLVM; then
+        : # Do nothing
+      else
+        echo "The ${arch} in ${object_file} is NOT bitcode-enabled."
+        all_objects_have_bitcode=false
+        break
+      fi
+    done
+    if [[ "$all_objects_have_bitcode" = "true" ]]; then
+      echo "The ${arch} in ${executable_file} is fully bitcode-enabled."
+      xcrun ld -r -bitcode_bundle -exported_symbols_list \
+        "${WHITELIST_FILE_PATH}" \
+        $LD_DEBUGGABLE_FLAGS \
+        -filelist "${objects_file_list}" -o "${arch_file}_processed.o"
+    else
+      echo "The ${arch} in ${executable_file} is NOT fully bitcode-enabled."
+      xcrun ld -r -exported_symbols_list \
+        "${WHITELIST_FILE_PATH}" \
+        $LD_DEBUGGABLE_FLAGS \
+        -filelist "${objects_file_list}" -o "${arch_file}_processed.o"
+    fi
+
+    output_object="${framework}/${arch}"
+
+    mv "${arch_file}_processed.o" "${output_object}"
+    rm -rf "${archdir}"
+    rm "${objects_file_list}"
+    merge_cmd+=(-arch "${arch}" "${output_object}")
+done
+
+# Repackages the processed object files.
+unzip "${INPUT_FRAMEWORK}"
+merge_cmd+=(-create -output "${BUNDLE_NAME}")
+"${merge_cmd[@]}"
+
+chmod +x "${BUNDLE_NAME}"
+rm "${executable_file}"
+mv "${BUNDLE_NAME}" "${executable_file}"
+( TZ=UTC find "${BUNDLE_NAME}.framework/" -exec touch -h -t 198001010000 {} \+ )
+zip --compression-method store --symlinks --recurse-paths --quiet "${OUTPUT}" "${BUNDLE_NAME}.framework/"
diff --git a/tensorflow/lite/experimental/ios/ios.bzl b/tensorflow/lite/experimental/ios/ios.bzl
index 976c6b09a97..3181b587e72 100644
--- a/tensorflow/lite/experimental/ios/ios.bzl
+++ b/tensorflow/lite/experimental/ios/ios.bzl
@@ -1,5 +1,8 @@
 """TensorFlow Lite Build Configurations for iOS"""
 
+# Placeholder for Google-internal load statements.
+load("@build_bazel_rules_apple//apple:ios.bzl", "ios_static_framework")
+
 TFL_MINIMUM_OS_VERSION = "9.0"
 
 # Default tags for filtering iOS targets. Targets are restricted to Apple platforms.
@@ -13,3 +16,58 @@ TFL_DISABLED_SANITIZER_TAGS = [
     "nomsan",
     "notsan",
 ]
+
+# iOS static framework with symbol whitelist. Exported C++ symbbols might cause
+# symbol collision with other libraries. List of symbols to whitelist can be
+# generated by running `nm -m -g FRAMEWORK_LIBRARY | grep _TfLite` for framework
+# built with `ios_static_framework` rule.
+def tflite_ios_static_framework(
+        name,
+        bundle_name,
+        whitelist_symbols_file,
+        exclude_resources = True,
+        **kwargs):
+    """TFLite variant of ios_static_framework with symbol hiding.
+
+    Args:
+      name: The name of the target.
+      bundle_name: The name to give to the framework bundle, without the
+          ".framework" extension. If omitted, the target's name will be used.
+      whitelist_symbols_file: a file including a list of whitelisted symbols,
+          one symbol per line.
+      exclude_resources: Indicates whether resources should be excluded from the
+          bundle. This can be used to avoid unnecessarily bundling resources if
+          the static framework is being distributed in a different fashion, such
+          as a Cocoapod.
+      **kwargs: Pass-through arguments.
+    """
+
+    preprocessed_name = "Preprocessed_" + name
+    ios_static_framework(
+        name = preprocessed_name,
+        bundle_name = bundle_name,
+        exclude_resources = exclude_resources,
+        **kwargs
+    )
+
+    framework_target = ":{}.zip".format(preprocessed_name)
+
+    srcs = [
+        framework_target,
+        whitelist_symbols_file,
+    ]
+    cmd = ("INPUT_FRAMEWORK=\"$(location " + framework_target + ")\" " +
+           "BUNDLE_NAME=\"" + bundle_name + "\" " +
+           "WHITELIST_FILE_PATH=\"$(location " + whitelist_symbols_file + ")\" " +
+           "OUTPUT=\"$(OUTS)\" " +
+           "\"$(location //tensorflow/lite/experimental/ios:hide_symbols_with_whitelist)\"")
+
+    native.genrule(
+        name = name,
+        srcs = srcs,
+        outs = [name + ".zip"],
+        cmd = cmd,
+        tools = [
+            "//tensorflow/lite/experimental/ios:hide_symbols_with_whitelist",
+        ],
+    )
diff --git a/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteC.txt b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteC.txt
new file mode 100644
index 00000000000..e8ae288ea8f
--- /dev/null
+++ b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteC.txt
@@ -0,0 +1 @@
+_TfLite*
diff --git a/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCCoreML.txt b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCCoreML.txt
new file mode 100644
index 00000000000..817b4a7f2ec
--- /dev/null
+++ b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCCoreML.txt
@@ -0,0 +1,2 @@
+_TfLiteCoreMlDelegateCreate
+_TfLiteCoreMlDelegateDelete
diff --git a/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCMetal.txt b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCMetal.txt
new file mode 100644
index 00000000000..b66b059eef0
--- /dev/null
+++ b/tensorflow/lite/experimental/ios/whitelist_TensorFlowLiteCMetal.txt
@@ -0,0 +1,2 @@
+_TFLGpuDelegateCreate
+_TFLGpuDelegateDelete

From 6128ffea4615c5de807beb515bb528422ba88782 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Wed, 17 Jun 2020 02:55:28 -0700
Subject: [PATCH 0387/1390] Properly configure the block and grid dimensions
 when launching generated kernels.

PiperOrigin-RevId: 316853965
Change-Id: I1fb68ef970f9d4b9bf1ae66fa84a2183aa3e8186
---
 tensorflow/core/kernels/cubin_headers/BUILD   |  1 +
 .../mlir_generated_cwise_op_gpu_tanh.cu.cc    | 69 ++++++++++++++-----
 2 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/tensorflow/core/kernels/cubin_headers/BUILD b/tensorflow/core/kernels/cubin_headers/BUILD
index b8ba164fbc3..ec8b44050db 100644
--- a/tensorflow/core/kernels/cubin_headers/BUILD
+++ b/tensorflow/core/kernels/cubin_headers/BUILD
@@ -37,4 +37,5 @@ gen_kernel_library(
         "f32",
         "f64",
     ],
+    unroll_factors = "4",
 )
diff --git a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
index 40dd7c7e49e..76d1a46aedd 100644
--- a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
+++ b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
@@ -45,9 +45,40 @@ Status CreateKernel(absl::string_view kernel_name, uint64_t num_args,
   return stream_exec->GetKernel(loader_spec, kernel_base.get());
 }
 
-class MlirGenerateTanhOp : public OpKernel {
+struct LaunchConfig {
+  se::BlockDim blockDim;
+  se::ThreadDim threadDim;
+}
+
+LaunchConfig
+GetLaunchConfiguration(std::vector tile_sizes,
+                       std::vector unrolling_factors,
+                       std::vector shape) {
+  LaunchConfig result;
+  // Ensure the vectors are length 3 and pad with ones.
+  tile_sizes.resize(3, 1);
+  unrolling_factors.resize(3, 1);
+  shape.resize(3, 1);
+  // We know that the kernel was generated by mapping the three outer-most
+  // dimensions to x,y,z dimensions. So we only need to compute those.
+  for (int i = 0; i < 3; ++i) {
+    // The number of threads is given by the tiling size.
+    result.threadDim[i] = tile_sizes[i];
+    // Compute the number of grids. We use ceildiv here as we have to allocate
+    // an extra thread/block if the division is not even. The kernel contains
+    // code to handle the boundaries.
+    int number_of_threads =
+        (shape[i] + unrolling_factors[i] - 1) / unrolling_factors[i];
+    int number_of_grids =
+        (number_of_threads + tile_sizes[i] - 1) / tile_sizes[i];
+    result.blockDim[i] = number_of_grids;
+  }
+  return result;
+}
+
+class MlirGeneratedTanhOp : public OpKernel {
  public:
-  explicit MlirGenerateTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  explicit MlirGeneratedTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
     auto* stream = ctx->op_device_context()->stream();
@@ -88,11 +119,13 @@ class MlirGenerateTanhOp : public OpKernel {
     args.add_argument(inp.NumElements());
     args.add_argument(1);
 
-    // TODO(b/158649746): Choose block size and thread dim according to the
-    // number of input elements. For now, this supports at most 1024 elements.
+    // This has to be aligned with the configuration that was used when building
+    // the kernels. See the corresponding build rules in `cubin_headers/BUILD`.
+    LaunchCondig config =
+        GetLaunchConfiguration({256}, {4}, {inp.getNumElements()});
     OP_REQUIRES_OK(
-        ctx, stream->parent()->Launch(stream, se::ThreadDim(inp.NumElements()),
-                                      se::BlockDim(1), *kernel, args));
+        ctx, stream->parent()->Launch(stream, config.threadDim, config.blockDim,
+                                      *kernel, args));
   }
 
  protected:
@@ -103,26 +136,26 @@ class MlirGenerateTanhOp : public OpKernel {
   std::mutex mu_;
 };
 
-class MlirGenerateTanhF16Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF16Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF16Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF16Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF16Kernel;
   }
 };
 
-class MlirGenerateTanhF32Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF32Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF32Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF32Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF32Kernel;
   }
 };
 
-class MlirGenerateTanhF64Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF64Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF64Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF64Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF64Kernel;
   }
 };
@@ -130,11 +163,11 @@ class MlirGenerateTanhF64Op : public MlirGenerateTanhOp {
 
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF16Op);
+    MlirGeneratedTanhF16Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF32Op);
+    MlirGeneratedTanhF32Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF64Op);
+    MlirGeneratedTanhF64Op);
 }  // namespace tensorflow

From 13b840aa7299c77e744a2eed5a366cbf98988fb0 Mon Sep 17 00:00:00 2001
From: Agoniii <815244047@qq.com>
Date: Wed, 17 Jun 2020 19:10:45 +0800
Subject: [PATCH 0388/1390] modify Log to gry

---
 .../optimizers/auto_mixed_precision_test.cc        | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc b/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
index 951279d37cd..3fa8260739c 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
@@ -287,10 +287,10 @@ TEST_F(AutoMixedPrecisionTest, Simple) {
   Output clr2 = ops::Relu(s.WithOpName("clr2"), gry1);
   Output wht1 = ops::MatMul(s.WithOpName("wht1"), clr2, clr2);
   Output clr3 = ops::Relu(s.WithOpName("clr3"), wht1);
-  Output blk2 = ops::Log(s.WithOpName("blk2"), clr3);
-  Output clr4 = ops::Relu(s.WithOpName("clr4"), blk2);
-  Output blk3 = ops::SparseMatMul(s.WithOpName("blk3"), clr4, clr4);
-  Output clr5 = ops::Relu(s.WithOpName("clr5"), blk3);
+  Output gry2 = ops::Log(s.WithOpName("gry2"), clr3);
+  Output clr4 = ops::Relu(s.WithOpName("clr4"), gry2);
+  Output blk2 = ops::SparseMatMul(s.WithOpName("blk2"), clr4, clr4);
+  Output clr5 = ops::Relu(s.WithOpName("clr5"), blk2);
   Output fetch = ops::Identity(s.WithOpName("fetch"), clr5);
 
   GrapplerItem item;
@@ -313,10 +313,10 @@ TEST_F(AutoMixedPrecisionTest, Simple) {
   EXPECT_EQ(output_view.GetNode("clr2")->attr().at("T").type(), DT_HALF);
   EXPECT_EQ(output_view.GetNode("wht1")->attr().at("T").type(), DT_HALF);
   EXPECT_EQ(output_view.GetNode("clr3")->attr().at("T").type(), DT_HALF);
-  EXPECT_EQ(output_view.GetNode("blk2")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("gry2")->attr().at("T").type(), DT_FLOAT);
   EXPECT_EQ(output_view.GetNode("clr4")->attr().at("T").type(), DT_FLOAT);
-  EXPECT_EQ(output_view.GetNode("blk3")->attr().at("Ta").type(), DT_FLOAT);
-  EXPECT_EQ(output_view.GetNode("blk3")->attr().at("Tb").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("blk2")->attr().at("Ta").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("blk2")->attr().at("Tb").type(), DT_FLOAT);
   EXPECT_EQ(output_view.GetNode("clr5")->attr().at("T").type(), DT_FLOAT);
 
   auto tensors = EvaluateNodes(output, item.fetch);

From 286cd7fc6839bb2fc999fd16fb1801f6b30656b8 Mon Sep 17 00:00:00 2001
From: Elena Zhelezina 
Date: Wed, 17 Jun 2020 12:22:53 +0100
Subject: [PATCH 0389/1390] Addressed reviewer's comments.

Change-Id: I4b849e60540879ca89483ede675c63631bc9417b
---
 tensorflow/lite/python/lite.py | 25 +++++++------------------
 1 file changed, 7 insertions(+), 18 deletions(-)

diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py
index 06796ba820b..cb2f1853619 100644
--- a/tensorflow/lite/python/lite.py
+++ b/tensorflow/lite/python/lite.py
@@ -299,32 +299,21 @@ class QuantizationMode(object):
 
     inference_input_type = input_ty if input_ty else constants.FLOAT
     inference_output_type = output_ty if output_ty else constants.FLOAT
-    if self.post_training_int8_no_float():
+
+    if self.post_training_int8_no_float() \
+      or self.post_training_int16x8_no_float():
       return True, {
           "inference_input_type": inference_input_type,
           "inference_output_type": inference_output_type,
-          "activations_type": constants.INT8,
+          "activations_type": self.activations_type(),
           "allow_float": False
       }
-    elif self.post_training_int8_allow_float():
+    elif self.post_training_int8_allow_float() \
+      or self.post_training_int16x8_allow_float():
       return True, {
           "inference_input_type": inference_input_type,
           "inference_output_type": inference_output_type,
-          "activations_type": constants.INT8,
-          "allow_float": True
-      }
-    elif self.post_training_int16x8_no_float():
-      return True, {
-          "inference_input_type": inference_input_type,
-          "inference_output_type": inference_output_type,
-          "activations_type": constants.INT16,
-          "allow_float": False
-      }
-    elif self.post_training_int16x8_allow_float():
-      return True, {
-          "inference_input_type": inference_input_type,
-          "inference_output_type": inference_output_type,
-          "activations_type": constants.INT16,
+          "activations_type": self.activations_type(),
           "allow_float": True
       }
     else:

From 9eafb72689fb2e99cd06cb09ec84de96e2f2a509 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Wed, 17 Jun 2020 05:06:09 -0700
Subject: [PATCH 0390/1390] Properly configure the block and grid dimensions
 when launching generated kernels.

PiperOrigin-RevId: 316867392
Change-Id: I975a9d1e29e954760532a82985dd016707aa6d02
---
 tensorflow/core/kernels/cubin_headers/BUILD   |  1 -
 .../mlir_generated_cwise_op_gpu_tanh.cu.cc    | 69 +++++--------------
 2 files changed, 18 insertions(+), 52 deletions(-)

diff --git a/tensorflow/core/kernels/cubin_headers/BUILD b/tensorflow/core/kernels/cubin_headers/BUILD
index ec8b44050db..b8ba164fbc3 100644
--- a/tensorflow/core/kernels/cubin_headers/BUILD
+++ b/tensorflow/core/kernels/cubin_headers/BUILD
@@ -37,5 +37,4 @@ gen_kernel_library(
         "f32",
         "f64",
     ],
-    unroll_factors = "4",
 )
diff --git a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
index 76d1a46aedd..40dd7c7e49e 100644
--- a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
+++ b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
@@ -45,40 +45,9 @@ Status CreateKernel(absl::string_view kernel_name, uint64_t num_args,
   return stream_exec->GetKernel(loader_spec, kernel_base.get());
 }
 
-struct LaunchConfig {
-  se::BlockDim blockDim;
-  se::ThreadDim threadDim;
-}
-
-LaunchConfig
-GetLaunchConfiguration(std::vector tile_sizes,
-                       std::vector unrolling_factors,
-                       std::vector shape) {
-  LaunchConfig result;
-  // Ensure the vectors are length 3 and pad with ones.
-  tile_sizes.resize(3, 1);
-  unrolling_factors.resize(3, 1);
-  shape.resize(3, 1);
-  // We know that the kernel was generated by mapping the three outer-most
-  // dimensions to x,y,z dimensions. So we only need to compute those.
-  for (int i = 0; i < 3; ++i) {
-    // The number of threads is given by the tiling size.
-    result.threadDim[i] = tile_sizes[i];
-    // Compute the number of grids. We use ceildiv here as we have to allocate
-    // an extra thread/block if the division is not even. The kernel contains
-    // code to handle the boundaries.
-    int number_of_threads =
-        (shape[i] + unrolling_factors[i] - 1) / unrolling_factors[i];
-    int number_of_grids =
-        (number_of_threads + tile_sizes[i] - 1) / tile_sizes[i];
-    result.blockDim[i] = number_of_grids;
-  }
-  return result;
-}
-
-class MlirGeneratedTanhOp : public OpKernel {
+class MlirGenerateTanhOp : public OpKernel {
  public:
-  explicit MlirGeneratedTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  explicit MlirGenerateTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
     auto* stream = ctx->op_device_context()->stream();
@@ -119,13 +88,11 @@ class MlirGeneratedTanhOp : public OpKernel {
     args.add_argument(inp.NumElements());
     args.add_argument(1);
 
-    // This has to be aligned with the configuration that was used when building
-    // the kernels. See the corresponding build rules in `cubin_headers/BUILD`.
-    LaunchCondig config =
-        GetLaunchConfiguration({256}, {4}, {inp.getNumElements()});
+    // TODO(b/158649746): Choose block size and thread dim according to the
+    // number of input elements. For now, this supports at most 1024 elements.
     OP_REQUIRES_OK(
-        ctx, stream->parent()->Launch(stream, config.threadDim, config.blockDim,
-                                      *kernel, args));
+        ctx, stream->parent()->Launch(stream, se::ThreadDim(inp.NumElements()),
+                                      se::BlockDim(1), *kernel, args));
   }
 
  protected:
@@ -136,26 +103,26 @@ class MlirGeneratedTanhOp : public OpKernel {
   std::mutex mu_;
 };
 
-class MlirGeneratedTanhF16Op : public MlirGeneratedTanhOp {
+class MlirGenerateTanhF16Op : public MlirGenerateTanhOp {
  public:
-  explicit MlirGeneratedTanhF16Op(OpKernelConstruction* ctx)
-      : MlirGeneratedTanhOp(ctx) {
+  explicit MlirGenerateTanhF16Op(OpKernelConstruction* ctx)
+      : MlirGenerateTanhOp(ctx) {
     cubin_data_ = kTanhF16Kernel;
   }
 };
 
-class MlirGeneratedTanhF32Op : public MlirGeneratedTanhOp {
+class MlirGenerateTanhF32Op : public MlirGenerateTanhOp {
  public:
-  explicit MlirGeneratedTanhF32Op(OpKernelConstruction* ctx)
-      : MlirGeneratedTanhOp(ctx) {
+  explicit MlirGenerateTanhF32Op(OpKernelConstruction* ctx)
+      : MlirGenerateTanhOp(ctx) {
     cubin_data_ = kTanhF32Kernel;
   }
 };
 
-class MlirGeneratedTanhF64Op : public MlirGeneratedTanhOp {
+class MlirGenerateTanhF64Op : public MlirGenerateTanhOp {
  public:
-  explicit MlirGeneratedTanhF64Op(OpKernelConstruction* ctx)
-      : MlirGeneratedTanhOp(ctx) {
+  explicit MlirGenerateTanhF64Op(OpKernelConstruction* ctx)
+      : MlirGenerateTanhOp(ctx) {
     cubin_data_ = kTanhF64Kernel;
   }
 };
@@ -163,11 +130,11 @@ class MlirGeneratedTanhF64Op : public MlirGeneratedTanhOp {
 
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGeneratedTanhF16Op);
+    MlirGenerateTanhF16Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGeneratedTanhF32Op);
+    MlirGenerateTanhF32Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGeneratedTanhF64Op);
+    MlirGenerateTanhF64Op);
 }  // namespace tensorflow

From 401dad16ea5a4b0fc13e373c347404b98f2b030a Mon Sep 17 00:00:00 2001
From: Nick Kreeger 
Date: Wed, 17 Jun 2020 06:43:11 -0700
Subject: [PATCH 0391/1390] Add default values in MicroInterpreter
 constructors.

The MicroInterpreter uses a few values to check state - there is a scenario where these values are not always defaulted to internal states. This can cause an exception when the interpreter tries to run. To ensure things work properly, default values.

I also updated the MicroInterpreter test to use the new RecordingMicroAllocator. Two new tests have been added:

1.) Ensure that the interpreter fails to allocate with too small an arena at Invoke() (insured by recording allocation APIs)

2.) Ensure that the interpreter does not allocate anything at construction time - only at Invoke() (or manually with AllocateTensors()).

This will give us better coverage when we add more tenant use cases.

PiperOrigin-RevId: 316877994
Change-Id: I0582080a1fb649276076371be991a13392324801
---
 tensorflow/lite/micro/BUILD                   |   1 +
 tensorflow/lite/micro/micro_interpreter.cc    |   4 +
 .../lite/micro/micro_interpreter_test.cc      | 111 ++++++++++++++++++
 3 files changed, 116 insertions(+)

diff --git a/tensorflow/lite/micro/BUILD b/tensorflow/lite/micro/BUILD
index 32d7271734e..f63d9778634 100644
--- a/tensorflow/lite/micro/BUILD
+++ b/tensorflow/lite/micro/BUILD
@@ -268,6 +268,7 @@ tflite_micro_cc_test(
         ":micro_framework",
         ":micro_utils",
         ":op_resolvers",
+        ":recording_allocators",
         ":test_helpers",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/micro/testing:micro_test",
diff --git a/tensorflow/lite/micro/micro_interpreter.cc b/tensorflow/lite/micro/micro_interpreter.cc
index c20eb1f0984..6b17a5ffe84 100644
--- a/tensorflow/lite/micro/micro_interpreter.cc
+++ b/tensorflow/lite/micro/micro_interpreter.cc
@@ -83,6 +83,8 @@ MicroInterpreter::MicroInterpreter(const Model* model,
       error_reporter_(error_reporter),
       allocator_(*MicroAllocator::Create(tensor_arena, tensor_arena_size,
                                          error_reporter)),
+      tensors_allocated_(false),
+      initialization_status_(kTfLiteError),
       context_helper_(error_reporter_, &allocator_) {
   Init(profiler);
 }
@@ -96,6 +98,8 @@ MicroInterpreter::MicroInterpreter(const Model* model,
       op_resolver_(*op_resolver),
       error_reporter_(error_reporter),
       allocator_(*allocator),
+      tensors_allocated_(false),
+      initialization_status_(kTfLiteError),
       context_helper_(error_reporter_, &allocator_) {
   Init(profiler);
 }
diff --git a/tensorflow/lite/micro/micro_interpreter_test.cc b/tensorflow/lite/micro/micro_interpreter_test.cc
index 079e23d33eb..93d095d3c68 100644
--- a/tensorflow/lite/micro/micro_interpreter_test.cc
+++ b/tensorflow/lite/micro/micro_interpreter_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/lite/micro/micro_mutable_op_resolver.h"
 #include "tensorflow/lite/micro/micro_optional_debug_tools.h"
 #include "tensorflow/lite/micro/micro_utils.h"
+#include "tensorflow/lite/micro/recording_micro_allocator.h"
 #include "tensorflow/lite/micro/test_helpers.h"
 #include "tensorflow/lite/micro/testing/micro_test.h"
 
@@ -244,6 +245,7 @@ TF_LITE_MICRO_TEST(TestIncompleteInitialization) {
   tflite::testing::MockOpResolver mock_resolver;
   constexpr size_t allocator_buffer_size = 2048;
   uint8_t allocator_buffer[allocator_buffer_size];
+
   tflite::MicroInterpreter interpreter(model, mock_resolver, allocator_buffer,
                                        allocator_buffer_size,
                                        micro_test::reporter);
@@ -276,4 +278,113 @@ TF_LITE_MICRO_TEST(InterpreterWithProfilerShouldProfileOps) {
 #endif
 }
 
+TF_LITE_MICRO_TEST(TestIncompleteInitializationAllocationsWithSmallArena) {
+  const tflite::Model* model = tflite::testing::GetComplexMockModel();
+  TF_LITE_MICRO_EXPECT_NE(nullptr, model);
+
+  tflite::testing::MockOpResolver mock_resolver;
+  // 1kb is too small for the ComplexMockModel:
+  constexpr size_t allocator_buffer_size = 1048;
+  uint8_t allocator_buffer[allocator_buffer_size];
+
+  tflite::RecordingMicroAllocator* allocator =
+      tflite::RecordingMicroAllocator::Create(
+          allocator_buffer, allocator_buffer_size, micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_NE(nullptr, allocator);
+
+  tflite::MicroInterpreter interpreter(model, &mock_resolver, allocator,
+                                       micro_test::reporter);
+
+  // Interpreter fails because arena is too small:
+  TF_LITE_MICRO_EXPECT_EQ(interpreter.Invoke(), kTfLiteError);
+
+  // Ensure allocations are zero (ignore tail since some internal structs are
+  // initialized with this space):
+  TF_LITE_MICRO_EXPECT_EQ(
+      0, allocator->GetSimpleMemoryAllocator()->GetHeadUsedBytes());
+  TF_LITE_MICRO_EXPECT_EQ(
+      0, allocator
+             ->GetRecordedAllocation(
+                 tflite::RecordedAllocationType::kTfLiteTensorArray)
+             .used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(
+      0, allocator
+             ->GetRecordedAllocation(tflite::RecordedAllocationType::
+                                         kTfLiteTensorArrayQuantizationData)
+             .used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(
+      0,
+      allocator
+          ->GetRecordedAllocation(
+              tflite::RecordedAllocationType::kTfLiteTensorVariableBufferData)
+          .used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(
+      0,
+      allocator->GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
+          .used_bytes);
+}
+
+TF_LITE_MICRO_TEST(TestInterpreterDoesNotAllocateUntilInvoke) {
+  const tflite::Model* model = tflite::testing::GetComplexMockModel();
+  TF_LITE_MICRO_EXPECT_NE(nullptr, model);
+
+  tflite::testing::MockOpResolver mock_resolver;
+  constexpr size_t allocator_buffer_size = 1024 * 4;
+  uint8_t allocator_buffer[allocator_buffer_size];
+
+  tflite::RecordingMicroAllocator* allocator =
+      tflite::RecordingMicroAllocator::Create(
+          allocator_buffer, allocator_buffer_size, micro_test::reporter);
+  TF_LITE_MICRO_EXPECT_NE(nullptr, allocator);
+
+  tflite::MicroInterpreter interpreter(model, &mock_resolver, allocator,
+                                       micro_test::reporter);
+
+  // Ensure allocations are zero (ignore tail since some internal structs are
+  // initialized with this space):
+  TF_LITE_MICRO_EXPECT_EQ(
+      0, allocator->GetSimpleMemoryAllocator()->GetHeadUsedBytes());
+  TF_LITE_MICRO_EXPECT_EQ(
+      0, allocator
+             ->GetRecordedAllocation(
+                 tflite::RecordedAllocationType::kTfLiteTensorArray)
+             .used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(
+      0,
+      allocator
+          ->GetRecordedAllocation(
+              tflite::RecordedAllocationType::kTfLiteTensorVariableBufferData)
+          .used_bytes);
+  TF_LITE_MICRO_EXPECT_EQ(
+      0,
+      allocator->GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
+          .used_bytes);
+
+  TF_LITE_MICRO_EXPECT_EQ(interpreter.Invoke(), kTfLiteOk);
+  allocator->PrintAllocations();
+
+  // Allocation sizes vary based on platform - check that allocations are now
+  // non-zero:
+  TF_LITE_MICRO_EXPECT_GT(
+      allocator->GetSimpleMemoryAllocator()->GetHeadUsedBytes(), 0);
+  TF_LITE_MICRO_EXPECT_GT(
+      allocator
+          ->GetRecordedAllocation(
+              tflite::RecordedAllocationType::kTfLiteTensorArray)
+          .used_bytes,
+      0);
+  TF_LITE_MICRO_EXPECT_GT(
+
+      allocator
+          ->GetRecordedAllocation(
+              tflite::RecordedAllocationType::kTfLiteTensorVariableBufferData)
+          .used_bytes,
+      0);
+  TF_LITE_MICRO_EXPECT_GT(
+
+      allocator->GetRecordedAllocation(tflite::RecordedAllocationType::kOpData)
+          .used_bytes,
+      0);
+}
+
 TF_LITE_MICRO_TESTS_END

From f44b07ed4e5f2b338d3b3d9f17875b86b8636a92 Mon Sep 17 00:00:00 2001
From: Mark Daoust 
Date: Wed, 17 Jun 2020 07:20:46 -0700
Subject: [PATCH 0392/1390] Fix formatting

Some of these indents are triggering markdown's (horrible) "4-space indent is a code block" feature

PiperOrigin-RevId: 316883013
Change-Id: If2b53a6788d3179b868a62fb6b4caeeb08caa4bf
---
 tensorflow/python/ops/map_fn.py | 94 ++++++++++++++++-----------------
 1 file changed, 46 insertions(+), 48 deletions(-)

diff --git a/tensorflow/python/ops/map_fn.py b/tensorflow/python/ops/map_fn.py
index 96810805c18..516f427ad08 100644
--- a/tensorflow/python/ops/map_fn.py
+++ b/tensorflow/python/ops/map_fn.py
@@ -108,31 +108,29 @@ def map_fn(fn,
 
   `fn_output_signature` can be specified using any of the following:
 
-    * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
-    * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
-    * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
-    * A (possibly nested) tuple, list, or dict containing the above types.
+  * A `tf.DType` or `tf.TensorSpec` (to describe a `tf.Tensor`)
+  * A `tf.RaggedTensorSpec` (to describe a `tf.RaggedTensor`)
+  * A `tf.SparseTensorSpec` (to describe a `tf.sparse.SparseTensor`)
+  * A (possibly nested) tuple, list, or dict containing the above types.
 
   #### RaggedTensors
 
   `map_fn` supports `tf.RaggedTensor` inputs and outputs.  In particular:
 
-    * If `elems` is a `RaggedTensor`, then `fn` will be called with each
-      row of that ragged tensor.
+  * If `elems` is a `RaggedTensor`, then `fn` will be called with each
+    row of that ragged tensor.
+    * If `elems` has only one ragged dimension, then the values passed to
+      `fn` will be `tf.Tensor`s.
+    * If `elems` has multiple ragged dimensions, then the values passed to
+      `fn` will be `tf.RaggedTensor`s with one fewer ragged dimension.
 
-      * If `elems` has only one ragged dimension, then the values passed to
-        `fn` will be `tf.Tensor`s.
-      * If `elems` has multiple ragged dimensions, then the values passed to
-        `fn` will be `tf.RaggedTensor`s with one fewer ragged dimension.
-
-    * If the result of `map_fn` should be a `RaggedTensor`, then use a
-      `tf.RaggedTensorSpec` to specify `fn_output_signature`.
-
-      * If `fn` returns `tf.Tensor`s with varying sizes, then use a
-        `tf.RaggedTensorSpec` with `ragged_rank=0` to combine them into a
-        single ragged tensor (which will have ragged_rank=1).
-      * If `fn` returns `tf.RaggedTensor`s, then use a `tf.RaggedTensorSpec`
-        with the same `ragged_rank`.
+  * If the result of `map_fn` should be a `RaggedTensor`, then use a
+    `tf.RaggedTensorSpec` to specify `fn_output_signature`.
+    * If `fn` returns `tf.Tensor`s with varying sizes, then use a
+      `tf.RaggedTensorSpec` with `ragged_rank=0` to combine them into a
+      single ragged tensor (which will have ragged_rank=1).
+    * If `fn` returns `tf.RaggedTensor`s, then use a `tf.RaggedTensorSpec`
+      with the same `ragged_rank`.
 
   >>> # Example: RaggedTensor input
   >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
@@ -150,10 +148,10 @@ def map_fn(fn,
   *rows* of a `RaggedTensor`.  If you wish to map a function over the
   individual values, then you should use:
 
-    * `tf.ragged.map_flat_values(fn, rt)`
-      (if fn is expressible as TensorFlow ops)
-    * `rt.with_flat_values(map_fn(fn, rt.flat_values))`
-      (otherwise)
+  * `tf.ragged.map_flat_values(fn, rt)`
+    (if fn is expressible as TensorFlow ops)
+  * `rt.with_flat_values(map_fn(fn, rt.flat_values))`
+    (otherwise)
 
   E.g.:
 
@@ -165,14 +163,14 @@ def map_fn(fn,
 
   `map_fn` supports `tf.sparse.SparseTensor` inputs and outputs.  In particular:
 
-    * If `elems` is a `SparseTensor`, then `fn` will be called with each row
-      of that sparse tensor. In particular, the value passed to `fn` will be a
-      `tf.sparse.SparseTensor` with one fewer dimension than `elems`.
+  * If `elems` is a `SparseTensor`, then `fn` will be called with each row
+    of that sparse tensor. In particular, the value passed to `fn` will be a
+    `tf.sparse.SparseTensor` with one fewer dimension than `elems`.
 
-    * If the result of `map_fn` should be a `SparseTensor`, then use a
-      `tf.SparseTensorSpec` to specify `fn_output_signature`.  The individual
-      `SparseTensor`s returned by `fn` will be stacked into a single
-      `SparseTensor` with one more dimension.
+  * If the result of `map_fn` should be a `SparseTensor`, then use a
+    `tf.SparseTensorSpec` to specify `fn_output_signature`.  The individual
+    `SparseTensor`s returned by `fn` will be stacked into a single
+    `SparseTensor` with one more dimension.
 
   >>> # Example: SparseTensor input
   >>> st = tf.sparse.SparseTensor([[0, 0], [2, 0], [2, 1]], [2, 3, 4], [4, 4])
@@ -195,15 +193,15 @@ def map_fn(fn,
   *rows* of a `SparseTensor`.  If you wish to map a function over the nonzero
   values, then you should use:
 
-    * If the function is expressible as TensorFlow ops, use:
-      ```python
-      tf.sparse.SparseTensor(st.indices, fn(st.values), st.dense_shape)
-      ```
-    * Otherwise, use:
-      ```python
-      tf.sparse.SparseTensor(st.indices, tf.map_fn(fn, st.values),
-                             st.dense_shape)
-      ```
+  * If the function is expressible as TensorFlow ops, use:
+    ```python
+    tf.sparse.SparseTensor(st.indices, fn(st.values), st.dense_shape)
+    ```
+  * Otherwise, use:
+    ```python
+    tf.sparse.SparseTensor(st.indices, tf.map_fn(fn, st.values),
+                           st.dense_shape)
+    ```
 
   #### `map_fn` vs. vectorized operations
 
@@ -215,14 +213,14 @@ def map_fn(fn,
 
   `map_fn` should typically only be used if one of the following is true:
 
-    * It is difficult or expensive to express the desired transform with
-      vectorized operations.
-    * `fn` creates large intermediate values, so an equivalent vectorized
-      transform would take too much memory.
-    * Processing elements in parallel is more efficient than an equivalent
-      vectorized transform.
-    * Efficiency of the transform is not critical, and using `map_fn` is
-      more readable.
+  * It is difficult or expensive to express the desired transform with
+    vectorized operations.
+  * `fn` creates large intermediate values, so an equivalent vectorized
+    transform would take too much memory.
+  * Processing elements in parallel is more efficient than an equivalent
+    vectorized transform.
+  * Efficiency of the transform is not critical, and using `map_fn` is
+    more readable.
 
   E.g., the example given above that maps `fn=lambda t: tf.range(t, t + 3)`
   across `elems` could be rewritten more efficiently using vectorized ops:
@@ -255,7 +253,7 @@ def map_fn(fn,
            [2, 3, 4]], dtype=int32)>
 
 
-  Note that if you use the `tf.function` decorator, any non-TensorFlow Python
+  Note: if you use the `tf.function` decorator, any non-TensorFlow Python
   code that you may have written in your function won't get executed. See
   `tf.function` for more  details. The recommendation would be to debug without
   `tf.function` but switch to it to get performance benefits of running `map_fn`

From 3ebc53c394e34dfd780e6bf59ad3c96bd9b3fa79 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Thu, 11 Jun 2020 15:59:21 -0500
Subject: [PATCH 0393/1390] Relax stub include version checking.

Remove upper bound on version check for latest inc files.
---
 tensorflow/stream_executor/cuda/cublas_stub.cc   | 10 ++++------
 tensorflow/stream_executor/cuda/cuda_stub.cc     | 10 ++++------
 tensorflow/stream_executor/cuda/cudart_stub.cc   | 10 ++++------
 tensorflow/stream_executor/cuda/cusolver_stub.cc | 10 ++++------
 tensorflow/stream_executor/cuda/cusparse_stub.cc | 10 ++++------
 5 files changed, 20 insertions(+), 30 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cublas_stub.cc b/tensorflow/stream_executor/cuda/cublas_stub.cc
index 1cbfd51316c..76f3d9b134e 100644
--- a/tensorflow/stream_executor/cuda/cublas_stub.cc
+++ b/tensorflow/stream_executor/cuda/cublas_stub.cc
@@ -63,14 +63,12 @@ typedef enum {} cublasMath_t;
 
 #if CUDA_VERSION < 10000
 #include "tensorflow/stream_executor/cuda/cublas_9_0.inc"
-#elif CUDA_VERSION == 10000
+#elif CUDA_VERSION < 10010
 #include "tensorflow/stream_executor/cuda/cublas_10_0.inc"
-#elif CUDA_VERSION == 10010
+#elif CUDA_VERSION < 10020
 #include "tensorflow/stream_executor/cuda/cublas_10_1.inc"
-#elif CUDA_VERSION == 10020
+#elif CUDA_VERSION < 11000
 #include "tensorflow/stream_executor/cuda/cublas_10_2.inc"
-#elif CUBLAS_VER_MAJOR == 11 && CUBLAS_VER_MINOR == 0
-#include "tensorflow/stream_executor/cuda/cublas_11_0.inc"
 #else
-#error "We have no wrapper for this version."
+#include "tensorflow/stream_executor/cuda/cublas_11_0.inc"
 #endif
diff --git a/tensorflow/stream_executor/cuda/cuda_stub.cc b/tensorflow/stream_executor/cuda/cuda_stub.cc
index ce02be89c22..58c898a54ee 100644
--- a/tensorflow/stream_executor/cuda/cuda_stub.cc
+++ b/tensorflow/stream_executor/cuda/cuda_stub.cc
@@ -95,14 +95,12 @@ typedef void(CUDA_CB* CUhostFn)(void* userData);
 
 #if CUDA_VERSION < 10000
 #include "tensorflow/stream_executor/cuda/cuda_9_0.inc"
-#elif CUDA_VERSION == 10000
+#elif CUDA_VERSION < 10010
 #include "tensorflow/stream_executor/cuda/cuda_10_0.inc"
-#elif CUDA_VERSION <= 10010
+#elif CUDA_VERSION < 10020
 #include "tensorflow/stream_executor/cuda/cuda_10_1.inc"
-#elif CUDA_VERSION <= 10020
+#elif CUDA_VERSION < 11000
 #include "tensorflow/stream_executor/cuda/cuda_10_2.inc"
-#elif CUDA_VERSION <= 11000
-#include "tensorflow/stream_executor/cuda/cuda_11_0.inc"
 #else
-#error "We have no wrapper for this version."
+#include "tensorflow/stream_executor/cuda/cuda_11_0.inc"
 #endif
diff --git a/tensorflow/stream_executor/cuda/cudart_stub.cc b/tensorflow/stream_executor/cuda/cudart_stub.cc
index 3b9e0f2937b..2ab9d142e3c 100644
--- a/tensorflow/stream_executor/cuda/cudart_stub.cc
+++ b/tensorflow/stream_executor/cuda/cudart_stub.cc
@@ -53,16 +53,14 @@ cudaError_t GetSymbolNotFoundError() {
 // A bunch of new symbols were introduced in version 10
 #if CUDART_VERSION < 10000
 #include "tensorflow/stream_executor/cuda/cuda_runtime_9_0.inc"
-#elif CUDART_VERSION == 10000
+#elif CUDART_VERSION < 10010
 #include "tensorflow/stream_executor/cuda/cuda_runtime_10_0.inc"
-#elif CUDART_VERSION == 10010
+#elif CUDART_VERSION < 10020
 #include "tensorflow/stream_executor/cuda/cuda_runtime_10_1.inc"
-#elif CUDART_VERSION == 10020
+#elif CUDART_VERSION < 11000
 #include "tensorflow/stream_executor/cuda/cuda_runtime_10_2.inc"
-#elif CUDART_VERSION == 11000
-#include "tensorflow/stream_executor/cuda/cuda_runtime_11_0.inc"
 #else
-#error "We have no wrapper for this version."
+#include "tensorflow/stream_executor/cuda/cuda_runtime_11_0.inc"
 #endif
 #undef __dv
 #undef __CUDA_DEPRECATED
diff --git a/tensorflow/stream_executor/cuda/cusolver_stub.cc b/tensorflow/stream_executor/cuda/cusolver_stub.cc
index a4b9cc37f9b..edf87c3dc0b 100644
--- a/tensorflow/stream_executor/cuda/cusolver_stub.cc
+++ b/tensorflow/stream_executor/cuda/cusolver_stub.cc
@@ -53,14 +53,12 @@ cusolverStatus_t GetSymbolNotFoundError() {
 
 #if CUDA_VERSION < 10000
 #include "tensorflow/stream_executor/cuda/cusolver_dense_9_0.inc"
-#elif CUDA_VERSION == 10000
+#elif CUDA_VERSION < 10010
 #include "tensorflow/stream_executor/cuda/cusolver_dense_10_0.inc"
-#elif CUDA_VERSION == 10010
+#elif CUDA_VERSION < 10020
 #include "tensorflow/stream_executor/cuda/cusolver_dense_10_1.inc"
-#elif CUDA_VERSION == 10020
+#elif CUDA_VERSION < 11000
 #include "tensorflow/stream_executor/cuda/cusolver_dense_10_2.inc"
-#elif CUDA_VERSION == 11000
-#include "tensorflow/stream_executor/cuda/cusolver_dense_11_0.inc"
 #else
-#error "We don't have a wrapper for this version."
+#include "tensorflow/stream_executor/cuda/cusolver_dense_11_0.inc"
 #endif
diff --git a/tensorflow/stream_executor/cuda/cusparse_stub.cc b/tensorflow/stream_executor/cuda/cusparse_stub.cc
index ae56402fbc3..caed4d1008e 100644
--- a/tensorflow/stream_executor/cuda/cusparse_stub.cc
+++ b/tensorflow/stream_executor/cuda/cusparse_stub.cc
@@ -53,14 +53,12 @@ cusparseStatus_t GetSymbolNotFoundError() {
 
 #if CUDA_VERSION < 10000
 #include "tensorflow/stream_executor/cuda/cusparse_9_0.inc"
-#elif CUDA_VERSION == 10000
+#elif CUDA_VERSION < 10010
 #include "tensorflow/stream_executor/cuda/cusparse_10_0.inc"
-#elif CUDA_VERSION == 10010
+#elif CUDA_VERSION < 10020
 #include "tensorflow/stream_executor/cuda/cusparse_10_1.inc"
-#elif CUDA_VERSION == 10020
+#elif CUDA_VERSION < 11000
 #include "tensorflow/stream_executor/cuda/cusparse_10_2.inc"
-#elif CUSPARSE_VER_MAJOR == 11 && CUSPARSE_VER_MINOR == 0
-#include "tensorflow/stream_executor/cuda/cusparse_11_0.inc"
 #else
-#error "We don't have a wrapper for this version."
+#include "tensorflow/stream_executor/cuda/cusparse_11_0.inc"
 #endif

From 3d4ca5a00a58aae4273c682f05dcb6665def8d02 Mon Sep 17 00:00:00 2001
From: Michael Banfield 
Date: Wed, 17 Jun 2020 08:02:15 -0700
Subject: [PATCH 0394/1390] Support fetching configured runtime version without
 using Cloud TPU API.

PiperOrigin-RevId: 316888815
Change-Id: I9f840076122e220c80b7b301f2290a6d4f595f1a
---
 tensorflow/python/tpu/client/client.py      | 20 +++++++++++++++++++-
 tensorflow/python/tpu/client/client_test.py | 16 ++++++++++++++++
 tensorflow/python/tpu/client/version.py     |  2 +-
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/tpu/client/client.py b/tensorflow/python/tpu/client/client.py
index 972f02437d8..c834a57c153 100644
--- a/tensorflow/python/tpu/client/client.py
+++ b/tensorflow/python/tpu/client/client.py
@@ -20,6 +20,7 @@ from __future__ import division
 from __future__ import print_function
 
 import datetime
+import json
 import logging
 import os
 import time
@@ -48,6 +49,7 @@ _DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL'
 _GCE_METADATA_ENDPOINT = 'http://metadata.google.internal'
 _DEFAULT_ENDPOINT_PORT = '8470'
 _OOM_EVENT_COOL_TIME_SEC = 90
+_VERSION_SWITCHER_ENDPOINT = 'http://{}:8475/requestversion'
 
 
 def _utcnow():
@@ -277,6 +279,22 @@ class Client(object):
 
   def runtime_version(self):
     """Return runtime version of the TPU."""
+
+    if not self._use_api:
+      # Fallback on getting version directly from TPU.
+      url = _VERSION_SWITCHER_ENDPOINT.format(
+          self.network_endpoints()[0]['ipAddress'])
+      try:
+        req = request.Request(url)
+        resp = request.urlopen(req)
+        version_details = json.loads(resp.read())
+        return version_details.get('currentVersion')
+      except HTTPError as e:
+        status_code = e.code
+        if status_code == 404:
+          return None
+        else:
+          raise e
     return self._get_tpu_property('tensorflowVersion')
 
   def accelerator_type(self):
@@ -350,7 +368,7 @@ class Client(object):
           be sent.
       """
       ip_address = worker['ipAddress']
-      url = 'http://{}:8475/requestversion/{}?restartType={}'.format(
+      url = (_VERSION_SWITCHER_ENDPOINT + '/{}?restartType={}').format(
           ip_address, version, restart_type)
       req = request.Request(url, data=b'')
       try:
diff --git a/tensorflow/python/tpu/client/client_test.py b/tensorflow/python/tpu/client/client_test.py
index 9d7f29ad476..f53f09cd3d5 100644
--- a/tensorflow/python/tpu/client/client_test.py
+++ b/tensorflow/python/tpu/client/client_test.py
@@ -630,6 +630,22 @@ class CloudTpuClientTest(test.TestCase):
         'http://5.6.7.8:8475/requestversion/1.15?restartType=ifNeeded'
     ], sorted(paths))
 
+  @mock.patch.object(request, 'urlopen')
+  def testGetTpuVersion(self, urlopen):
+    c = client.Client(
+        tpu='grpc://1.2.3.4:8470')
+    resp = mock.Mock()
+    resp.read.side_effect = ['{}', '{"currentVersion": "someVersion"}']
+    urlopen.return_value = resp
+    self.assertIsNone(c.runtime_version(), 'Missing key should be handled.')
+    self.assertEqual(
+        'someVersion', c.runtime_version(), 'Should return configured version.')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion',
+        'http://1.2.3.4:8475/requestversion',
+    ], sorted(paths))
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/tpu/client/version.py b/tensorflow/python/tpu/client/version.py
index 001059a91da..a91586640fc 100644
--- a/tensorflow/python/tpu/client/version.py
+++ b/tensorflow/python/tpu/client/version.py
@@ -18,4 +18,4 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-__version__ = "0.9"
+__version__ = "0.10"

From 23be4f5d44f28734620fa508d9807d9aca2ce074 Mon Sep 17 00:00:00 2001
From: Rahul Joshi 
Date: Wed, 17 Jun 2020 08:32:24 -0700
Subject: [PATCH 0395/1390] - Fix
 PropagatePotentiallyWrittenWithinUnhandledOp() to mark resource uses within  
 regions as potentially written when a resource operand is seen. This also
 fixes   the case when multiple resources are used as operands in the same
 unhandled op - Add test case to demonstrate the issue

PiperOrigin-RevId: 316893534
Change-Id: I9e688a90155efd990eb5ef835c23933825bcbdd0
---
 ...f_saved_model_optimize_global_tensors.mlir | 88 ++++++++++++++++++-
 .../transforms/optimize_global_tensors.cc     | 22 ++---
 2 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
index 9d8911d306d..0c68cf0cf64 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
@@ -1,7 +1,7 @@
 // RUN: tf-opt -tf-saved-model-optimize-global-tensors -split-input-file %s | FileCheck %s
 
 //===----------------------------------------------------------------------===//
-// Freezing.
+// Immutability.
 //===----------------------------------------------------------------------===//
 
 module attributes {tf_saved_model.semantics} {
@@ -142,3 +142,89 @@ module attributes {tf_saved_model.semantics} {
 // Test running the pass on a module that does not have
 // tf_saved_model.semantics.
 module {}
+
+// -----
+
+// Test use as an input in unhandled op
+module attributes {tf_saved_model.semantics} {
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor, value = dense<42.> : tensor } : () -> ()
+
+  func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v})
+  attributes {tf_saved_model.exported_names = ["f"]} {
+    "tf.unhandled_op"(%arg0) : (tensor>>) -> ()
+    return
+  }
+}
+
+
+// -----
+
+// Test use as a region capture in an unhandled op
+module attributes {tf_saved_model.semantics} {
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor, value = dense<42.> : tensor } : () -> ()
+
+  func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v})
+  attributes {tf_saved_model.exported_names = ["f"]} {
+    "tf.unhandled"() ({
+      %val = "tf.ReadVariableOp"(%arg0) : (tensor>>) -> tensor
+      "tf.unhandled_terminator"() : () -> ()
+    }) : () -> ()
+    return
+  }
+}
+
+// -----
+
+// Test use as region capture as well as input in an unhandled op
+// to the unhandled op.
+module attributes {tf_saved_model.semantics} {
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor, value = dense<42.> : tensor } : () -> ()
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "u", type = tensor, value = dense<22.> : tensor } : () -> ()
+
+  func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v}, %arg1: tensor>> {tf_saved_model.bound_input = @u})
+  attributes {tf_saved_model.exported_names = ["f"]} {
+    %0 = "tf.unhandled"(%arg0) ({
+      %val = "tf.ReadVariableOp"(%arg1) : (tensor>>) -> tensor
+      "tf.unhandled_terminator"() : () -> ()
+    }) : (tensor>>) -> (tensor>>)
+    return
+  }
+}
+
+// -----
+
+// Test multiple global tensors uses as operands for an unhandled op.
+module attributes {tf_saved_model.semantics} {
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor, value = dense<42.> : tensor } : () -> ()
+
+  // CHECK: "tf_saved_model.global_tensor"() {
+  // CHECK-SAME: is_mutable
+  // CHECK-SAME: } : () -> ()
+  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "u", type = tensor, value = dense<22.> : tensor } : () -> ()
+
+  func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v}, %arg1: tensor>> {tf_saved_model.bound_input = @u})
+  attributes {tf_saved_model.exported_names = ["f"]} {
+    "tf.unhandled"(%arg0, %arg1) : (tensor>>, tensor>>) -> ()
+    return
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc b/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
index cd8f988fd5f..07cc6203cbd 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
@@ -56,14 +56,14 @@ struct GlobalTensorUse {
 using GlobalTensorUsesMap =
     std::map>;
 
-static bool IsResourceType(Type type) {
+bool IsResourceType(Type type) {
   if (auto tensor_type = type.dyn_cast()) {
     return tensor_type.getElementType().isa();
   }
   return false;
 }
 
-static bool IsResource(Value value) { return IsResourceType(value.getType()); }
+bool IsResource(Value value) { return IsResourceType(value.getType()); }
 
 class ResourceAnalyzer {
  public:
@@ -129,30 +129,24 @@ class ResourceAnalyzer {
       // this errs on the side of being conservative. We should improve
       // this by using either a property or a trait that clearly
       // identifies ops with resource mutating behavior.
-      if (PropagatePotentiallyWrittenWithinUnhandledOp(op)) {
-        return;
-      }
+      PropagatePotentiallyWrittenWithinUnhandledOp(op);
     });
     return success();
   }
 
   // If an op is not one of the handled ones, we assume all resource usages
   // within its purview are mutating in nature.
-  bool PropagatePotentiallyWrittenWithinUnhandledOp(Operation* op) {
+  void PropagatePotentiallyWrittenWithinUnhandledOp(Operation* op) {
     for (auto operand : op->getOperands()) {
       if (IsResource(operand)) {
         SetPotentiallyWritten(operand);
-        return true;
       }
     }
-    bool uses_resources = false;
     visitUsedValuesDefinedAbove(op->getRegions(), [&](OpOperand* operand) {
       if (IsResource(operand->get())) {
         SetPotentiallyWritten(operand->get());
-        uses_resources = true;
       }
     });
-    return uses_resources;
   }
 
   // Given a funcOp associated with the callee and operands from the
@@ -212,7 +206,7 @@ bool IsImmutable(GlobalTensorOp global_tensor,
   return true;
 }
 
-static GlobalTensorUsesMap CreateGlobalTensorUsesMap(ModuleOp module) {
+GlobalTensorUsesMap CreateGlobalTensorUsesMap(ModuleOp module) {
   GlobalTensorUsesMap global_tensor_uses;
 
   SymbolTable symbol_table(module);
@@ -293,13 +287,13 @@ void OptimizeGlobalTensorsPass::runOnOperation() {
   EraseUnusedGlobalTensors(module, global_tensor_uses);
 }
 
-}  // namespace
-
 // For "opt" to pick up this pass.
-static PassRegistration pass(
+PassRegistration pass(
     "tf-saved-model-optimize-global-tensors",
     "Optimize tf_saved_model.global_tensor's.");
 
+}  // namespace
+
 std::unique_ptr> CreateOptimizeGlobalTensorsPass() {
   return std::make_unique();
 }

From fc296acdc1d454596d9e0e531656858f3b0acca6 Mon Sep 17 00:00:00 2001
From: Mark Daoust 
Date: Wed, 17 Jun 2020 09:00:01 -0700
Subject: [PATCH 0396/1390] Make axis handling for Normalization more robust.

PiperOrigin-RevId: 316898233
Change-Id: I6888216ed21c4d2a482772fb2a314160750185b6
---
 .../layers/preprocessing/normalization.py     | 67 ++++++++++++-------
 .../preprocessing/normalization_test.py       | 43 ++++++++++++
 2 files changed, 86 insertions(+), 24 deletions(-)

diff --git a/tensorflow/python/keras/layers/preprocessing/normalization.py b/tensorflow/python/keras/layers/preprocessing/normalization.py
index 09564cbb064..ba2f7eaae89 100644
--- a/tensorflow/python/keras/layers/preprocessing/normalization.py
+++ b/tensorflow/python/keras/layers/preprocessing/normalization.py
@@ -53,10 +53,13 @@ class Normalization(CombinerPreprocessingLayer):
 
   Attributes:
       axis: Integer or tuple of integers, the axis or axes that should be
-        normalized (typically the features axis). We will normalize each element
-        in the specified axis. If set to 'None', the layer will perform scalar
-        normalization (diving the input by a single scalar value). 0 (the batch
-        axis) is not allowed.
+        "kept". These axes are not be summed over when calculating the
+        normalization statistics. By default the last axis, the `features` axis
+        is kept and any `space` or `time` axes are summed. Each element in the
+        the axes that are kept is normalized independently. If `axis` is set to
+        'None', the layer will perform scalar normalization (diving the input
+        by a single scalar value). The `batch` axis, 0, is always summed over
+        (`axis=0` is not allowed).
 
   Examples:
 
@@ -78,10 +81,18 @@ class Normalization(CombinerPreprocessingLayer):
     # time, the dtype value will change to reflect it.
     dtype = dtype or K.floatx()
 
+    # Standardize `axis` to a tuple.
+    if axis is None:
+      axis = ()
+    elif isinstance(axis, int):
+      axis = (axis,)
+    else:
+      axis = tuple(axis)
+
     super(Normalization, self).__init__(
         combiner=_NormalizingCombiner(axis), dtype=dtype, **kwargs)
 
-    if axis == 0:
+    if 0 in axis:
       raise ValueError('The argument \'axis\' may not be 0.')
 
     self.axis = axis
@@ -90,18 +101,27 @@ class Normalization(CombinerPreprocessingLayer):
     input_shape = tensor_shape.TensorShape(input_shape).as_list()
     if len(input_shape) == 1:
       input_shape = input_shape + [1]
+
+    ndim = len(input_shape)
+
+    # Sort `self.axis` to avoid transposing `mean_and_var_shape`.
+    # Negative axes are not sortable until you know the number of dimensions.
+    original_axis = self.axis
+    self.axis = tuple(sorted(self.axis,
+                             key=lambda a: a if a >= 0 else ndim + a))
+
+    if any(a < 1-ndim for a in self.axis) or any(a >= ndim for a in self.axis):
+      raise ValueError('All `axis` values must be in '
+                       'the range [1-ndim, ndim-1].\n'
+                       'Got:\n'
+                       '    ndim: {}\n'
+                       '    axis: {}'.format(ndim, original_axis))
+
     self._broadcast_shape = [1 for _ in range(len(input_shape))]
-    if isinstance(self.axis, (tuple, list)):
-      mean_and_var_shape = []
-      for i in self.axis:
-        mean_and_var_shape.append(input_shape[i])
-        self._broadcast_shape[i] = input_shape[i]
-    else:
-      if self.axis is None:
-        mean_and_var_shape = ()
-      else:
-        mean_and_var_shape = input_shape[self.axis]
-        self._broadcast_shape[self.axis] = input_shape[self.axis]
+    mean_and_var_shape = []
+    for i in self.axis:
+      mean_and_var_shape.append(input_shape[i])
+      self._broadcast_shape[i] = input_shape[i]
 
     # count is not used in this class's call() method, but is used to re-create
     # the accumulator during multiple calls to 'adapt'.
@@ -179,11 +199,13 @@ class _NormalizingCombiner(Combiner):
     if values.ndim == 1:
       values = np.expand_dims(values, 1)
 
+    # `np.delete` ignores negative indexes, so use a mask to delete items.
+    axis_mask = np.ones([values.ndim], dtype=bool)
+    axis_mask[np.array(self.axis, dtype=np.int32)] = False
+
     # This is the shape of all reduced axes (not specified in 'axis').
-    if self.axis is None:
-      reduction_counts = values.shape
-    else:
-      reduction_counts = np.delete(values.shape, self.axis)
+
+    reduction_counts = np.array(values.shape)[axis_mask]
     # We get the number of elements that will be reduced by multiplying all
     # values of 'shape' corresponding to the reduced axes.
     count = np.prod(reduction_counts, dtype=np.int64)
@@ -191,10 +213,7 @@ class _NormalizingCombiner(Combiner):
     # We want to reduce across dimensions except those specified in 'axis'
     # when using np.mean or np.variance; create the tuple of axes to reduce
     # over here.
-    if self.axis is None:
-      reduction_axes = None
-    else:
-      reduction_axes = tuple(np.delete(range(values.ndim), self.axis))
+    reduction_axes = tuple(np.arange(values.ndim)[axis_mask])
 
     mean = np.mean(values, axis=reduction_axes, dtype=np.float64)
     variance = np.var(values, axis=reduction_axes, dtype=np.float64)
diff --git a/tensorflow/python/keras/layers/preprocessing/normalization_test.py b/tensorflow/python/keras/layers/preprocessing/normalization_test.py
index 75ef9370899..f5f68d9c51a 100644
--- a/tensorflow/python/keras/layers/preprocessing/normalization_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/normalization_test.py
@@ -275,6 +275,49 @@ class NormalizationTest(keras_parameterized.TestCase,
     if context.executing_eagerly():
       self.assertAllClose(output.numpy(), [[-1], [1], [-1], [1]])
 
+  @parameterized.parameters(
+      {"axis": 0},
+      {"axis": (-1, 0)},
+  )
+  def test_zeros_fail_init(self, axis):
+    cls = get_layer_class()
+    with self.assertRaisesRegex(ValueError,
+                                "The argument 'axis' may not be 0."):
+      cls(axis=axis)
+
+  @parameterized.parameters(
+      # Out of bounds
+      {"axis": 3},
+      {"axis": -3},
+      # In a tuple
+      {"axis": (1, 3)},
+      {"axis": (1, -3)},
+  )
+  def test_bad_axis_fail_build(self, axis):
+    cls = get_layer_class()
+    layer = cls(axis=axis)
+    with self.assertRaisesRegex(ValueError,
+                                r"in the range \[1-ndim, ndim-1\]."):
+      layer.build([None, 2, 3])
+
+  @parameterized.parameters(
+      # Results should be identical no matter how the axes are specified (3d).
+      {"axis": (1, 2)},
+      {"axis": (2, 1)},
+      {"axis": (1, -1)},
+      {"axis": (-1, 1)},
+  )
+  def test_axis_permutations(self, axis):
+    cls = get_layer_class()
+    layer = cls(axis=axis)
+    # data.shape = [2, 2, 3]
+    data = np.array([[[0., 1., 2.], [0., 2., 6.]],
+                     [[2., 3., 4.], [3., 6., 10.]]])
+    expect = np.array([[[-1., -1., -1.], [-1., -1., -1.]],
+                       [[1., 1., 1.], [1., 1., 1.]]])
+    layer.adapt(data)
+    self.assertAllClose(expect, layer(data))
+
 
 if __name__ == "__main__":
   test.main()

From cb264418f6be7a8ce4fcfc6ee19a60404e428162 Mon Sep 17 00:00:00 2001
From: Edward Loper 
Date: Wed, 17 Jun 2020 09:19:04 -0700
Subject: [PATCH 0397/1390] Add tf.strings.format() and tf.print() to support
 RaggedTensors.

PiperOrigin-RevId: 316901841
Change-Id: I5fa78acc118557fcf43d1b805149172f8547c0e1
---
 tensorflow/python/ops/ragged/BUILD            |  28 ++-
 .../python/ops/ragged/ragged_dispatch.py      |  11 +-
 .../python/ops/ragged/ragged_dispatch_test.py |  22 +-
 .../python/ops/ragged/ragged_print_op_test.py | 195 ++++++++++++++++++
 .../python/ops/ragged/ragged_string_ops.py    | 117 ++++++++++-
 5 files changed, 359 insertions(+), 14 deletions(-)
 create mode 100644 tensorflow/python/ops/ragged/ragged_print_op_test.py

diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index b2a02b82454..95e5602a246 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -1264,9 +1264,35 @@ py_test(
     srcs_version = "PY2AND3",
     deps = [
         ":ragged_array_ops",
-        "//tensorflow/python:constant_op",
+        ":ragged_factory_ops",
+        ":ragged_tensor",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:dtypes",
+        "//tensorflow/python:errors",
         "//tensorflow/python:framework_test_lib",
         "//tensorflow/python:platform_test",
+        "//tensorflow/python:tensor_shape",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_test(
+    name = "ragged_print_op_test",
+    srcs = ["ragged_print_op_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    deps = [
+        ":ragged",  # fixdeps: keep
+        ":ragged_factory_ops",
+        ":ragged_string_ops",
+        ":ragged_tensor",
+        "//tensorflow/python:constant_op",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:logging_ops",
+        "//tensorflow/python:platform_test",
+        "//tensorflow/python:sparse_ops",
+        "//tensorflow/python/eager:def_function",
         "@absl_py//absl/testing:parameterized",
     ],
 )
diff --git a/tensorflow/python/ops/ragged/ragged_dispatch.py b/tensorflow/python/ops/ragged/ragged_dispatch.py
index f13bed07ba0..5c9388b8677 100644
--- a/tensorflow/python/ops/ragged/ragged_dispatch.py
+++ b/tensorflow/python/ops/ragged/ragged_dispatch.py
@@ -29,6 +29,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import clip_ops
 from tensorflow.python.ops import data_flow_ops
 from tensorflow.python.ops import gen_bitwise_ops
+from tensorflow.python.ops import logging_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import parsing_ops
@@ -510,6 +511,7 @@ _RAGGED_DISPATCH_OPS = [
      ['data', 'segment_ids']),
     (math_ops.unsorted_segment_sqrt_n, ragged_math_ops.segment_sqrt_n,
      ['data', 'segment_ids']),
+    (string_ops.string_format, ragged_string_ops.string_format, ['[inputs]']),
     (string_ops.reduce_join_v2, ragged_string_ops.reduce_join, ['inputs']),
     (math_ops.reduce_sum, ragged_math_ops.reduce_sum, ['input_tensor']),
     (math_ops.reduce_prod, ragged_math_ops.reduce_prod, ['input_tensor']),
@@ -549,7 +551,7 @@ def register_dispatchers():
     RaggedDispatcher(original_op, ragged_op, args).register(original_op)
 
 
-def _ragged_op_signature(op, ragged_args):
+def _ragged_op_signature(op, ragged_args, ragged_varargs=False):
   """Returns a signature for the given op, marking ragged args in bold."""
   op_name = tf_export.get_canonical_name_for_symbol(op)
   argspec = tf_inspect.getfullargspec(op)
@@ -566,7 +568,10 @@ def _ragged_op_signature(op, ragged_args):
 
   # Add varargs and keyword args
   if argspec.varargs:
-    arg_names.append('*' + argspec.varargs)
+    if ragged_varargs:
+      arg_names.append('***' + argspec.varargs + '**')
+    else:
+      arg_names.append('*' + argspec.varargs)
   if argspec.varkw:
     arg_names.append('**' + argspec.varkw)
 
@@ -597,6 +602,8 @@ def ragged_op_list(tf_version=1):
       arginfos = _get_arg_infos(op, ragged_args)
       ragged_args = [arginfo.position for arginfo in arginfos]
       lines.append(_ragged_op_signature(op, ragged_args))
+  lines.append(
+      _ragged_op_signature(logging_ops.print_v2, [], ragged_varargs=True))
   return ('\n\n### Additional ops that support `RaggedTensor`\n\n'
           'Arguments that accept `RaggedTensor`s are marked in **bold**.\n\n' +
           '\n'.join(sorted(lines)) + 'n')
diff --git a/tensorflow/python/ops/ragged/ragged_dispatch_test.py b/tensorflow/python/ops/ragged/ragged_dispatch_test.py
index 60d9f6c8713..193e329e18a 100644
--- a/tensorflow/python/ops/ragged/ragged_dispatch_test.py
+++ b/tensorflow/python/ops/ragged/ragged_dispatch_test.py
@@ -142,8 +142,7 @@ BINARY_INT_OPS = [
 
 # pylint: disable=g-complex-comprehension
 @test_util.run_all_in_graph_and_eager_modes
-class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase,
-                               parameterized.TestCase):
+class RaggedDispatchTest(test_util.TensorFlowTestCase, parameterized.TestCase):
 
   def assertSameShape(self, x, y):
     """Checks that x and y have the same shape (including ragged shapes)."""
@@ -763,7 +762,12 @@ class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase,
               'tensor': ragged_factory_ops.constant_value([[1, 2, 3], [4, 5]]),
               'axis': [0, -1]
           },
-          expected=ragged_factory_ops.constant_value([[5, 4], [3, 2, 1]]))
+          expected=ragged_factory_ops.constant_value([[5, 4], [3, 2, 1]])),
+      dict(
+          op=string_ops.string_format,
+          kwargs={'template': 'Hi {}',
+                  'inputs': [ragged_factory_ops.constant_value([[1, 2], [3]])]},
+          expected='Hi [[1, 2], [3]]'),
   ])
   def testRaggedDispatch(self, op, expected, args=(), result_is_list=False,
                          kwargs=None):
@@ -819,14 +823,14 @@ class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase,
         'math.unsorted_segment_mean', 'math.unsorted_segment_min',
         'math.unsorted_segment_prod', 'math.unsorted_segment_sqrt_n',
         'math.unsorted_segment_sum', 'one_hot', 'ones_like', 'rank', 'realdiv',
-        'reduce_all', 'size', 'squeeze', 'stack', 'strings.as_string',
+        'math.reduce_all', 'size', 'squeeze', 'stack', 'strings.as_string',
         'strings.join', 'strings.length', 'strings.reduce_join',
         'strings.regex_full_match', 'strings.regex_replace', 'strings.strip',
         'strings.substr', 'strings.to_hash_bucket_fast',
         'strings.to_hash_bucket_strong', 'strings.to_hash_bucket',
         'strings.to_number', 'strings.unicode_script', 'tile', 'truncatediv',
         'truncatemod', 'zeros_like', 'dynamic_partition', 'reverse',
-        'nn.dropout',
+        'nn.dropout', 'strings.format', 'print'
     ]
 
     # Ops that should be listed as supported in v1 only.
@@ -838,15 +842,15 @@ class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase,
 
     v1_ragged_ops = ragged_dispatch.ragged_op_list(tf_version=1)
     for element in supported_ops + supported_ops_v1:
-      self.assertIn(element, v1_ragged_ops)
+      self.assertIn('`tf.' + element + '`', v1_ragged_ops)
     for element in supported_ops_v2:
-      self.assertNotIn(element, v1_ragged_ops)
+      self.assertNotIn('`tf.' + element + '`', v1_ragged_ops)
 
     v2_ragged_ops = ragged_dispatch.ragged_op_list(tf_version=2)
     for element in supported_ops + supported_ops_v2:
-      self.assertIn(element, v2_ragged_ops)
+      self.assertIn('`tf.' + element + '`', v2_ragged_ops)
     for element in supported_ops_v1:
-      self.assertNotIn(element, v2_ragged_ops)
+      self.assertNotIn('`tf.' + element + '`', v2_ragged_ops)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/ops/ragged/ragged_print_op_test.py b/tensorflow/python/ops/ragged/ragged_print_op_test.py
new file mode 100644
index 00000000000..2b612d463d0
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_print_op_test.py
@@ -0,0 +1,195 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for tf.print with ragged tensors.
+
+Note: ragged support for tf.print is implemented by RaggedPrintV2Dispatcher in
+ragged_dispatch.py.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os.path
+import tempfile
+from absl.testing import parameterized
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedPrintV2Test(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  # pylint: disable=g-long-lambda
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='2d_int_values',
+          inputs=lambda: [ragged_factory_ops.constant([[1, 2], [3]])],
+          expected='[[1, 2], [3]]\n'),
+      dict(
+          testcase_name='3d_int_values',
+          inputs=lambda: [ragged_factory_ops.constant([[[1, 2], [3]], [[4]]])],
+          expected='[[[1, 2], [3]], [[4]]]\n'),
+      dict(
+          testcase_name='2d_str_values',
+          inputs=lambda: [ragged_factory_ops.constant([['a', 'b'], ['c']])],
+          expected="[['a', 'b'], ['c']]\n"),
+      dict(
+          testcase_name='2d_str_values_with_escaping',
+          inputs=lambda: [ragged_factory_ops.constant([["a'b"], ['c"d']])],
+          expected="[['a\\'b'], ['c\"d']]\n"),
+      dict(
+          testcase_name='two_ragged_values',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              ragged_factory_ops.constant([[5], [], [6, 7, 8]])
+          ],
+          expected='[[1, 2], [3]] [[5], [], [6, 7, 8]]\n'),
+      dict(
+          testcase_name='ragged_value_and_non_tensor_values',
+          inputs=lambda:
+          ['a', 5, True,
+           ragged_factory_ops.constant([[1, 2], [3]]), 'c'],
+          expected='a 5 True [[1, 2], [3]] c\n'),
+      dict(
+          testcase_name='ragged_value_and_dense_value',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              constant_op.constant([[1, 2], [3, 4]])
+          ],
+          expected='[[1, 2], [3]] [[1 2]\n [3 4]]\n'),
+      dict(
+          testcase_name='ragged_value_and_sparse_value',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              sparse_ops.from_dense([[1]])
+          ],
+          expected=(
+              '[[1, 2], [3]] '
+              "'SparseTensor(indices=[[0 0]], values=[1], shape=[1 1])'\n")),
+      dict(
+          testcase_name='summarize_default',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          expected=('[[1, 2, 3, ..., 7, 8, 9], [10], [], '
+                    '..., '
+                    '[], [], [11, 12]]\n')),
+      dict(
+          testcase_name='summarize_2',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          summarize=2,
+          expected='[[1, 2, ..., 8, 9], [10], ..., [], [11, 12]]\n'),
+      dict(
+          testcase_name='summarize_neg1',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          summarize=-1,
+          expected=('[[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], '
+                    '[], [], [], [], [11, 12]]\n')),
+  ])
+  def testRaggedPrint(self, inputs, expected, summarize=None):
+    if callable(inputs):
+      inputs = inputs()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      path = os.path.join(tmpdirname, 'print_output')
+      kwargs = {'output_stream': 'file://{}'.format(path)}
+      if summarize is not None:
+        kwargs.update(summarize=summarize)
+      self.evaluate(logging_ops.print_v2(*inputs, **kwargs))
+      actual = open(path, 'r').read()
+      self.assertEqual(repr(actual), repr(expected))
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedToStringTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      ('2d_int', [[1, 2], [], [3, 4, 5]], '[[1, 2], [], [3, 4, 5]]'),
+      ('2d_str', [['a'], ['b'], ['c', 'd']], "[['a'], ['b'], ['c', 'd']]"),
+      ('3d_int', [[[1, 2], []], [[3, 4, 5]]], '[[[1, 2], []], [[3, 4, 5]]]'),
+      ('escape', [["a'b"], [r'c\d']], r"[['a\'b'], ['c\\d']]"),
+      dict(testcase_name='2d_empty', rt=[], ragged_rank=1, expected='[]'),
+      dict(testcase_name='3d_empty', rt=[], ragged_rank=2, expected='[]'),
+      dict(
+          testcase_name='3d_rrank1',
+          rt=[[[1, 2], [3, 4]], [], [[5, 6]]],
+          ragged_rank=1,
+          expected='[[[1, 2], [3, 4]], [], [[5, 6]]]'),
+      dict(
+          testcase_name='2d_empty_row', rt=[[]], ragged_rank=1,
+          expected='[[]]'),
+      dict(
+          testcase_name='3d_empty_row', rt=[[]], ragged_rank=2,
+          expected='[[]]'),
+      dict(
+          testcase_name='summarize_1',
+          rt=[[1, 2, 3, 4, 5], [], [6], [7], [8, 9]],
+          summarize=1,
+          expected='[[1, ..., 5], ..., [8, 9]]'),
+      dict(
+          testcase_name='summarize_2',
+          rt=[[1, 2, 3, 4, 5], [], [6], [7], [8, 9]],
+          summarize=2,
+          expected='[[1, 2, ..., 4, 5], [], ..., [7], [8, 9]]'),
+  ])
+  def testRaggedToString(self, rt, expected, summarize=None, ragged_rank=None):
+    rt = ragged_factory_ops.constant(rt, ragged_rank=ragged_rank)
+    actual = ragged_string_ops.ragged_tensor_to_string(rt, summarize=summarize)
+    self.assertAllEqual(actual, expected)
+
+  @parameterized.named_parameters([
+      ('maxelts_BadType', [[1]], "Expected summarize .*, got 'foo'", 'foo'),
+      ('maxelts_0', [[1]], 'Expected summarize to be .*, got 0', 0),
+      ('maxelts_Neg2', [[1]], 'Expected summarize to be .*, got -2', -2),
+  ])
+  def testRaggedToStringErrors(self,
+                               rt,
+                               error,
+                               summarize=None,
+                               exception=ValueError):
+    rt = ragged_factory_ops.constant(rt)
+    with self.assertRaisesRegex(exception, error):
+      self.evaluate(
+          ragged_string_ops.ragged_tensor_to_string(rt, summarize=summarize))
+
+  def testRaggedToStringUnknownRank(self):
+
+    @def_function.function(
+        input_signature=[ragged_tensor.RaggedTensorSpec(ragged_rank=1)])
+    def f(rt):
+      return ragged_string_ops.ragged_tensor_to_string(rt)
+
+    with self.assertRaisesRegex(
+        ValueError, 'RaggedTensor to_string requires '
+        'that rt.shape.rank is not None'):
+      f(ragged_factory_ops.constant([[1, 2], [3]]))
+
+
+if __name__ == '__main__':
+  googletest.main()
diff --git a/tensorflow/python/ops/ragged/ragged_string_ops.py b/tensorflow/python/ops/ragged/ragged_string_ops.py
index 0d9c4d506f3..0ac23c298ba 100755
--- a/tensorflow/python/ops/ragged/ragged_string_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_string_ops.py
@@ -18,10 +18,13 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import gen_string_ops
 from tensorflow.python.ops import string_ops
 from tensorflow.python.ops.ragged import ragged_array_ops
@@ -30,9 +33,14 @@ from tensorflow.python.ops.ragged import ragged_tensor
 from tensorflow.python.util import compat as util_compat
 from tensorflow.python.util import deprecation
 from tensorflow.python.util import dispatch
+from tensorflow.python.util.lazy_loader import LazyLoader
 from tensorflow.python.util.tf_export import tf_export
 
 
+map_fn_lib = LazyLoader("map_fn_lib", globals(),
+                        "tensorflow.python.ops.map_fn")
+
+
 @tf_export("strings.bytes_split")
 @dispatch.add_dispatch_support
 def string_bytes_split(input, name=None):  # pylint: disable=redefined-builtin
@@ -640,7 +648,7 @@ def strings_split_v1(input=None, sep=None, maxsplit=-1,  # pylint: disable=redef
         input, dtype=dtypes.string, name="input")
 
     if input.shape.rank == 0:
-      input = gen_array_ops.expand_dims(input, 0)
+      input = array_ops.expand_dims(input, 0)
 
     if result_type == "SparseTensor":
       if input.shape.rank == 1:
@@ -813,3 +821,108 @@ def ngrams(data,
         values=output, row_splits=output_splits, validate=False)
     return array_ops.reshape(output.flat_values,
                              dense_shape) if to_tensor else output
+
+
+def string_format(template, inputs, placeholder="{}", summarize=3, name=None):
+  """Version of tf.strings.format that handles RaggedTensors."""
+  if tensor_util.is_tensor(inputs) or ragged_tensor.is_ragged(inputs):
+    inputs = [inputs]
+
+  split_template = template.split(placeholder)
+  if len(inputs) != len(split_template) - 1:
+    raise ValueError("num placeholders in template and num inputs must match"
+                     ": {} vs {}".format(len(split_template) - 1, len(inputs)))
+
+  with ops.name_scope(name, "StringFormat", [inputs]):
+    output_pieces = [constant_op.constant(split_template[0])]
+    for i, input in enumerate(inputs):
+      if ragged_tensor.is_ragged(input):
+        output_pieces.append(ragged_tensor_to_string(input, summarize))
+      else:
+        output_pieces.append(string_ops.string_format(
+            "{}", [input], summarize=summarize))
+      output_pieces.append(constant_op.constant(split_template[i + 1]))
+    if len(output_pieces) == 1:
+      return output_pieces[0]
+    else:
+      return string_ops.reduce_join(output_pieces)
+
+
+def ragged_tensor_to_string(rt, summarize=None):
+  """Returns a scalar string tensor with the contents of a RaggedTensor.
+
+  Requires that `rt.shape.rank` is not `None`.
+
+  Note: this converts the entire `RaggedTensor` into a single string scalar.
+  If you want to convert individual elements, use `tf.strings.as_string(rt)`.
+
+  >>> rt1 = tf.ragged.constant([[1, 2, 3], [4, 5]])
+  >>> ragged_tensor_to_string(rt1).numpy()
+  b'[[1, 2, 3], [4, 5]]'
+
+  >>> rt2 = tf.ragged.constant([[['a'], ['b', 'c']], [['d', 'e', 'f'], []]])
+  >>> ragged_tensor_to_string(rt2).numpy()
+  b"[[['a'], ['b', 'c']], [['d', 'e', 'f'], []]]"
+
+  >>> rt3 = tf.ragged.constant([[1], [2, 3, 4, 5, 6], [], [], [7], [8, 9]])
+  >>> ragged_tensor_to_string(rt3, summarize=2).numpy()
+  b'[[1], [2, 3, ..., 5, 6], ..., [7], [8, 9]]'
+
+  Args:
+    rt: The RaggedTensor that should be converted to a string.
+    summarize: If specified, then only the first and last `summarize` elements
+      within each dimension are included in the string. If `-1` or `None`, then
+      all elements are included.
+  """
+  if (summarize is not None and summarize != -1 and
+      not (isinstance(summarize, int) and summarize > 0)):
+    raise ValueError("Expected summarize to be -1 or a positive int, got %r" %
+                     summarize)
+  with ops.name_scope(None, "AsString", [rt]):
+    rt = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt)
+    if rt.shape.rank is None:
+      raise ValueError("RaggedTensor to_string requires that rt.shape.rank "
+                       "is not None.")
+    # Convert all elements of `rt` to strings.
+    if rt.dtype == dtypes.string:
+      escaped = string_ops.regex_replace(rt.flat_values, r"(['\\])", r"\\\1")
+      str_t = rt.with_flat_values("'" + escaped + "'")
+    else:
+      str_t = rt.with_flat_values(string_ops.as_string(rt.flat_values))
+
+    return _ragged_tensor_to_string(str_t, summarize)
+
+
+def _ragged_tensor_to_string(string_tensor, summarize):
+  """Returns a scalar string tensor with the contents of `string_tensor`.
+
+  Args:
+    string_tensor: A potentially ragged tensor with dtype=string.
+    summarize: Include only the first and last `summarize` elements of each
+      dimension.  If `-1` or `None`, then include all elements.
+
+  Returns:
+    A scalar string Tensor.
+  """
+  if string_tensor.shape.rank == 1:
+    pieces = string_tensor
+  else:
+    pieces = map_fn_lib.map_fn(
+        lambda s: _ragged_tensor_to_string(s, summarize),
+        string_tensor,
+        fn_output_signature=tensor_spec.TensorSpec(None, dtypes.string))
+  if summarize not in (-1, None):
+    pieces = control_flow_ops.cond(
+        _nrows(string_tensor) <= 2 * summarize,
+        lambda: pieces,
+        lambda: array_ops.concat(  # pylint: disable=g-long-lambda
+            [pieces[:summarize], ["..."], pieces[-summarize:]],
+            axis=0))
+  return "[" + string_ops.reduce_join(pieces, separator=", ") + "]"
+
+
+def _nrows(tensor, out_type=dtypes.int32):
+  if isinstance(tensor, ragged_tensor.RaggedTensor):
+    return tensor.nrows(out_type=out_type)
+  else:
+    return array_ops.shape(tensor, out_type=out_type)[0]

From fa6074832921268ffe81dfcca6b2133d19721b29 Mon Sep 17 00:00:00 2001
From: Marat Dukhan 
Date: Wed, 17 Jun 2020 09:27:58 -0700
Subject: [PATCH 0398/1390] Recognize "armeabi" cpu as Linux/ARM in XNNPACK
 backend

PiperOrigin-RevId: 316903436
Change-Id: I22ea414026df1f78da63df6ff6c872c65e9bcf5c
---
 tensorflow/workspace.bzl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index cb1ea721fb0..354c4c353b9 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -164,11 +164,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
 
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "bd5fd63a09222cd092f0c058b576cf044fb4074f2c4ce8a6fc32fc43d155f9c7",
-        strip_prefix = "XNNPACK-ae046f5a5127084bfe41090afdf1c1d4c9874b77",
+        sha256 = "714d650828b1409e88ccb2a62b36a47827bdcddd875bfcfd3b321fe1b7b1c106",
+        strip_prefix = "XNNPACK-b8e7b076a0c2e7356a69b8478fcd76498d357b45",
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/ae046f5a5127084bfe41090afdf1c1d4c9874b77.zip",
-            "https://github.com/google/XNNPACK/archive/ae046f5a5127084bfe41090afdf1c1d4c9874b77.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/b8e7b076a0c2e7356a69b8478fcd76498d357b45.zip",
+            "https://github.com/google/XNNPACK/archive/b8e7b076a0c2e7356a69b8478fcd76498d357b45.zip",
         ],
     )
 
@@ -184,11 +184,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
 
     tf_http_archive(
         name = "pthreadpool",
-        sha256 = "c4d4a16053ec0e5125dbd8ae1d6d2ba99601d6fdcf8601a0d51d02a048c40348",
-        strip_prefix = "pthreadpool-e1642461b3b0217d23d6664d839a060f54e4e652",
+        sha256 = "03312bd7d8d9e379d685258963ee8820767158b5946cdd00336ff17dae851001",
+        strip_prefix = "pthreadpool-029c88620802e1361ccf41d1970bd5b07fd6b7bb",
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/Maratyszcza/pthreadpool/archive/e1642461b3b0217d23d6664d839a060f54e4e652.zip",
-            "https://github.com/Maratyszcza/pthreadpool/archive/e1642461b3b0217d23d6664d839a060f54e4e652.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/Maratyszcza/pthreadpool/archive/029c88620802e1361ccf41d1970bd5b07fd6b7bb.zip",
+            "https://github.com/Maratyszcza/pthreadpool/archive/029c88620802e1361ccf41d1970bd5b07fd6b7bb.zip",
         ],
     )
 

From 225bdf60f3c4f51ab5568a53d31b0799369a1b89 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 09:30:39 -0700
Subject: [PATCH 0399/1390] Added new way of scalar reading in
 ParseMultiplyScalar.

PiperOrigin-RevId: 316903891
Change-Id: I6a160c0a1a83cff9ee6b1df9c3f59ea0477d6c30
---
 tensorflow/lite/delegates/gpu/common/model_builder.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/gpu/common/model_builder.cc b/tensorflow/lite/delegates/gpu/common/model_builder.cc
index dc671a47691..01f94c94888 100644
--- a/tensorflow/lite/delegates/gpu/common/model_builder.cc
+++ b/tensorflow/lite/delegates/gpu/common/model_builder.cc
@@ -1271,7 +1271,7 @@ class MulOperationParser : public TFLiteOperationParser {
                                    GraphFloat32* graph, ObjectReader* reader) {
     RETURN_IF_ERROR(reader->AddInput(node, runtime_tensor));
     MultiplyAttributes attr;
-    if (constant_dims->size <= 0) {
+    if (constant_dims->size <= 0 || NumElements(constant_dims) == 1) {
       Tensor tensor;
       RETURN_IF_ERROR(reader->ReadTensor(constant_tensor, &tensor));
       attr.param = tensor.data[0];

From 51ccd6911b8bab58df2e8be4f31ced43b04cff96 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 09:31:04 -0700
Subject: [PATCH 0400/1390] Choosing better setting in convolution for Intel.

PiperOrigin-RevId: 316903965
Change-Id: I9ff6c2a5026059011b5ccf7beddb8111b419ff8d
---
 tensorflow/lite/delegates/gpu/cl/cl_device.cc |  6 +++++
 tensorflow/lite/delegates/gpu/cl/cl_device.h  |  1 +
 .../delegates/gpu/cl/kernels/conv_powervr.cc  | 23 +++++++++++++++----
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.cc b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
index aea81d5e659..64e07428515 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_device.cc
+++ b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
@@ -495,6 +495,12 @@ std::string CLDevice::GetPlatformVersion() const {
   return GetPlatformInfo(platform_id_, CL_PLATFORM_VERSION);
 }
 
+bool CLDevice::IsCL20OrHigher() const {
+  return info_.cl_version != OpenCLVersion::CL_1_0 &&
+         info_.cl_version != OpenCLVersion::CL_1_1 &&
+         info_.cl_version != OpenCLVersion::CL_1_2;
+}
+
 bool CLDevice::IsAdreno() const { return info_.vendor == Vendor::QUALCOMM; }
 
 bool CLDevice::IsAdreno3xx() const {
diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.h b/tensorflow/lite/delegates/gpu/cl/cl_device.h
index 1df16aa3bad..ae6a1d11af6 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_device.h
+++ b/tensorflow/lite/delegates/gpu/cl/cl_device.h
@@ -178,6 +178,7 @@ class CLDevice {
   bool SupportsExtension(const std::string& extension) const;
   bool SupportsFP32RTN() const;
   bool SupportsFP16RTN() const;
+  bool IsCL20OrHigher() const;
   bool IsAdreno() const;
   bool IsAdreno3xx() const;
   bool IsAdreno4xx() const;
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
index 363a0157420..bd694e7cc4f 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
@@ -859,12 +859,27 @@ ConvPowerVR::ConvParams ConvPowerVR::GuessBestParams(
     conv_params.src_depth_loop_size = 1;
     conv_params.weights_upload_type = WeightsUploadType::GLOBAL_MEM;
   } else if (device.IsIntel()) {
+    if (different_weights_for_height) {
+      conv_params.work_group_size = int3(16, 1, 1);
+      conv_params.work_group_launch_order = int3(0, 1, 2);
+      conv_params.fixed_work_group_size = true;
+    } else {
+      conv_params.linear_hw = true;
+      conv_params.work_group_size = int3(16, 1, 1);
+      conv_params.work_group_launch_order = int3(0, 1, 2);
+      conv_params.fixed_work_group_size = true;
+    }
     conv_params.block_size = int3(1, 1, 4);
-    conv_params.work_group_size = int3(8, 2, 1);
-    conv_params.work_group_launch_order = int3(0, 1, 2);
-    conv_params.fixed_work_group_size = true;
     conv_params.src_depth_loop_size = 1;
-    conv_params.weights_upload_type = WeightsUploadType::LOCAL_MEM_BY_THREADS;
+    if (definition.precision != CalculationsPrecision::F32_F16 &&
+        device.SupportsExtension("cl_khr_subgroups") &&
+        device.SupportsExtension("cl_intel_required_subgroup_size") &&
+        device.IsCL20OrHigher()) {
+      conv_params.weights_upload_type =
+          WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST;
+    } else {
+      conv_params.weights_upload_type = WeightsUploadType::LOCAL_MEM_BY_THREADS;
+    }
     if (dst_depth % 4 == 0 || dst_depth >= 8) {
       conv_params.block_size.z = 4;
     } else if (dst_depth % 2 == 0 || dst_depth >= 4) {

From 8b7af2c77091b4e4286aa3ef9efda324646fbf99 Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Wed, 17 Jun 2020 23:41:36 +0700
Subject: [PATCH 0401/1390] Use char** instead of string

---
 tensorflow/c/env.cc              | 10 ++--------
 tensorflow/c/env.h               |  5 +----
 tensorflow/core/platform/path.cc |  2 --
 3 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc
index 43879a18359..6d8528bc42f 100644
--- a/tensorflow/c/env.cc
+++ b/tensorflow/c/env.cc
@@ -147,14 +147,8 @@ TF_StringStream* TF_GetLocalTempDirectories() {
   return list;
 }
 
-void TF_GetTempFileName(const char* extension, std::string* name,
-                        TF_Status* status) {
-  *name = ::tensorflow::io::GetTempFilename(extension);
-  if (name->length() == 0) {
-    TF_SetStatus(status, TF_INTERNAL, "Can not get temp file name");
-  } else {
-    TF_SetStatus(status, TF_OK, "");
-  }
+void TF_GetTempFileName(const char* extension, char** name) {
+  *name = strdup(::tensorflow::io::GetTempFilename(extension).c_str());
 }
 
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) {
diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h
index 273a3b5e142..29ec417a75e 100644
--- a/tensorflow/c/env.h
+++ b/tensorflow/c/env.h
@@ -20,8 +20,6 @@ limitations under the License.
 #include 
 #include 
 
-#include 
-
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/c/tf_file_statistics.h"
 
@@ -157,8 +155,7 @@ TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void);
 // Creates a temporary file name with an extension.
 // The caller is responsible for freeing the returned pointer.
 TF_CAPI_EXPORT extern void TF_GetTempFileName(const char* extension,
-                                              std::string* name,
-                                              TF_Status* status);
+                                              char** name);
 
 // Returns the number of nanoseconds since the Unix epoch.
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void);
diff --git a/tensorflow/core/platform/path.cc b/tensorflow/core/platform/path.cc
index f9442ccba0f..1e88328aace 100644
--- a/tensorflow/core/platform/path.cc
+++ b/tensorflow/core/platform/path.cc
@@ -327,8 +327,6 @@ string GetTempFilename(const string& extension) {
   }
   LOG(FATAL) << "No temp directory found.";
 #endif
-  // Return an empty string to indicate that we can not create temp file name.
-  return "";
 }
 
 bool GetTestUndeclaredOutputsDir(string* dir) {

From b780ee931b1d4c14aae7ac534c937a88204d52c2 Mon Sep 17 00:00:00 2001
From: Edward Loper 
Date: Wed, 17 Jun 2020 09:33:05 -0700
Subject: [PATCH 0402/1390] Replace the mechanism used to register & look up
 Python types from c code in tensorflow/python/util.h with one that supports
 non-type symbols as well.

PiperOrigin-RevId: 316904361
Change-Id: I4ec98c861742efddcebd140ff9e1a6ff567cc94c
---
 tensorflow/python/util/util.cc                | 56 +++++++++++++------
 tensorflow/python/util/util.h                 | 15 ++++-
 tensorflow/python/util/util_wrapper.cc        |  4 ++
 .../tools/def_file_filter/symbols_pybind.txt  |  1 +
 4 files changed, 56 insertions(+), 20 deletions(-)

diff --git a/tensorflow/python/util/util.cc b/tensorflow/python/util/util.cc
index 1d0dd695d74..cf8581443e7 100644
--- a/tensorflow/python/util/util.cc
+++ b/tensorflow/python/util/util.cc
@@ -29,15 +29,25 @@ limitations under the License.
 namespace tensorflow {
 namespace swig {
 
-std::unordered_map* PythonTypesMap() {
+namespace {
+string PyObjectToString(PyObject* o);
+}  // namespace
+
+std::unordered_map* RegisteredPyObjectMap() {
   static auto* m = new std::unordered_map();
   return m;
 }
 
-PyObject* GetRegisteredType(const string& key) {
-  auto* m = PythonTypesMap();
-  auto it = m->find(key);
-  if (it == m->end()) return nullptr;
+PyObject* GetRegisteredPyObject(const string& name) {
+  const auto* m = RegisteredPyObjectMap();
+  auto it = m->find(name);
+  if (it == m->end()) {
+    PyErr_SetString(PyExc_TypeError,
+                    tensorflow::strings::StrCat("No object with name ", name,
+                                                " has been registered.")
+                        .c_str());
+    return nullptr;
+  }
   return it->second;
 }
 
@@ -49,26 +59,35 @@ PyObject* RegisterType(PyObject* type_name, PyObject* type) {
                         .c_str());
     return nullptr;
   }
+  return RegisterPyObject(type_name, type);
+}
 
+PyObject* RegisterPyObject(PyObject* name, PyObject* value) {
   string key;
-  if (PyBytes_Check(type_name)) {
-    key = PyBytes_AsString(type_name);
-  }
+  if (PyBytes_Check(name)) {
+    key = PyBytes_AsString(name);
 #if PY_MAJOR_VERSION >= 3
-  if (PyUnicode_Check(type_name)) {
-    key = PyUnicode_AsUTF8(type_name);
-  }
+  } else if (PyUnicode_Check(name)) {
+    key = PyUnicode_AsUTF8(name);
 #endif
-
-  if (PythonTypesMap()->find(key) != PythonTypesMap()->end()) {
+  } else {
     PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat(
-                                         "Type already registered for ", key)
+                                         "Expected name to be a str, got",
+                                         PyObjectToString(name))
                                          .c_str());
     return nullptr;
   }
 
-  Py_INCREF(type);
-  PythonTypesMap()->emplace(key, type);
+  auto* m = RegisteredPyObjectMap();
+  if (m->find(key) != m->end()) {
+    PyErr_SetString(PyExc_TypeError, tensorflow::strings::StrCat(
+                                         "Value already registered for ", key)
+                                         .c_str());
+    return nullptr;
+  }
+
+  Py_INCREF(value);
+  m->emplace(key, value);
 
   Py_RETURN_NONE;
 }
@@ -196,7 +215,7 @@ class CachedTypeCheck {
 // Returns 0 otherwise.
 // Returns -1 if an error occurred (e.g., if 'type_name' is not registered.)
 int IsInstanceOfRegisteredType(PyObject* obj, const char* type_name) {
-  PyObject* type_obj = GetRegisteredType(type_name);
+  PyObject* type_obj = GetRegisteredPyObject(type_name);
   if (TF_PREDICT_FALSE(type_obj == nullptr)) {
     PyErr_SetString(PyExc_RuntimeError,
                     tensorflow::strings::StrCat(
@@ -513,7 +532,8 @@ class AttrsValueIterator : public ValueIterator {
 };
 
 bool IsSparseTensorValueType(PyObject* o) {
-  PyObject* sparse_tensor_value_type = GetRegisteredType("SparseTensorValue");
+  PyObject* sparse_tensor_value_type =
+      GetRegisteredPyObject("SparseTensorValue");
   if (TF_PREDICT_FALSE(sparse_tensor_value_type == nullptr)) {
     return false;
   }
diff --git a/tensorflow/python/util/util.h b/tensorflow/python/util/util.h
index 23438b43c53..fc0b864416e 100644
--- a/tensorflow/python/util/util.h
+++ b/tensorflow/python/util/util.h
@@ -19,6 +19,8 @@ limitations under the License.
 
 #include 
 
+#include 
+
 namespace tensorflow {
 namespace swig {
 
@@ -270,11 +272,20 @@ PyObject* FlattenForData(PyObject* nested);
 PyObject* AssertSameStructureForData(PyObject* o1, PyObject* o2,
                                      bool check_types);
 
-// RegisterType is used to pass PyTypeObject (which is defined in python) for an
-// arbitrary identifier `type_name` into C++.
+// Registers a Python object so it can be looked up from c++.  The set of
+// valid names, and the expected values for those names, are listed in
+// the documentation for `RegisteredPyObjects`.  Returns PyNone.
+PyObject* RegisterPyObject(PyObject* name, PyObject* value);
+
+// Variant of RegisterPyObject that requires the object's value to be a type.
 PyObject* RegisterType(PyObject* type_name, PyObject* type);
 
 }  // namespace swig
+
+// Returns a borrowed reference to an object that was registered with
+// RegisterPyObject.  (Do not call PY_DECREF on the result).
+PyObject* GetRegisteredPyObject(const std::string& name);
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_PYTHON_UTIL_UTIL_H_
diff --git a/tensorflow/python/util/util_wrapper.cc b/tensorflow/python/util/util_wrapper.cc
index dd74306413c..63c70d785cc 100644
--- a/tensorflow/python/util/util_wrapper.cc
+++ b/tensorflow/python/util/util_wrapper.cc
@@ -30,6 +30,10 @@ PYBIND11_MODULE(_pywrap_utils, m) {
           return tensorflow::PyoOrThrow(
               tensorflow::swig::RegisterType(type_name.ptr(), type.ptr()));
         });
+  m.def("RegisterPyObject", [](const py::handle& name, const py::handle& type) {
+    return tensorflow::PyoOrThrow(
+        tensorflow::swig::RegisterPyObject(name.ptr(), type.ptr()));
+  });
   m.def(
       "IsTensor",
       [](const py::handle& o) {
diff --git a/tensorflow/tools/def_file_filter/symbols_pybind.txt b/tensorflow/tools/def_file_filter/symbols_pybind.txt
index e72ef973ff2..07f5906aa08 100644
--- a/tensorflow/tools/def_file_filter/symbols_pybind.txt
+++ b/tensorflow/tools/def_file_filter/symbols_pybind.txt
@@ -17,6 +17,7 @@ tensorflow::swig::Flatten
 tensorflow::swig::IsSequenceForData
 tensorflow::swig::FlattenForData
 tensorflow::swig::AssertSameStructureForData
+tensorflow::swig::RegisterPyObject
 tensorflow::swig::RegisterType
 tensorflow::swig::IsEagerTensorSlow
 

From d8e0beacd934a3a4e7ee3c14240a9e398a929a27 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 09:33:37 -0700
Subject: [PATCH 0403/1390] Reshape&Reshapex4 converted to new style.

PiperOrigin-RevId: 316904479
Change-Id: I7c1fb0ca5a31fc1f82545d70cfcdcfb7d63bcd6a
---
 .../lite/delegates/gpu/cl/kernels/reshape.cc  | 185 ++++++++----------
 .../delegates/gpu/cl/kernels/reshapex4.cc     | 127 ++++++------
 2 files changed, 140 insertions(+), 172 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc b/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
index e1589e9d682..a99fff0a1da 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
@@ -25,92 +25,24 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetReshapeBatchedCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHSBPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHSBPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
+std::string GetReshapeBatchedCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
 
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
-  c += "    int src_channels,          \n";
-  c += "    int dst_channels           \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int linear_id = get_global_id(0);\n";
-  c += "  int X = linear_id / dst_size.w;\n";
-  c += "  int B = linear_id % dst_size.w;\n";
+  c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+  c += "  int B = linear_id % args.dst_tensor.Batch();\n";
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z || B >= "
-       "dst_size.w) return;\n";
-  c += "  FLT temps[4];\n";
-  c += "  temps[0] = (FLT)(0.0f);\n";
-  c += "  temps[1] = (FLT)(0.0f);\n";
-  c += "  temps[2] = (FLT)(0.0f);\n";
-  c += "  temps[3] = (FLT)(0.0f);\n";
-  c += "  int base = ((B * dst_size.y + Y)* dst_size.x + X)* dst_channels + Z "
-       "* 4;\n";
-  c += "  for (int i = 0; i < 4; ++i) {\n";
-  c += "    int dst_channel = Z * 4 + i;\n";
-  c += "    if (dst_channel < dst_channels) {;\n";
-  c += "      int p = base + i;\n";
-  c += "      int src_c = p % src_channels;\n";
-  c += "      p = p / src_channels;\n";
-  c += "      int src_x = p % src_size.x;\n";
-  c += "      p = p / src_size.x;\n";
-  c += "      int src_y = p % src_size.y;\n";
-  c += "      int src_b = p / src_size.y;\n";
-  c += "      int src_z = src_c / 4;\n";
-  c += "      int src_sub_ch = src_c % 4;\n";
-  c += "      FLT4 t =" +
-       src_tensor.ReadWHSB("src_x", "src_y", "src_z", "src_b") + ";\n";
-  c += "      FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
-  c += "      temps[i] = t_ar[src_sub_ch];\n";
-  c += "    }\n";
-  c += "  }\n";
-  c += "  FLT4 result = (FLT4)(temps[0], temps[1], temps[2], temps[3]);\n";
-  const LinkingContext context{"result", "X * dst_size.w + B", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHSB("result", "X", "Y", "Z", "B");
-  c += "}\n";
-  return c;
-}
-
-std::string GetReshapeCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
-
-  std::string c = GetCommonDefines(op_def.precision);
-  c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
-  c += "    int src_channels,          \n";
-  c += "    int dst_channels           \n";
-  c += ") {\n";
-  c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) { \n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
   c += "    return; \n";
   c += "  } \n";
   c += "  FLT temps[4];\n";
@@ -118,25 +50,73 @@ std::string GetReshapeCode(
   c += "  temps[1] = (FLT)(0.0f);\n";
   c += "  temps[2] = (FLT)(0.0f);\n";
   c += "  temps[3] = (FLT)(0.0f);\n";
+  c += "  int base = ((B * args.dst_tensor.Height() + Y) * "
+       "args.dst_tensor.Width() + X) * args.dst_tensor.Channels() + Z * 4;\n";
   c += "  for (int i = 0; i < 4; ++i) {\n";
   c += "    int dst_channel = Z * 4 + i;\n";
-  c += "    if (dst_channel < dst_channels) {;\n";
-  c += "      int p = dst_channel + dst_channels * (X + dst_size.x * Y);\n";
-  c += "      int src_c = p % src_channels;\n";
-  c += "      p = p / src_channels;\n";
-  c += "      int src_x = p % src_size.x;\n";
-  c += "      int src_y = p / src_size.x;\n";
+  c += "    if (dst_channel < args.dst_tensor.Channels()) {;\n";
+  c += "      int p = base + i;\n";
+  c += "      int src_c = p % args.src_tensor.Channels();\n";
+  c += "      p = p / args.src_tensor.Channels();\n";
+  c += "      int src_x = p % args.src_tensor.Width();\n";
+  c += "      p = p / args.src_tensor.Width();\n";
+  c += "      int src_y = p % args.src_tensor.Height();\n";
+  c += "      int src_b = p / args.src_tensor.Height();\n";
   c += "      int src_z = src_c / 4;\n";
   c += "      int src_sub_ch = src_c % 4;\n";
-  c += "      FLT4 t =" + src_tensor.ReadWHS("src_x", "src_y", "src_z") + ";\n";
+  c += "      FLT4 t = args.src_tensor.Read(src_x, src_y, src_z, src_b);\n";
   c += "      FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
   c += "      temps[i] = t_ar[src_sub_ch];\n";
   c += "    }\n";
   c += "  }\n";
   c += "  FLT4 result = (FLT4)(temps[0], temps[1], temps[2], temps[3]);\n";
-  const LinkingContext context{"result", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("result", "X", "Y", "Z");
+  c += "  args.dst_tensor.Write(result, X, Y, Z, B);\n";
+  c += "}\n";
+  return c;
+}
+
+std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+
+  std::string c = GetCommonDefines(op_def.precision);
+  c += "__kernel void main_function(\n";
+  c += "$0) {\n";
+  c += "  int X = get_global_id(0);\n";
+  c += "  int Y = get_global_id(1);\n";
+  c += "  int Z = get_global_id(2);\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  FLT temps[4];\n";
+  c += "  temps[0] = (FLT)(0.0f);\n";
+  c += "  temps[1] = (FLT)(0.0f);\n";
+  c += "  temps[2] = (FLT)(0.0f);\n";
+  c += "  temps[3] = (FLT)(0.0f);\n";
+  c += "  int base = (Y * args.dst_tensor.Width() + X) * "
+       "args.dst_tensor.Channels() + Z * 4;\n";
+  c += "  for (int i = 0; i < 4; ++i) {\n";
+  c += "    int dst_channel = Z * 4 + i;\n";
+  c += "    if (dst_channel < args.dst_tensor.Channels()) {;\n";
+  c += "      int p = base + i;\n";
+  c += "      int src_c = p % args.src_tensor.Channels();\n";
+  c += "      p = p / args.src_tensor.Channels();\n";
+  c += "      int src_x = p % args.src_tensor.Width();\n";
+  c += "      int src_y = p / args.src_tensor.Width();\n";
+  c += "      int src_z = src_c / 4;\n";
+  c += "      int src_sub_ch = src_c % 4;\n";
+  c += "      FLT4 t = args.src_tensor.Read(src_x, src_y, src_z);\n";
+  c += "      FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
+  c += "      temps[i] = t_ar[src_sub_ch];\n";
+  c += "    }\n";
+  c += "  }\n";
+  c += "  FLT4 result = (FLT4)(temps[0], temps[1], temps[2], temps[3]);\n";
+  c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
   c += "}\n";
   return c;
 }
@@ -157,24 +137,25 @@ Reshape& Reshape::operator=(Reshape&& operation) {
 }
 
 absl::Status Reshape::Compile(const CreationContext& creation_context) {
-  const auto code = definition_.IsBatchSupported()
-                        ? GetReshapeBatchedCode(definition_, linked_operations_)
-                        : GetReshapeCode(definition_, linked_operations_);
+  std::string code = definition_.IsBatchSupported()
+                         ? GetReshapeBatchedCode(definition_, &args_)
+                         : GetReshapeCode(definition_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Reshape::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->Channels()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->Channels()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Reshape::GetGridSize() const {
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc b/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
index de6813e741f..0847fce5836 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
@@ -25,82 +25,66 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetReshapeBatchedCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHSBPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHSBPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
+std::string GetReshapeBatchedCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
 
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size              \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int linear_id = get_global_id(0);\n";
-  c += "  int X = linear_id / dst_size.w;\n";
-  c += "  int B = linear_id % dst_size.w;\n";
+  c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+  c += "  int B = linear_id % args.dst_tensor.Batch();\n";
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z || B >= "
-       "dst_size.w) return;\n";
-  c += "  int dst_bhwc4 = ((B * dst_size.y + Y) * dst_size.x + X) * dst_size.z "
-       "+ Z;\n";
-  c += "  int src_z = dst_bhwc4 % src_size.z;\n";
-  c += "  dst_bhwc4 = dst_bhwc4 / src_size.z;\n";
-  c += "  int src_x = dst_bhwc4 % src_size.x;\n";
-  c += "  dst_bhwc4 = dst_bhwc4 / src_size.x;\n";
-  c += "  int src_y = dst_bhwc4 % src_size.y;\n";
-  c += "  int src_b = dst_bhwc4 / src_size.y;\n";
-  c += "  FLT4 result =" +
-       src_tensor.ReadWHSB("src_x", "src_y", "src_z", "src_b") + ";\n";
-  const LinkingContext context{"result", "X * dst_size.w + B", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHSB("result", "X", "Y", "Z", "B");
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  int dst_bhwc4 = ((B * args.dst_tensor.Height() + Y) * "
+       "args.dst_tensor.Width() + X) * args.dst_tensor.Slices() + Z;\n";
+  c += "  int src_z = dst_bhwc4 % args.src_tensor.Slices();\n";
+  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Slices();\n";
+  c += "  int src_x = dst_bhwc4 % args.src_tensor.Width();\n";
+  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Width();\n";
+  c += "  int src_y = dst_bhwc4 % args.src_tensor.Height();\n";
+  c += "  int src_b = dst_bhwc4 / args.src_tensor.Height();\n";
+  c += "  FLT4 result = args.src_tensor.Read(src_x, src_y, src_z, src_b);\n";
+  c += "  args.dst_tensor.Write(result, X, Y, Z, B);\n";
   c += "}\n";
   return c;
 }
 
-std::string GetReshapeCode(
-    const OperationDef& op_def,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
+  args->AddObjectRef(
+      "src_tensor", AccessType::READ,
+      absl::make_unique(op_def.src_tensors[0]));
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
 
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size              \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
-  c += "  int dst_hwc4 = (Y * dst_size.x + X) * dst_size.z + Z;\n";
-  c += "  int src_z = dst_hwc4 % src_size.z;\n";
-  c += "  dst_hwc4 = dst_hwc4 / src_size.z;\n";
-  c += "  int src_x = dst_hwc4 % src_size.x;\n";
-  c += "  int src_y = dst_hwc4 / src_size.x;\n";
-  c +=
-      "  FLT4 result =" + src_tensor.ReadWHS("src_x", "src_y", "src_z") + ";\n";
-  const LinkingContext context{"result", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("result", "X", "Y", "Z");
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  int dst_hwc4 = (Y * args.dst_tensor.Width() + X) * "
+       "args.dst_tensor.Slices() + Z;\n";
+  c += "  int src_z = dst_hwc4 % args.src_tensor.Slices();\n";
+  c += "  dst_hwc4 = dst_hwc4 / args.src_tensor.Slices();\n";
+  c += "  int src_x = dst_hwc4 % args.src_tensor.Width();\n";
+  c += "  int src_y = dst_hwc4 / args.src_tensor.Width();\n";
+  c += "  FLT4 result = args.src_tensor.Read(src_x, src_y, src_z);\n";
+  c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
   c += "}\n";
   return c;
 }
@@ -121,22 +105,25 @@ Reshapex4& Reshapex4::operator=(Reshapex4&& operation) {
 }
 
 absl::Status Reshapex4::Compile(const CreationContext& creation_context) {
-  const auto code = definition_.IsBatchSupported()
-                        ? GetReshapeBatchedCode(definition_, linked_operations_)
-                        : GetReshapeCode(definition_, linked_operations_);
+  std::string code = definition_.IsBatchSupported()
+                         ? GetReshapeBatchedCode(definition_, &args_)
+                         : GetReshapeCode(definition_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Reshapex4::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWHSB()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Reshapex4::GetGridSize() const {

From c870b9f9203c989c7e2940c7f60f2e288e42138c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 09:50:14 -0700
Subject: [PATCH 0404/1390] Qualify uses of std::string

PiperOrigin-RevId: 316907769
Change-Id: I01c3ca5f0fb1b1b2284ae5969da968cec7493583
---
 .../lite/toco/allocate_transient_arrays.cc    |  35 ++--
 tensorflow/lite/toco/args.cc                  |  26 +--
 tensorflow/lite/toco/args.h                   |  66 +++----
 tensorflow/lite/toco/dump_graphviz.cc         | 103 +++++-----
 tensorflow/lite/toco/dump_graphviz.h          |   4 +-
 tensorflow/lite/toco/export_tensorflow.cc     | 183 +++++++++---------
 tensorflow/lite/toco/export_tensorflow.h      |   3 +-
 tensorflow/lite/toco/format_port.h            |  10 +-
 tensorflow/lite/toco/import_tensorflow.cc     | 142 +++++++-------
 tensorflow/lite/toco/import_tensorflow.h      |   2 +-
 .../lite/toco/import_tensorflow_test.cc       |   6 +-
 tensorflow/lite/toco/model.h                  |  38 ++--
 tensorflow/lite/toco/model_cmdline_flags.cc   |  30 +--
 tensorflow/lite/toco/model_cmdline_flags.h    |   2 +-
 .../lite/toco/model_cmdline_flags_test.cc     |   6 +-
 tensorflow/lite/toco/tensorflow_util.cc       |   6 +-
 tensorflow/lite/toco/tensorflow_util.h        |   2 +-
 tensorflow/lite/toco/toco.cc                  |   2 +-
 tensorflow/lite/toco/toco_cmdline_flags.cc    |   6 +-
 tensorflow/lite/toco/toco_cmdline_flags.h     |   3 +-
 .../lite/toco/toco_cmdline_flags_test.cc      |   4 +-
 tensorflow/lite/toco/toco_convert.cc          |  14 +-
 tensorflow/lite/toco/toco_convert.h           |   4 +-
 tensorflow/lite/toco/toco_convert_test.cc     |  24 +--
 tensorflow/lite/toco/toco_port.cc             |  16 +-
 tensorflow/lite/toco/toco_port.h              |  18 +-
 tensorflow/lite/toco/toco_tooling.cc          |   9 +-
 tensorflow/lite/toco/toco_tooling.h           |   7 +-
 tensorflow/lite/toco/tooling_util.cc          | 163 ++++++++--------
 tensorflow/lite/toco/tooling_util.h           |  74 +++----
 30 files changed, 523 insertions(+), 485 deletions(-)

diff --git a/tensorflow/lite/toco/allocate_transient_arrays.cc b/tensorflow/lite/toco/allocate_transient_arrays.cc
index 3ec53c9c2d6..76279f6e62d 100644
--- a/tensorflow/lite/toco/allocate_transient_arrays.cc
+++ b/tensorflow/lite/toco/allocate_transient_arrays.cc
@@ -55,8 +55,8 @@ bool EndsAt(const ArrayLifespan& lifespan, std::size_t op_index) {
 // Helper function for ComputeArrayLifespans: updates one ArrayLifespan for
 // one array for one op.
 void UpdateArrayLifespan(
-    const string& array_name, std::size_t op_index,
-    std::unordered_map* array_lifespans) {
+    const std::string& array_name, std::size_t op_index,
+    std::unordered_map* array_lifespans) {
   if (array_lifespans->count(array_name)) {
     auto& lifespan = array_lifespans->at(array_name);
     if (!lifespan.persistent) {
@@ -74,7 +74,7 @@ void UpdateArrayLifespan(
 // Computes the ArrayLifespan for each array.
 void ComputeArrayLifespans(
     const Model& model,
-    std::unordered_map* array_lifespans) {
+    std::unordered_map* array_lifespans) {
   CHECK(array_lifespans->empty());
   for (const auto& rnn_state : model.flags.rnn_states()) {
     ArrayLifespan lifespan;
@@ -159,7 +159,8 @@ class Allocator {
 
 // Returns the required transient allocation size (in bytes) for a given array,
 // or 0 if it's not a transient array.
-std::size_t TransientArraySize(const Model& model, const string& array_name,
+std::size_t TransientArraySize(const Model& model,
+                               const std::string& array_name,
                                std::size_t transient_data_alignment) {
   if (!IsAllocatableTransientArray(model, array_name)) {
     return 0;
@@ -191,7 +192,7 @@ std::size_t TransientArraySize(const Model& model, const string& array_name,
 
 // Allocates an array: call this for every array just before the first
 // op where it is used.
-void AllocateTransientArray(const Model& model, const string& array_name,
+void AllocateTransientArray(const Model& model, const std::string& array_name,
                             Allocator* allocator,
                             std::size_t transient_data_alignment) {
   if (!IsAllocatableTransientArray(model, array_name)) {
@@ -206,7 +207,7 @@ void AllocateTransientArray(const Model& model, const string& array_name,
 
 // Deallocates an array: call this for every array just after the last
 // op where it is used.
-void DeallocateTransientArray(const Model& model, const string& array_name,
+void DeallocateTransientArray(const Model& model, const std::string& array_name,
                               Allocator* allocator) {
   if (!IsAllocatableTransientArray(model, array_name)) {
     return;
@@ -216,7 +217,7 @@ void DeallocateTransientArray(const Model& model, const string& array_name,
   allocator->Deallocate(*array->alloc);
 }
 
-void PushBackIfNotFound(const string& s, std::vector* v) {
+void PushBackIfNotFound(const std::string& s, std::vector* v) {
   if (std::find(v->begin(), v->end(), s) == v->end()) {
     v->push_back(s);
   }
@@ -227,7 +228,7 @@ void PushBackIfNotFound(const string& s, std::vector* v) {
 void AllocateTransientArrays(Model* model,
                              std::size_t transient_data_alignment) {
   // Precompute the lifespans for all arrays.
-  std::unordered_map array_lifespans;
+  std::unordered_map array_lifespans;
   ComputeArrayLifespans(*model, &array_lifespans);
 
   // In case of variable batch, our convention will be to compute the
@@ -250,7 +251,7 @@ void AllocateTransientArrays(Model* model,
 
   // Construct a sorted map of array names, so that other layout engines can
   // match exactly.
-  std::map ordered_arrays_map;
+  std::map ordered_arrays_map;
   for (const auto& pair : model->GetArrayMap()) {
     ordered_arrays_map[pair.first] = pair.second.get();
   }
@@ -258,7 +259,7 @@ void AllocateTransientArrays(Model* model,
   // Allocate persistent arrays (like RNN states). For them, 'transient'
   // is a misnormer, should read 'workspace'.
   for (const auto& array_pair : ordered_arrays_map) {
-    const string& array_name = array_pair.first;
+    const std::string& array_name = array_pair.first;
     auto it = array_lifespans.find(array_name);
     if (it != array_lifespans.end() && it->second.persistent) {
       AllocateTransientArray(*model, array_name, &allocator,
@@ -270,7 +271,7 @@ void AllocateTransientArrays(Model* model,
        op_index++) {
     const auto& op = model->operators[op_index];
     // Allocate those arrays whose lifespan starts exactly here.
-    std::vector arrays_to_allocate;
+    std::vector arrays_to_allocate;
     for (const auto& input : op->inputs) {
       if (StartsAt(array_lifespans[input], op_index)) {
         PushBackIfNotFound(input, &arrays_to_allocate);
@@ -281,13 +282,13 @@ void AllocateTransientArrays(Model* model,
         PushBackIfNotFound(output, &arrays_to_allocate);
       }
     }
-    for (const string& array : arrays_to_allocate) {
+    for (const std::string& array : arrays_to_allocate) {
       AllocateTransientArray(*model, array, &allocator,
                              transient_data_alignment);
     }
 
     // Deallocate those arrays whose lifespan ends exactly here.
-    std::vector arrays_to_deallocate;
+    std::vector arrays_to_deallocate;
     for (const auto& input : op->inputs) {
       if (EndsAt(array_lifespans[input], op_index)) {
         PushBackIfNotFound(input, &arrays_to_deallocate);
@@ -298,7 +299,7 @@ void AllocateTransientArrays(Model* model,
         PushBackIfNotFound(output, &arrays_to_deallocate);
       }
     }
-    for (const string& array : arrays_to_deallocate) {
+    for (const std::string& array : arrays_to_deallocate) {
       DeallocateTransientArray(*model, array, &allocator);
     }
   }
@@ -309,7 +310,7 @@ void AllocateTransientArrays(Model* model,
   std::size_t optimal_transient_alloc_size = 0;
   std::size_t persistent_alloc_size = 0;
   for (const auto& array_pair : ordered_arrays_map) {
-    const string& array_name = array_pair.first;
+    const std::string& array_name = array_pair.first;
     auto it = array_lifespans.find(array_name);
     if (it != array_lifespans.end() && it->second.persistent) {
       persistent_alloc_size +=
@@ -320,7 +321,7 @@ void AllocateTransientArrays(Model* model,
     // for each operator, compute the sum of the sizes of the array that must
     // be live during the execution of this operator, plus the size of
     // persistent arrays that must be live at all times.
-    std::vector non_persistent_edges;
+    std::vector non_persistent_edges;
     for (const auto& input : op->inputs) {
       if (!array_lifespans[input].persistent) {
         PushBackIfNotFound(input, &non_persistent_edges);
@@ -332,7 +333,7 @@ void AllocateTransientArrays(Model* model,
       }
     }
     std::size_t size = persistent_alloc_size;
-    for (const string& edge : non_persistent_edges) {
+    for (const std::string& edge : non_persistent_edges) {
       size += TransientArraySize(*model, edge, transient_data_alignment);
     }
     // The optimal total size is the maximum of all operator-specific sizes.
diff --git a/tensorflow/lite/toco/args.cc b/tensorflow/lite/toco/args.cc
index ce67de900d7..c30b98ce516 100644
--- a/tensorflow/lite/toco/args.cc
+++ b/tensorflow/lite/toco/args.cc
@@ -94,14 +94,16 @@ bool SplitStructuredLine(absl::string_view line, char delimiter,
 }
 
 inline bool TryStripPrefixString(absl::string_view str,
-                                 absl::string_view prefix, string* result) {
+                                 absl::string_view prefix,
+                                 std::string* result) {
   bool res = absl::ConsumePrefix(&str, prefix);
   result->assign(str.begin(), str.end());
   return res;
 }
 
 inline bool TryStripSuffixString(absl::string_view str,
-                                 absl::string_view suffix, string* result) {
+                                 absl::string_view suffix,
+                                 std::string* result) {
   bool res = absl::ConsumeSuffix(&str, suffix);
   result->assign(str.begin(), str.end());
   return res;
@@ -109,7 +111,7 @@ inline bool TryStripSuffixString(absl::string_view str,
 
 }  // namespace
 
-bool Arg::Parse(string text) {
+bool Arg::Parse(std::string text) {
   parsed_value_.elements.clear();
   specified_ = true;
   // strings::Split("") produces {""}, but we need {} on empty input.
@@ -125,7 +127,7 @@ bool Arg::Parse(string text) {
   return true;
 }
 
-bool Arg::Parse(string text) {
+bool Arg::Parse(std::string text) {
   parsed_value_.elements.clear();
   specified_ = true;
 
@@ -138,24 +140,24 @@ bool Arg::Parse(string text) {
   // TODO(aselle): Change argument parsing when absl supports structuredline.
   SplitStructuredLine(text_disposable_copy, ',', "{}", &outer_vector);
   for (const absl::string_view& outer_member_stringpiece : outer_vector) {
-    string outer_member(outer_member_stringpiece);
+    std::string outer_member(outer_member_stringpiece);
     if (outer_member.empty()) {
       continue;
     }
-    string outer_member_copy = outer_member;
+    std::string outer_member_copy = outer_member;
     absl::StripAsciiWhitespace(&outer_member);
     if (!TryStripPrefixString(outer_member, "{", &outer_member)) return false;
     if (!TryStripSuffixString(outer_member, "}", &outer_member)) return false;
-    const std::vector inner_fields_vector =
+    const std::vector inner_fields_vector =
         absl::StrSplit(outer_member, ',');
 
-    std::unordered_map element;
-    for (const string& member_field : inner_fields_vector) {
-      std::vector outer_member_key_value =
+    std::unordered_map element;
+    for (const std::string& member_field : inner_fields_vector) {
+      std::vector outer_member_key_value =
           absl::StrSplit(member_field, ':');
       if (outer_member_key_value.size() != 2) return false;
-      string& key = outer_member_key_value[0];
-      string& value = outer_member_key_value[1];
+      std::string& key = outer_member_key_value[0];
+      std::string& value = outer_member_key_value[1];
       absl::StripAsciiWhitespace(&key);
       absl::StripAsciiWhitespace(&value);
       if (element.count(key) != 0) return false;
diff --git a/tensorflow/lite/toco/args.h b/tensorflow/lite/toco/args.h
index 20fa5ecc20c..e1fe209062e 100644
--- a/tensorflow/lite/toco/args.h
+++ b/tensorflow/lite/toco/args.h
@@ -35,7 +35,7 @@ struct IntList {
   std::vector elements;
 };
 struct StringMapList {
-  std::vector> elements;
+  std::vector> elements;
 };
 
 // command_line_flags.h don't track whether or not a flag is specified. Arg
@@ -82,13 +82,13 @@ template <>
 class Arg final {
  public:
   // Provide default_value() to arg list
-  string default_value() const { return ""; }
+  std::string default_value() const { return ""; }
   // Return true if the command line argument was specified on the command line.
   bool specified() const { return specified_; }
   // Bind the parse member function so tensorflow::Flags can call it.
-  bool Parse(string text);
+  bool Parse(std::string text);
 
-  std::function bind() {
+  std::function bind() {
     return std::bind(&Arg::Parse, this, std::placeholders::_1);
   }
 
@@ -103,14 +103,14 @@ template <>
 class Arg final {
  public:
   // Provide default_value() to StringMapList
-  string default_value() const { return ""; }
+  std::string default_value() const { return ""; }
   // Return true if the command line argument was specified on the command line.
   bool specified() const { return specified_; }
   // Bind the parse member function so tensorflow::Flags can call it.
 
-  bool Parse(string text);
+  bool Parse(std::string text);
 
-  std::function bind() {
+  std::function bind() {
     return std::bind(&Arg::Parse, this, std::placeholders::_1);
   }
 
@@ -123,18 +123,18 @@ class Arg final {
 
 // Flags that describe a model. See model_cmdline_flags.cc for details.
 struct ParsedModelFlags {
-  Arg input_array;
-  Arg input_arrays;
-  Arg output_array;
-  Arg output_arrays;
-  Arg input_shapes;
+  Arg input_array;
+  Arg input_arrays;
+  Arg output_array;
+  Arg output_arrays;
+  Arg input_shapes;
   Arg batch_size = Arg(1);
   Arg mean_value = Arg(0.f);
-  Arg mean_values;
+  Arg mean_values;
   Arg std_value = Arg(1.f);
-  Arg std_values;
-  Arg input_data_type;
-  Arg input_data_types;
+  Arg std_values;
+  Arg input_data_type;
+  Arg input_data_types;
   Arg variable_batch = Arg(false);
   Arg input_shape;
   Arg rnn_states;
@@ -142,44 +142,44 @@ struct ParsedModelFlags {
   Arg change_concat_input_ranges = Arg(true);
   // Debugging output options.
   // TODO(benoitjacob): these shouldn't be ModelFlags.
-  Arg graphviz_first_array;
-  Arg graphviz_last_array;
-  Arg dump_graphviz;
+  Arg graphviz_first_array;
+  Arg graphviz_last_array;
+  Arg dump_graphviz;
   Arg dump_graphviz_video = Arg(false);
-  Arg conversion_summary_dir;
+  Arg conversion_summary_dir;
   Arg allow_nonexistent_arrays = Arg(false);
   Arg allow_nonascii_arrays = Arg(false);
-  Arg arrays_extra_info_file;
-  Arg model_flags_file;
+  Arg arrays_extra_info_file;
+  Arg model_flags_file;
 };
 
 // Flags that describe the operation you would like to do (what conversion
 // you want). See toco_cmdline_flags.cc for details.
 struct ParsedTocoFlags {
-  Arg input_file;
-  Arg savedmodel_directory;
-  Arg output_file;
-  Arg input_format = Arg("TENSORFLOW_GRAPHDEF");
-  Arg output_format = Arg("TFLITE");
-  Arg savedmodel_tagset;
+  Arg input_file;
+  Arg savedmodel_directory;
+  Arg output_file;
+  Arg input_format = Arg("TENSORFLOW_GRAPHDEF");
+  Arg output_format = Arg("TFLITE");
+  Arg savedmodel_tagset;
   // TODO(aselle): command_line_flags  doesn't support doubles
   Arg default_ranges_min = Arg(0.);
   Arg default_ranges_max = Arg(0.);
   Arg default_int16_ranges_min = Arg(0.);
   Arg default_int16_ranges_max = Arg(0.);
-  Arg inference_type;
-  Arg inference_input_type;
+  Arg inference_type;
+  Arg inference_input_type;
   Arg drop_fake_quant = Arg(false);
   Arg reorder_across_fake_quant = Arg(false);
   Arg allow_custom_ops = Arg(false);
   Arg allow_dynamic_tensors = Arg(true);
-  Arg custom_opdefs;
+  Arg custom_opdefs;
   Arg post_training_quantize = Arg(false);
   Arg quantize_to_float16 = Arg(false);
   // Deprecated flags
   Arg quantize_weights = Arg(false);
-  Arg input_type;
-  Arg input_types;
+  Arg input_type;
+  Arg input_types;
   Arg debug_disable_recurrent_cell_fusion = Arg(false);
   Arg drop_control_dependency = Arg(false);
   Arg propagate_fake_quant_num_bits = Arg(false);
diff --git a/tensorflow/lite/toco/dump_graphviz.cc b/tensorflow/lite/toco/dump_graphviz.cc
index 68d3b957129..006d5546c60 100644
--- a/tensorflow/lite/toco/dump_graphviz.cc
+++ b/tensorflow/lite/toco/dump_graphviz.cc
@@ -77,7 +77,9 @@ class Color {
 
   // Returns the string serialization of this color in graphviz format,
   // for use as 'fillcolor' in boxes.
-  string AsHexString() const { return StringF("#%.2X%.2X%.2X", r_, g_, b_); }
+  std::string AsHexString() const {
+    return StringF("#%.2X%.2X%.2X", r_, g_, b_);
+  }
   // The color to use for this node; will be used as 'fillcolor'
   // for its box. See Color::AsHexString. A suitable, different
   // color will be chosen for the 'fontcolor' for the inside text
@@ -85,7 +87,7 @@ class Color {
   // Returns the serialization in graphviz format of a suitable color to use
   // 'fontcolor' in the same boxes. It should black or white, whichever offers
   // the better contrast from AsHexString().
-  string TextColorString() const {
+  std::string TextColorString() const {
     // https://en.wikipedia.org/wiki/Relative_luminance
     const float luminance = 0.2126f * r_ + 0.7152f * g_ + 0.0722f * b_;
     const uint8 l = luminance > 128.f ? 0 : 255;
@@ -96,7 +98,7 @@ class Color {
   uint8 r_ = 0, g_ = 0, b_ = 0;
 };
 
-Color HashStringToColor(string s) {
+Color HashStringToColor(std::string s) {
   // Return a unique color for a name.
   //
   // This function removes Tensorflow anti-collision suffixes (eg "_2"), hashes
@@ -120,8 +122,8 @@ Color HashStringToColor(string s) {
   return Color(color_word);
 }
 
-void GetArrayColorAndShape(const Model& model, const string& array_name,
-                           Color* color, string* shape) {
+void GetArrayColorAndShape(const Model& model, const std::string& array_name,
+                           Color* color, std::string* shape) {
   // All colors in this file are from:
   // https://material.io/guidelines/style/color.html
   // Arrays involved in RNN back-edges have a different color
@@ -167,7 +169,8 @@ void GetArrayColorAndShape(const Model& model, const string& array_name,
   *shape = "box";
 }
 
-string GetArrayCompassPt(const Model& model, const string& array_name) {
+std::string GetArrayCompassPt(const Model& model,
+                              const std::string& array_name) {
   // The "compass point" is the point on the node where edge connections are
   // made. For most arrays we don't care, but input's and outputs look better
   // connected at the tip of the "house" and "invhouse" shapes used. So we
@@ -191,7 +194,7 @@ string GetArrayCompassPt(const Model& model, const string& array_name) {
   return "";
 }
 
-void AppendArrayVal(string* string, Array const& array, int index) {
+void AppendArrayVal(std::string* string, Array const& array, int index) {
   if (array.buffer->type == ArrayDataType::kFloat) {
     const auto& data = array.GetBuffer().data;
     if (index >= data.size()) {
@@ -231,10 +234,10 @@ void AppendArrayVal(string* string, Array const& array, int index) {
   }
 }
 
-typedef std::map Attributes;
+typedef std::map Attributes;
 
-string AttributesToHtml(Attributes attributes) {
-  string html;
+std::string AttributesToHtml(Attributes attributes) {
+  std::string html;
   for (const auto& attr : attributes) {
     html += R"CODE()CODE";
     html += attr.first;
@@ -245,8 +248,8 @@ string AttributesToHtml(Attributes attributes) {
   return html;
 }
 
-string GetArrayLabel(const Model& model, const string& array_id) {
-  string html;
+std::string GetArrayLabel(const Model& model, const std::string& array_id) {
+  std::string html;
 
   // Use HTML-like labels (http://www.graphviz.org/doc/info/shapes.html#html)
   html += "<";
@@ -265,7 +268,7 @@ string GetArrayLabel(const Model& model, const string& array_id) {
   html += R"CODE()CODE";
   html += R"CODE()CODE";
   AppendF(&html, R"CODE(%s)CODE",
-          std::vector(absl::StrSplit(array_id, '/')).back());
+          std::vector(absl::StrSplit(array_id, '/')).back());
   html += R"CODE()CODE";
   html += "";
 
@@ -371,7 +374,7 @@ Attributes GetOpAttributes(const Model& model, const Operator& op) {
   switch (op.type) {
     case OperatorType::kConv: {
       const auto& conv_op = static_cast(op);
-      string stride;
+      std::string stride;
       AppendF(&stride, "%d", conv_op.stride_width);
       stride += kUnicodeMult;
       AppendF(&stride, "%d", conv_op.stride_height);
@@ -382,7 +385,7 @@ Attributes GetOpAttributes(const Model& model, const Operator& op) {
     }
     case OperatorType::kDepthwiseConv: {
       const auto& depthconv_op = static_cast(op);
-      string stride;
+      std::string stride;
       AppendF(&stride, "%d", depthconv_op.stride_width);
       stride += kUnicodeMult;
       AppendF(&stride, "%d", depthconv_op.stride_height);
@@ -426,9 +429,9 @@ Color GetOpColor(const Operator& op) {
   }
 }
 
-string GetOpLabel(const Model& model, const Operator& op) {
+std::string GetOpLabel(const Model& model, const Operator& op) {
   // Use HTML-like labels (http://www.graphviz.org/doc/info/shapes.html#html)
-  string html;
+  std::string html;
   html += "<";
 
   // Begin Table
@@ -462,7 +465,8 @@ string GetOpLabel(const Model& model, const Operator& op) {
   if (op.type == OperatorType::kUnsupported) {
     html += static_cast(op).tensorflow_op;
   } else {
-    html += string(absl::StripPrefix(OperatorTypeName(op.type), "TensorFlow"));
+    html +=
+        std::string(absl::StripPrefix(OperatorTypeName(op.type), "TensorFlow"));
   }
   html += R"CODE()CODE";
   html += "";
@@ -498,7 +502,7 @@ string GetOpLabel(const Model& model, const Operator& op) {
   return html;
 }
 
-float GetLog2BufferSize(const Model& model, const string& array_id) {
+float GetLog2BufferSize(const Model& model, const std::string& array_id) {
   auto& array = model.GetArray(array_id);
   if (array.has_shape()) {
     int buffer_size = 0;
@@ -510,22 +514,23 @@ float GetLog2BufferSize(const Model& model, const string& array_id) {
   return 0.0f;
 }
 
-string GetOpId(int op_index) { return StringF("op%05d", op_index); }
+std::string GetOpId(int op_index) { return StringF("op%05d", op_index); }
 
-void DumpOperator(const Model& model, string* output_file, int op_index) {
+void DumpOperator(const Model& model, std::string* output_file, int op_index) {
   // Dump node for operator.
   const Operator& op = *model.operators[op_index];
   Color color = GetOpColor(op);
-  string label = GetOpLabel(model, op);
-  string op_id = GetOpId(op_index);
+  std::string label = GetOpLabel(model, op);
+  std::string op_id = GetOpId(op_index);
   AppendF(output_file, kOpNodeFmt, op_id, label, color.AsHexString(),
           color.TextColorString());
 }
 
-void DumpOperatorEdges(const Model& model, string* output_file, int op_index) {
+void DumpOperatorEdges(const Model& model, std::string* output_file,
+                       int op_index) {
   // Inputs
   const Operator& op = *model.operators[op_index];
-  string op_id = GetOpId(op_index);
+  std::string op_id = GetOpId(op_index);
   for (int i = 0; i < op.inputs.size(); i++) {
     const auto& input = op.inputs[i];
     if (!model.HasArray(input)) {
@@ -546,7 +551,7 @@ void DumpOperatorEdges(const Model& model, string* output_file, int op_index) {
       // would otherwise skew the layout.
       weight = 1.0f;
     }
-    string compass_pt = GetArrayCompassPt(model, input);
+    std::string compass_pt = GetArrayCompassPt(model, input);
     AppendF(output_file, kInputEdgeFmt, input, compass_pt, op_id, i, line_width,
             weight);
   }
@@ -563,7 +568,7 @@ void DumpOperatorEdges(const Model& model, string* output_file, int op_index) {
     if (!IsArrayConsumed(model, output)) {
       weight = 1.0f;
     }
-    string compass_pt = GetArrayCompassPt(model, output);
+    std::string compass_pt = GetArrayCompassPt(model, output);
     AppendF(output_file, kOutputEdgeFmt, op_id, i, output, compass_pt,
             line_width, weight);
   }
@@ -572,19 +577,19 @@ void DumpOperatorEdges(const Model& model, string* output_file, int op_index) {
 struct Node {
   Node() : math_ops(0) {}
   // Name used as a key in the model's array map
-  string array_id;
+  std::string array_id;
 
   // Estimated number of math ops incurred by this node (the sum of the op
   // with this array as 1st output, plus all children nodes).
   int64 math_ops;
 
   // A map of child nodes keyed by name.
-  std::map> children;
+  std::map> children;
 };
 
-string GetSubgraphLabel(Node const& node, const string& subgraph) {
+std::string GetSubgraphLabel(Node const& node, const std::string& subgraph) {
   // Use HTML-like labels (http://www.graphviz.org/doc/info/shapes.html#html)
-  string html;
+  std::string html;
   html += "<";
 
   // Begin Table
@@ -613,19 +618,19 @@ string GetSubgraphLabel(Node const& node, const string& subgraph) {
   return html;
 }
 
-void DumpSubgraphHeader(string* output_file, Node const& node,
-                        const string& node_name) {
+void DumpSubgraphHeader(std::string* output_file, Node const& node,
+                        const std::string& node_name) {
   Color color = HashStringToColor(node_name);
-  string label = GetSubgraphLabel(node, node_name);
+  std::string label = GetSubgraphLabel(node, node_name);
   AppendF(output_file, kSubgraphFmt, node_name, color.AsHexString(), label);
 }
 
-void DumpArray(const Model& model, string* output_file,
-               const string& array_id) {
+void DumpArray(const Model& model, std::string* output_file,
+               const std::string& array_id) {
   Color color;
-  string shape;
+  std::string shape;
   GetArrayColorAndShape(model, array_id, &color, &shape);
-  string label = GetArrayLabel(model, array_id);
+  std::string label = GetArrayLabel(model, array_id);
   AppendF(output_file, kArrayNodeFmt, array_id, label, array_id, shape,
           color.AsHexString(), color.TextColorString());
 
@@ -638,8 +643,8 @@ void DumpArray(const Model& model, string* output_file,
   }
 }
 
-void DumpNode(const Model& model, string* output_file, const string& node_name,
-              Node const& node) {
+void DumpNode(const Model& model, std::string* output_file,
+              const std::string& node_name, Node const& node) {
   bool not_root = !node_name.empty();
   if (not_root) {
     DumpSubgraphHeader(output_file, node, node_name);
@@ -662,7 +667,7 @@ void DumpNode(const Model& model, string* output_file, const string& node_name,
   }
 }
 
-int64 GetArithmeticOpsCount(const Model& model, const string& array_id) {
+int64 GetArithmeticOpsCount(const Model& model, const std::string& array_id) {
   for (const auto& op : model.operators) {
     if (!op->outputs.empty() && op->outputs[0] == array_id) {
       int64 count;
@@ -676,15 +681,15 @@ int64 GetArithmeticOpsCount(const Model& model, const string& array_id) {
   return 0;
 }
 
-void InsertNode(const Model& model, const string& array_id, Node* node,
-                std::vector prefixes, int64* math_ops) {
+void InsertNode(const Model& model, const std::string& array_id, Node* node,
+                std::vector prefixes, int64* math_ops) {
   if (prefixes.empty()) {
     // Base case: store array in this node.
     node->array_id = array_id;
     *math_ops = GetArithmeticOpsCount(model, array_id);
   } else {
     // Insert into the sub-tree for that prefix.
-    string prefix = prefixes.back();
+    std::string prefix = prefixes.back();
     prefixes.pop_back();
     if (node->children.count(prefix) == 0) {
       // Create a new node if this prefix is unseen.
@@ -700,16 +705,16 @@ void InsertNode(const Model& model, const string& array_id, Node* node,
 void BuildArrayTree(const Model& model, Node* tree) {
   // Delimit array names by path "/", then place into a tree based on this path.
   for (const auto& array_id : model.GetArrayMap()) {
-    std::vector prefixes = absl::StrSplit(array_id.first, '/');
+    std::vector prefixes = absl::StrSplit(array_id.first, '/');
     std::reverse(prefixes.begin(), prefixes.end());
     int64 math_ops;  // Temporary storage for math ops used during recursion.
     InsertNode(model, array_id.first, tree, prefixes, &math_ops);
   }
 }
 
-string GetGraphLabel(const Model& model, const string& graph_name) {
+std::string GetGraphLabel(const Model& model, const std::string& graph_name) {
   // Use HTML-like labels (http://www.graphviz.org/doc/info/shapes.html#html)
-  string html;
+  std::string html;
   html += "<";
 
   // Begin Table
@@ -753,8 +758,8 @@ string GetGraphLabel(const Model& model, const string& graph_name) {
 }
 }  // namespace
 
-void DumpGraphviz(const Model& model, string* output_file,
-                  const string& graph_name) {
+void DumpGraphviz(const Model& model, std::string* output_file,
+                  const std::string& graph_name) {
   // Start graphviz format
   AppendF(output_file, kGraphFmt, GetGraphLabel(model, graph_name));
 
diff --git a/tensorflow/lite/toco/dump_graphviz.h b/tensorflow/lite/toco/dump_graphviz.h
index 9bb74dac3f8..0e847896552 100644
--- a/tensorflow/lite/toco/dump_graphviz.h
+++ b/tensorflow/lite/toco/dump_graphviz.h
@@ -21,8 +21,8 @@ limitations under the License.
 
 namespace toco {
 
-void DumpGraphviz(const Model& model, string* output_file_contents,
-                  const string& graph_name);
+void DumpGraphviz(const Model& model, std::string* output_file_contents,
+                  const std::string& graph_name);
 
 }  // namespace toco
 
diff --git a/tensorflow/lite/toco/export_tensorflow.cc b/tensorflow/lite/toco/export_tensorflow.cc
index ec3fb386d10..7ecf6cc7d44 100644
--- a/tensorflow/lite/toco/export_tensorflow.cc
+++ b/tensorflow/lite/toco/export_tensorflow.cc
@@ -49,7 +49,7 @@ namespace toco {
 namespace {
 
 tensorflow::DataType GetTensorFlowDataType(ArrayDataType data_type,
-                                           const string& error_location) {
+                                           const std::string& error_location) {
   switch (data_type) {
     case ArrayDataType::kBool:
       return tensorflow::DT_BOOL;
@@ -74,12 +74,12 @@ tensorflow::DataType GetTensorFlowDataType(ArrayDataType data_type,
 }
 
 tensorflow::DataType GetTensorFlowDataTypeForOp(ArrayDataType data_type,
-                                                const string& op_name) {
+                                                const std::string& op_name) {
   return GetTensorFlowDataType(data_type, "op '" + op_name + "'");
 }
 
 tensorflow::DataType GetTensorFlowDataType(const Model& model,
-                                           const string& array_name) {
+                                           const std::string& array_name) {
   return GetTensorFlowDataType(model.GetArray(array_name).data_type,
                                "array '" + array_name + "'");
 }
@@ -113,8 +113,8 @@ void ExportFloatArray(const Shape& input_shape, const float* input_data,
     }
   }
   output_tensor->set_tensor_content(
-      string(reinterpret_cast(input_data),
-             sizeof(*input_data) * input_flat_size));
+      std::string(reinterpret_cast(input_data),
+                  sizeof(*input_data) * input_flat_size));
 }
 
 void ExportFloatArray(AxesOrder input_axes_order, const Shape& input_shape,
@@ -137,7 +137,7 @@ void ExportFloatArray(AxesOrder input_axes_order, const Shape& input_shape,
                    legacy_scalar_policy);
 }
 
-bool HasAlreadyExportedConst(const string& name,
+bool HasAlreadyExportedConst(const std::string& name,
                              const GraphDef& tensorflow_graph) {
   for (const auto& node : tensorflow_graph.node()) {
     if (node.op() == "Const" && node.name() == name) {
@@ -147,7 +147,7 @@ bool HasAlreadyExportedConst(const string& name,
   return false;
 }
 
-void ConvertFloatTensorConst(const string& name, const Shape& input_shape,
+void ConvertFloatTensorConst(const std::string& name, const Shape& input_shape,
                              const float* input_data,
                              AxesOrder input_axes_order,
                              AxesOrder output_axes_order,
@@ -165,7 +165,7 @@ void ConvertFloatTensorConst(const string& name, const Shape& input_shape,
                    tensor, legacy_scalar_policy);
 }
 
-void ConvertFloatTensorConst(const string& name, const Shape& input_shape,
+void ConvertFloatTensorConst(const std::string& name, const Shape& input_shape,
                              const float* input_data,
                              AxesOrder input_axes_order,
                              AxesOrder output_axes_order,
@@ -175,7 +175,7 @@ void ConvertFloatTensorConst(const string& name, const Shape& input_shape,
                           LegacyScalarPolicy::kAvoidLegacyScalars);
 }
 
-void ConvertFloatTensorConst(const Model& model, const string& name,
+void ConvertFloatTensorConst(const Model& model, const std::string& name,
                              AxesOrder input_axes_order,
                              AxesOrder output_axes_order,
                              GraphDef* tensorflow_graph) {
@@ -193,7 +193,7 @@ void ConvertFloatTensorConst(const Model& model, const string& name,
                           output_axes_order, tensorflow_graph);
 }
 
-void ConvertFloatTensorConst(const Model& model, const string& name,
+void ConvertFloatTensorConst(const Model& model, const std::string& name,
                              GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -214,7 +214,7 @@ void ConvertFloatTensorConst(const Model& model, const string& name,
                    LegacyScalarPolicy::kAvoidLegacyScalars);
 }
 
-void ConvertBoolTensorConst(const Model& model, const string& name,
+void ConvertBoolTensorConst(const Model& model, const std::string& name,
                             GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -238,7 +238,7 @@ void ConvertBoolTensorConst(const Model& model, const string& name,
   }
 }
 
-void ConvertIntTensorConst(const Model& model, const string& name,
+void ConvertIntTensorConst(const Model& model, const std::string& name,
                            GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -262,7 +262,8 @@ void ConvertIntTensorConst(const Model& model, const string& name,
   }
 }
 
-void CreateIntTensorConst(const string& name, const std::vector& data,
+void CreateIntTensorConst(const std::string& name,
+                          const std::vector& data,
                           const std::vector& shape,
                           GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
@@ -286,7 +287,7 @@ void CreateIntTensorConst(const string& name, const std::vector& data,
   CHECK_EQ(num_elements, data.size());
 }
 
-void ConvertComplex64TensorConst(const Model& model, const string& name,
+void ConvertComplex64TensorConst(const Model& model, const std::string& name,
                                  GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -311,7 +312,7 @@ void ConvertComplex64TensorConst(const Model& model, const string& name,
   }
 }
 
-void CreateMatrixShapeTensorConst(const string& name, int rows, int cols,
+void CreateMatrixShapeTensorConst(const std::string& name, int rows, int cols,
                                   GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -324,12 +325,12 @@ void CreateMatrixShapeTensorConst(const string& name, int rows, int cols,
   tensor->set_dtype(DT_INT32);
   const int32 data[2] = {cols, rows};
   tensor->set_tensor_content(
-      string(reinterpret_cast(data), sizeof(data)));
+      std::string(reinterpret_cast(data), sizeof(data)));
   auto* shape = tensor->mutable_tensor_shape();
   shape->add_dim()->set_size(2);
 }
 
-void CreateDummyConcatDimTensorConst(const string& name, int dim,
+void CreateDummyConcatDimTensorConst(const std::string& name, int dim,
                                      GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
     return;
@@ -343,7 +344,7 @@ void CreateDummyConcatDimTensorConst(const string& name, int dim,
   tensor->add_int_val(dim);
 }
 
-void CreateReshapeShapeTensorConst(const string& name,
+void CreateReshapeShapeTensorConst(const std::string& name,
                                    const std::vector& shape,
                                    GraphDef* tensorflow_graph) {
   if (HasAlreadyExportedConst(name, *tensorflow_graph)) {
@@ -367,7 +368,7 @@ void CreateReshapeShapeTensorConst(const string& name,
   }
 }
 
-string WalkUpToConstantArray(const Model& model, const string& name) {
+std::string WalkUpToConstantArray(const Model& model, const std::string& name) {
   const Array& original_array = model.GetArray(name);
   if (original_array.buffer) {
     return name;
@@ -375,7 +376,7 @@ string WalkUpToConstantArray(const Model& model, const string& name) {
   const auto* op = GetOpWithOutput(model, name);
   CHECK(op);
   CHECK(op->type == OperatorType::kFakeQuant);
-  const string& input_of_fakequant_name = op->inputs[0];
+  const std::string& input_of_fakequant_name = op->inputs[0];
   const Array& input_of_fakequant = model.GetArray(input_of_fakequant_name);
   CHECK(input_of_fakequant.buffer);
   return input_of_fakequant_name;
@@ -384,7 +385,7 @@ string WalkUpToConstantArray(const Model& model, const string& name) {
 void ConvertConvOperator(const Model& model, const ConvOperator& src_op,
                          GraphDef* tensorflow_graph) {
   const bool has_bias = src_op.inputs.size() >= 3;
-  string conv_output = src_op.outputs[0];
+  std::string conv_output = src_op.outputs[0];
   if (has_bias) {
     conv_output += "/conv";
   }
@@ -395,7 +396,7 @@ void ConvertConvOperator(const Model& model, const ConvOperator& src_op,
   *conv2d_op->add_input() = src_op.inputs[0];
   *conv2d_op->add_input() = src_op.inputs[1];
   (*conv2d_op->mutable_attr())["T"].set_type(DT_FLOAT);
-  const string& weights_array_name =
+  const std::string& weights_array_name =
       WalkUpToConstantArray(model, src_op.inputs[1]);
   const auto& weights_array = model.GetArray(weights_array_name);
   CHECK(weights_array.buffer->type == ArrayDataType::kFloat);
@@ -414,7 +415,7 @@ void ConvertConvOperator(const Model& model, const ConvOperator& src_op,
     dilations.mutable_list()->add_i(src_op.dilation_width_factor);
     dilations.mutable_list()->add_i(1);
   }
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -432,7 +433,7 @@ void ConvertConvOperator(const Model& model, const ConvOperator& src_op,
     biasadd_op->add_input(src_op.inputs[2]);
     (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT);
     CHECK(model.HasArray(src_op.inputs[2]));
-    const string& bias_array_name =
+    const std::string& bias_array_name =
         WalkUpToConstantArray(model, src_op.inputs[2]);
     const auto& bias_array = model.GetArray(bias_array_name);
     // TODO(b/62904716) Bias arrays should be 1-D, and used directly.
@@ -452,7 +453,7 @@ void ConvertDepthwiseConvOperator(const Model& model,
                                   const DepthwiseConvOperator& src_op,
                                   GraphDef* tensorflow_graph) {
   const bool has_bias = src_op.inputs.size() >= 3;
-  string conv_output = src_op.outputs[0];
+  std::string conv_output = src_op.outputs[0];
   if (has_bias) {
     conv_output += "/conv";
   }
@@ -469,7 +470,7 @@ void ConvertDepthwiseConvOperator(const Model& model,
   // That's only a matter of constructing a Dims object; the actual
   // array layout is the same.
   CHECK(model.HasArray(src_op.inputs[1]));
-  const string& src_weights_name =
+  const std::string& src_weights_name =
       WalkUpToConstantArray(model, src_op.inputs[1]);
   const auto& src_weights_array = model.GetArray(src_weights_name);
   const auto& src_weights_shape = src_weights_array.shape();
@@ -505,7 +506,7 @@ void ConvertDepthwiseConvOperator(const Model& model,
     dilations.mutable_list()->add_i(src_op.dilation_width_factor);
     dilations.mutable_list()->add_i(1);
   }
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -523,7 +524,8 @@ void ConvertDepthwiseConvOperator(const Model& model,
     biasadd_op->add_input(src_op.inputs[2]);
     (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT);
     CHECK(model.HasArray(src_op.inputs[2]));
-    const string& bias_name = WalkUpToConstantArray(model, src_op.inputs[2]);
+    const std::string& bias_name =
+        WalkUpToConstantArray(model, src_op.inputs[2]);
     const auto& bias_array = model.GetArray(bias_name);
     // TODO(b/62904716) Bias arrays should be 1-D, and used directly.
     Shape bias_shape_1d = bias_array.shape();
@@ -548,7 +550,7 @@ void ConvertTransposeConvOperator(const Model& model,
   *conv2d_op->add_input() = src_op.inputs[1];
   *conv2d_op->add_input() = src_op.inputs[2];
   (*conv2d_op->mutable_attr())["T"].set_type(DT_FLOAT);
-  const string& weights_array_name = WalkUpToConstantArray(
+  const std::string& weights_array_name = WalkUpToConstantArray(
       model, src_op.inputs[TransposeConvOperator::WEIGHTS]);
   const auto& weights_array = model.GetArray(weights_array_name);
   CHECK(weights_array.buffer->type == ArrayDataType::kFloat);
@@ -559,7 +561,7 @@ void ConvertTransposeConvOperator(const Model& model,
   strides.mutable_list()->add_i(src_op.stride_height);
   strides.mutable_list()->add_i(src_op.stride_width);
   strides.mutable_list()->add_i(1);
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -596,9 +598,9 @@ void ConvertFullyConnectedOperator(const Model& model,
                                    const FullyConnectedOperator& src_op,
                                    GraphDef* tensorflow_graph) {
   // Reshape input activations to have the shape expected by the MatMul.
-  const string reshape_output =
+  const std::string reshape_output =
       AvailableArrayName(model, src_op.outputs[0] + "/reshape");
-  const string reshape_shape =
+  const std::string reshape_shape =
       AvailableArrayName(model, reshape_output + "/shape");
   const auto& fc_weights_array = model.GetArray(src_op.inputs[1]);
   const auto& fc_weights_shape = fc_weights_array.shape();
@@ -614,7 +616,7 @@ void ConvertFullyConnectedOperator(const Model& model,
       GetTensorFlowDataType(model, src_op.inputs[0]));
 
   const bool has_bias = src_op.inputs.size() >= 3;
-  string matmul_output = src_op.outputs[0];
+  std::string matmul_output = src_op.outputs[0];
   if (has_bias) {
     matmul_output += "/matmul";
   }
@@ -622,9 +624,9 @@ void ConvertFullyConnectedOperator(const Model& model,
   // Transpose the RHS input from column-major to row-major to match TensorFlow
   // expectations. This is the inverse of the transpose we do during
   // ResolveTensorFlowMatMul.
-  const string transpose_output =
+  const std::string transpose_output =
       AvailableArrayName(model, matmul_output + "/transpose_weights");
-  const string transpose_perm =
+  const std::string transpose_perm =
       AvailableArrayName(model, transpose_output + "/perm");
   CreateIntTensorConst(transpose_perm, {1, 0}, {2}, tensorflow_graph);
   tensorflow::NodeDef* transpose_op = tensorflow_graph->add_node();
@@ -733,9 +735,9 @@ void ConvertReluOperator(const Model& model, const ReluOperator& src_op,
 
 void ConvertRelu1Operator(const Relu1Operator& src_op,
                           GraphDef* tensorflow_graph) {
-  const string max_bounds = src_op.outputs[0] + "/max_bounds";
-  const string min_bounds = src_op.outputs[0] + "/min_bounds";
-  const string max_output = src_op.outputs[0] + "/max_output";
+  const std::string max_bounds = src_op.outputs[0] + "/max_bounds";
+  const std::string min_bounds = src_op.outputs[0] + "/min_bounds";
+  const std::string max_output = src_op.outputs[0] + "/max_output";
 
   tensorflow::NodeDef* max_bounds_const_op = tensorflow_graph->add_node();
   max_bounds_const_op->set_op("Const");
@@ -808,15 +810,16 @@ void ConvertTanhOperator(const TanhOperator& src_op,
 
 void ConvertSoftmaxOperator(const Model& model, const SoftmaxOperator& src_op,
                             GraphDef* tensorflow_graph) {
-  string softmax_input;
+  std::string softmax_input;
   Operator* providing_op = GetOpWithOutput(model, src_op.inputs[0]);
   if (providing_op != nullptr && providing_op->type == OperatorType::kReshape) {
     softmax_input = src_op.inputs[0];
   } else {
     // Insert a reshape operator that reduces the dimensions down to the 2 that
     // are required for TensorFlow Logits.
-    const string reshape_output = src_op.outputs[0] + "/softmax_insert_reshape";
-    const string softmax_size = src_op.outputs[0] + "/softmax_insert_size";
+    const std::string reshape_output =
+        src_op.outputs[0] + "/softmax_insert_reshape";
+    const std::string softmax_size = src_op.outputs[0] + "/softmax_insert_size";
     softmax_input = reshape_output;
 
     tensorflow::NodeDef* reshape_op = tensorflow_graph->add_node();
@@ -848,16 +851,17 @@ void ConvertSoftmaxOperator(const Model& model, const SoftmaxOperator& src_op,
 void ConvertLogSoftmaxOperator(const Model& model,
                                const LogSoftmaxOperator& src_op,
                                GraphDef* tensorflow_graph) {
-  string softmax_input;
+  std::string softmax_input;
   Operator* providing_op = GetOpWithOutput(model, src_op.inputs[0]);
   if (providing_op != nullptr && providing_op->type == OperatorType::kReshape) {
     softmax_input = src_op.inputs[0];
   } else {
     // Insert a reshape operator that reduces the dimensions down to the 2 that
     // are required for TensorFlow Logits.
-    const string reshape_output =
+    const std::string reshape_output =
         src_op.outputs[0] + "/log_softmax_insert_reshape";
-    const string softmax_size = src_op.outputs[0] + "/log_softmax_insert_size";
+    const std::string softmax_size =
+        src_op.outputs[0] + "/log_softmax_insert_size";
     softmax_input = reshape_output;
 
     tensorflow::NodeDef* reshape_op = tensorflow_graph->add_node();
@@ -886,11 +890,12 @@ void ConvertLogSoftmaxOperator(const Model& model,
 
 void ConvertL2NormalizationOperator(const L2NormalizationOperator& src_op,
                                     GraphDef* tensorflow_graph) {
-  const string square_output = src_op.outputs[0] + "/square";
-  const string sum_reduction_indices = src_op.outputs[0] + "/reduction_indices";
-  const string sum_output = src_op.outputs[0] + "/sum";
-  const string rsqrt_output = src_op.outputs[0] + "/rsqrt";
-  const string rsqrt_tiled_output = src_op.outputs[0] + "/rsqrt_tiled";
+  const std::string square_output = src_op.outputs[0] + "/square";
+  const std::string sum_reduction_indices =
+      src_op.outputs[0] + "/reduction_indices";
+  const std::string sum_output = src_op.outputs[0] + "/sum";
+  const std::string rsqrt_output = src_op.outputs[0] + "/rsqrt";
+  const std::string rsqrt_tiled_output = src_op.outputs[0] + "/rsqrt_tiled";
 
   tensorflow::NodeDef* sum_reduction_indices_op = tensorflow_graph->add_node();
   sum_reduction_indices_op->set_op("Const");
@@ -975,7 +980,7 @@ void ConvertMaxPoolOperator(const MaxPoolOperator& src_op,
   strides.mutable_list()->add_i(src_op.stride_height);
   strides.mutable_list()->add_i(src_op.stride_width);
   strides.mutable_list()->add_i(1);
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -1003,7 +1008,7 @@ void ConvertAveragePoolOperator(const AveragePoolOperator& src_op,
   strides.mutable_list()->add_i(src_op.stride_height);
   strides.mutable_list()->add_i(src_op.stride_width);
   strides.mutable_list()->add_i(1);
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -1026,7 +1031,7 @@ void ConvertConcatenationOperator(const Model& model,
   tensorflow::NodeDef* dc_op = tensorflow_graph->add_node();
   dc_op->set_op("ConcatV2");
   dc_op->set_name(src_op.outputs[0]);
-  const string dummy_axis = src_op.outputs[0] + "/axis";
+  const std::string dummy_axis = src_op.outputs[0] + "/axis";
   CreateDummyConcatDimTensorConst(dummy_axis, src_op.axis, tensorflow_graph);
   for (const auto& input : src_op.inputs) {
     *dc_op->add_input() = input;
@@ -1060,8 +1065,8 @@ void ConvertTensorFlowReshapeOperator(const Model& model,
 
 void ConvertL2PoolOperator(const L2PoolOperator& src_op,
                            GraphDef* tensorflow_graph) {
-  const string square_output = src_op.outputs[0] + "/square";
-  const string avgpool_output = src_op.outputs[0] + "/avgpool";
+  const std::string square_output = src_op.outputs[0] + "/square";
+  const std::string avgpool_output = src_op.outputs[0] + "/avgpool";
 
   tensorflow::NodeDef* square_op = tensorflow_graph->add_node();
   square_op->set_op("Square");
@@ -1069,7 +1074,7 @@ void ConvertL2PoolOperator(const L2PoolOperator& src_op,
   *square_op->add_input() = src_op.inputs[0];
   (*square_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  string padding;
+  std::string padding;
   if (src_op.padding.type == PaddingType::kSame) {
     padding = "SAME";
   } else if (src_op.padding.type == PaddingType::kValid) {
@@ -1235,7 +1240,7 @@ void ConvertGatherOperator(const Model& model, const GatherOperator& src_op,
   } else {
     // Constant axis.
     CHECK_EQ(src_op.inputs.size(), 2);
-    const string gather_axis =
+    const std::string gather_axis =
         AvailableArrayName(model, gather_op->name() + "/axis");
     CreateIntTensorConst(gather_axis, {src_op.axis.value()}, {},
                          tensorflow_graph);
@@ -1454,8 +1459,8 @@ absl::string_view FindLongestCommonPrefix(absl::string_view a,
 
   const char* pa = a.data();
   const char* pb = b.data();
-  string::difference_type count = 0;
-  const string::difference_type limit = std::min(a.size(), b.size());
+  std::string::difference_type count = 0;
+  const std::string::difference_type limit = std::min(a.size(), b.size());
   while (count < limit && *pa == *pb) {
     ++pa;
     ++pb;
@@ -1469,12 +1474,12 @@ absl::string_view FindLongestCommonPrefix(absl::string_view a,
 void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
                              GraphDef* tensorflow_graph) {
   // Find the base name
-  const string base(
+  const std::string base(
       FindLongestCommonPrefix(src_op.outputs[LstmCellOperator::STATE_OUTPUT],
                               src_op.outputs[LstmCellOperator::ACTIV_OUTPUT]));
 
   // Concatenate inputs
-  const string concat_output = base + "basic_lstm_cell/concat";
+  const std::string concat_output = base + "basic_lstm_cell/concat";
   // Op names have been chosen to match the tf.slim LSTM naming
   // as closely as possible.
   const int axis =
@@ -1484,7 +1489,7 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
       1;
   // Note that DATA_INPUT may have extra size 1 dimensions, but TF concat
   // works the same since the tensor has the same underlying data layout.
-  const string axis_output = concat_output + "/axis";
+  const std::string axis_output = concat_output + "/axis";
   CreateDummyConcatDimTensorConst(axis_output, axis, tensorflow_graph);
   tensorflow::NodeDef* concat_op = tensorflow_graph->add_node();
   concat_op->set_op("ConcatV2");
@@ -1497,9 +1502,9 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   (*concat_op->mutable_attr())["N"].set_i(2);  // Number of inputs
 
   // Write weights
-  const string weights_output = base + "weights";
+  const std::string weights_output = base + "weights";
   CHECK(model.HasArray(src_op.inputs[LstmCellOperator::WEIGHTS_INPUT]));
-  const string weights_name = WalkUpToConstantArray(
+  const std::string weights_name = WalkUpToConstantArray(
       model, src_op.inputs[LstmCellOperator::WEIGHTS_INPUT]);
   const auto& weights_array = model.GetArray(weights_name);
   // Convert 4D FullyConnected weights into 2D matrix
@@ -1513,7 +1518,7 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
                           AxesOrder::kCR, AxesOrder::kRC, tensorflow_graph);
 
   // Fully connected matrix multiply
-  const string matmul_output = base + "MatMul";
+  const std::string matmul_output = base + "MatMul";
   tensorflow::NodeDef* matmul_op = tensorflow_graph->add_node();
   matmul_op->set_op("MatMul");
   matmul_op->set_name(matmul_output);
@@ -1524,9 +1529,9 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   (*matmul_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
   // Write biases
-  const string biases_output = base + "biases";
+  const std::string biases_output = base + "biases";
   CHECK(model.HasArray(src_op.inputs[LstmCellOperator::BIASES_INPUT]));
-  const string bias_name = WalkUpToConstantArray(
+  const std::string bias_name = WalkUpToConstantArray(
       model, src_op.inputs[LstmCellOperator::BIASES_INPUT]);
   const auto& bias_array = model.GetArray(bias_name);
   // TODO(b/62904716) Bias arrays should be 1-D, and used directly.
@@ -1542,7 +1547,7 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
                           LegacyScalarPolicy::kDoCreateLegacyScalars);
 
   // Add biases
-  string biasadd_output = base + "BiasAdd";
+  std::string biasadd_output = base + "BiasAdd";
   tensorflow::NodeDef* biasadd_op = tensorflow_graph->add_node();
   biasadd_op->set_op("BiasAdd");
   biasadd_op->set_name(biasadd_output);
@@ -1552,10 +1557,10 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   (*biasadd_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
   // Split
-  string split_dim_output = base + "split/split_dim";
+  std::string split_dim_output = base + "split/split_dim";
   // The dimension is the same as the concatenation dimension
   CreateDummyConcatDimTensorConst(split_dim_output, axis, tensorflow_graph);
-  string split_output = base + "split";
+  std::string split_output = base + "split";
   tensorflow::NodeDef* split_op = tensorflow_graph->add_node();
   split_op->set_op("Split");
   split_op->set_name(split_output);
@@ -1565,21 +1570,21 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   (*split_op->mutable_attr())["num_split"].set_i(4);  // Split into four outputs
 
   // Activation functions and memory computations
-  const string tanh_0_output = base + "Tanh";
+  const std::string tanh_0_output = base + "Tanh";
   tensorflow::NodeDef* tanh_0_op = tensorflow_graph->add_node();
   tanh_0_op->set_op("Tanh");
   tanh_0_op->set_name(tanh_0_output);
   *tanh_0_op->add_input() = split_output + ":1";
   (*tanh_0_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string sigmoid_1_output = base + "Sigmoid_1";
+  const std::string sigmoid_1_output = base + "Sigmoid_1";
   tensorflow::NodeDef* logistic_1_op = tensorflow_graph->add_node();
   logistic_1_op->set_op("Sigmoid");
   logistic_1_op->set_name(sigmoid_1_output);
   *logistic_1_op->add_input() = split_output;
   (*logistic_1_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string mul_1_output = base + "mul_1";
+  const std::string mul_1_output = base + "mul_1";
   tensorflow::NodeDef* mul_1_op = tensorflow_graph->add_node();
   mul_1_op->set_op("Mul");
   mul_1_op->set_name(mul_1_output);
@@ -1587,21 +1592,21 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   *mul_1_op->add_input() = tanh_0_output;
   (*mul_1_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string sigmoid_0_output = base + "Sigmoid";
+  const std::string sigmoid_0_output = base + "Sigmoid";
   tensorflow::NodeDef* logistic_2_op = tensorflow_graph->add_node();
   logistic_2_op->set_op("Sigmoid");
   logistic_2_op->set_name(sigmoid_0_output);
   *logistic_2_op->add_input() = split_output + ":2";
   (*logistic_2_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string sigmoid_2_output = base + "Sigmoid_2";
+  const std::string sigmoid_2_output = base + "Sigmoid_2";
   tensorflow::NodeDef* logistic_3_op = tensorflow_graph->add_node();
   logistic_3_op->set_op("Sigmoid");
   logistic_3_op->set_name(sigmoid_2_output);
   *logistic_3_op->add_input() = split_output + ":3";
   (*logistic_3_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string mul_0_output = base + "mul";
+  const std::string mul_0_output = base + "mul";
   tensorflow::NodeDef* mul_0_op = tensorflow_graph->add_node();
   mul_0_op->set_op("Mul");
   mul_0_op->set_name(mul_0_output);
@@ -1609,7 +1614,8 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   *mul_0_op->add_input() = sigmoid_0_output;
   (*mul_0_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string add_1_output = src_op.outputs[LstmCellOperator::STATE_OUTPUT];
+  const std::string add_1_output =
+      src_op.outputs[LstmCellOperator::STATE_OUTPUT];
   tensorflow::NodeDef* add_1_op = tensorflow_graph->add_node();
   add_1_op->set_op("Add");
   add_1_op->set_name(add_1_output);
@@ -1617,14 +1623,15 @@ void ConvertLstmCellOperator(const Model& model, const LstmCellOperator& src_op,
   *add_1_op->add_input() = mul_1_output;
   (*add_1_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string tanh_1_output = base + "Tanh_1";
+  const std::string tanh_1_output = base + "Tanh_1";
   tensorflow::NodeDef* tanh_1_op = tensorflow_graph->add_node();
   tanh_1_op->set_op("Tanh");
   tanh_1_op->set_name(tanh_1_output);
   *tanh_1_op->add_input() = add_1_output;
   (*tanh_1_op->mutable_attr())["T"].set_type(DT_FLOAT);
 
-  const string mul_2_output = src_op.outputs[LstmCellOperator::ACTIV_OUTPUT];
+  const std::string mul_2_output =
+      src_op.outputs[LstmCellOperator::ACTIV_OUTPUT];
   tensorflow::NodeDef* mul_2_op = tensorflow_graph->add_node();
   mul_2_op->set_op("Mul");
   mul_2_op->set_name(mul_2_output);
@@ -1730,7 +1737,8 @@ void ConvertPadV2Operator(const Model& model, const PadV2Operator& src_op,
   shape->add_dim()->set_size(2);
 }
 
-void CreateSliceInput(const string& input_name, const std::vector& values,
+void CreateSliceInput(const std::string& input_name,
+                      const std::vector& values,
                       GraphDef* tensorflow_graph) {
   tensorflow::NodeDef* params_op = tensorflow_graph->add_node();
   params_op->set_op("Const");
@@ -1797,7 +1805,8 @@ void ConvertSliceOperator(const Model& model, const SliceOperator& src_op,
 
 template 
 void ConvertReduceOperator(const Model& model, const T& src_op,
-                           GraphDef* tensorflow_graph, const string& op_name) {
+                           GraphDef* tensorflow_graph,
+                           const std::string& op_name) {
   tensorflow::NodeDef* new_op = tensorflow_graph->add_node();
   new_op->set_op(op_name);
   new_op->set_name(src_op.outputs[0]);
@@ -2412,7 +2421,7 @@ void ConvertOperator(const Model& model, const Operator& src_op,
   }
 }
 
-void AddPlaceholder(const string& name, ArrayDataType type,
+void AddPlaceholder(const std::string& name, ArrayDataType type,
                     GraphDef* tensorflow_graph) {
   tensorflow::NodeDef* placeholder = tensorflow_graph->add_node();
   placeholder->set_op("Placeholder");
@@ -2444,8 +2453,8 @@ void AddPlaceholder(const string& name, ArrayDataType type,
   placeholder->set_name(name);
 }
 
-void AddPlaceholderForRNNState(const Model& model, const string& name, int size,
-                               GraphDef* tensorflow_graph) {
+void AddPlaceholderForRNNState(const Model& model, const std::string& name,
+                               int size, GraphDef* tensorflow_graph) {
   tensorflow::NodeDef* placeholder = tensorflow_graph->add_node();
   placeholder->set_op("Placeholder");
   placeholder->set_name(name);
@@ -2484,7 +2493,7 @@ void ExportTensorFlowGraphDefImplementation(const Model& model,
   // after, as some operators need to export arrays that they reference
   // in a specific way, rather than in the generic way done below.
   for (const auto& array_pair : model.GetArrayMap()) {
-    const string& array_name = array_pair.first;
+    const std::string& array_name = array_pair.first;
     const auto& array = *array_pair.second;
     if (array.buffer) {
       switch (array.data_type) {
@@ -2510,12 +2519,12 @@ void ExportTensorFlowGraphDefImplementation(const Model& model,
 
 void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model) {
   for (const auto& array_kv : model->GetArrayMap()) {
-    const string& array_name = array_kv.first;
+    const std::string& array_name = array_kv.first;
     Array& array = *array_kv.second;
     if (!array.buffer || !array.minmax) {
       continue;
     }
-    const string& wrapped_array_name =
+    const std::string& wrapped_array_name =
         AvailableArrayName(*model, array_name + "/data");
     Array& wrapped_array = model->GetOrCreateArray(wrapped_array_name);
     wrapped_array.data_type = array.data_type;
@@ -2533,7 +2542,7 @@ void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model) {
 }
 
 void ExportTensorFlowGraphDef(const Model& model,
-                              string* output_file_contents) {
+                              std::string* output_file_contents) {
   CHECK(output_file_contents->empty());
   GraphDef tensorflow_graph;
   ExportTensorFlowGraphDefImplementation(model, &tensorflow_graph);
diff --git a/tensorflow/lite/toco/export_tensorflow.h b/tensorflow/lite/toco/export_tensorflow.h
index 09c966ded62..bc7ccd8d875 100644
--- a/tensorflow/lite/toco/export_tensorflow.h
+++ b/tensorflow/lite/toco/export_tensorflow.h
@@ -20,7 +20,8 @@ limitations under the License.
 
 namespace toco {
 
-void ExportTensorFlowGraphDef(const Model& model, string* output_file_contents);
+void ExportTensorFlowGraphDef(const Model& model,
+                              std::string* output_file_contents);
 
 void EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(Model* model);
 
diff --git a/tensorflow/lite/toco/format_port.h b/tensorflow/lite/toco/format_port.h
index 3c154e5ad48..47d39069c9f 100644
--- a/tensorflow/lite/toco/format_port.h
+++ b/tensorflow/lite/toco/format_port.h
@@ -38,13 +38,13 @@ inline const char* IdentityOrConvertStringToRaw(const std::string& foo) {
 
 // Delegate to TensorFlow Appendf function until absl has an equivalent.
 template 
-inline void AppendFHelper(string* destination, const char* fmt,
+inline void AppendFHelper(std::string* destination, const char* fmt,
                           Args&&... args) {
   tensorflow::strings::Appendf(destination, fmt, args...);
 }
 
 // Specialization for no argument format string (avoid security bug).
-inline void AppendFHelper(string* destination, const char* fmt) {
+inline void AppendFHelper(std::string* destination, const char* fmt) {
   tensorflow::strings::Appendf(destination, "%s", fmt);
 }
 
@@ -52,15 +52,15 @@ inline void AppendFHelper(string* destination, const char* fmt) {
 // pointed to by destination. fmt follows C printf semantics.
 // One departure is that %s can be driven by a std::string or string.
 template 
-inline void AppendF(string* destination, const char* fmt, Args&&... args) {
+inline void AppendF(std::string* destination, const char* fmt, Args&&... args) {
   AppendFHelper(destination, fmt, IdentityOrConvertStringToRaw(args)...);
 }
 
 // Return formatted string (with format fmt and args args). fmt follows C printf
 // semantics. One departure is that %s can be driven by a std::string or string.
 template 
-inline string StringF(const char* fmt, Args&&... args) {
-  string result;
+inline std::string StringF(const char* fmt, Args&&... args) {
+  std::string result;
   AppendFHelper(&result, fmt, IdentityOrConvertStringToRaw(args)...);
   return result;
 }
diff --git a/tensorflow/lite/toco/import_tensorflow.cc b/tensorflow/lite/toco/import_tensorflow.cc
index 3124133047e..2adfe838c3d 100644
--- a/tensorflow/lite/toco/import_tensorflow.cc
+++ b/tensorflow/lite/toco/import_tensorflow.cc
@@ -67,7 +67,7 @@ using tensorflow::TensorShapeProto;
 namespace toco {
 
 namespace {
-bool HasAttr(const NodeDef& node, const string& attr_name) {
+bool HasAttr(const NodeDef& node, const std::string& attr_name) {
   return node.attr().count(attr_name) > 0;
 }
 
@@ -78,14 +78,15 @@ bool HasWildcardDimension(const TensorShapeProto& shape) {
   return false;
 }
 
-const string& GetStringAttr(const NodeDef& node, const string& attr_name) {
+const std::string& GetStringAttr(const NodeDef& node,
+                                 const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kS);
   return attr.s();
 }
 
-int64 GetIntAttr(const NodeDef& node, const string& attr_name) {
+int64 GetIntAttr(const NodeDef& node, const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name)) << attr_name << " not found in:\n"
                                   << node.DebugString();
   const auto& attr = node.attr().at(attr_name);
@@ -93,14 +94,14 @@ int64 GetIntAttr(const NodeDef& node, const string& attr_name) {
   return attr.i();
 }
 
-float GetFloatAttr(const NodeDef& node, const string& attr_name) {
+float GetFloatAttr(const NodeDef& node, const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kF);
   return attr.f();
 }
 
-bool GetBoolAttr(const NodeDef& node, const string& attr_name) {
+bool GetBoolAttr(const NodeDef& node, const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kB);
@@ -108,7 +109,7 @@ bool GetBoolAttr(const NodeDef& node, const string& attr_name) {
 }
 
 tensorflow::DataType GetDataTypeAttr(const NodeDef& node,
-                                     const string& attr_name) {
+                                     const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kType);
@@ -116,14 +117,15 @@ tensorflow::DataType GetDataTypeAttr(const NodeDef& node,
 }
 
 const TensorShapeProto& GetShapeAttr(const NodeDef& node,
-                                     const string& attr_name) {
+                                     const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kShape);
   return attr.shape();
 }
 
-const TensorProto& GetTensorAttr(const NodeDef& node, const string& attr_name) {
+const TensorProto& GetTensorAttr(const NodeDef& node,
+                                 const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name)) << "No attr named '" << attr_name << "'";
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kTensor);
@@ -131,7 +133,7 @@ const TensorProto& GetTensorAttr(const NodeDef& node, const string& attr_name) {
 }
 
 const AttrValue::ListValue& GetListAttr(const NodeDef& node,
-                                        const string& attr_name) {
+                                        const std::string& attr_name) {
   CHECK(HasAttr(node, attr_name));
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kList);
@@ -139,10 +141,10 @@ const AttrValue::ListValue& GetListAttr(const NodeDef& node,
 }
 
 tensorflow::Status CheckOptionalAttr(const NodeDef& node,
-                                     const string& attr_name,
-                                     const string& expected_value) {
+                                     const std::string& attr_name,
+                                     const std::string& expected_value) {
   if (HasAttr(node, attr_name)) {
-    const string& value = GetStringAttr(node, attr_name);
+    const std::string& value = GetStringAttr(node, attr_name);
     if (value != expected_value) {
       return tensorflow::errors::InvalidArgument(
           "Unexpected value for attribute '" + attr_name + "'. Expected '" +
@@ -153,7 +155,7 @@ tensorflow::Status CheckOptionalAttr(const NodeDef& node,
 }
 
 tensorflow::Status CheckOptionalAttr(
-    const NodeDef& node, const string& attr_name,
+    const NodeDef& node, const std::string& attr_name,
     const tensorflow::DataType& expected_value) {
   if (HasAttr(node, attr_name)) {
     const tensorflow::DataType& value = GetDataTypeAttr(node, attr_name);
@@ -168,7 +170,7 @@ tensorflow::Status CheckOptionalAttr(
 
 template 
 tensorflow::Status ExpectValue(const T1& v1, const T2& v2,
-                               const string& description) {
+                               const std::string& description) {
   if (v1 == v2) return tensorflow::Status::OK();
   return tensorflow::errors::InvalidArgument(absl::StrCat(
       "Unexpected ", description, ": got ", v1, ", expected ", v2));
@@ -244,8 +246,8 @@ template <>
 struct TensorTraits {
   static int size(const TensorProto& p) { return p.float_val_size(); }
   static float get(const TensorProto& p, int i) { return p.float_val(i); }
-  static string accessor_name() { return "float_val"; }
-  static string type_name() { return "float"; }
+  static std::string accessor_name() { return "float_val"; }
+  static std::string type_name() { return "float"; }
   static void CopyFromContent(const TensorProto& p, std::vector* data) {
     toco::port::CopyToBuffer(p.tensor_content(),
                              reinterpret_cast(data->data()));
@@ -256,8 +258,8 @@ template <>
 struct TensorTraits {
   static int size(const TensorProto& p) { return p.int_val_size(); }
   static uint8_t get(const TensorProto& p, int i) { return p.int_val(i); }
-  static string accessor_name() { return "int_val"; }
-  static string type_name() { return "uint8"; }
+  static std::string accessor_name() { return "int_val"; }
+  static std::string type_name() { return "uint8"; }
   static void CopyFromContent(const TensorProto& p,
                               std::vector* data) {
     toco::port::CopyToBuffer(p.tensor_content(),
@@ -272,8 +274,8 @@ struct TensorTraits> {
     return std::complex(p.scomplex_val(2 * i),
                                p.scomplex_val(2 * i + 1));
   }
-  static string accessor_name() { return "scomplex_val"; }
-  static string type_name() { return "complex64"; }
+  static std::string accessor_name() { return "scomplex_val"; }
+  static std::string type_name() { return "complex64"; }
   static void CopyFromContent(const TensorProto& p,
                               std::vector>* data) {
     toco::port::CopyToBuffer(p.tensor_content(),
@@ -285,8 +287,8 @@ template <>
 struct TensorTraits {
   static int size(const TensorProto& p) { return p.int_val_size(); }
   static int32 get(const TensorProto& p, int i) { return p.int_val(i); }
-  static string accessor_name() { return "int_val"; }
-  static string type_name() { return "int32"; }
+  static std::string accessor_name() { return "int_val"; }
+  static std::string type_name() { return "int32"; }
   static void CopyFromContent(const TensorProto& p, std::vector* data) {
     toco::port::CopyToBuffer(p.tensor_content(),
                              reinterpret_cast(data->data()));
@@ -297,8 +299,8 @@ template <>
 struct TensorTraits {
   static int size(const TensorProto& p) { return p.int64_val_size(); }
   static int64 get(const TensorProto& p, int i) { return p.int64_val(i); }
-  static string accessor_name() { return "int64_val"; }
-  static string type_name() { return "int64"; }
+  static std::string accessor_name() { return "int64_val"; }
+  static std::string type_name() { return "int64"; }
   static void CopyFromContent(const TensorProto& p, std::vector* data) {
     toco::port::CopyToBuffer(p.tensor_content(),
                              reinterpret_cast(data->data()));
@@ -309,8 +311,8 @@ template <>
 struct TensorTraits {
   static int size(const TensorProto& p) { return p.bool_val_size(); }
   static bool get(const TensorProto& p, int i) { return p.bool_val(i); }
-  static string accessor_name() { return "bool_val"; }
-  static string type_name() { return "bool"; }
+  static std::string accessor_name() { return "bool_val"; }
+  static std::string type_name() { return "bool"; }
   static void CopyFromContent(const TensorProto& p, std::vector* data) {
     std::vector buf(p.tensor_content().size());
     toco::port::CopyToBuffer(p.tensor_content(), buf.data());
@@ -348,8 +350,8 @@ tensorflow::Status ImportTensorData(const TensorProto& input_tensor,
       (*output_data)[i] = last;
     }
   } else {
-    string accessor_name = TensorTraits::accessor_name();
-    string type_name = TensorTraits::type_name();
+    std::string accessor_name = TensorTraits::accessor_name();
+    std::string type_name = TensorTraits::type_name();
     return tensorflow::errors::InvalidArgument(
         absl::StrCat("Neither input_content (",
                      input_tensor.tensor_content().size() / sizeof(T), ") nor ",
@@ -527,10 +529,11 @@ tensorflow::Status CheckInputsCount(
 }
 
 template 
-string CreateConstArray(Model* model, string const& name,
-                        std::vector > const& data) {
+std::string CreateConstArray(
+    Model* model, std::string const& name,
+    std::vector> const& data) {
   // Utility function to create a const 1D array, useful for input parameters.
-  string array_name = toco::AvailableArrayName(*model, name);
+  std::string array_name = toco::AvailableArrayName(*model, name);
   auto& array = model->GetOrCreateArray(array_name);
   array.data_type = T;
   array.mutable_shape()->mutable_dims()->emplace_back(
@@ -576,7 +579,7 @@ void GetOutputNamesFromNodeDef(const NodeDef& node,
     ++next_output;
   };
   for (int i = 0; i < op_def.output_arg_size(); ++i) {
-    string multiples = op_def.output_arg(i).number_attr();
+    std::string multiples = op_def.output_arg(i).number_attr();
     if (!multiples.empty()) {
       CHECK(HasAttr(node, multiples)) << "No attr named " << multiples;
       int num_outputs = GetIntAttr(node, multiples);
@@ -584,7 +587,7 @@ void GetOutputNamesFromNodeDef(const NodeDef& node,
         add_output();
       }
     } else {
-      string list = op_def.output_arg(i).type_list_attr();
+      std::string list = op_def.output_arg(i).type_list_attr();
       if (!list.empty()) {
         CHECK(HasAttr(node, list)) << "No attr named " << list;
         const AttrValue::ListValue& list_value = GetListAttr(node, list);
@@ -624,7 +627,7 @@ void GetOutputTypesFromNodeDef(const NodeDef& node,
   };
 
   for (int i = 0; i < op_def.output_arg_size(); ++i) {
-    string multiples = op_def.output_arg(i).number_attr();
+    std::string multiples = op_def.output_arg(i).number_attr();
     if (!multiples.empty()) {
       CHECK(HasAttr(node, multiples)) << "No attr named " << multiples;
       int num_outputs = GetIntAttr(node, multiples);
@@ -633,7 +636,7 @@ void GetOutputTypesFromNodeDef(const NodeDef& node,
         add_type(type);
       }
     } else {
-      string list = op_def.output_arg(i).type_list_attr();
+      std::string list = op_def.output_arg(i).type_list_attr();
       if (!list.empty()) {
         CHECK(HasAttr(node, list)) << "No attr named " << list;
         const AttrValue::ListValue& list_value = GetListAttr(node, list);
@@ -1057,7 +1060,7 @@ tensorflow::Status ConvertIdentityNOperator(
   for (int i = 0; i < node.input_size(); ++i) {
     auto* op = new TensorFlowIdentityOperator;
     const auto& input_name = node.input(i);
-    string output_name = node.name();
+    std::string output_name = node.name();
     if (i > 0) {
       output_name = output_name + ":" + std::to_string(i);
     }
@@ -1756,13 +1759,13 @@ tensorflow::Status ConvertBatchNormWithGlobalNormalizationOperator(
   // to the input, before feeding it into TensorFlowRsqrtOperator.
   // CHECK_EQ(GetFloatAttr(node, "variance_epsilon"), 0.001f);
 
-  string multiplier = node.name() + "_mul";
+  std::string multiplier = node.name() + "_mul";
   if (GetBoolAttr(node, "scale_after_normalization")) {
     // Create graph:
     //   v -> RSQRT ->
     //                 MUL  -> multiplier
     //   gamma  ----->
-    string rsqrt = node.name() + "_rsqrt";
+    std::string rsqrt = node.name() + "_rsqrt";
 
     auto* rsqrt_op = new TensorFlowRsqrtOperator;
     rsqrt_op->inputs.push_back(node.input(2));
@@ -1803,17 +1806,19 @@ tensorflow::Status ConvertFusedBatchNormOperator(
   TF_QCHECK_OK(CheckInputsCount(node, tf_import_flags, 5));
 
   // Declare shortcuts for the inputs.
-  const string& gamma_input = node.input(1);
-  const string& beta_input = node.input(2);
-  const string& moving_mean_input = node.input(3);
-  const string& moving_variance_input = node.input(4);
+  const std::string& gamma_input = node.input(1);
+  const std::string& beta_input = node.input(2);
+  const std::string& moving_mean_input = node.input(3);
+  const std::string& moving_variance_input = node.input(4);
 
   // Create an array holding the epsilon value (typically, 0.001).
-  const string epsilon_array_name = CreateConstArray(
-      model, node.name() + "_epsilon_array", {GetFloatAttr(node, "epsilon")});
+  const std::string epsilon_array_name =
+      CreateConstArray(model,
+                                              node.name() + "_epsilon_array",
+                                              {GetFloatAttr(node, "epsilon")});
 
   // Add epsilon to the moving variance.
-  const string epsilon_add_op_name = node.name() + "_epsilon";
+  const std::string epsilon_add_op_name = node.name() + "_epsilon";
   auto* epsilon_add_op = new AddOperator;
   epsilon_add_op->inputs.push_back(moving_variance_input);
   epsilon_add_op->inputs.push_back(epsilon_array_name);
@@ -1821,14 +1826,14 @@ tensorflow::Status ConvertFusedBatchNormOperator(
   model->operators.emplace_back(epsilon_add_op);
 
   // Take the inverse square root of the (variance + epsilon).
-  const string rsqrt_op_name = node.name() + "_rsqrt";
+  const std::string rsqrt_op_name = node.name() + "_rsqrt";
   auto* rsqrt_op = new TensorFlowRsqrtOperator;
   rsqrt_op->inputs.push_back(epsilon_add_op_name);
   rsqrt_op->outputs.push_back(rsqrt_op_name);
   model->operators.emplace_back(rsqrt_op);
 
   // Multiply the result by gamma.
-  const string multiplier = node.name() + "_mul";
+  const std::string multiplier = node.name() + "_mul";
   auto* mul_op = new MulOperator;
   mul_op->inputs.push_back(rsqrt_op_name);
   mul_op->inputs.push_back(gamma_input);
@@ -1966,8 +1971,8 @@ tensorflow::Status ConvertTransposeConvOperator(
         << "].";
   }
 
-  const string& weights_name = node.input(TransposeConvOperator::WEIGHTS);
-  const string& transposed_weights_name = weights_name + "_transposed";
+  const std::string& weights_name = node.input(TransposeConvOperator::WEIGHTS);
+  const std::string& transposed_weights_name = weights_name + "_transposed";
   // Check if a TransposeOperator was already created for these weights
   // (can happen when multiple layers share the same weights).
   const Operator* existing_transpose =
@@ -1980,7 +1985,7 @@ tensorflow::Status ConvertTransposeConvOperator(
     // because they consider this a backward conv, inverting the sense of
     // input/output.)
     TransposeOperator* transpose = new TransposeOperator;
-    string perm_array = CreateConstArray(
+    std::string perm_array = CreateConstArray(
         model, node.name() + "_transpose_perm", {2, 0, 1, 3});
     transpose->inputs = {weights_name, perm_array};
     transpose->outputs = {transposed_weights_name};
@@ -2137,10 +2142,10 @@ tensorflow::Status ConvertReverseSequenceOperator(
 void StripCaretFromArrayNames(Model* model) {
   for (auto& op : model->operators) {
     for (auto& input : op->inputs) {
-      input = string(absl::StripPrefix(input, "^"));
+      input = std::string(absl::StripPrefix(input, "^"));
     }
     for (auto& output : op->outputs) {
-      output = string(absl::StripPrefix(output, "^"));
+      output = std::string(absl::StripPrefix(output, "^"));
     }
   }
   for (auto& array : model->GetArrayMap()) {
@@ -2152,7 +2157,7 @@ void StripCaretFromArrayNames(Model* model) {
 
 void StripZeroOutputIndexFromInputs(NodeDef* node) {
   for (auto& input : *node->mutable_input()) {
-    input = string(absl::StripSuffix(input, ":0"));
+    input = std::string(absl::StripSuffix(input, ":0"));
   }
 }
 
@@ -2170,15 +2175,15 @@ void StripZeroOutputIndexFromInputs(NodeDef* node) {
 // all nodes, we can use that information.
 void AddExtraOutputs(Model* model) {
   // Construct the list of all arrays consumed by anything in the graph.
-  std::vector consumed_arrays;
+  std::vector consumed_arrays;
   // Add arrays consumed by an op.
   for (const auto& consumer_op : model->operators) {
-    for (const string& input : consumer_op->inputs) {
+    for (const std::string& input : consumer_op->inputs) {
       consumed_arrays.push_back(input);
     }
   }
   // Add global outputs of the model.
-  for (const string& output_array : model->flags.output_arrays()) {
+  for (const std::string& output_array : model->flags.output_arrays()) {
     consumed_arrays.push_back(output_array);
   }
   // Add arrays consumed by a RNN back-edge.
@@ -2187,7 +2192,7 @@ void AddExtraOutputs(Model* model) {
   }
   // Now add operator outputs so that all arrays that are consumed,
   // are produced.
-  for (const string& consumed_array : consumed_arrays) {
+  for (const std::string& consumed_array : consumed_arrays) {
     // Test if consumed_array is already the output of some op.
     // This has occurred in a model where separate nodes had names of the form
     // foo:$i with the same base name foo.
@@ -2195,7 +2200,7 @@ void AddExtraOutputs(Model* model) {
       continue;
     }
     // Split the consumed array name into the form name:output_index.
-    const std::vector& split = absl::StrSplit(consumed_array, ':');
+    const std::vector& split = absl::StrSplit(consumed_array, ':');
     // If not of the form name:output_index, then this is not an additional
     // output of a node with multiple outputs, so nothing to do here.
     if (split.size() != 2) {
@@ -2288,7 +2293,7 @@ tensorflow::Status ConvertTopKV2Operator(
   op->inputs.push_back(node.input(0));
   // K can be encoded as attr (TopK) convert it to a const.
   if (HasAttr(node, "k")) {
-    string k_array = CreateConstArray(
+    std::string k_array = CreateConstArray(
         model, node.name() + "k", {static_cast(GetIntAttr(node, "k"))});
     op->inputs.push_back(k_array);
   } else {
@@ -2346,7 +2351,7 @@ tensorflow::Status ConvertSparseToDenseOperator(
   TF_QCHECK_OK(CheckInputsCount(node, tf_import_flags, 4));
 
   auto* op = new SparseToDenseOperator;
-  for (const string& input : node.input()) {
+  for (const std::string& input : node.input()) {
     op->inputs.push_back(input);
   }
   op->outputs.push_back(node.name());
@@ -2371,7 +2376,7 @@ tensorflow::Status ConvertOneHotOperator(
 
   auto op = absl::make_unique();
   op->axis = HasAttr(node, "axis") ? GetIntAttr(node, "axis") : -1;
-  for (const string& input : node.input()) {
+  for (const std::string& input : node.input()) {
     op->inputs.push_back(input);
   }
   op->outputs.push_back(node.name());
@@ -2386,7 +2391,7 @@ tensorflow::Status ConvertCTCBeamSearchDecoderOperator(
   TF_QCHECK_OK(CheckInputsCount(node, tf_import_flags, 2));
 
   auto* op = new CTCBeamSearchDecoderOperator;
-  for (const string& input : node.input()) {
+  for (const std::string& input : node.input()) {
     op->inputs.push_back(input);
   }
 
@@ -2434,7 +2439,7 @@ tensorflow::Status ConvertUnidirectionalSequenceLstm(
         count++;
       } else {
         // Optional input.
-        string optional_name = node.name() + "_" + std::to_string(idx);
+        std::string optional_name = node.name() + "_" + std::to_string(idx);
         model->CreateOptionalArray(optional_name);
         op->inputs[idx] = optional_name;
       }
@@ -2442,7 +2447,7 @@ tensorflow::Status ConvertUnidirectionalSequenceLstm(
   } else {  // Legacy version.
     std::vector done(kInputsSize);
     int idx = 0;
-    for (const string& input : node.input()) {
+    for (const std::string& input : node.input()) {
       int real_index = indices.i(idx);
       op->inputs[real_index] = (input);
       done[real_index] = true;
@@ -2451,7 +2456,7 @@ tensorflow::Status ConvertUnidirectionalSequenceLstm(
 
     for (int idx = 0; idx < done.size(); idx++) {
       if (!done[idx]) {
-        string optional_name = node.name() + "_" + std::to_string(idx);
+        std::string optional_name = node.name() + "_" + std::to_string(idx);
         model->CreateOptionalArray(optional_name);
         op->inputs[idx] = optional_name;
       }
@@ -2491,7 +2496,7 @@ tensorflow::Status ConvertUnidirectionalSequenceRnn(
   }
 
   auto* op = new UnidirectionalSequenceRnnOperator();
-  for (const string& input : node.input()) {
+  for (const std::string& input : node.input()) {
     op->inputs.push_back(input);
   }
   // Only use the last one as input.
@@ -2703,7 +2708,8 @@ std::unique_ptr ImportTensorFlowGraphDef(
         << "Unsupported explicit zero output index: "
         << specified_input_array.name();
   }
-  for (const string& specified_output_array : model_flags.output_arrays()) {
+  for (const std::string& specified_output_array :
+       model_flags.output_arrays()) {
     CHECK(!absl::EndsWith(specified_output_array, ":0"))
         << "Unsupported explicit zero output index: " << specified_output_array;
   }
@@ -2746,7 +2752,7 @@ std::unique_ptr ImportTensorFlowGraphDef(
 
 std::unique_ptr ImportTensorFlowGraphDef(
     const ModelFlags& model_flags, const TensorFlowImportFlags& tf_import_flags,
-    const string& input_file_contents) {
+    const std::string& input_file_contents) {
   std::unique_ptr tf_graph(new GraphDef);
   CHECK(ParseFromStringEitherTextOrBinary(input_file_contents, tf_graph.get()));
 
diff --git a/tensorflow/lite/toco/import_tensorflow.h b/tensorflow/lite/toco/import_tensorflow.h
index 4ada25e2fbe..a95cfee2e75 100644
--- a/tensorflow/lite/toco/import_tensorflow.h
+++ b/tensorflow/lite/toco/import_tensorflow.h
@@ -43,7 +43,7 @@ std::unique_ptr ImportTensorFlowGraphDef(
 // flags.
 std::unique_ptr ImportTensorFlowGraphDef(
     const ModelFlags& model_flags, const TensorFlowImportFlags& tf_import_flags,
-    const string& input_file_contents);
+    const std::string& input_file_contents);
 
 // Gets a list of supported ops by their names.
 std::vector GetPotentiallySupportedOps();
diff --git a/tensorflow/lite/toco/import_tensorflow_test.cc b/tensorflow/lite/toco/import_tensorflow_test.cc
index eb6ed3fdd74..98ce18bf38e 100644
--- a/tensorflow/lite/toco/import_tensorflow_test.cc
+++ b/tensorflow/lite/toco/import_tensorflow_test.cc
@@ -163,7 +163,7 @@ void BuildConstNode(std::initializer_list shape,
 TEST(FlexImportTest, ConditionalConst) {
   Model model;
   auto build_and_import_node =
-      [&model](const string& name, std::initializer_list shape,
+      [&model](const std::string& name, std::initializer_list shape,
                tensorflow::DataType dtype, int64_t num_elements) {
         NodeDef node;
         BuildConstNode(shape, dtype, num_elements, &node);
@@ -486,8 +486,8 @@ class TensorContentTest : public ::testing::Test {
         break;
     }
     t.set_tensor_content(
-        string(reinterpret_cast(allocated_content.get()),
-               num_elements * sizeof(T)));
+        std::string(reinterpret_cast(allocated_content.get()),
+                    num_elements * sizeof(T)));
 
     AttrValue value_attr;
     SetAttrValue(t, &value_attr);
diff --git a/tensorflow/lite/toco/model.h b/tensorflow/lite/toco/model.h
index 89ea9d997f9..58397f5a3eb 100644
--- a/tensorflow/lite/toco/model.h
+++ b/tensorflow/lite/toco/model.h
@@ -287,7 +287,7 @@ struct DataTypeImpl {
 };
 template <>
 struct DataTypeImpl {
-  typedef string Type;
+  typedef std::string Type;
 };
 template <>
 struct DataTypeImpl {
@@ -398,10 +398,10 @@ struct Operator {
   // names to addresses is given by the Model, which owns both Operator's and
   // Array's. Thus, an Operator on its own doesn't contain much information,
   // it is meant to be used in conjunction with the Model that owns it.
-  std::vector inputs;
+  std::vector inputs;
 
   // Output activation arrays. Same comments as for inputs apply here too.
-  std::vector outputs;
+  std::vector outputs;
 
   // If true, the operator has more outputs than are listed in the 'outputs'
   // member. These need to be resolved by some graph transformation.
@@ -415,7 +415,7 @@ struct Operator {
   // It's guaranteed to be filled for `TensorFlowUnsupportedOperator`.
   // It's not guaranteed to be filled for other ops. Ops created by graph
   // transformations won't have TensorFlow NodeDef.
-  string tensorflow_node_def;
+  std::string tensorflow_node_def;
 
  protected:
   // Constructor used by subclasses for specific OperatorType's.
@@ -1693,7 +1693,7 @@ struct TensorFlowUnsupportedOperator : Operator {
   TensorFlowUnsupportedOperator() : Operator(OperatorType::kUnsupported) {}
 
   // The original TF operation type. Used for diagnostic purposes.
-  string tensorflow_op;
+  std::string tensorflow_op;
   // A boolean indicating if the unsupported op should be treated as quantized.
   bool quantized = false;
   // A boolean indicating if the unsupported op output should allow float values
@@ -2393,14 +2393,16 @@ struct Array {
 // Owns everything.
 class Model {
  public:
-  using ArrayMap = std::unordered_map>;
+  using ArrayMap = std::unordered_map>;
 
-  bool HasArray(const string& name) const { return arrays.count(name) > 0; }
-  Array& GetArray(const string& name) const {
+  bool HasArray(const std::string& name) const {
+    return arrays.count(name) > 0;
+  }
+  Array& GetArray(const std::string& name) const {
     DCHECK(HasArray(name)) << "Array not found: " << name;
     return *arrays.at(name);
   }
-  Array& GetOrCreateArray(const string& name) {
+  Array& GetOrCreateArray(const std::string& name) {
     // Make sure name is not used by an optional array
     DCHECK(!optional_arrays.count(name));
     if (!HasArray(name)) {
@@ -2410,17 +2412,17 @@ class Model {
     Array& result = GetArray(name);
     return result;
   }
-  void CreateOptionalArray(const string& name) {
+  void CreateOptionalArray(const std::string& name) {
     DCHECK(!arrays.count(name) && !optional_arrays.count(name));
     optional_arrays.insert(name);
   }
-  bool IsOptionalArray(const string& name) const {
+  bool IsOptionalArray(const std::string& name) const {
     return optional_arrays.count(name);
   }
 
   // Note that this invalidates all array iterators.
-  void EraseArray(const string& name) { arrays.erase(name); }
-  void EraseArrays(std::function discardable) {
+  void EraseArray(const std::string& name) { arrays.erase(name); }
+  void EraseArrays(std::function discardable) {
     for (auto it = arrays.begin(); it != arrays.end();) {
       if (discardable(it->first)) {
         it = arrays.erase(it);
@@ -2434,17 +2436,17 @@ class Model {
 
   int64 ArithmeticOpsCount() const { return ops_count; }
 
-  void AddInvalidInputArray(string invalid_input_array) {
+  void AddInvalidInputArray(std::string invalid_input_array) {
     invalid_input_arrays_.insert(invalid_input_array);
   }
 
-  const std::unordered_set& GetInvalidInputArrays() const {
+  const std::unordered_set& GetInvalidInputArrays() const {
     return invalid_input_arrays_;
   }
 
   // Optional arrays are used for optional tensors,
   // these tensors do not have data, but with reserved names as op inputs.
-  std::set optional_arrays;
+  std::set optional_arrays;
 
   // The list of operators. Notice how it's a list of unique_ptr's, implying
   // that the Model is what owns Operator's and keeps them alive.
@@ -2467,10 +2469,10 @@ class Model {
   // that the Model is what owns Array's and keeps them alive.
   // The Operator's refer to these Array's by their name strings, not by their
   // addresses. See Operator::inputs, Operator::outputs.
-  std::unordered_map> arrays;
+  std::unordered_map> arrays;
 
   // Invalid input arrays.
-  std::unordered_set invalid_input_arrays_;
+  std::unordered_set invalid_input_arrays_;
 };
 
 // OperatorSignature contains the information required to making versioning
diff --git a/tensorflow/lite/toco/model_cmdline_flags.cc b/tensorflow/lite/toco/model_cmdline_flags.cc
index 86a1cedd612..d3c48aee2fe 100644
--- a/tensorflow/lite/toco/model_cmdline_flags.cc
+++ b/tensorflow/lite/toco/model_cmdline_flags.cc
@@ -36,7 +36,7 @@ limitations under the License.
 namespace toco {
 
 bool ParseModelFlagsFromCommandLineFlags(
-    int* argc, char* argv[], string* msg,
+    int* argc, char* argv[], std::string* msg,
     ParsedModelFlags* parsed_model_flags_ptr) {
   ParsedModelFlags& parsed_flags = *parsed_model_flags_ptr;
   using tensorflow::Flag;
@@ -188,7 +188,7 @@ void ReadModelFlagsFromCommandLineFlags(
   // Load proto containing the initial model flags.
   // Additional flags specified on the command line will overwrite the values.
   if (parsed_model_flags.model_flags_file.specified()) {
-    string model_flags_file_contents;
+    std::string model_flags_file_contents;
     QCHECK(port::file::GetContents(parsed_model_flags.model_flags_file.value(),
                                    &model_flags_file_contents,
                                    port::file::Defaults())
@@ -217,9 +217,9 @@ void ReadModelFlagsFromCommandLineFlags(
   }
 
   if (parsed_model_flags.output_arrays.specified()) {
-    std::vector output_arrays =
+    std::vector output_arrays =
         absl::StrSplit(parsed_model_flags.output_arrays.value(), ',');
-    for (const string& output_array : output_arrays) {
+    for (const std::string& output_array : output_arrays) {
       model_flags->add_output_arrays(output_array);
     }
   }
@@ -251,7 +251,7 @@ void ReadModelFlagsFromCommandLineFlags(
     QCHECK(uses_multi_input_flags);
     for (const auto& input_array :
          absl::StrSplit(parsed_model_flags.input_arrays.value(), ',')) {
-      model_flags->add_input_arrays()->set_name(string(input_array));
+      model_flags->add_input_arrays()->set_name(std::string(input_array));
     }
   }
   if (parsed_model_flags.mean_value.specified()) {
@@ -261,7 +261,7 @@ void ReadModelFlagsFromCommandLineFlags(
   }
   if (parsed_model_flags.mean_values.specified()) {
     QCHECK(uses_multi_input_flags);
-    std::vector mean_values =
+    std::vector mean_values =
         absl::StrSplit(parsed_model_flags.mean_values.value(), ',');
     QCHECK(mean_values.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < mean_values.size(); ++i) {
@@ -278,7 +278,7 @@ void ReadModelFlagsFromCommandLineFlags(
   }
   if (parsed_model_flags.std_values.specified()) {
     QCHECK(uses_multi_input_flags);
-    std::vector std_values =
+    std::vector std_values =
         absl::StrSplit(parsed_model_flags.std_values.value(), ',');
     QCHECK(std_values.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < std_values.size(); ++i) {
@@ -296,7 +296,7 @@ void ReadModelFlagsFromCommandLineFlags(
   }
   if (parsed_model_flags.input_data_types.specified()) {
     QCHECK(uses_multi_input_flags);
-    std::vector input_data_types =
+    std::vector input_data_types =
         absl::StrSplit(parsed_model_flags.input_data_types.value(), ',');
     QCHECK(input_data_types.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < input_data_types.size(); ++i) {
@@ -319,7 +319,7 @@ void ReadModelFlagsFromCommandLineFlags(
   }
   if (parsed_model_flags.input_shapes.specified()) {
     QCHECK(uses_multi_input_flags);
-    std::vector input_shapes =
+    std::vector input_shapes =
         absl::StrSplit(parsed_model_flags.input_shapes.value(), ':');
     QCHECK(input_shapes.size() == model_flags->input_arrays_size());
     for (size_t i = 0; i < input_shapes.size(); ++i) {
@@ -352,8 +352,8 @@ void ReadModelFlagsFromCommandLineFlags(
   for (const auto& element : parsed_model_flags.rnn_states.value().elements) {
     auto* rnn_state_proto = model_flags->add_rnn_states();
     for (const auto& kv_pair : element) {
-      const string& key = kv_pair.first;
-      const string& value = kv_pair.second;
+      const std::string& key = kv_pair.first;
+      const std::string& value = kv_pair.second;
       if (key == "state_array") {
         rnn_state_proto->set_state_array(value);
       } else if (key == "back_edge_source_array") {
@@ -377,8 +377,8 @@ void ReadModelFlagsFromCommandLineFlags(
   for (const auto& element : parsed_model_flags.model_checks.value().elements) {
     auto* model_check_proto = model_flags->add_model_checks();
     for (const auto& kv_pair : element) {
-      const string& key = kv_pair.first;
-      const string& value = kv_pair.second;
+      const std::string& key = kv_pair.first;
+      const std::string& value = kv_pair.second;
       if (key == "count_type") {
         model_check_proto->set_count_type(value);
       } else if (key == "count_min") {
@@ -411,7 +411,7 @@ void ReadModelFlagsFromCommandLineFlags(
   }
 
   if (parsed_model_flags.arrays_extra_info_file.specified()) {
-    string arrays_extra_info_file_contents;
+    std::string arrays_extra_info_file_contents;
     CHECK(port::file::GetContents(
               parsed_model_flags.arrays_extra_info_file.value(),
               &arrays_extra_info_file_contents, port::file::Defaults())
@@ -443,7 +443,7 @@ void ParseModelFlagsOrDie(int* argc, char* argv[]) {
   // TODO(aselle): in the future allow Google version to use
   // flags, and only use this mechanism for open source
   auto* flags = UncheckedGlobalParsedModelFlags(false);
-  string msg;
+  std::string msg;
   bool model_success =
       toco::ParseModelFlagsFromCommandLineFlags(argc, argv, &msg, flags);
   if (!model_success || !msg.empty()) {
diff --git a/tensorflow/lite/toco/model_cmdline_flags.h b/tensorflow/lite/toco/model_cmdline_flags.h
index 1642e053199..23e79e62047 100644
--- a/tensorflow/lite/toco/model_cmdline_flags.h
+++ b/tensorflow/lite/toco/model_cmdline_flags.h
@@ -28,7 +28,7 @@ namespace toco {
 // is successful. msg has the usage string if there was an error or
 // "--help" was specified
 bool ParseModelFlagsFromCommandLineFlags(
-    int* argc, char* argv[], string* msg,
+    int* argc, char* argv[], std::string* msg,
     ParsedModelFlags* parsed_model_flags_ptr);
 // Populate the ModelFlags proto with model data.
 void ReadModelFlagsFromCommandLineFlags(
diff --git a/tensorflow/lite/toco/model_cmdline_flags_test.cc b/tensorflow/lite/toco/model_cmdline_flags_test.cc
index bff8e4843a0..b87e200095c 100644
--- a/tensorflow/lite/toco/model_cmdline_flags_test.cc
+++ b/tensorflow/lite/toco/model_cmdline_flags_test.cc
@@ -35,8 +35,8 @@ TEST(ModelCmdlineFlagsTest, ParseArgsStringMapList) {
       "back_edge_source_array:rnn/basic_lstm_cell/Mul_2,size:4}",
       nullptr};
 
-  string expected_input_arrays = "input_1";
-  std::vector> expected_rnn_states;
+  std::string expected_input_arrays = "input_1";
+  std::vector> expected_rnn_states;
   expected_rnn_states.push_back(
       {{"state_array", "rnn/BasicLSTMCellZeroState/zeros"},
        {"back_edge_source_array", "rnn/basic_lstm_cell/Add_1"},
@@ -46,7 +46,7 @@ TEST(ModelCmdlineFlagsTest, ParseArgsStringMapList) {
        {"back_edge_source_array", "rnn/basic_lstm_cell/Mul_2"},
        {"size", "4"}});
 
-  string message;
+  std::string message;
   ParsedModelFlags result_flags;
 
   EXPECT_TRUE(ParseModelFlagsFromCommandLineFlags(
diff --git a/tensorflow/lite/toco/tensorflow_util.cc b/tensorflow/lite/toco/tensorflow_util.cc
index db9388b040c..bf5d8016857 100644
--- a/tensorflow/lite/toco/tensorflow_util.cc
+++ b/tensorflow/lite/toco/tensorflow_util.cc
@@ -37,16 +37,16 @@ namespace toco {
 using tensorflow::AttrValue;
 using tensorflow::GraphDef;
 
-void LogDumpGraphDef(int log_level, const string& message,
+void LogDumpGraphDef(int log_level, const std::string& message,
                      const GraphDef& tf_graph) {
   if (!VLOG_IS_ON(log_level)) {
     return;
   }
-  std::set ops;
+  std::set ops;
   for (const auto& node : tf_graph.node()) {
     ops.insert(node.op());
   }
-  string dump;
+  std::string dump;
   toco::port::AppendF(&dump, R"MSG(
 BEGIN DUMP OF TENSORFLOW GRAPHDEF (%s)
 There are %d nodes.
diff --git a/tensorflow/lite/toco/tensorflow_util.h b/tensorflow/lite/toco/tensorflow_util.h
index 010fbe88b21..6abad52b9cb 100644
--- a/tensorflow/lite/toco/tensorflow_util.h
+++ b/tensorflow/lite/toco/tensorflow_util.h
@@ -24,7 +24,7 @@ limitations under the License.
 
 namespace toco {
 
-void LogDumpGraphDef(int log_level, const string& message,
+void LogDumpGraphDef(int log_level, const std::string& message,
                      const tensorflow::GraphDef& tf_graph);
 
 }  // namespace toco
diff --git a/tensorflow/lite/toco/toco.cc b/tensorflow/lite/toco/toco.cc
index aa7e43350ca..18800c7b726 100644
--- a/tensorflow/lite/toco/toco.cc
+++ b/tensorflow/lite/toco/toco.cc
@@ -21,7 +21,7 @@ limitations under the License.
 #include "tensorflow/lite/toco/toco_convert.h"
 
 int main(int argc, char** argv) {
-  toco::string msg;
+  std::string msg;
   toco::ParsedTocoFlags parsed_toco_flags;
   toco::ParsedModelFlags parsed_model_flags;
 
diff --git a/tensorflow/lite/toco/toco_cmdline_flags.cc b/tensorflow/lite/toco/toco_cmdline_flags.cc
index c133db8f2a4..223cfd40775 100644
--- a/tensorflow/lite/toco/toco_cmdline_flags.cc
+++ b/tensorflow/lite/toco/toco_cmdline_flags.cc
@@ -29,7 +29,7 @@ limitations under the License.
 namespace toco {
 
 bool ParseTocoFlagsFromCommandLineFlags(
-    int* argc, char* argv[], string* msg,
+    int* argc, char* argv[], std::string* msg,
     ParsedTocoFlags* parsed_toco_flags_ptr) {
   using tensorflow::Flag;
   ParsedTocoFlags& parsed_flags = *parsed_toco_flags_ptr;
@@ -212,7 +212,7 @@ enum class FlagRequirement {
 
 // Enforces the FlagRequirements are met for a given flag.
 template 
-void EnforceFlagRequirement(const T& flag, const string& flag_name,
+void EnforceFlagRequirement(const T& flag, const std::string& flag_name,
                             FlagRequirement requirement) {
   if (requirement == FlagRequirement::kMustBeSpecified) {
     QCHECK(flag.specified()) << "Missing required flag " << flag_name;
@@ -317,7 +317,7 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags,
            "type of input arrays, use --input_data_type. If you are trying to "
            "control the quantization/dequantization of real-numbers input "
            "arrays in the output file, use --inference_input_type.";
-    std::vector input_types =
+    std::vector input_types =
         absl::StrSplit(parsed_toco_flags.input_types.value(), ',');
     QCHECK(!input_types.empty());
     for (int i = 1; i < input_types.size(); i++) {
diff --git a/tensorflow/lite/toco/toco_cmdline_flags.h b/tensorflow/lite/toco/toco_cmdline_flags.h
index cf57055abc2..278c49d25e3 100644
--- a/tensorflow/lite/toco/toco_cmdline_flags.h
+++ b/tensorflow/lite/toco/toco_cmdline_flags.h
@@ -25,7 +25,8 @@ namespace toco {
 // Parse and remove arguments handled from toco. Returns true if parsing
 // is successful. msg has the usage string if there was an error or
 // "--help" was specified
-bool ParseTocoFlagsFromCommandLineFlags(int* argc, char* argv[], string* msg,
+bool ParseTocoFlagsFromCommandLineFlags(int* argc, char* argv[],
+                                        std::string* msg,
                                         ParsedTocoFlags* parsed_toco_flags_ptr);
 // Populate the TocoFlags proto with parsed_toco_flags data.
 void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags,
diff --git a/tensorflow/lite/toco/toco_cmdline_flags_test.cc b/tensorflow/lite/toco/toco_cmdline_flags_test.cc
index a1066e063bc..c1d0e2f7e9b 100644
--- a/tensorflow/lite/toco/toco_cmdline_flags_test.cc
+++ b/tensorflow/lite/toco/toco_cmdline_flags_test.cc
@@ -29,7 +29,7 @@ TEST(TocoCmdlineFlagsTest, DefaultValue) {
   // TF flag parsing lib is relaying on this invariant.
   const char* args[] = {"toco", nullptr};
 
-  string message;
+  std::string message;
   ParsedTocoFlags result_flags;
 
   EXPECT_TRUE(ParseTocoFlagsFromCommandLineFlags(
@@ -41,7 +41,7 @@ TEST(TocoCmdlineFlagsTest, ParseFlags) {
   int argc = 2;
   const char* args[] = {"toco", "--allow_dynamic_tensors=false", nullptr};
 
-  string message;
+  std::string message;
   ParsedTocoFlags result_flags;
 
   EXPECT_TRUE(ParseTocoFlagsFromCommandLineFlags(
diff --git a/tensorflow/lite/toco/toco_convert.cc b/tensorflow/lite/toco/toco_convert.cc
index 62dacef0b60..0e04b9ce143 100644
--- a/tensorflow/lite/toco/toco_convert.cc
+++ b/tensorflow/lite/toco/toco_convert.cc
@@ -32,7 +32,7 @@ namespace toco {
 namespace {
 
 // Checks the permissions of the output file to ensure it is writeable.
-void CheckOutputFilePermissions(const Arg& output_file) {
+void CheckOutputFilePermissions(const Arg& output_file) {
   QCHECK(output_file.specified()) << "Missing required flag --output_file.\n";
   QCHECK(port::file::Writable(output_file.value()).ok())
       << "Specified output_file is not writable: " << output_file.value()
@@ -40,7 +40,7 @@ void CheckOutputFilePermissions(const Arg& output_file) {
 }
 
 // Checks the permissions of the frozen model file.
-void CheckFrozenModelPermissions(const Arg& input_file) {
+void CheckFrozenModelPermissions(const Arg& input_file) {
   QCHECK(input_file.specified()) << "Missing required flag --input_file.\n";
   QCHECK(port::file::Exists(input_file.value(), port::file::Defaults()).ok())
       << "Specified input_file does not exist: " << input_file.value() << ".\n";
@@ -55,7 +55,7 @@ void CheckFrozenModelPermissions(const Arg& input_file) {
 void ReadInputData(const ParsedTocoFlags& parsed_toco_flags,
                    const ParsedModelFlags& parsed_model_flags,
                    TocoFlags* toco_flags, ModelFlags* model_flags,
-                   string* graph_def_contents) {
+                   std::string* graph_def_contents) {
   port::CheckInitGoogleIsDone("InitGoogle is not done yet.\n");
 
   // Ensure savedmodel_directory is not set.
@@ -71,10 +71,10 @@ void ReadInputData(const ParsedTocoFlags& parsed_toco_flags,
 }
 }  // namespace
 
-tensorflow::Status Convert(const string& graph_def_contents,
+tensorflow::Status Convert(const std::string& graph_def_contents,
                            const TocoFlags& toco_flags,
                            const ModelFlags& model_flags,
-                           string* output_file_contents,
+                           std::string* output_file_contents,
                            int64* arithmetic_ops_count = nullptr) {
   std::unique_ptr model =
       Import(toco_flags, model_flags, graph_def_contents);
@@ -95,12 +95,12 @@ tensorflow::Status Convert(const ParsedTocoFlags& parsed_toco_flags,
   TocoFlags toco_flags;
   ReadTocoFlagsFromCommandLineFlags(parsed_toco_flags, &toco_flags);
 
-  string graph_def_contents;
+  std::string graph_def_contents;
   ReadInputData(parsed_toco_flags, parsed_model_flags, &toco_flags,
                 &model_flags, &graph_def_contents);
   CheckOutputFilePermissions(parsed_toco_flags.output_file);
 
-  string output_file_contents;
+  std::string output_file_contents;
   TF_RETURN_IF_ERROR(Convert(graph_def_contents, toco_flags, model_flags,
                              &output_file_contents));
 
diff --git a/tensorflow/lite/toco/toco_convert.h b/tensorflow/lite/toco/toco_convert.h
index 4e3ffe5119b..85abcfcc3bb 100644
--- a/tensorflow/lite/toco/toco_convert.h
+++ b/tensorflow/lite/toco/toco_convert.h
@@ -22,10 +22,10 @@ limitations under the License.
 
 namespace toco {
 
-tensorflow::Status Convert(const string& graph_def_contents,
+tensorflow::Status Convert(const std::string& graph_def_contents,
                            const TocoFlags& toco_flags,
                            const ModelFlags& model_flags,
-                           string* output_file_contents,
+                           std::string* output_file_contents,
                            int64* arithmetic_ops_count = nullptr);
 
 tensorflow::Status Convert(const ParsedTocoFlags& parsed_toco_flags,
diff --git a/tensorflow/lite/toco/toco_convert_test.cc b/tensorflow/lite/toco/toco_convert_test.cc
index b02c1043f2b..bd6a8e79b45 100644
--- a/tensorflow/lite/toco/toco_convert_test.cc
+++ b/tensorflow/lite/toco/toco_convert_test.cc
@@ -32,8 +32,8 @@ TEST(TocoTest, BadInputFormat) {
   TocoFlags toco_flags;
   ModelFlags model_flags;
 
-  string input;
-  string output;
+  std::string input;
+  std::string output;
 
   EXPECT_DEATH(Convert(input, toco_flags, model_flags, &output).ok(),
                "Unhandled input_format='FILE_FORMAT_UNKNOWN'");
@@ -44,8 +44,8 @@ TEST(TocoTest, MissingOutputArrays) {
   ModelFlags model_flags;
 
   toco_flags.set_input_format(TENSORFLOW_GRAPHDEF);
-  string input;
-  string output;
+  std::string input;
+  std::string output;
 
   EXPECT_DEATH(Convert(input, toco_flags, model_flags, &output).ok(),
                "This model does not define output arrays, so a --output_arrays "
@@ -58,8 +58,8 @@ TEST(TocoTest, BadOutputArray) {
 
   toco_flags.set_input_format(TENSORFLOW_GRAPHDEF);
   model_flags.add_output_arrays("output1");
-  string input;
-  string output;
+  std::string input;
+  std::string output;
 
   EXPECT_DEATH(Convert(input, toco_flags, model_flags, &output).ok(),
                "Specified output array .output1. is not produced by any op "
@@ -72,7 +72,7 @@ TEST(TocoTest, BadOutputFormat) {
 
   toco_flags.set_input_format(TENSORFLOW_GRAPHDEF);
   model_flags.add_output_arrays("output1");
-  string input = R"GraphDef(
+  std::string input = R"GraphDef(
     node {
       name: "output1"
       input: "input1"
@@ -82,7 +82,7 @@ TEST(TocoTest, BadOutputFormat) {
     }
   )GraphDef";
 
-  string output;
+  std::string output;
 
   EXPECT_DEATH(Convert(input, toco_flags, model_flags, &output).ok(),
                "Unhandled output_format='FILE_FORMAT_UNKNOWN'");
@@ -97,7 +97,7 @@ TEST(TocoTest, SimpleFloatModel) {
 
   // Inputs are automatically selected (but that might not be a good idea).
   model_flags.add_output_arrays("output1");
-  string input = R"GraphDef(
+  std::string input = R"GraphDef(
     node {
       name: "input1"
       op: "Placeholder"
@@ -117,7 +117,7 @@ TEST(TocoTest, SimpleFloatModel) {
     }
   )GraphDef";
 
-  string output;
+  std::string output;
   EXPECT_TRUE(Convert(input, toco_flags, model_flags, &output).ok());
   EXPECT_TRUE(!output.empty());
 }
@@ -139,7 +139,7 @@ TEST(TocoTest, TransientStringTensors) {
   indices_1->set_name("indices1");
 
   model_flags.add_output_arrays("output1");
-  string input = R"GraphDef(
+  std::string input = R"GraphDef(
     node {
       name: "input1"
       op: "Placeholder"
@@ -169,7 +169,7 @@ TEST(TocoTest, TransientStringTensors) {
     }
   )GraphDef";
 
-  string output;
+  std::string output;
 
   EXPECT_TRUE(Convert(input, toco_flags, model_flags, &output).ok());
   EXPECT_TRUE(!output.empty());
diff --git a/tensorflow/lite/toco/toco_port.cc b/tensorflow/lite/toco/toco_port.cc
index d2f1d102c5a..8352e0fd9f2 100644
--- a/tensorflow/lite/toco/toco_port.cc
+++ b/tensorflow/lite/toco/toco_port.cc
@@ -28,7 +28,7 @@ double round(double x) { return ::round(x); }
 
 namespace toco {
 namespace port {
-void CopyToBuffer(const string& src, char* dest) {
+void CopyToBuffer(const std::string& src, char* dest) {
   memcpy(dest, src.data(), src.size());
 }
 
@@ -84,7 +84,7 @@ toco::port::file::Options ToOptions(const ::file::Options& options) {
   return Options();
 }
 
-tensorflow::Status Writable(const string& filename) {
+tensorflow::Status Writable(const std::string& filename) {
   File* f = nullptr;
   const auto status = ::file::Open(filename, "w", &f, ::file::Defaults());
   if (f) {
@@ -93,28 +93,30 @@ tensorflow::Status Writable(const string& filename) {
   return ToStatus(status);
 }
 
-tensorflow::Status Readable(const string& filename,
+tensorflow::Status Readable(const std::string& filename,
                             const file::Options& options) {
   return ToStatus(::file::Readable(filename, ::file::Defaults()));
 }
 
-tensorflow::Status Exists(const string& filename,
+tensorflow::Status Exists(const std::string& filename,
                           const file::Options& options) {
   auto status = ::file::Exists(filename, ::file::Defaults());
   return ToStatus(status);
 }
 
-tensorflow::Status GetContents(const string& filename, string* contents,
+tensorflow::Status GetContents(const std::string& filename,
+                               std::string* contents,
                                const file::Options& options) {
   return ToStatus(::file::GetContents(filename, contents, ::file::Defaults()));
 }
 
-tensorflow::Status SetContents(const string& filename, const string& contents,
+tensorflow::Status SetContents(const std::string& filename,
+                               const std::string& contents,
                                const file::Options& options) {
   return ToStatus(::file::SetContents(filename, contents, ::file::Defaults()));
 }
 
-string JoinPath(const string& a, const string& b) {
+std::string JoinPath(const std::string& a, const std::string& b) {
   return ::file::JoinPath(a, b);
 }
 
diff --git a/tensorflow/lite/toco/toco_port.h b/tensorflow/lite/toco/toco_port.h
index 5a80d29b72a..e57420fba4f 100644
--- a/tensorflow/lite/toco/toco_port.h
+++ b/tensorflow/lite/toco/toco_port.h
@@ -68,21 +68,23 @@ inline Options Defaults() {
   Options o;
   return o;
 }
-tensorflow::Status GetContents(const string& filename, string* contents,
+tensorflow::Status GetContents(const std::string& filename,
+                               std::string* contents, const Options& options);
+tensorflow::Status SetContents(const std::string& filename,
+                               const std::string& contents,
                                const Options& options);
-tensorflow::Status SetContents(const string& filename, const string& contents,
-                               const Options& options);
-string JoinPath(const string& base, const string& filename);
-tensorflow::Status Writable(const string& filename);
-tensorflow::Status Readable(const string& filename, const Options& options);
-tensorflow::Status Exists(const string& filename, const Options& options);
+std::string JoinPath(const std::string& base, const std::string& filename);
+tensorflow::Status Writable(const std::string& filename);
+tensorflow::Status Readable(const std::string& filename,
+                            const Options& options);
+tensorflow::Status Exists(const std::string& filename, const Options& options);
 }  // namespace file
 
 // Copy `src` string to `dest`. User must ensure `dest` has enough space.
 #if defined(PLATFORM_GOOGLE)
 void CopyToBuffer(const ::absl::Cord& src, char* dest);
 #endif  // PLATFORM_GOOGLE
-void CopyToBuffer(const string& src, char* dest);
+void CopyToBuffer(const std::string& src, char* dest);
 
 inline uint32 ReverseBits32(uint32 n) {
   n = ((n >> 1) & 0x55555555) | ((n & 0x55555555) << 1);
diff --git a/tensorflow/lite/toco/toco_tooling.cc b/tensorflow/lite/toco/toco_tooling.cc
index da0915f9739..25b48b54135 100644
--- a/tensorflow/lite/toco/toco_tooling.cc
+++ b/tensorflow/lite/toco/toco_tooling.cc
@@ -37,7 +37,7 @@ namespace toco {
 namespace {
 // CHECK-fails if the model contains a kUnsupported operation.
 void CheckUnsupportedOperations(const Model& model) {
-  std::set unsupported_ops;
+  std::set unsupported_ops;
   for (auto& op : model.operators) {
     if (op->type == OperatorType::kUnsupported) {
       unsupported_ops.insert(
@@ -172,7 +172,7 @@ void SetFinalDataTypeOnInputs(const TocoFlags& toco_flags, Model* model) {
   }
 
   for (int i = 0; i < model->flags.input_arrays_size(); i++) {
-    string const& array_name = model->flags.input_arrays(i).name();
+    std::string const& array_name = model->flags.input_arrays(i).name();
     auto* array = &model->GetArray(array_name);
     // Note that the notion of changing data types only applies to real-numbers
     // arrays (see the documentation for inference_input_type).
@@ -209,7 +209,7 @@ void SetFinalDataTypeOnInputs(const TocoFlags& toco_flags, Model* model) {
 
 std::unique_ptr Import(const TocoFlags& toco_flags,
                               const ModelFlags& model_flags,
-                              const string& input_file_contents) {
+                              const std::string& input_file_contents) {
   std::unique_ptr model;
   switch (toco_flags.input_format()) {
     case TENSORFLOW_GRAPHDEF: {
@@ -473,7 +473,8 @@ tensorflow::Status TransformWithStatus(const TocoFlags& toco_flags,
 }
 
 tensorflow::Status Export(const TocoFlags& toco_flags, const Model& model,
-                          bool allow_custom_ops, string* output_file_contents) {
+                          bool allow_custom_ops,
+                          std::string* output_file_contents) {
   switch (toco_flags.output_format()) {
     case TENSORFLOW_GRAPHDEF:
       ExportTensorFlowGraphDef(model, output_file_contents);
diff --git a/tensorflow/lite/toco/toco_tooling.h b/tensorflow/lite/toco/toco_tooling.h
index 36996151949..581df4b14fd 100644
--- a/tensorflow/lite/toco/toco_tooling.h
+++ b/tensorflow/lite/toco/toco_tooling.h
@@ -27,7 +27,7 @@ namespace toco {
 // Imports the input file into a Model object.
 std::unique_ptr Import(const TocoFlags& toco_flags,
                               const ModelFlags& model_flags,
-                              const string& input_file_contents);
+                              const std::string& input_file_contents);
 
 // Transforms a Model. The resulting Model is ready to be passed
 // to Export with the exact same toco_flags.
@@ -42,11 +42,12 @@ inline void Transform(const TocoFlags& toco_flags, Model* model) {
 // Transform, to a file of the format given by
 // toco_flags.output_format().
 tensorflow::Status Export(const TocoFlags& toco_flags, const Model& model,
-                          bool allow_custom_ops, string* output_file_contents);
+                          bool allow_custom_ops,
+                          std::string* output_file_contents);
 
 // This if for backward-compatibility with internal tools.
 inline void Export(const TocoFlags& toco_flags, const Model& model,
-                   string* output_file_contents) {
+                   std::string* output_file_contents) {
   auto status = Export(toco_flags, model, true, output_file_contents);
   if (!status.ok()) {
     LOG(QFATAL) << status.error_message();
diff --git a/tensorflow/lite/toco/tooling_util.cc b/tensorflow/lite/toco/tooling_util.cc
index 82ef4445a84..be4cda8aa3d 100644
--- a/tensorflow/lite/toco/tooling_util.cc
+++ b/tensorflow/lite/toco/tooling_util.cc
@@ -53,8 +53,8 @@ absl::string_view FindLongestCommonPrefix(absl::string_view a,
   return absl::string_view(a.data(), count);
 }
 
-string LogName(const Operator& op) {
-  const string& opname = HelpfulOperatorTypeName(op);
+std::string LogName(const Operator& op) {
+  const std::string& opname = HelpfulOperatorTypeName(op);
   if (op.outputs.empty()) {
     return toco::port::StringF("{%s operator}", opname);
   } else {
@@ -63,7 +63,7 @@ string LogName(const Operator& op) {
   }
 }
 
-string ArrayDataTypeName(ArrayDataType data_type) {
+std::string ArrayDataTypeName(ArrayDataType data_type) {
   switch (data_type) {
     case ArrayDataType::kFloat:
       return "float";
@@ -96,7 +96,7 @@ string ArrayDataTypeName(ArrayDataType data_type) {
   }
 }
 
-bool IsInputArray(const Model& model, const string& array_name) {
+bool IsInputArray(const Model& model, const std::string& array_name) {
   for (const auto& input_array : model.flags.input_arrays()) {
     if (array_name == input_array.name()) {
       return true;
@@ -105,7 +105,7 @@ bool IsInputArray(const Model& model, const string& array_name) {
   return false;
 }
 
-bool IsOutputArray(const Model& model, const string& array_name) {
+bool IsOutputArray(const Model& model, const std::string& array_name) {
   for (const auto& output_array : model.flags.output_arrays()) {
     if (array_name == output_array) {
       return true;
@@ -114,7 +114,7 @@ bool IsOutputArray(const Model& model, const string& array_name) {
   return false;
 }
 
-bool IsArrayConsumed(const Model& model, const string& name) {
+bool IsArrayConsumed(const Model& model, const std::string& name) {
   if (GetOpWithInput(model, name)) {
     return true;
   }
@@ -131,7 +131,7 @@ bool IsArrayConsumed(const Model& model, const string& name) {
 
 int CountTrueOutputs(const Model& model, const Operator& op) {
   int count = 0;
-  for (const string& output : op.outputs) {
+  for (const std::string& output : op.outputs) {
     if (IsArrayConsumed(model, output)) {
       ++count;
     }
@@ -139,7 +139,7 @@ int CountTrueOutputs(const Model& model, const Operator& op) {
   return count;
 }
 
-int CountOpsWithInput(const Model& model, const string& array_name) {
+int CountOpsWithInput(const Model& model, const std::string& array_name) {
   int count = 0;
   for (const auto& op : model.operators) {
     for (auto& input : op->inputs) {
@@ -155,7 +155,7 @@ int CountOpsWithInput(const Model& model, const string& array_name) {
   return count;
 }
 
-bool DeleteArrayIfUnused(const string& array_name, Model* model) {
+bool DeleteArrayIfUnused(const std::string& array_name, Model* model) {
   if (IsDiscardableArray(*model, array_name) &&
       CountOpsWithInput(*model, array_name) == 0 &&
       GetOpWithOutput(*model, array_name) == nullptr) {
@@ -165,7 +165,7 @@ bool DeleteArrayIfUnused(const string& array_name, Model* model) {
   return false;
 }
 
-bool DeleteArrayIfUnusedOutsideOfOp(const string& array_name,
+bool DeleteArrayIfUnusedOutsideOfOp(const std::string& array_name,
                                     const Operator* op, Model* model) {
   if (!IsDiscardableArray(*model, array_name)) {
     return false;
@@ -187,10 +187,10 @@ bool DeleteArrayIfUnusedOutsideOfOp(const string& array_name,
 }
 
 void DeleteOpAndArrays(Model* model, const Operator* op) {
-  for (const string& array_name : op->inputs) {
+  for (const std::string& array_name : op->inputs) {
     DeleteArrayIfUnusedOutsideOfOp(array_name, op, model);
   }
-  for (const string& array_name : op->outputs) {
+  for (const std::string& array_name : op->outputs) {
     DeleteArrayIfUnusedOutsideOfOp(array_name, op, model);
   }
   auto op_it = FindOp(*model, op);
@@ -199,7 +199,7 @@ void DeleteOpAndArrays(Model* model, const Operator* op) {
 }
 
 std::vector>::const_iterator FindOpWithOutput(
-    const Model& model, const string& array_name) {
+    const Model& model, const std::string& array_name) {
   for (auto it = model.operators.begin(); it != model.operators.end(); ++it) {
     for (auto& output : it->get()->outputs) {
       if (output == array_name) {
@@ -211,7 +211,7 @@ std::vector>::const_iterator FindOpWithOutput(
 }
 
 std::vector>::iterator FindOpWithOutput(
-    Model& model, const string& array_name) {
+    Model& model, const std::string& array_name) {
   for (auto it = model.operators.begin(); it != model.operators.end(); ++it) {
     for (auto& output : it->get()->outputs) {
       if (output == array_name) {
@@ -222,14 +222,14 @@ std::vector>::iterator FindOpWithOutput(
   return model.operators.end();
 }
 
-Operator* GetOpWithOutput(const Model& model, const string& array_name) {
+Operator* GetOpWithOutput(const Model& model, const std::string& array_name) {
   auto it = FindOpWithOutput(model, array_name);
   return it == model.operators.end() ? nullptr : it->get();
 }
 
 // GetFirstOpWithInput assumes that this finds the first op.
 std::vector>::const_iterator FindOpWithInput(
-    const Model& model, const string& array_name) {
+    const Model& model, const std::string& array_name) {
   for (auto it = model.operators.begin(); it != model.operators.end(); ++it) {
     for (auto& input : it->get()->inputs) {
       if (input == array_name) {
@@ -241,7 +241,7 @@ std::vector>::const_iterator FindOpWithInput(
 }
 
 std::vector>::iterator FindOpWithInput(
-    Model& model, const string& array_name) {
+    Model& model, const std::string& array_name) {
   for (auto it = model.operators.begin(); it != model.operators.end(); ++it) {
     for (auto& input : it->get()->inputs) {
       if (input == array_name) {
@@ -272,18 +272,19 @@ std::vector>::iterator FindOp(Model& model,
   return model.operators.end();
 }
 
-Operator* GetOpWithInput(const Model& model, const string& array_name) {
+Operator* GetOpWithInput(const Model& model, const std::string& array_name) {
   auto it = FindOpWithInput(model, array_name);
   return it == model.operators.end() ? nullptr : it->get();
 }
 
-Operator* GetFirstOpWithInput(const Model& model, const string& array_name) {
+Operator* GetFirstOpWithInput(const Model& model,
+                              const std::string& array_name) {
   auto it = FindOpWithInput(model, array_name);
   return it == model.operators.end() ? nullptr : it->get();
 }
 
-void ReplaceArrayUsage(Model* model, const string& old_array_name,
-                       const string& new_array_name) {
+void ReplaceArrayUsage(Model* model, const std::string& old_array_name,
+                       const std::string& new_array_name) {
   for (auto& op_it : model->operators) {
     Operator* op = op_it.get();
     for (size_t i = 0; i < op->inputs.size(); ++i) {
@@ -299,11 +300,12 @@ void ReplaceArrayUsage(Model* model, const string& old_array_name,
   }
 }
 
-string FormatArraysList(const Model& model, const std::vector& list) {
+std::string FormatArraysList(const Model& model,
+                             const std::vector& list) {
   if (list.empty()) {
     return "[]";
   }
-  string result = "";
+  std::string result = "";
   if (list.size() > 1) {
     result += "[ ";
   }
@@ -459,7 +461,7 @@ const char* OperatorTypeName(OperatorType type) {
   }
 }
 
-string HelpfulOperatorTypeName(const Operator& op) {
+std::string HelpfulOperatorTypeName(const Operator& op) {
   if (op.type == OperatorType::kUnsupported) {
     return toco::port::StringF(
         "(Unsupported TensorFlow op: %s)",
@@ -503,7 +505,7 @@ void LogSummary(int log_level, const Model& model) {
   }
 }
 
-void LogArray(int log_level, const Model& model, const string& name) {
+void LogArray(int log_level, const Model& model, const std::string& name) {
   VLOG(log_level) << "Array: " << name;
   if (!model.HasArray(name)) {
     VLOG(log_level) << "  DOES NOT EXIST";
@@ -524,7 +526,7 @@ void LogArray(int log_level, const Model& model, const string& name) {
     if (array_shape.dimensions_count() == 0) {
       VLOG(log_level) << "  (Zero dimensions)";
     } else {
-      string message = "  Dims: ";
+      std::string message = "  Dims: ";
       bool first = true;
       for (const int dim : array_shape.dims()) {
         if (!first) {
@@ -568,10 +570,10 @@ void DumpGraphvizVideoFrame(const Model& model) {
   // this new video-dumping feature.
   static int dump_id = 0;
   static std::unordered_set dump_hashes;
-  string graphviz_dump;
+  std::string graphviz_dump;
   DumpGraphviz(model, &graphviz_dump,
                toco::port::StringF("VIDEO frame:%05d", dump_id));
-  std::size_t hash = std::hash{}(graphviz_dump);
+  std::size_t hash = std::hash{}(graphviz_dump);
   if (!dump_hashes.count(hash)) {
     LOG(INFO) << "DUMPING GRAPHVIZ VIDEO FRAME: " << dump_id;
     dump_hashes.insert(hash);
@@ -585,13 +587,13 @@ void DumpGraphvizVideoFrame(const Model& model) {
   }
 }
 
-void LogDump(int log_level, const string& message, const Model& model) {
+void LogDump(int log_level, const std::string& message, const Model& model) {
   namespace port = toco::port;
   const auto& dump_options = *GraphVizDumpOptions::singleton();
 
   DumpGraphvizVideoFrame(model);
   if (!dump_options.dump_graphviz.empty()) {
-    string graphviz_dump;
+    std::string graphviz_dump;
 
     DumpGraphviz(model, &graphviz_dump, message);
     const auto result = port::file::SetContents(
@@ -608,7 +610,7 @@ void LogDump(int log_level, const string& message, const Model& model) {
   }
   VLOG(log_level) << "BEGIN DUMP OF TOCO MODEL (" << message << ")";
   LogSummary(log_level, model);
-  std::unordered_set already_printed_arrays;
+  std::unordered_set already_printed_arrays;
   for (const auto& op : model.operators) {
     for (const auto& input : op->inputs) {
       if (!already_printed_arrays.count(input)) {
@@ -759,7 +761,7 @@ int RequiredBufferSizeForShape(const Shape& shape) {
   return max_offset;
 }
 
-bool IsConstantParameterArray(const Model& model, const string& name) {
+bool IsConstantParameterArray(const Model& model, const std::string& name) {
   if (!model.HasArray(name)) {
     return false;
   }
@@ -858,7 +860,7 @@ bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array) {
 namespace {
 // Take an array name, which may be something like "name:3_5" and make it
 // acceptable as a TF node name, say "name_3_5";
-string SanitizeNameForTFNode(const string& array_name) {
+std::string SanitizeNameForTFNode(const std::string& array_name) {
   auto node_name = array_name;
   std::replace(node_name.begin(), node_name.end(), ':', '_');
   return node_name;
@@ -866,7 +868,7 @@ string SanitizeNameForTFNode(const string& array_name) {
 
 void CheckInputArraysAreNotOutputArrays(const ModelFlags& model_flags) {
   for (const auto& input_array : model_flags.input_arrays()) {
-    for (const string& output_array : model_flags.output_arrays()) {
+    for (const std::string& output_array : model_flags.output_arrays()) {
       QCHECK_NE(input_array.name(), output_array)
           << "The array " << output_array
           << " is listed in both --input_arrays and --output_arrays.";
@@ -874,7 +876,7 @@ void CheckInputArraysAreNotOutputArrays(const ModelFlags& model_flags) {
   }
 }
 
-bool IsAsciiPrintable(const string& name) {
+bool IsAsciiPrintable(const std::string& name) {
   for (char c : name) {
     if (!absl::ascii_isprint(c)) {
       return false;
@@ -883,8 +885,8 @@ bool IsAsciiPrintable(const string& name) {
   return true;
 }
 
-string DumpAscii(const string& name) {
-  string result;
+std::string DumpAscii(const std::string& name) {
+  std::string result;
   port::AppendF(&result, "ASCII | Hex\n");
   port::AppendF(&result, "------+----\n");
   for (char c : name) {
@@ -909,7 +911,7 @@ void CheckNonAsciiIOArrays(const ModelFlags& model_flags) {
         << "Here is a dump of the string:\n\n"
         << DumpAscii(input_array.name());
   }
-  for (const string& output_array : model_flags.output_arrays()) {
+  for (const std::string& output_array : model_flags.output_arrays()) {
     QCHECK(IsAsciiPrintable(output_array))
         << "Non-ASCII-printable character found in --output_arrays: "
         << output_array << ". Pass --allow_nonascii_arrays to allow that. "
@@ -932,7 +934,7 @@ void CheckNonExistentIOArrays(const Model& model) {
       "Is it a typo? This should not happen. If you trigger this error "
       "please send a bug report (with code to reproduce this error), to the "
       "TensorFlow Lite team.";
-  for (const string& output_array : model.flags.output_arrays()) {
+  for (const std::string& output_array : model.flags.output_arrays()) {
     if (IsConstantParameterArray(model, output_array)) {
       continue;  // It is OK to request that a constant be an output.
     }
@@ -984,7 +986,7 @@ void FixNoMissingArray(Model* model) {
     }
   }
   if (model->flags.allow_nonexistent_arrays()) {
-    for (const string& output_array : model->flags.output_arrays()) {
+    for (const std::string& output_array : model->flags.output_arrays()) {
       model->GetOrCreateArray(output_array);
     }
     for (const auto& rnn_state : model->flags.rnn_states()) {
@@ -995,7 +997,7 @@ void FixNoMissingArray(Model* model) {
 }
 
 void CheckNoOrphanedArray(const Model& model) {
-  std::unordered_set arrays_without_known_use;
+  std::unordered_set arrays_without_known_use;
   for (const auto& array : model.GetArrayMap()) {
     if (IsDiscardableArray(model, array.first)) {
       arrays_without_known_use.insert(array.first);
@@ -1022,7 +1024,7 @@ void CheckNoOrphanedArray(const Model& model) {
 }
 
 void FixNoOrphanedArray(Model* model) {
-  std::unordered_set arrays_without_known_use;
+  std::unordered_set arrays_without_known_use;
   for (const auto& array : model->GetArrayMap()) {
     arrays_without_known_use.insert(array.first);
   }
@@ -1071,11 +1073,11 @@ void CheckEachArray(const Model& model) {
 
     // Check name.  Either "name_with_suffix_8", "name_with_port:3", but not
     // "name_with_both:3_8".
-    const string& name = array_entry.first;
+    const std::string& name = array_entry.first;
     auto colon_pos = name.find_first_of(":");
-    if (colon_pos != string::npos) {
+    if (colon_pos != std::string::npos) {
       CHECK_EQ(name.substr(colon_pos + 1).find_first_not_of("0123456789"),
-               string::npos)
+               std::string::npos)
           << "Array '" << name << "' has non-digit characters after colon.";
     }
     CHECK_GT(colon_pos, 0) << "Array '" << name
@@ -1084,7 +1086,7 @@ void CheckEachArray(const Model& model) {
 }
 
 void CheckOperatorOrdering(const Model& model) {
-  std::unordered_set arrays_behind_us;
+  std::unordered_set arrays_behind_us;
   for (const auto& array_entry : model.GetArrayMap()) {
     if (!GetOpWithOutput(model, array_entry.first)) {
       arrays_behind_us.insert(array_entry.first);
@@ -1103,13 +1105,13 @@ void CheckOperatorOrdering(const Model& model) {
       arrays_behind_us.insert(output);
     }
   }
-  for (const string& output_array : model.flags.output_arrays()) {
+  for (const std::string& output_array : model.flags.output_arrays()) {
     CHECK(arrays_behind_us.count(output_array));
   }
 }
 
 void FixOperatorOrdering(Model* model) {
-  std::unordered_set arrays_behind_us;
+  std::unordered_set arrays_behind_us;
   for (const auto& array_entry : model->GetArrayMap()) {
     if (!GetOpWithOutput(*model, array_entry.first)) {
       arrays_behind_us.insert(array_entry.first);
@@ -1123,7 +1125,7 @@ void FixOperatorOrdering(Model* model) {
   for (std::size_t i = 0; i < old_operators.size(); i++) {
     remaining.insert(i);
   }
-  std::unordered_map reason_why_leftover;
+  std::unordered_map reason_why_leftover;
   while (true) {
     bool inserted_something = false;
     for (const auto& i : remaining) {
@@ -1133,7 +1135,7 @@ void FixOperatorOrdering(Model* model) {
       for (const auto& input : op->inputs) {
         if (!IsConstantParameterArray(*model, input) &&
             !arrays_behind_us.count(input)) {
-          for (const string& output : op->outputs) {
+          for (const std::string& output : op->outputs) {
             reason_why_leftover[output] = input;
           }
           can_insert = false;
@@ -1166,15 +1168,15 @@ void FixOperatorOrdering(Model* model) {
     LOG(ERROR) << "BEGIN TRACE OF OPERATOR WITH BAD INPUT";
     LOG(ERROR) << "Here is the first-encountered operator with a bad input: ";
     const Operator* bad_op = old_operators[*remaining.begin()].get();
-    std::unordered_set bad_inputs_already_traced;
+    std::unordered_set bad_inputs_already_traced;
     // The following while(true) loop should always end with a LOG(FATAL).
     while (true) {
       LOG(ERROR) << HelpfulOperatorTypeName(*bad_op) << " : "
                  << FormatArraysList(*model, bad_op->inputs) << " -> "
                  << FormatArraysList(*model, bad_op->outputs);
       bool found_bad_output = false;
-      string bad_output;
-      for (const string& output : bad_op->outputs) {
+      std::string bad_output;
+      for (const std::string& output : bad_op->outputs) {
         if (reason_why_leftover.count(output)) {
           found_bad_output = true;
           bad_output = output;
@@ -1182,7 +1184,7 @@ void FixOperatorOrdering(Model* model) {
         }
       }
       CHECK(found_bad_output);
-      const string& bad_input = reason_why_leftover[bad_output];
+      const std::string& bad_input = reason_why_leftover[bad_output];
       LOG(ERROR) << "The bad input here is: " << bad_input;
       if (bad_inputs_already_traced.count(bad_input)) {
         LOG(FATAL)
@@ -1198,7 +1200,7 @@ void FixOperatorOrdering(Model* model) {
       bad_op = nullptr;
       for (const auto& i : remaining) {
         const Operator* op = old_operators[i].get();
-        for (const string& output : op->outputs) {
+        for (const std::string& output : op->outputs) {
           if (bad_input == output) {
             bad_op = op;
             break;
@@ -1233,7 +1235,7 @@ void CheckInvariants(const Model& model) {
 }
 
 void CheckCountInRange(const ::toco::ModelFlags::ModelCheck& model_check,
-                       const int count, const string& count_description) {
+                       const int count, const std::string& count_description) {
   if (model_check.count_min() >= 0) {
     CHECK_GE(count, model_check.count_min())
         << "Mismatch in " << count_description << ": count  was " << count
@@ -1251,7 +1253,7 @@ void CheckCountInRange(const ::toco::ModelFlags::ModelCheck& model_check,
 
 void CheckModelCounts(const Model& model) {
   std::unordered_multiset ops_by_type;
-  std::unordered_map op_type_by_name;
+  std::unordered_map op_type_by_name;
   if (model.flags.model_checks_size() == 0) {
     return;
   }
@@ -1261,7 +1263,7 @@ void CheckModelCounts(const Model& model) {
     op_type_by_name[OperatorTypeName(op->type)] = op->type;
   }
   for (const auto& model_check : model.flags.model_checks()) {
-    string count_type = model_check.count_type();
+    std::string count_type = model_check.count_type();
     if (count_type == "None") {
       continue;
     } else if (count_type == "Arrays") {
@@ -1284,12 +1286,12 @@ void CheckModelCounts(const Model& model) {
 }
 
 void FixEdgeArrays(Model* model) {
-  for (const string& output_array_name : model->flags.output_arrays()) {
+  for (const std::string& output_array_name : model->flags.output_arrays()) {
     if (!GetOpWithOutput(*model, output_array_name)) {
       // Output has no operator producing it. Change that by inserting a copy.
       LOG(WARNING) << "Fixing constant output array " << output_array_name
                    << " by inserting a copy. This is not optimal.";
-      string intermediate_array_name =
+      std::string intermediate_array_name =
           AvailableArrayName(*model, output_array_name + "_copy");
       CloneArray(model, output_array_name, intermediate_array_name);
       InsertCopyOperator(model, intermediate_array_name, output_array_name);
@@ -1378,8 +1380,8 @@ void CopyArrayAttribs(const Array& source_array, Array* target_array) {
 }
 }  // namespace
 
-void InsertCopyOperator(Model* model, const string& source_array_name,
-                        const string& target_array_name) {
+void InsertCopyOperator(Model* model, const std::string& source_array_name,
+                        const std::string& target_array_name) {
   // Reshape to the same size. This should be a no-op.
   const Array& source_array = model->GetArray(source_array_name);
   std::vector shape = source_array.shape().dims();
@@ -1404,8 +1406,8 @@ void InsertCopyOperator(Model* model, const string& source_array_name,
   model->operators.emplace_back(copy_op);
 }
 
-void CloneArray(Model* model, const string& source_array_name,
-                const string& target_array_name) {
+void CloneArray(Model* model, const std::string& source_array_name,
+                const std::string& target_array_name) {
   CHECK(!model->HasArray(target_array_name));
   const Array& source_array = model->GetArray(source_array_name);
   Array& target_array = model->GetOrCreateArray(target_array_name);
@@ -1479,7 +1481,7 @@ void MakeArrayDims(int num_dims, int batch, int height, int width, int depth,
   }
 }
 
-void CreateOrCheckRnnStateArray(const string& name, int size,
+void CreateOrCheckRnnStateArray(const std::string& name, int size,
                                 int state_num_dims, Model* model) {
   int batch = 1;
   int num_dims = -1;
@@ -1781,7 +1783,7 @@ int ElementSize(ArrayDataType data_type) {
   }
 }
 
-void DropMinMax(Model* model, const string& array_name) {
+void DropMinMax(Model* model, const std::string& array_name) {
   auto& array = model->GetArray(array_name);
   if (!!array.minmax) {
     LOG(WARNING) << "Dropping MinMax information in array " << array_name
@@ -1790,7 +1792,8 @@ void DropMinMax(Model* model, const string& array_name) {
   }
 }
 
-bool IsAllocatableTransientArray(const Model& model, const string& array_name) {
+bool IsAllocatableTransientArray(const Model& model,
+                                 const std::string& array_name) {
   // Optional array is not transient
   if (model.IsOptionalArray(array_name)) return false;
   // The model's input and output arrays are externally allocated.
@@ -1818,15 +1821,15 @@ bool IsAllocatableTransientArray(const Model& model, const string& array_name) {
   return true;
 }
 
-string AvailableArrayName(const Model& model, const string& name) {
-  string sanitized_name = SanitizeNameForTFNode(name);
+std::string AvailableArrayName(const Model& model, const std::string& name) {
+  std::string sanitized_name = SanitizeNameForTFNode(name);
   if (!model.HasArray(sanitized_name) &&
       !model.IsOptionalArray(sanitized_name)) {
     return sanitized_name;
   }
   const int kNumSuffixesToTry = 1000;
   for (int i = 0; i < kNumSuffixesToTry; i++) {
-    const string& name_with_suffix =
+    const std::string& name_with_suffix =
         toco::port::StringF("%s_%d", sanitized_name, i);
     if (!model.HasArray(name_with_suffix) &&
         !model.IsOptionalArray(name_with_suffix)) {
@@ -1839,7 +1842,7 @@ string AvailableArrayName(const Model& model, const string& name) {
   return "";
 }
 
-string ShapeToString(const Shape& shape) {
+std::string ShapeToString(const Shape& shape) {
   if (shape.dimensions_count() == 0) {
     return "[]";
   }
@@ -1847,7 +1850,7 @@ string ShapeToString(const Shape& shape) {
   return absl::StrCat("[ ", absl::StrJoin(shape.dims(), ", "), " ]");
 }
 
-void PrintArrayShape(Model* model, const string& name) {
+void PrintArrayShape(Model* model, const std::string& name) {
   if (!model->GetArray(name).has_shape()) {
     LOG(INFO) << name << " has no shape";
     return;
@@ -1856,7 +1859,7 @@ void PrintArrayShape(Model* model, const string& name) {
             << " has shape: " << ShapeToString(model->GetArray(name).shape());
 }
 
-bool IsArrayFullyConnectedWeights(const Model& model, const string& name) {
+bool IsArrayFullyConnectedWeights(const Model& model, const std::string& name) {
   bool is_fc_weights = false;
   bool is_something_else = false;
   for (const auto& op : model.operators) {
@@ -1874,8 +1877,8 @@ bool IsArrayFullyConnectedWeights(const Model& model, const string& name) {
   return is_fc_weights;
 }
 
-string CreateInt32Array(Model* model, const string& param_name,
-                        const std::vector& value) {
+std::string CreateInt32Array(Model* model, const std::string& param_name,
+                             const std::vector& value) {
   auto param_array_name = AvailableArrayName(*model, param_name);
   auto& param_array = model->GetOrCreateArray(param_array_name);
   param_array.mutable_shape()->ReplaceDims({static_cast(value.size())});
@@ -2031,7 +2034,7 @@ bool EstimateArithmeticOpsCount(const Model& model, int64* result) {
   return true;
 }
 
-string FormattedNumber(int64 x) {
+std::string FormattedNumber(int64 x) {
   const int64 million = 1000000;
   const int64 billion = 1000000000;
   if (x < 10000) {
@@ -2222,7 +2225,7 @@ int AxesCount(AxesOrder axes_order) {
   }
 }
 
-bool IsDiscardableArray(const Model& model, const string& array_name) {
+bool IsDiscardableArray(const Model& model, const std::string& array_name) {
   if (IsInputArray(model, array_name) || IsOutputArray(model, array_name)) {
     return false;
   }
@@ -2338,9 +2341,9 @@ void FinishBuildingRNNStates(Model* model) {
 
 // Returns the array names that match the ArraysExtraInfo's name and
 // name_regexp. The regexp match is for a full match.
-std::unordered_set ScanArrayNames(
+std::unordered_set ScanArrayNames(
     const Model& model, const toco::ArraysExtraInfo_Entry& entry) {
-  std::unordered_set matches;
+  std::unordered_set matches;
   if (model.HasArray(entry.name())) {
     matches.insert(entry.name());
   }
@@ -2409,7 +2412,7 @@ void UndoWeightsShuffling(Model* model) {
     if (fc_op.weights_format == FullyConnectedWeightsFormat::kDefault) {
       continue;
     }
-    const string& weights_name = fc_op.inputs[1];
+    const std::string& weights_name = fc_op.inputs[1];
     QCHECK_EQ(CountOpsWithInput(*model, weights_name), 1);
     auto& weights_array = model->GetArray(weights_name);
     QCHECK(weights_array.data_type == ArrayDataType::kUint8);
diff --git a/tensorflow/lite/toco/tooling_util.h b/tensorflow/lite/toco/tooling_util.h
index 6fd13be182c..438ce19970d 100644
--- a/tensorflow/lite/toco/tooling_util.h
+++ b/tensorflow/lite/toco/tooling_util.h
@@ -54,44 +54,45 @@ constexpr int kLogLevelModelUnchanged = 2;
 
 absl::string_view FindLongestCommonPrefix(absl::string_view a,
                                           absl::string_view b);
-string LogName(const Operator& op);
+std::string LogName(const Operator& op);
 
-string ArrayDataTypeName(ArrayDataType data_type);
+std::string ArrayDataTypeName(ArrayDataType data_type);
 
 // Returns true if the given array is specified as a model input array.
-bool IsInputArray(const Model& model, const string& array_name);
+bool IsInputArray(const Model& model, const std::string& array_name);
 // Returns true if the given array is specified as a model output array.
-bool IsOutputArray(const Model& model, const string& array_name);
+bool IsOutputArray(const Model& model, const std::string& array_name);
 
-bool IsArrayConsumed(const Model& model, const string& name);
+bool IsArrayConsumed(const Model& model, const std::string& name);
 int CountTrueOutputs(const Model& model, const Operator& op);
 
-int CountOpsWithInput(const Model& model, const string& array_name);
-bool DeleteArrayIfUnused(const string& array_name, Model* model);
+int CountOpsWithInput(const Model& model, const std::string& array_name);
+bool DeleteArrayIfUnused(const std::string& array_name, Model* model);
 
 // Deletes the op and any of its input and output arrays if they are unused
 // after the op has been deleted.
 void DeleteOpAndArrays(Model* model, const Operator* op);
 
 std::vector>::const_iterator FindOpWithOutput(
-    const Model& model, const string& array_name);
-Operator* GetOpWithOutput(const Model& model, const string& array_name);
+    const Model& model, const std::string& array_name);
+Operator* GetOpWithOutput(const Model& model, const std::string& array_name);
 
 std::vector>::iterator FindOpWithOutput(
-    Model& model, const string& array_name);
+    Model& model, const std::string& array_name);
 
 std::vector>::const_iterator FindOpWithInput(
-    const Model& model, const string& array_name);
+    const Model& model, const std::string& array_name);
 
 std::vector>::iterator FindOpWithInput(
-    Model& model, const string& array_name);
+    Model& model, const std::string& array_name);
 
-Operator* GetOpWithInput(const Model& model, const string& array_name);
-Operator* GetFirstOpWithInput(const Model& model, const string& array_name);
+Operator* GetOpWithInput(const Model& model, const std::string& array_name);
+Operator* GetFirstOpWithInput(const Model& model,
+                              const std::string& array_name);
 
 // Replaces all uses of the |old_array_name| with the |new_array_name|.
-void ReplaceArrayUsage(Model* model, const string& old_array_name,
-                       const string& new_array_name);
+void ReplaceArrayUsage(Model* model, const std::string& old_array_name,
+                       const std::string& new_array_name);
 
 std::vector>::const_iterator FindOp(
     const Model& model, const Operator* op);
@@ -99,15 +100,15 @@ std::vector>::iterator FindOp(Model& model,
                                                         const Operator* op);
 
 const char* OperatorTypeName(OperatorType type);
-string HelpfulOperatorTypeName(const Operator& op);
+std::string HelpfulOperatorTypeName(const Operator& op);
 
 // Whether the operator can be fused with an activation function. Note that this
 // will return false by default for new operators; fusing support is opt-in.
 bool OperatorSupportsFusedActivation(OperatorType type);
 
 void DumpGraphvizVideoFrame(const Model& model);
-void LogDump(int log_level, const string& message, const Model& model);
-void LogSummary(int log_level, const string& message, const Model& model);
+void LogDump(int log_level, const std::string& message, const Model& model);
+void LogSummary(int log_level, const std::string& message, const Model& model);
 
 // TODO(b/36075966): Clean up when dims superseded by array shape.
 void ExtendShape(Shape* shape, int new_shape_size);
@@ -143,12 +144,12 @@ inline ::tflite::RuntimeShape ToRuntimeShape(const Shape& shape) {
   return ::tflite::RuntimeShape(shape.dimensions_count(), shape.dims().data());
 }
 
-bool IsArrayFullyConnectedWeights(const Model& model, const string& name);
+bool IsArrayFullyConnectedWeights(const Model& model, const std::string& name);
 
 // If there is a wildcard dimension (-1), this may return a negative value.
 int RequiredBufferSizeForShape(const Shape& shape);
 
-bool IsConstantParameterArray(const Model& model, const string& name);
+bool IsConstantParameterArray(const Model& model, const std::string& name);
 
 // Compares two constant parameter arrays for exact equality.
 bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array);
@@ -193,12 +194,12 @@ void CopyArrayBuffer(const Array& source_array, Array* target_array) {
 
 // Inserts a no-op reshape operator between the source array and the target
 // array. This effectively just copies the data.
-void InsertCopyOperator(Model* model, const string& source_array_name,
-                        const string& target_array_name);
+void InsertCopyOperator(Model* model, const std::string& source_array_name,
+                        const std::string& target_array_name);
 
 // Clones an array with all data and parameters.
-void CloneArray(Model* model, const string& source_array_name,
-                const string& target_array_name);
+void CloneArray(Model* model, const std::string& source_array_name,
+                const std::string& target_array_name);
 
 void ResolveModelFlags(const ModelFlags& model_flags, Model* model);
 
@@ -245,32 +246,33 @@ inline std::vector ReverseOffset(const Shape& shape, int index) {
 
 int ElementSize(ArrayDataType data_type);
 
-void DropMinMax(Model* model, const string& array_name);
+void DropMinMax(Model* model, const std::string& array_name);
 
-bool IsAllocatableTransientArray(const Model& model, const string& array_name);
+bool IsAllocatableTransientArray(const Model& model,
+                                 const std::string& array_name);
 
-void CreateOrCheckRnnStateArray(const string& name, int size,
+void CreateOrCheckRnnStateArray(const std::string& name, int size,
                                 int state_num_dims, Model* model);
 
-string AvailableArrayName(const Model& model, const string& name);
+std::string AvailableArrayName(const Model& model, const std::string& name);
 
 // Formats a shape as a string: [ dims(0), dims(1), ..., dims(num_dims-1) ].
-string ShapeToString(const Shape& shape);
+std::string ShapeToString(const Shape& shape);
 
-void PrintArrayShape(Model* model, const string& name);
+void PrintArrayShape(Model* model, const std::string& name);
 
 void MakeArrayDims(int num_dims, int batch, int height, int width, int depth,
                    std::vector* out_dims);
 
 // Defines a constant int32 array with the provided values formatted for use
 // as op parameters.
-string CreateInt32Array(Model* model, const string& param_name,
-                        const std::vector& value);
+std::string CreateInt32Array(Model* model, const std::string& param_name,
+                             const std::vector& value);
 
 bool EstimateArithmeticOpsCount(const Model& model, const Operator& op,
                                 int64* result);
 bool EstimateArithmeticOpsCount(const Model& model, int64* result);
-string FormattedNumber(int64 x);
+std::string FormattedNumber(int64 x);
 
 int AxesCount(AxesOrder axes_order);
 
@@ -297,7 +299,7 @@ void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order,
 // that array. The idea is that we can't ever discard arrays that are either
 // an input or an output of the whole graph, or that appear in RNN back-edges,
 // as that would undercut explicit flags that the user might pass.
-bool IsDiscardableArray(const Model& model, const string& array_name);
+bool IsDiscardableArray(const Model& model, const std::string& array_name);
 
 void CheckFinalDataTypesSatisfied(const Model& model);
 
@@ -362,7 +364,7 @@ void CopyMinMaxAndQuantizationRelatedFields(const Array& src, Array* dst);
 
 // Delete Array if it's discardable and not referenced as input or output array
 // by any other op than the specified op.
-bool DeleteArrayIfUnusedOutsideOfOp(const string& array_name,
+bool DeleteArrayIfUnusedOutsideOfOp(const std::string& array_name,
                                     const Operator* op, Model* model);
 
 }  // namespace toco

From 406d9b5521a6bd665a8773573a529fe31b95e5ce Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Wed, 17 Jun 2020 09:53:14 -0700
Subject: [PATCH 0405/1390] Make the names of "activation", "activation_state",
 "output_state" variables consistent: Use "output_state", used by the float
 implementation (LstmStepFloat).

Remove a few variables in implementation and test that were representing the same value (the index of these tensors, mostly).

Also rename "input cell state" to just "cell state".

PiperOrigin-RevId: 316908413
Change-Id: Icb64ecd31c90f45ef21cf7d48849fb2ec0975d3a
---
 tensorflow/lite/kernels/lstm.cc               | 169 +++++++++---------
 tensorflow/lite/kernels/lstm_eval.cc          |  82 +++++----
 tensorflow/lite/kernels/lstm_eval.h           |  15 +-
 tensorflow/lite/kernels/lstm_shared.h         |   4 +-
 tensorflow/lite/kernels/lstm_test.cc          |  76 ++++----
 .../kernels/unidirectional_sequence_lstm.cc   |  53 +++---
 .../unidirectional_sequence_lstm_test.cc      |  39 ++--
 .../calibration/builtin_logging_ops/lstm.cc   |  33 ++--
 8 files changed, 226 insertions(+), 245 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index 74caafbd0c7..0e0c1b9c0f0 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -68,19 +68,19 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   const float cell_clip = params->cell_clip;
   const float proj_clip = params->proj_clip;
 
-  const TfLiteTensor* cell_tensor =
-      GetVariableInput(context, node, kInputCellStateTensor);
-  TF_LITE_ENSURE(context, cell_tensor != nullptr);
+  const TfLiteTensor* cell_state =
+      GetVariableInput(context, node, kCellStateTensor);
+  TF_LITE_ENSURE(context, cell_state != nullptr);
   const TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
 
-  auto* cell_params =
-      static_cast(cell_tensor->quantization.params);
+  auto* cell_state_params =
+      static_cast(cell_state->quantization.params);
   auto* proj_params = static_cast(
       output_tensor->quantization.params);
   if (cell_clip > 0.0) {
-    integer_lstm_param->quantized_cell_clip = static_cast(
-        std::min(std::max(cell_clip / cell_params->scale->data[0], -32768.0f),
-                 32767.0f));
+    integer_lstm_param->quantized_cell_clip = static_cast(std::min(
+        std::max(cell_clip / cell_state_params->scale->data[0], -32768.0f),
+        32767.0f));
   } else {
     integer_lstm_param->quantized_cell_clip = 0;
   }
@@ -134,9 +134,9 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   const TfLiteTensor* projection_weights =
       GetOptionalInputTensor(context, node, kProjectionWeightsTensor);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
 
   // Since we have already checked that weights are all there or none, we can
   // check the existence of only one to get the condition.
@@ -187,7 +187,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   float layer_norm_forget_scale = default_scale;
   float layer_norm_cell_scale = default_scale;
   float layer_norm_output_scale = default_scale;
-  float activation_scale = default_scale;
+  float output_state_scale = default_scale;
   int cell_scale = 1;
 
   // Effective scales.
@@ -231,7 +231,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   if (use_projection) {
     proj_weight_scale = projection_weights->params.scale;
   }
-  activation_scale = activation_state->params.scale;
+  output_state_scale = output_state->params.scale;
 
   input_to_forget_weight_scale = input_to_forget_weights->params.scale;
   input_to_cell_weight_scale = input_to_cell_weights->params.scale;
@@ -240,12 +240,8 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   recurrent_to_cell_weight_scale = recurrent_to_cell_weights->params.scale;
   recurrent_to_output_weight_scale = recurrent_to_output_weights->params.scale;
 
-  // Get cell state.
-  TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
-  TF_LITE_ENSURE(context, cell_state != nullptr);
+  // Check cell state (already used above)
   TF_LITE_ENSURE(context, CheckedLog2(cell_state->params.scale, &cell_scale));
-
   TF_LITE_ENSURE(context, cell_scale <= -9);
   integer_lstm_param->cell_scale = cell_scale;
   input_scale = input->params.scale;
@@ -255,31 +251,32 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
     effective_input_to_input_scale =
         input_to_input_weight_scale * input_scale / intermediate_scale[0];
     effective_recurrent_to_input_scale = recurrent_to_input_weight_scale *
-                                         activation_scale /
+                                         output_state_scale /
                                          intermediate_scale[0];
   }
   effective_input_to_forget_scale =
       input_to_forget_weight_scale * input_scale / intermediate_scale[1];
   effective_recurrent_to_forget_scale = recurrent_to_forget_weight_scale *
-                                        activation_scale /
+                                        output_state_scale /
                                         intermediate_scale[1];
 
   effective_input_to_cell_scale =
       input_to_cell_weight_scale * input_scale / intermediate_scale[2];
-  effective_recurrent_to_cell_scale =
-      recurrent_to_cell_weight_scale * activation_scale / intermediate_scale[2];
+  effective_recurrent_to_cell_scale = recurrent_to_cell_weight_scale *
+                                      output_state_scale /
+                                      intermediate_scale[2];
 
   effective_input_to_output_scale =
       input_to_output_weight_scale * input_scale / intermediate_scale[3];
   effective_recurrent_to_output_scale = recurrent_to_output_weight_scale *
-                                        activation_scale /
+                                        output_state_scale /
                                         intermediate_scale[3];
 
   effective_hidden_scale =
       std::pow(2, -15) / intermediate_scale[4] * std::pow(2, -15);
 
   effective_proj_scale =
-      proj_weight_scale * intermediate_scale[4] / activation_scale;
+      proj_weight_scale * intermediate_scale[4] / output_state_scale;
 
   if (use_peephole) {
     if (!use_cifg) {
@@ -419,11 +416,10 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   const TfLiteTensor* projection_bias =
       GetOptionalInputTensor(context, node, kProjectionBiasTensor);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
-  TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
+  TfLiteTensor* cell_state = GetVariableInput(context, node, kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   // Since we have already checked that weights are all there or none, we can
@@ -456,7 +452,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   int32_t* output_bias_ptr = nullptr;
   int32_t* proj_bias_ptr = nullptr;
   int16_t* cell_ptr = nullptr;
-  int8_t* activation_ptr = nullptr;
+  int8_t* output_state_ptr = nullptr;
 
   // Scales.
   const float default_scale = 1.0;
@@ -477,7 +473,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   float layer_norm_forget_scale = default_scale;
   float layer_norm_cell_scale = default_scale;
   float layer_norm_output_scale = default_scale;
-  float activation_scale = default_scale;
+  float output_state_scale = default_scale;
 
   // Effective scales.
   float effective_input_to_input_scale = default_scale;
@@ -495,7 +491,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
 
   // Zero points
   int input_zp = 0;
-  int activation_zp = 0;
+  int output_state_zp = 0;
 
   // Populate all the values.
   if (!use_cifg) {
@@ -537,7 +533,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
       proj_bias_ptr = projection_bias->data.i32;
     }
   }
-  activation_scale = activation_state->params.scale;
+  output_state_scale = output_state->params.scale;
 
   input_to_forget_weight_ptr = input_to_forget_weights->data.int8;
   input_to_forget_weight_scale = input_to_forget_weights->params.scale;
@@ -554,11 +550,11 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   forget_bias_ptr = forget_gate_bias->data.i32;
   cell_bias_ptr = cell_bias->data.i32;
   output_bias_ptr = output_gate_bias->data.i32;
-  activation_ptr = activation_state->data.int8;
+  output_state_ptr = output_state->data.int8;
   cell_ptr = cell_state->data.i16;
   input_scale = input->params.scale;
   input_zp = input->params.zero_point;
-  activation_zp = activation_state->params.zero_point;
+  output_state_zp = output_state->params.zero_point;
 
   std::vector intermediate_scale;
   for (int i = 0; i < 12; ++i) {
@@ -575,27 +571,28 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
     effective_input_to_input_scale =
         input_to_input_weight_scale * input_scale / intermediate_scale[1];
     effective_recurrent_to_input_scale = recurrent_to_input_weight_scale *
-                                         activation_scale /
+                                         output_state_scale /
                                          intermediate_scale[2];
   }
   effective_input_to_forget_scale =
       input_to_forget_weight_scale * input_scale / intermediate_scale[4];
   effective_recurrent_to_forget_scale = recurrent_to_forget_weight_scale *
-                                        activation_scale /
+                                        output_state_scale /
                                         intermediate_scale[5];
 
   effective_input_to_cell_scale =
       input_to_cell_weight_scale * input_scale / intermediate_scale[7];
-  effective_recurrent_to_cell_scale =
-      recurrent_to_cell_weight_scale * activation_scale / intermediate_scale[8];
+  effective_recurrent_to_cell_scale = recurrent_to_cell_weight_scale *
+                                      output_state_scale /
+                                      intermediate_scale[8];
 
   effective_input_to_output_scale =
       input_to_output_weight_scale * input_scale / intermediate_scale[10];
   effective_recurrent_to_output_scale = recurrent_to_output_weight_scale *
-                                        activation_scale /
+                                        output_state_scale /
                                         intermediate_scale[11];
   effective_proj_scale =
-      proj_weight_scale * std::pow(2, -15) / activation_scale;
+      proj_weight_scale * std::pow(2, -15) / output_state_scale;
 
   if (use_peephole) {
     if (!use_cifg) {
@@ -698,18 +695,16 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   const float cell_clip = params->cell_clip;
   const float proj_clip = params->proj_clip;
 
-  const TfLiteTensor* cell_tensor =
-      GetInput(context, node, kInputCellStateTensor);
   const TfLiteTensor* output_tensor = GetOutput(context, node, kOutputTensor);
 
-  auto* cell_params = reinterpret_cast(
-      cell_tensor->quantization.params);
+  auto* cell_state_params = reinterpret_cast(
+      cell_state->quantization.params);
   auto* proj_params = reinterpret_cast(
       output_tensor->quantization.params);
-  TF_LITE_ENSURE_EQ(context, cell_params->scale->data[0], 1.0 / 32768);
+  TF_LITE_ENSURE_EQ(context, cell_state_params->scale->data[0], 1.0 / 32768);
   if (cell_clip > 0.0 && cell_clip < 1.0) {
     integer_lstm_param->quantized_cell_clip =
-        static_cast(cell_clip / cell_params->scale->data[0]);
+        static_cast(cell_clip / cell_state_params->scale->data[0]);
   } else {
     integer_lstm_param->quantized_cell_clip = 0;
   }
@@ -1026,12 +1021,12 @@ TfLiteStatus PopulatePrecomputedZPTimesWeightsWithBias(TfLiteContext* context,
                                                        OpData* op_data,
                                                        TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  const TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
+  const TfLiteTensor* output_state =
+      GetVariableInput(context, node, kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
 
   const int32_t input_zero_point = -input->params.zero_point;
-  const int32_t activation_zero_point = -activation_state->params.zero_point;
+  const int32_t output_state_zero_point = -output_state->params.zero_point;
 
   const TfLiteTensor* input_to_input_weights =
       GetOptionalInputTensor(context, node, kInputToInputWeightsTensor);
@@ -1083,8 +1078,8 @@ TfLiteStatus PopulatePrecomputedZPTimesWeightsWithBias(TfLiteContext* context,
   TF_LITE_ENSURE_OK(
       context,
       PrecomputeZeroPointTimesWeightWithBias(
-          context, activation_zero_point, recurrent_to_forget_weights, nullptr,
-          &(integer_lstm_params->recurrent_to_forget_effective_bias)));
+          context, output_state_zero_point, recurrent_to_forget_weights,
+          nullptr, &(integer_lstm_params->recurrent_to_forget_effective_bias)));
 
   // Modulation gate.
   const TfLiteTensor* cell_gate_bias =
@@ -1097,7 +1092,7 @@ TfLiteStatus PopulatePrecomputedZPTimesWeightsWithBias(TfLiteContext* context,
   TF_LITE_ENSURE_OK(
       context,
       PrecomputeZeroPointTimesWeightWithBias(
-          context, activation_zero_point, recurrent_to_cell_weights, nullptr,
+          context, output_state_zero_point, recurrent_to_cell_weights, nullptr,
           &(integer_lstm_params->recurrent_to_cell_effective_bias)));
 
   // Output gate.
@@ -1112,8 +1107,8 @@ TfLiteStatus PopulatePrecomputedZPTimesWeightsWithBias(TfLiteContext* context,
   TF_LITE_ENSURE_OK(
       context,
       PrecomputeZeroPointTimesWeightWithBias(
-          context, activation_zero_point, recurrent_to_output_weights, nullptr,
-          &(integer_lstm_params->recurrent_to_output_effective_bias)));
+          context, output_state_zero_point, recurrent_to_output_weights,
+          nullptr, &(integer_lstm_params->recurrent_to_output_effective_bias)));
 
   // Input gate. The calculation is only meaningful for non-cifg case.
   const TfLiteTensor* input_gate_bias =
@@ -1126,7 +1121,7 @@ TfLiteStatus PopulatePrecomputedZPTimesWeightsWithBias(TfLiteContext* context,
   TF_LITE_ENSURE_OK(
       context,
       PrecomputeZeroPointTimesWeightWithBias(
-          context, activation_zero_point, recurrent_to_input_weights, nullptr,
+          context, output_state_zero_point, recurrent_to_input_weights, nullptr,
           &(integer_lstm_params->recurrent_to_input_effective_bias)));
 
   // Projection bias. The calculation is only meaningful for with projection.
@@ -1198,20 +1193,19 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       context, CheckInputTensorDimensions(context, node, n_input, n_output,
                                           n_cell, use_layer_norm, is_integer));
 
-  // Get the pointer to output, activation_state and cell_state tensors.
+  // Get the pointer to output, output_state and cell_state tensors.
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
-  TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
+  TfLiteTensor* cell_state = GetVariableInput(context, node, kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   // Check the shape of input state tensors.
   // These tensor may be 1D or 2D. It's fine as long as the total size is
   // correct.
-  TF_LITE_ENSURE_EQ(context, NumElements(activation_state), n_batch * n_output);
+  TF_LITE_ENSURE_EQ(context, NumElements(output_state), n_batch * n_output);
   TF_LITE_ENSURE_EQ(context, NumElements(cell_state), n_batch * n_cell);
 
   // Resize the output tensors.
@@ -1275,7 +1269,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   if (is_hybrid_op) {
     op_data->compute_row_sums = true;
     // Allocate temporary tensors to store quantized values of input,
-    // activation_state and cell_state tensors.
+    // output_state and cell_state tensors.
     node->temporaries->data[1] = op_data->scratch_tensor_index + 1;
     TfLiteTensor* input_quantized = GetTemporary(context, node, /*index=*/1);
     input_quantized->type = input_to_output_weights->type;
@@ -1286,17 +1280,17 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                                                        input_quantized_size));
     }
     node->temporaries->data[2] = op_data->scratch_tensor_index + 2;
-    TfLiteTensor* activation_state_quantized =
+    TfLiteTensor* output_state_quantized =
         GetTemporary(context, node, /*index=*/2);
-    activation_state_quantized->type = input_to_output_weights->type;
-    activation_state_quantized->allocation_type = kTfLiteArenaRw;
-    if (!TfLiteIntArrayEqual(activation_state_quantized->dims,
-                             activation_state->dims)) {
-      TfLiteIntArray* activation_state_quantized_size =
-          TfLiteIntArrayCopy(activation_state->dims);
-      TF_LITE_ENSURE_OK(
-          context, context->ResizeTensor(context, activation_state_quantized,
-                                         activation_state_quantized_size));
+    output_state_quantized->type = input_to_output_weights->type;
+    output_state_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(output_state_quantized->dims,
+                             output_state->dims)) {
+      TfLiteIntArray* output_state_quantized_size =
+          TfLiteIntArrayCopy(output_state->dims);
+      TF_LITE_ENSURE_OK(context,
+                        context->ResizeTensor(context, output_state_quantized,
+                                              output_state_quantized_size));
     }
     node->temporaries->data[3] = op_data->scratch_tensor_index + 3;
     TfLiteTensor* cell_state_quantized =
@@ -1540,11 +1534,10 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* projection_bias =
       GetOptionalInputTensor(context, node, kProjectionBiasTensor);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
-  TfLiteTensor* cell_state =
-      GetVariableInput(context, node, kInputCellStateTensor);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
+  TfLiteTensor* cell_state = GetVariableInput(context, node, kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
@@ -1569,7 +1562,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
           projection_bias, params, /*forward_sequence=*/true,
           /*time_major=*/true,
-          /*output_offset=*/0, scratch_buffer, activation_state, cell_state,
+          /*output_offset=*/0, scratch_buffer, output_state, cell_state,
           output);
     }
     case kTfLiteUInt8:
@@ -1580,7 +1573,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
         TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
         TfLiteTensor* input_quantized =
             GetTemporary(context, node, /*index=*/1);
-        TfLiteTensor* activation_state_quantized =
+        TfLiteTensor* output_state_quantized =
             GetTemporary(context, node, /*index=*/2);
         TfLiteTensor* cell_state_quantized =
             GetTemporary(context, node, /*index=*/3);
@@ -1614,8 +1607,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
             /*time_major=*/true, /*output_offset=*/0, scratch_buffer,
             scaling_factors, prod_scaling_factors, recovered_cell_weights,
             input_quantized,
-            /*aux_input_quantized=*/nullptr, activation_state_quantized,
-            cell_state_quantized, activation_state, cell_state,
+            /*aux_input_quantized=*/nullptr, output_state_quantized,
+            cell_state_quantized, output_state, cell_state,
             output_scratch_buffer, output, zero_points, row_sums, row_sums_size,
             &op_data->compute_row_sums,
             CpuBackendContext::GetFromContext(context));
@@ -1638,9 +1631,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               forget_layer_norm_coefficients, cell_layer_norm_coefficients,
               output_layer_norm_coefficients, input_gate_bias, forget_gate_bias,
               cell_bias, output_gate_bias, projection_weights, projection_bias,
-              params, &op_data->integer_lstm_param, activation_state,
-              cell_state, output, scratch0, scratch1, scratch2, scratch3,
-              scratch4, scratch5, CpuBackendContext::GetFromContext(context));
+              params, &op_data->integer_lstm_param, output_state, cell_state,
+              output, scratch0, scratch1, scratch2, scratch3, scratch4,
+              scratch5, CpuBackendContext::GetFromContext(context));
         } else {
           TfLiteTensor* scratch0 = GetTemporary(context, node, /*index=*/0);
           TfLiteTensor* scratch1 = GetTemporary(context, node, /*index=*/1);
@@ -1660,7 +1653,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               forget_layer_norm_coefficients, cell_layer_norm_coefficients,
               output_layer_norm_coefficients, input_gate_bias, forget_gate_bias,
               cell_bias, output_gate_bias, projection_weights, projection_bias,
-              params, activation_state, cell_state, output,
+              params, output_state, cell_state, output,
               &op_data->integer_lstm_param, scratch0, scratch1, scratch2,
               scratch3, scratch4, scratch5, scratch6, scratch7);
           return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index b4d43414d89..65f68b34251 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -900,7 +900,7 @@ inline void LstmStepHybrid(
 
 // Fully quantized lstm kernel for 16 bit gate matmul output.
 //
-// Input activation of size n_batch * n_input:
+// Input tensor of size n_batch * n_input:
 //   input_ptr
 //
 // LSTM weights:
@@ -972,7 +972,7 @@ inline void LstmStepHybrid(
 //   cell_scale: the power of two scale for cell state.
 //
 // Zero points:
-//   activation_zp: zero point of activation
+//   output_state_zp: zero point of output state
 //   hidden_zp: zero point for hidden state.
 //
 // Temporary pre-allocated storage for the calculation. Each is of size n_cell *
@@ -1048,8 +1048,8 @@ inline void LstmStepInteger(
     const int32_t* input_to_input_effective_bias,
     const int32_t* recurrent_to_input_effective_bias,
     const int32_t* projection_effective_bias, int32 n_batch, int32 n_cell,
-    int32 n_input, int32 n_output, int8_t* activation_ptr,
-    int32_t activation_zp, int16_t* cell_ptr, int8_t* output_ptr,
+    int32 n_input, int32 n_output, int8_t* output_state_ptr,
+    int32_t output_state_zp, int16_t* cell_ptr, int8_t* output_ptr,
     int16_t* scratch_0_ptr, int16_t* scratch_1_ptr, int16_t* scratch_2_ptr,
     int16_t* scratch_3_ptr, int8_t* scratch_4_ptr, int32_t* scratch_5_ptr,
     CpuBackendContext* context) {
@@ -1088,7 +1088,7 @@ inline void LstmStepInteger(
       n_batch, n_input, n_cell, 0, scratch_5_ptr, scratch_1_ptr, context);
 
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-      activation_ptr, recurrent_to_forget_effective_bias,
+      output_state_ptr, recurrent_to_forget_effective_bias,
       recurrent_to_forget_weight_ptr, effective_recurrent_to_forget_scale_a,
       effective_recurrent_to_forget_scale_b, n_batch, n_output, n_cell, 0,
       scratch_5_ptr, scratch_1_ptr, context);
@@ -1115,7 +1115,7 @@ inline void LstmStepInteger(
       n_input, n_cell, 0, scratch_5_ptr, scratch_2_ptr, context);
 
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-      activation_ptr, recurrent_to_cell_effective_bias,
+      output_state_ptr, recurrent_to_cell_effective_bias,
       recurrent_to_cell_weight_ptr, effective_recurrent_to_cell_scale_a,
       effective_recurrent_to_cell_scale_b, n_batch, n_output, n_cell, 0,
       scratch_5_ptr, scratch_2_ptr, context);
@@ -1139,7 +1139,7 @@ inline void LstmStepInteger(
         n_batch, n_input, n_cell, 0, scratch_5_ptr, scratch_0_ptr, context);
 
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-        activation_ptr, recurrent_to_input_effective_bias,
+        output_state_ptr, recurrent_to_input_effective_bias,
         recurrent_to_input_weight_ptr, effective_recurrent_to_input_scale_a,
         effective_recurrent_to_input_scale_b, n_batch, n_output, n_cell, 0,
         scratch_5_ptr, scratch_0_ptr, context);
@@ -1180,7 +1180,7 @@ inline void LstmStepInteger(
       n_batch, n_input, n_cell, 0, scratch_5_ptr, scratch_3_ptr, context);
 
   tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-      activation_ptr, recurrent_to_output_effective_bias,
+      output_state_ptr, recurrent_to_output_effective_bias,
       recurrent_to_output_weight_ptr, effective_recurrent_to_output_scale_a,
       effective_recurrent_to_output_scale_b, n_batch, n_output, n_cell, 0,
       scratch_5_ptr, scratch_3_ptr, context);
@@ -1213,7 +1213,7 @@ inline void LstmStepInteger(
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
         scratch_4_ptr, projection_effective_bias, proj_weight_ptr,
         effective_proj_scale_a, effective_proj_scale_b, n_batch, n_cell,
-        n_output, activation_zp, scratch_5_ptr, output_ptr, context);
+        n_output, output_state_zp, scratch_5_ptr, output_ptr, context);
     if (quantized_proj_clip > 0) {
       tensor_utils::CwiseClipping(output_ptr, quantized_proj_clip, n_batch,
                                   n_output);
@@ -1221,12 +1221,12 @@ inline void LstmStepInteger(
   } else {
     std::copy_n(scratch_4_ptr, n_batch * n_output, output_ptr);
   }
-  std::copy_n(output_ptr, n_batch * n_output, activation_ptr);
+  std::copy_n(output_ptr, n_batch * n_output, output_state_ptr);
 }
 
 // Fully quantized lstm kernel for 8 bit gate matmul output.
 //
-// Input activation of size n_batch * n_input:
+// Input tensor of size n_batch * n_input:
 //   input_ptr
 //
 // LSTM weights:
@@ -1298,7 +1298,7 @@ inline void LstmStepInteger(
 //   cell_scale: the power of two scale for cell state.
 //
 // Zero points:
-//   activation_zp: zero point of activation
+//   output_state_zp: zero point of output state.
 //   hidden_zp: zero point for hidden state.
 //
 // Temporary pre-allocated storage for the calculation. Each is of size n_cell *
@@ -1367,8 +1367,8 @@ void LstmStepInteger(
     const int32_t* intermediate_scale_a, const int32_t* intermediate_scale_b,
     const int32_t* intermediate_zp, int32 quantized_cell_clip,
     int32 quantized_proj_clip, int32 n_batch, int32 n_cell, int32 n_input,
-    int32 n_output, int32 output_batch_leading_dim, int8_t* activation_ptr,
-    int32_t activation_zp, int16_t* cell_ptr, int8_t* output_ptr,
+    int32 n_output, int32 output_batch_leading_dim, int8_t* output_state_ptr,
+    int32_t output_state_zp, int16_t* cell_ptr, int8_t* output_ptr,
     int8_t* scratch0, int8_t* scratch1, int16_t* scratch2, int16_t* scratch3,
     int16_t* scratch4, int16_t* scratch5, int16_t* scratch6,
     int16_t* scratch7) {
@@ -1381,7 +1381,7 @@ void LstmStepInteger(
       n_batch, n_input, n_cell, scratch0, intermediate_zp[4]);
 
   tensor_utils::MatrixBatchVectorMultiply(
-      activation_ptr, activation_zp, recurrent_to_forget_weight_ptr,
+      output_state_ptr, output_state_zp, recurrent_to_forget_weight_ptr,
       effective_recurrent_to_forget_scale_a,
       effective_recurrent_to_forget_scale_b, n_batch, n_output, n_cell,
       scratch1, intermediate_zp[5]);
@@ -1408,7 +1408,7 @@ void LstmStepInteger(
       n_input, n_cell, scratch0, intermediate_zp[7]);
 
   tensor_utils::MatrixBatchVectorMultiply(
-      activation_ptr, activation_zp, recurrent_to_cell_weight_ptr,
+      output_state_ptr, output_state_zp, recurrent_to_cell_weight_ptr,
       effective_recurrent_to_cell_scale_a, effective_recurrent_to_cell_scale_b,
       n_batch, n_output, n_cell, scratch1, intermediate_zp[8]);
 
@@ -1434,7 +1434,7 @@ void LstmStepInteger(
       n_batch, n_input, n_cell, scratch0, intermediate_zp[10]);
 
   tensor_utils::MatrixBatchVectorMultiply(
-      activation_ptr, activation_zp, recurrent_to_output_weight_ptr,
+      output_state_ptr, output_state_zp, recurrent_to_output_weight_ptr,
       effective_recurrent_to_output_scale_a,
       effective_recurrent_to_output_scale_b, n_batch, n_output, n_cell,
       scratch1, intermediate_zp[11]);
@@ -1478,7 +1478,7 @@ void LstmStepInteger(
   // Projection.
   tensor_utils::MatrixBatchVectorMultiply(
       scratch3, proj_weight_ptr, effective_proj_scale_a, effective_proj_scale_b,
-      proj_bias_ptr, n_batch, n_cell, n_output, activation_zp, output_ptr);
+      proj_bias_ptr, n_batch, n_cell, n_output, output_state_zp, output_ptr);
 
   // Projection clipping.
   if (quantized_proj_clip > 0) {
@@ -1486,8 +1486,8 @@ void LstmStepInteger(
                                 n_output);
   }
 
-  // Copy output to activation.
-  std::copy_n(output_ptr, n_batch * n_output, activation_ptr);
+  // Copy output to output state.
+  std::copy_n(output_ptr, n_batch * n_output, output_state_ptr);
 }
 
 }  // namespace
@@ -1518,9 +1518,8 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
-    int output_offset, TfLiteTensor* scratch_buffer,
-    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
-    TfLiteTensor* output) {
+    int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
+    TfLiteTensor* cell_state, TfLiteTensor* output) {
   TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3);
   int max_time, n_batch;
   if (input->dims->size == 3) {
@@ -1604,10 +1603,9 @@ TfLiteStatus EvalFloat(
           GetTensorData(projection_weights),
           GetTensorData(projection_bias), params, n_batch, n_cell,
           n_input, aux_input_size, n_output, output_batch_leading_dim,
-          GetTensorData(activation_state),
-          GetTensorData(cell_state), input_gate_scratch,
-          forget_gate_scratch, cell_gate_scratch, output_gate_scratch,
-          output_ptr);
+          GetTensorData(output_state), GetTensorData(cell_state),
+          input_gate_scratch, forget_gate_scratch, cell_gate_scratch,
+          output_gate_scratch, output_ptr);
     }
   } else {
     for (int b = 0; b < n_batch; b++) {
@@ -1628,9 +1626,9 @@ TfLiteStatus EvalFloat(
         float* output_ptr = GetTensorData(output) +
                             time_offset * output_step + output_offset;
 
-        // Offset the {activation,cell}_state pointers to the right batch.
-        float* activation_state_ptr = GetTensorData(activation_state) +
-                                      b * output_batch_leading_dim;
+        // Offset the {output,cell}_state pointers to the right batch.
+        float* output_state_ptr =
+            GetTensorData(output_state) + b * output_batch_leading_dim;
         float* cell_state_ptr = GetTensorData(cell_state) + b * n_cell;
         // Offset the scratch pointers to the right batch.
         float* input_gate_scratch_ptr =
@@ -1666,7 +1664,7 @@ TfLiteStatus EvalFloat(
             GetTensorData(projection_weights),
             GetTensorData(projection_bias), params, /*n_batch=*/1,
             n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
-            activation_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
+            output_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
             forget_gate_scratch_ptr, cell_gate_scratch_ptr,
             output_gate_scratch_ptr, output_ptr);
       }
@@ -1939,10 +1937,10 @@ TfLiteStatus EvalInteger8x8_16(
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
-    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
-    TfLiteTensor* output, TfLiteTensor* scratch0, TfLiteTensor* scratch1,
-    TfLiteTensor* scratch2, TfLiteTensor* scratch3, TfLiteTensor* scratch4,
-    TfLiteTensor* scratch5, CpuBackendContext* context) {
+    TfLiteTensor* output_state, TfLiteTensor* cell_state, TfLiteTensor* output,
+    TfLiteTensor* scratch0, TfLiteTensor* scratch1, TfLiteTensor* scratch2,
+    TfLiteTensor* scratch3, TfLiteTensor* scratch4, TfLiteTensor* scratch5,
+    CpuBackendContext* context) {
   TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3);
   const int n_input = input->dims->data[input->dims->size - 1];
   int max_time, n_batch;
@@ -1959,7 +1957,7 @@ TfLiteStatus EvalInteger8x8_16(
   const int n_output = recurrent_to_output_weights->dims->data[1];
 
   // Activation zero point
-  int activation_zp = activation_state->params.zero_point;
+  int output_state_zp = output_state->params.zero_point;
 
   // Get params for time/batch/sequence.
   const int output_batch_leading_dim =
@@ -2042,8 +2040,8 @@ TfLiteStatus EvalInteger8x8_16(
         integer_lstm_param->input_to_input_effective_bias.get(),
         integer_lstm_param->recurrent_to_input_effective_bias.get(),
         integer_lstm_param->projection_effective_bias.get(), n_batch, n_cell,
-        n_input, n_output, GetTensorData(activation_state),
-        activation_zp, GetTensorData(cell_state), output_ptr,
+        n_input, n_output, GetTensorData(output_state), output_state_zp,
+        GetTensorData(cell_state), output_ptr,
         GetTensorData(scratch0), GetTensorData(scratch1),
         GetTensorData(scratch2), GetTensorData(scratch3),
         GetTensorData(scratch4), GetTensorData(scratch5),
@@ -2072,7 +2070,7 @@ TfLiteStatus EvalInteger8x8_8(
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
     const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
-    const TfLiteLSTMParams* params, TfLiteTensor* activation_state,
+    const TfLiteLSTMParams* params, TfLiteTensor* output_state,
     TfLiteTensor* cell_state, TfLiteTensor* output,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
     TfLiteTensor* scratch0, TfLiteTensor* scratch1, TfLiteTensor* scratch2,
@@ -2131,11 +2129,11 @@ TfLiteStatus EvalInteger8x8_8(
   const int32_t* output_bias_ptr = GetTensorData(output_gate_bias);
   const int32_t* proj_bias_ptr = GetTensorData(projection_bias);
   int16_t* cell_ptr = GetTensorData(cell_state);
-  int8_t* activation_ptr = GetTensorData(activation_state);
+  int8_t* output_state_ptr = GetTensorData(output_state);
   int8_t* output_ptr = nullptr;
 
   const int32 input_zp = input->params.zero_point;
-  const int32 activation_zp = activation_state->params.zero_point;
+  const int32 output_state_zp = output_state->params.zero_point;
 
   // Get params for time/batch/sequence.
   const int output_batch_leading_dim =
@@ -2222,7 +2220,7 @@ TfLiteStatus EvalInteger8x8_8(
         integer_lstm_param->intermediate_zp,
         integer_lstm_param->quantized_cell_clip,
         integer_lstm_param->quantized_proj_clip, n_batch, n_cell, n_input,
-        n_output, output_batch_leading_dim, activation_ptr, activation_zp,
+        n_output, output_batch_leading_dim, output_state_ptr, output_state_zp,
         cell_ptr, output_ptr, GetTensorData(scratch0),
         GetTensorData(scratch1), GetTensorData(scratch2),
         GetTensorData(scratch3), GetTensorData(scratch4),
diff --git a/tensorflow/lite/kernels/lstm_eval.h b/tensorflow/lite/kernels/lstm_eval.h
index 91f47b18df6..3c9b4bccf42 100644
--- a/tensorflow/lite/kernels/lstm_eval.h
+++ b/tensorflow/lite/kernels/lstm_eval.h
@@ -120,9 +120,8 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
-    int output_offset, TfLiteTensor* scratch_buffer,
-    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
-    TfLiteTensor* output);
+    int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
+    TfLiteTensor* cell_state, TfLiteTensor* output);
 
 TfLiteStatus EvalHybrid(
     const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights,
@@ -179,10 +178,10 @@ TfLiteStatus EvalInteger8x8_16(
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
-    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
-    TfLiteTensor* output, TfLiteTensor* scratch0, TfLiteTensor* scratch1,
-    TfLiteTensor* scratch2, TfLiteTensor* scratch3, TfLiteTensor* scratch4,
-    TfLiteTensor* scratch5, CpuBackendContext* context);
+    TfLiteTensor* output_state, TfLiteTensor* cell_state, TfLiteTensor* output,
+    TfLiteTensor* scratch0, TfLiteTensor* scratch1, TfLiteTensor* scratch2,
+    TfLiteTensor* scratch3, TfLiteTensor* scratch4, TfLiteTensor* scratch5,
+    CpuBackendContext* context);
 
 TfLiteStatus EvalInteger8x8_8(
     const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights,
@@ -203,7 +202,7 @@ TfLiteStatus EvalInteger8x8_8(
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
     const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
-    const TfLiteLSTMParams* params, TfLiteTensor* activation_state,
+    const TfLiteLSTMParams* params, TfLiteTensor* output_state,
     TfLiteTensor* cell_state, TfLiteTensor* output,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
     TfLiteTensor* scratch0, TfLiteTensor* scratch1, TfLiteTensor* scratch2,
diff --git a/tensorflow/lite/kernels/lstm_shared.h b/tensorflow/lite/kernels/lstm_shared.h
index 9e29650a3d8..0907be9094b 100644
--- a/tensorflow/lite/kernels/lstm_shared.h
+++ b/tensorflow/lite/kernels/lstm_shared.h
@@ -57,8 +57,8 @@ constexpr int kProjectionBiasTensor = 17;  // Optional
 
 // These state tensors are defined as variable tensors, and will be modified by
 // this op.
-constexpr int kInputActivationStateTensor = 18;
-constexpr int kInputCellStateTensor = 19;
+constexpr int kOutputStateTensor = 18;
+constexpr int kCellStateTensor = 19;
 
 // Layer norm coefficient tensors of size {n_cell}, representing a diagonal
 // matrix.
diff --git a/tensorflow/lite/kernels/lstm_test.cc b/tensorflow/lite/kernels/lstm_test.cc
index ba5ee6508cc..f8594f9adf0 100644
--- a/tensorflow/lite/kernels/lstm_test.cc
+++ b/tensorflow/lite/kernels/lstm_test.cc
@@ -104,10 +104,10 @@ class LSTMOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    // Adding the 2 input state tensors.
-    input_activation_state_ =
+    // Adding the 2 state tensors.
+    output_state_ =
         AddInput(TensorData{TensorType_FLOAT32, {n_batch_, n_output_}}, true);
-    input_cell_state_ =
+    cell_state_ =
         AddInput(TensorData{TensorType_FLOAT32, {n_batch_, n_cell_}}, true);
 
     // Layer norm weights.
@@ -266,13 +266,11 @@ class LSTMOpModel : public SingleOpModel {
 
   int projection_weights_;
   int projection_bias_;
-  int input_activation_state_;
-  int input_cell_state_;
-
-  int output_;
   int output_state_;
   int cell_state_;
 
+  int output_;
+
   int n_batch_;
   int n_input_;
   int n_cell_;
@@ -553,7 +551,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingOmittedLayerNormLstmTest,
                        {0, 0},  // projection_weight tensor
                        {0},     // projection_bias tensor
 
-                       {n_batch, n_output},  // activation_state tensor
+                       {n_batch, n_output},  // output_state tensor
                        {n_batch, n_cell},    // cell_state tensor
 
                        {0},  // input_layer_norm_coefficient tensor
@@ -1697,7 +1695,7 @@ TEST_F(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {n_cell},  // input_layer_norm_coefficient tensor
@@ -1768,7 +1766,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {n_cell},  // input_layer_norm_coefficient tensor
@@ -1841,7 +1839,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLayerNormLstmInt8Test,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {n_cell},  // input_layer_norm_coefficient tensor
@@ -1955,7 +1953,7 @@ TEST_F(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {0},       // input_layer_norm_coefficient tensor
@@ -2026,7 +2024,7 @@ TEST_P(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {0},       // input_layer_norm_coefficient tensor
@@ -2098,7 +2096,7 @@ TEST_P(CifgPeepholeProjectionNoClippingLayerNormLstmInt8Test,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {0},       // input_layer_norm_coefficient tensor
@@ -2216,13 +2214,13 @@ class LSTMIntegerOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    // Adding the 2 input state tensors.
-    input_activation_state_ = AddInput({TensorType_INT16, input_shapes[18],
-                                        ranges[18].first, ranges[18].second},
-                                       true);
-    input_cell_state_ = AddInput({TensorType_INT16, input_shapes[19],
-                                  ranges[19].first, ranges[19].second},
-                                 true);
+    // Adding the 2 state tensors.
+    output_state_ = AddInput({TensorType_INT16, input_shapes[18],
+                              ranges[18].first, ranges[18].second},
+                             true);
+    cell_state_ = AddInput({TensorType_INT16, input_shapes[19],
+                            ranges[19].first, ranges[19].second},
+                           true);
 
     // Layer norm weights.
     if (use_layer_norm) {
@@ -2386,8 +2384,6 @@ class LSTMIntegerOpModel : public SingleOpModel {
 
   int projection_weights_;
   int projection_bias_;
-  int input_activation_state_;
-  int input_cell_state_;
 
   int intermediates_[5];
 
@@ -2483,7 +2479,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionNoPeephole) {
       {n_output, n_cell},  // projection_weight tensor
       {0},                 // projection_bias tensor
 
-      {n_batch, n_output},  // activation_state tensor
+      {n_batch, n_output},  // output_state tensor
       {n_batch, n_cell},    // cell_state tensor
 
       {n_cell},  // input_layer_norm_coefficient tensor
@@ -2517,14 +2513,14 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionNoPeephole) {
       {-0.5, 0.5},  // projection_weight tensor
       {-1, 1},      // projection_bias tensor
 
-      {-1.0, 32767.0 / 32768},  // activation_state tensor
+      {-1.0, 32767.0 / 32768},  // output_state tensor
       {-1, 1},                  // cell_state tensor
 
       {-1.00001, 1.0},  // input_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // forget_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // cell_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // output_layer_norm_coefficient tensor
-      // Output scale is the same as input activation scale and only activation
+      // Output scale is the same as output_state scale and only output_state
       // scale is used in the op, so this is only provided for clarity.
       {-1.0, 32767.0 / 32768},  // output tensor.
   };
@@ -2685,7 +2681,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionYesPeephole) {
       {n_output, n_cell},  // projection_weight tensor
       {0},                 // projection_bias tensor
 
-      {n_batch, n_output},  // activation_state tensor
+      {n_batch, n_output},  // output_state tensor
       {n_batch, n_cell},    // cell_state tensor
 
       {n_cell},  // input_layer_norm_coefficient tensor
@@ -2719,14 +2715,14 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionYesPeephole) {
       {-0.5, 0.5},  // projection_weight tensor
       {-1, 1},      // projection_bias tensor
 
-      {-1.0, 32767.0 / 32768},  // activation_state tensor
+      {-1.0, 32767.0 / 32768},  // output_state tensor
       {-1, 1},                  // cell_state tensor
 
       {-0.5, 0.5},  // input_layer_norm_coefficient tensor
       {-0.5, 0.5},  // forget_layer_norm_coefficient tensor
       {-1.0, 1.0},  // cell_layer_norm_coefficient tensor
       {-1.0, 1.0},  // output_layer_norm_coefficient tensor
-      // Output scale is the same as input activation scale and only activation
+      // Output scale is the same as output_state scale and only output_state
       // scale is used in the op, so this is only provided for clarity.
       {-1.0, 32767.0 / 32768},  // output tensor.
   };
@@ -2892,13 +2888,13 @@ class LSTMIntegerOpModel8x8_8 : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    // Adding the 2 input state tensors.
-    input_activation_state_ = AddInput({TensorType_INT16, input_shapes[18],
-                                        ranges[18].first, ranges[18].second},
-                                       true);
-    input_cell_state_ = AddInput({TensorType_INT16, input_shapes[19],
-                                  ranges[19].first, ranges[19].second},
-                                 true);
+    // Adding the 2 state tensors.
+    output_state_ = AddInput({TensorType_INT16, input_shapes[18],
+                              ranges[18].first, ranges[18].second},
+                             true);
+    cell_state_ = AddInput({TensorType_INT16, input_shapes[19],
+                            ranges[19].first, ranges[19].second},
+                           true);
 
     // Layer norm weights.
     if (use_layer_norm) {
@@ -3062,8 +3058,6 @@ class LSTMIntegerOpModel8x8_8 : public SingleOpModel {
 
   int projection_weights_;
   int projection_bias_;
-  int input_activation_state_;
-  int input_cell_state_;
 
   int intermediates_[12];
 
@@ -3160,7 +3154,7 @@ TEST(LSTMIntegerOpModel8x8_8, CifgYesLayerNormNoYesProjectionNoPeephole) {
       {n_output, n_cell},  // projection_weight tensor
       {n_output},          // projection_bias tensor
 
-      {n_batch, n_output},  // activation_state tensor
+      {n_batch, n_output},  // output_state tensor
       {n_batch, n_cell},    // cell_state tensor
 
       {0},       // input_layer_norm_coefficient tensor
@@ -3194,14 +3188,14 @@ TEST(LSTMIntegerOpModel8x8_8, CifgYesLayerNormNoYesProjectionNoPeephole) {
       {-0.5, 0.5},  // projection_weight tensor
       {-1, 1},      // projection_bias tensor
 
-      {-1.0, 32767.0 / 32768},  // activation_state tensor
+      {-1.0, 32767.0 / 32768},  // output_state tensor
       {-1.0, 32767.0 / 32768},  // cell_state tensor
 
       {-1.00001, 1.0},  // input_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // forget_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // cell_layer_norm_coefficient tensor
       {-1.00001, 1.0},  // output_layer_norm_coefficient tensor
-      // Output scale is the same as input activation scale and only activation
+      // Output scale is the same as output_state scale and only output_state
       // scale is used in the op, so this is only provided for clarity.
       {-1.0, 32767.0 / 32768},  // output tensor.
   };
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
index 95864196f18..f1c0f9d42a6 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
@@ -317,20 +317,20 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                     CheckInputTensorDimensions(context, node, n_input, n_output,
                                                n_cell, is_layer_norm_lstm));
 
-  // Get the pointer to output, activation_state and cell_state buffer tensors.
+  // Get the pointer to output, output_state and cell_state buffer tensors.
   TfLiteTensor* output = GetOutput(context, node, lstm::full::kOutputTensor);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, lstm::full::kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, lstm::full::kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
   TfLiteTensor* cell_state =
-      GetVariableInput(context, node, lstm::full::kInputCellStateTensor);
+      GetVariableInput(context, node, lstm::full::kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   // Check the shape of input state tensors.
   // These tensor may be 1D or 2D. It's fine as long as the total size is
   // correct.
-  TF_LITE_ENSURE_EQ(context, NumElements(activation_state), n_batch * n_output);
+  TF_LITE_ENSURE_EQ(context, NumElements(output_state), n_batch * n_output);
   TF_LITE_ENSURE_EQ(context, NumElements(cell_state), n_batch * n_cell);
 
   // Resize the output tensors.
@@ -370,7 +370,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   if (IsHybridOp(input, input_to_output_weights)) {
     op_data->compute_row_sums = true;
     // Allocate temporary tensors to store quantized values of input,
-    // activation_state and cell_state tensors.
+    // output_state and cell_state tensors.
     node->temporaries->data[kInputQuantized] =
         scratch_tensor_index + kInputQuantized;
     TfLiteTensor* input_quantized =
@@ -384,17 +384,17 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     }
     node->temporaries->data[kOutputStateQuantized] =
         scratch_tensor_index + kOutputStateQuantized;
-    TfLiteTensor* activation_state_quantized =
+    TfLiteTensor* output_state_quantized =
         GetTemporary(context, node, kOutputStateQuantized);
-    activation_state_quantized->type = input_to_output_weights->type;
-    activation_state_quantized->allocation_type = kTfLiteArenaRw;
-    if (!TfLiteIntArrayEqual(activation_state_quantized->dims,
-                             activation_state->dims)) {
-      TfLiteIntArray* activation_state_quantized_size =
-          TfLiteIntArrayCopy(activation_state->dims);
-      TF_LITE_ENSURE_OK(
-          context, context->ResizeTensor(context, activation_state_quantized,
-                                         activation_state_quantized_size));
+    output_state_quantized->type = input_to_output_weights->type;
+    output_state_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(output_state_quantized->dims,
+                             output_state->dims)) {
+      TfLiteIntArray* output_state_quantized_size =
+          TfLiteIntArrayCopy(output_state->dims);
+      TF_LITE_ENSURE_OK(context,
+                        context->ResizeTensor(context, output_state_quantized,
+                                              output_state_quantized_size));
     }
     node->temporaries->data[kCellStateQuantized] =
         scratch_tensor_index + kCellStateQuantized;
@@ -559,11 +559,11 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   // Index the scratch buffers pointers to the global scratch buffer.
   TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
 
-  TfLiteTensor* activation_state =
-      GetVariableInput(context, node, lstm::full::kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
+  TfLiteTensor* output_state =
+      GetVariableInput(context, node, lstm::full::kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
   TfLiteTensor* cell_state =
-      GetVariableInput(context, node, lstm::full::kInputCellStateTensor);
+      GetVariableInput(context, node, lstm::full::kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   const TfLiteTensor* input_layer_norm_coefficients =
@@ -613,14 +613,14 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
           forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
           projection_bias, &lstm_params, /*forward_sequence=*/true, time_major,
-          /*output_offset=*/0, scratch_buffer, activation_state, cell_state,
+          /*output_offset=*/0, scratch_buffer, output_state, cell_state,
           output);
     }
     case kTfLiteUInt8:
     case kTfLiteInt8: {
       OpData* op_data = reinterpret_cast(node->user_data);
       TfLiteTensor* input_quantized = GetTemporary(context, node, /*index=*/1);
-      TfLiteTensor* activation_state_quantized =
+      TfLiteTensor* output_state_quantized =
           GetTemporary(context, node, /*index=*/2);
       TfLiteTensor* cell_state_quantized =
           GetTemporary(context, node, /*index=*/3);
@@ -652,10 +652,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           projection_bias, &lstm_params, /*forward_sequence=*/true, time_major,
           /*output_offset=*/0, scratch_buffer, scaling_factors,
           prod_scaling_factors, recovered_cell_weights, input_quantized,
-          /*aux_input_quantized=*/nullptr, activation_state_quantized,
-          cell_state_quantized, activation_state, cell_state, accum_scratch,
-          output, zero_points, row_sums, row_sums_size,
-          &op_data->compute_row_sums,
+          /*aux_input_quantized=*/nullptr, output_state_quantized,
+          cell_state_quantized, output_state, cell_state, accum_scratch, output,
+          zero_points, row_sums, row_sums_size, &op_data->compute_row_sums,
           CpuBackendContext::GetFromContext(context));
     }
     default:
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
index 43cc75f894b..ec20d76ae2e 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
@@ -100,13 +100,12 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    // Adding the 2 input state tensors.
-    input_activation_state_ =
+    // Adding the 2 state tensors.
+    output_state_ =
         AddInput(TensorData{TensorType_FLOAT32, {n_output_ * n_batch_}},
                  /*is_variable=*/true);
-    input_cell_state_ =
-        AddInput(TensorData{TensorType_FLOAT32, {n_cell_ * n_batch_}},
-                 /*is_variable=*/true);
+    cell_state_ = AddInput(TensorData{TensorType_FLOAT32, {n_cell_ * n_batch_}},
+                           /*is_variable=*/true);
 
     // Layer norm weights.
     if (is_layer_norm) {
@@ -256,8 +255,8 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
   int projection_weights_;
   int projection_bias_;
 
-  int input_activation_state_;
-  int input_cell_state_;
+  int output_state_;
+  int cell_state_;
 
   int input_layer_norm_coefficients_;
   int forget_layer_norm_coefficients_;
@@ -537,7 +536,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       });
 
@@ -599,7 +598,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       });
 
@@ -665,7 +664,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_UINT8, GetParam());
@@ -728,7 +727,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_INT8, GetParam());
@@ -840,7 +839,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       });
 
@@ -901,7 +900,7 @@ TEST_P(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_UINT8, GetParam());
@@ -964,7 +963,7 @@ TEST_P(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_INT8, GetParam());
@@ -1626,7 +1625,7 @@ TEST_F(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       });
 
@@ -1695,7 +1694,7 @@ TEST_P(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_UINT8, GetParam());
@@ -1766,7 +1765,7 @@ TEST_P(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {0},                 // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       },
       TensorType_INT8, GetParam());
@@ -2437,7 +2436,7 @@ TEST_F(NoCifgPeepholeProjectionAndBiasClippingUnidirectionalLstmTest,
           {n_output, n_cell},  // projection_weight tensor
           {n_output},          // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
       });
 
@@ -2643,7 +2642,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingLayerNormUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {0},       // input_layer_norm_coefficient tensor
@@ -2714,7 +2713,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
           {0, 0},  // projection_weight tensor
           {0},     // projection_bias tensor
 
-          {n_batch, n_output},  // activation_state tensor
+          {n_batch, n_output},  // output_state tensor
           {n_batch, n_cell},    // cell_state tensor
 
           {0},  // input_layer_norm_coefficient tensor
diff --git a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
index 0d4c614511d..88ea7c1d591 100644
--- a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
+++ b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
@@ -302,9 +302,8 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
-    int output_offset, TfLiteTensor* scratch_buffer,
-    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
-    TfLiteTensor* output, Logger* logger,
+    int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
+    TfLiteTensor* cell_state, TfLiteTensor* output, Logger* logger,
     const std::vector& intermediate_tensor_indexes,
     ErrorReporter* error_reporter) {
   TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3);
@@ -390,10 +389,10 @@ TfLiteStatus EvalFloat(
           GetTensorData(projection_weights),
           GetTensorData(projection_bias), params, n_batch, n_cell,
           n_input, aux_input_size, n_output, output_batch_leading_dim,
-          GetTensorData(activation_state),
-          GetTensorData(cell_state), input_gate_scratch,
-          forget_gate_scratch, cell_gate_scratch, output_gate_scratch,
-          output_ptr_time, logger, intermediate_tensor_indexes, error_reporter);
+          GetTensorData(output_state), GetTensorData(cell_state),
+          input_gate_scratch, forget_gate_scratch, cell_gate_scratch,
+          output_gate_scratch, output_ptr_time, logger,
+          intermediate_tensor_indexes, error_reporter);
     }
   } else {
     for (int b = 0; b < n_batch; b++) {
@@ -414,9 +413,9 @@ TfLiteStatus EvalFloat(
         float* output_ptr = GetTensorData(output) +
                             time_offset * output_step + output_offset;
 
-        // Offset the {activation,cell}_state pointers to the right batch.
-        float* activation_state_ptr = GetTensorData(activation_state) +
-                                      b * output_batch_leading_dim;
+        // Offset the {output,cell}_state pointers to the right batch.
+        float* output_state_ptr =
+            GetTensorData(output_state) + b * output_batch_leading_dim;
         float* cell_state_ptr = GetTensorData(cell_state) + b * n_cell;
         // Offset the scratch pointers to the right batch.
         float* input_gate_scratch_ptr =
@@ -452,7 +451,7 @@ TfLiteStatus EvalFloat(
             GetTensorData(projection_weights),
             GetTensorData(projection_bias), params, /*n_batch=*/1,
             n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
-            activation_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
+            output_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
             forget_gate_scratch_ptr, cell_gate_scratch_ptr,
             output_gate_scratch_ptr, output_ptr, logger,
             intermediate_tensor_indexes, error_reporter);
@@ -541,11 +540,11 @@ TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
   // Index the scratch buffers pointers to the global scratch buffer.
   TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
 
-  TfLiteTensor* activation_state = GetVariableInput(
-      context, node, ops::builtin::lstm::full::kInputActivationStateTensor);
-  TF_LITE_ENSURE(context, activation_state != nullptr);
+  TfLiteTensor* output_state = GetVariableInput(
+      context, node, ops::builtin::lstm::full::kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
   TfLiteTensor* cell_state = GetVariableInput(
-      context, node, ops::builtin::lstm::full::kInputCellStateTensor);
+      context, node, ops::builtin::lstm::full::kCellStateTensor);
   TF_LITE_ENSURE(context, cell_state != nullptr);
 
   TfLiteTensor* output =
@@ -574,8 +573,8 @@ TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
           forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
           projection_bias, params, /*forward_sequence=*/true,
           /*time_major=*/true,
-          /*output_offset=*/0, scratch_buffer, activation_state, cell_state,
-          output, logger, intermediate_tensor_indexes, error_reporter);
+          /*output_offset=*/0, scratch_buffer, output_state, cell_state, output,
+          logger, intermediate_tensor_indexes, error_reporter);
     }
     case kTfLiteUInt8:
     case kTfLiteInt8:

From 23d482eaa2efe2bb38de7eb4f89539be9e3aa32a Mon Sep 17 00:00:00 2001
From: Jared Duke 
Date: Wed, 17 Jun 2020 09:57:30 -0700
Subject: [PATCH 0406/1390] Add flag for using optimized TFLite CPU kernels on
 Android

Add an experimental flag which allows opting in to a set of highly
optimized floating point kernels provided via the XNNPACK delegate.
This is offered as a preview, with the plan to enable these kernels
by default in a future release. The flag can be enabled via:

  Interpreter.Options options =
      new Interpreter.Options().setUseXNNPACK(true);

See tensorflow/lite/delegates/xnnpack/README.md for more details about
these kernels and the associated delegate functionality.

PiperOrigin-RevId: 316909226
Change-Id: Ib60cf259225b8a48a9830ccbb24ec10534b038ce
---
 tensorflow/lite/delegates/xnnpack/BUILD       |  2 +
 .../delegates/xnnpack/xnnpack_delegate.cc     |  3 +
 tensorflow/lite/java/BUILD                    |  1 +
 .../java/org/tensorflow/lite/Interpreter.java | 27 +++++++++
 .../lite/NativeInterpreterWrapper.java        |  7 +++
 tensorflow/lite/java/src/main/native/BUILD    |  1 +
 .../native/nativeinterpreterwrapper_jni.cc    | 55 +++++++++++++++++++
 .../lite/InterpreterMobileNetTest.java        | 16 ++++++
 .../org/tensorflow/lite/InterpreterTest.java  | 32 +++++++++++
 9 files changed, 144 insertions(+)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index 5736a2995b1..97e6aea2a6b 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -21,6 +21,7 @@ cc_library(
     linkstatic = True,
     deps = [
         "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite:util",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/schema:schema_fbs",
@@ -47,6 +48,7 @@ cc_library(
     linkstatic = True,
     deps = [
         "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite:util",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/schema:schema_fbs",
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index c4c95b6b295..739e45f62e4 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/lite/builtin_ops.h"
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/minimal_logging.h"
 #include "tensorflow/lite/tools/optimize/sparsity/format_converter.h"
 
 namespace tflite {
@@ -52,6 +53,8 @@ class Delegate {
           pthreadpool_create(static_cast(options->num_threads)));
     }
 #endif
+    TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
+                         "Created TensorFlow Lite XNNPACK delegate for CPU.");
   }
 
   TfLiteIntArray* PrepareOpsToDelegate(TfLiteContext* context);
diff --git a/tensorflow/lite/java/BUILD b/tensorflow/lite/java/BUILD
index 101e98e3dd1..d0331bca3e5 100644
--- a/tensorflow/lite/java/BUILD
+++ b/tensorflow/lite/java/BUILD
@@ -408,6 +408,7 @@ tflite_jni_binary(
         "//tensorflow/lite/c:c_api",
         "//tensorflow/lite/c:c_api_experimental",
         "//tensorflow/lite/delegates/nnapi/java/src/main/native",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
         "//tensorflow/lite/java/src/main/native",
     ],
 )
diff --git a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
index 7c9c5644f47..5993ee7a037 100644
--- a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
+++ b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
@@ -137,10 +137,37 @@ public final class Interpreter implements AutoCloseable {
       return this;
     }
 
+    /**
+     * Experimental: Enable an optimized set of floating point CPU kernels (provided by XNNPACK).
+     *
+     * 
Enabling this flag will enable use of a new, highly optimized set of CPU kernels provided
+     * via the XNNPACK delegate. Currently, this is restricted to a subset of floating point
+     * operations. Eventually, we plan to enable this by default, as it can provide significant
+     * peformance benefits for many classes of floating point models. See
+     * https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md
+     * for more details.
+     *
+     * 
Things to keep in mind when enabling this flag:
+     *
+     * 
    
+     *   
  • Startup time and resize time may increase.
+     *   
  • Baseline memory consumption may increase.
+     *   
  • Compatibility with other delegates (e.g., GPU) has not been fully validated.
+     *   
  • Quantized models will not see any benefit.
+     * 

+     *
+     * 
WARNING: This is an experimental interface that is subject to change.
+     */
+    public Options setUseXNNPACK(boolean useXNNPACK) {
+      this.useXNNPACK = useXNNPACK;
+      return this;
+    }
+
     int numThreads = -1;
     Boolean useNNAPI;
     Boolean allowFp16PrecisionForFp32;
     Boolean allowBufferHandleOutput;
+    Boolean useXNNPACK;
     final List delegates = new ArrayList<>();
   }
 
diff --git a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
index 8eb3c66f3b5..5e9a6eecf00 100644
--- a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
+++ b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
@@ -80,6 +80,10 @@ final class NativeInterpreterWrapper implements AutoCloseable {
       allowBufferHandleOutput(interpreterHandle, options.allowBufferHandleOutput.booleanValue());
     }
     applyDelegates(options);
+    if (options.useXNNPACK != null) {
+      useXNNPACK(
+          interpreterHandle, errorHandle, options.useXNNPACK.booleanValue(), options.numThreads);
+    }
     allocateTensors(interpreterHandle, errorHandle);
     this.isMemoryAllocated = true;
   }
@@ -438,6 +442,9 @@ final class NativeInterpreterWrapper implements AutoCloseable {
 
   private static native void allowBufferHandleOutput(long interpreterHandle, boolean allow);
 
+  private static native void useXNNPACK(
+      long interpreterHandle, long errorHandle, boolean state, int numThreads);
+
   private static native long createErrorReporter(int size);
 
   private static native long createModel(String modelPathOrBuffer, long errorHandle);
diff --git a/tensorflow/lite/java/src/main/native/BUILD b/tensorflow/lite/java/src/main/native/BUILD
index fdbbc9dc72c..52f79615a9f 100644
--- a/tensorflow/lite/java/src/main/native/BUILD
+++ b/tensorflow/lite/java/src/main/native/BUILD
@@ -31,6 +31,7 @@ cc_library(
         "//tensorflow/lite:string_util",
         "//tensorflow/lite:util",
         "//tensorflow/lite/c:common",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate_hdrs_only",
         "//tensorflow/lite/experimental/tflite_api_dispatcher:tflite_api_dispatcher_with_kernels",
         "//tensorflow/lite/java/jni",
     ],
diff --git a/tensorflow/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc b/tensorflow/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
index 690b58ac1f4..7abe0f518f0 100644
--- a/tensorflow/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
+++ b/tensorflow/lite/java/src/main/native/nativeinterpreterwrapper_jni.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
 #include 
 #include 
 #include 
@@ -20,6 +21,7 @@ limitations under the License.
 #include 
 
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
 #include "tensorflow/lite/experimental/tflite_api_dispatcher/tflite_api_dispatcher.h"
 #include "tensorflow/lite/java/src/main/native/jni_utils.h"
 #include "tensorflow/lite/util.h"
@@ -323,6 +325,59 @@ Java_org_tensorflow_lite_NativeInterpreterWrapper_allowBufferHandleOutput(
   interpreter->SetAllowBufferHandleOutput(allow);
 }
 
+JNIEXPORT void JNICALL
+Java_org_tensorflow_lite_NativeInterpreterWrapper_useXNNPACK(
+    JNIEnv* env, jclass clazz, jlong handle, jlong error_handle, jboolean state,
+    jint num_threads) {
+  // If not using xnnpack, simply don't apply the delegate.
+  if (!state) {
+    return;
+  }
+
+  tflite_api_dispatcher::Interpreter* interpreter =
+      convertLongToInterpreter(env, handle);
+  if (interpreter == nullptr) {
+    return;
+  }
+
+  BufferErrorReporter* error_reporter =
+      convertLongToErrorReporter(env, error_handle);
+  if (error_reporter == nullptr) {
+    return;
+  }
+
+  // We use dynamic loading to avoid taking a hard dependency on XNNPack.
+  // This allows clients that use trimmed builds to save on binary size.
+  auto xnnpack_options_default =
+      reinterpret_cast(
+          dlsym(RTLD_DEFAULT, "TfLiteXNNPackDelegateOptionsDefault"));
+  auto xnnpack_create =
+      reinterpret_cast(
+          dlsym(RTLD_DEFAULT, "TfLiteXNNPackDelegateCreate"));
+  auto xnnpack_delete =
+      reinterpret_cast(
+          dlsym(RTLD_DEFAULT, "TfLiteXNNPackDelegateDelete"));
+
+  if (xnnpack_options_default && xnnpack_create && xnnpack_delete) {
+    TfLiteXNNPackDelegateOptions options = xnnpack_options_default();
+    if (num_threads > 0) {
+      options.num_threads = num_threads;
+    }
+    tflite_api_dispatcher::Interpreter::TfLiteDelegatePtr delegate(
+        xnnpack_create(&options), xnnpack_delete);
+    if (interpreter->ModifyGraphWithDelegate(std::move(delegate)) !=
+        kTfLiteOk) {
+      ThrowException(env, kIllegalArgumentException,
+                     "Internal error: Failed to apply XNNPACK delegate: %s",
+                     error_reporter->CachedErrorMessage());
+    }
+  } else {
+    ThrowException(env, kIllegalArgumentException,
+                   "Failed to load XNNPACK delegate from current runtime. "
+                   "Have you added the necessary dependencies?");
+  }
+}
+
 JNIEXPORT void JNICALL
 Java_org_tensorflow_lite_NativeInterpreterWrapper_numThreads(JNIEnv* env,
                                                              jclass clazz,
diff --git a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterMobileNetTest.java b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterMobileNetTest.java
index 446cf5f7b02..80b3bf3cab9 100644
--- a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterMobileNetTest.java
+++ b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterMobileNetTest.java
@@ -54,6 +54,16 @@ public final class InterpreterMobileNetTest {
     runMobileNetFloatTest(new Interpreter.Options().setNumThreads(2));
   }
 
+  @Test
+  public void testMobileNetEnhancedCpuKernels() {
+    runMobileNetFloatTest(new Interpreter.Options().setUseXNNPACK(true));
+  }
+
+  @Test
+  public void testMobileNetEnhancedCpuKernelsMultithreaded() {
+    runMobileNetFloatTest(new Interpreter.Options().setUseXNNPACK(true).setNumThreads(2));
+  }
+
   @Test
   public void testMobileNetQuantized() {
     runMobileNetQuantizedTest(new Interpreter.Options());
@@ -64,6 +74,12 @@ public final class InterpreterMobileNetTest {
     runMobileNetQuantizedTest(new Interpreter.Options().setNumThreads(2));
   }
 
+  @Test
+  public void testMobileNetQuantizedEnhancedCpu() {
+    // The "enhanced CPU flag" should only impact float models, this is a sanity test to confirm.
+    runMobileNetQuantizedTest(new Interpreter.Options().setUseXNNPACK(true));
+  }
+
   private static void runMobileNetFloatTest(Interpreter.Options options) {
     ByteBuffer img =
         TestUtils.getTestImageAsFloatByteBuffer(
diff --git a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
index 3daa9fe0766..f1d4ff147b1 100644
--- a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
+++ b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
@@ -409,6 +409,38 @@ public final class InterpreterTest {
     interpreter.close();
   }
 
+  @Test
+  public void testUseXNNPACK() throws Exception {
+    Interpreter interpreter =
+        new Interpreter(MODEL_BUFFER, new Interpreter.Options().setUseXNNPACK(true));
+    float[] oneD = {1.23f, 6.54f, 7.81f};
+    float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD};
+    float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD};
+    float[][][][] fourD = {threeD, threeD};
+    float[][][][] parsedOutputs = new float[2][8][8][3];
+    interpreter.run(fourD, parsedOutputs);
+    float[] outputOneD = parsedOutputs[0][0][0];
+    float[] expected = {3.69f, 19.62f, 23.43f};
+    assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder();
+    interpreter.close();
+  }
+
+  @Test
+  public void testResizeWithEnhancedCpuKernels() throws Exception {
+    Interpreter interpreter =
+        new Interpreter(MODEL_BUFFER, new Interpreter.Options().setUseXNNPACK(true));
+    float[] input = {1.f};
+    float[] output = new float[1];
+    interpreter.run(input, output);
+    assertThat(output).usingTolerance(0.1f).containsExactly(new float[] {3.f}).inOrder();
+
+    // The new input shape should trigger a resize. Inference should still work properly.
+    float[] input2 = {1.f, 2.f};
+    float[] output2 = new float[2];
+    interpreter.run(input2, output2);
+    assertThat(output2).usingTolerance(0.1f).containsExactly(new float[] {3.f, 6.f}).inOrder();
+  }
+
   @Test
   public void testRedundantClose() throws Exception {
     Interpreter interpreter = new Interpreter(MODEL_BUFFER);

From 076f1474faa95cb70277d96000c2803ea1f8852b Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Thu, 18 Jun 2020 00:24:43 +0700
Subject: [PATCH 0407/1390] Return instead of using out parameter

---
 tensorflow/c/env.cc | 4 ++--
 tensorflow/c/env.h  | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc
index 6d8528bc42f..ce715c43acb 100644
--- a/tensorflow/c/env.cc
+++ b/tensorflow/c/env.cc
@@ -147,8 +147,8 @@ TF_StringStream* TF_GetLocalTempDirectories() {
   return list;
 }
 
-void TF_GetTempFileName(const char* extension, char** name) {
-  *name = strdup(::tensorflow::io::GetTempFilename(extension).c_str());
+char* TF_GetTempFileName(const char* extension) {
+  return strdup(::tensorflow::io::GetTempFilename(extension).c_str());
 }
 
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) {
diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h
index 29ec417a75e..7dc7ac32f08 100644
--- a/tensorflow/c/env.h
+++ b/tensorflow/c/env.h
@@ -154,8 +154,7 @@ TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void);
 
 // Creates a temporary file name with an extension.
 // The caller is responsible for freeing the returned pointer.
-TF_CAPI_EXPORT extern void TF_GetTempFileName(const char* extension,
-                                              char** name);
+TF_CAPI_EXPORT extern char* TF_GetTempFileName(const char* extension);
 
 // Returns the number of nanoseconds since the Unix epoch.
 TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void);

From c1ae0ef7ce398891118797be3de3e8437f306c7b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 10:22:03 -0700
Subject: [PATCH 0408/1390] Qualify uses of std::string

PiperOrigin-RevId: 316914367
Change-Id: Iae32a48b4db10d313f9e9b72f56eb6a6ac64c55f
---
 tensorflow/lite/toco/python/toco_python_api.cc | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/toco/python/toco_python_api.cc b/tensorflow/lite/toco/python/toco_python_api.cc
index 3f3d301a40d..0f21d0854ae 100644
--- a/tensorflow/lite/toco/python/toco_python_api.cc
+++ b/tensorflow/lite/toco/python/toco_python_api.cc
@@ -47,9 +47,9 @@ namespace toco {
 
 void PopulateConversionLogHelper(const toco::ModelFlags& model_flags,
                                  toco::TocoFlags* toco_flags,
-                                 const string& input_contents_txt,
-                                 const string& output_file_contents_txt,
-                                 const string& error_message,
+                                 const std::string& input_contents_txt,
+                                 const std::string& output_file_contents_txt,
+                                 const std::string& error_message,
                                  GraphVizDumpOptions* dump_options) {
   // Make sure the graphviz file will be dumped under the same folder.
   dump_options->dump_graphviz = toco_flags->conversion_summary_dir();
@@ -167,7 +167,7 @@ PyObject* TocoConvert(PyObject* model_flags_proto_txt_raw,
     dump_options.dump_graphviz_video = toco_flags.dump_graphviz_include_video();
   }
 
-  string output_file_contents_txt;
+  std::string output_file_contents_txt;
   tensorflow::Status status;
   int64 arithmetic_ops_count;
 
@@ -221,7 +221,7 @@ PyObject* TocoGetPotentiallySupportedOps() {
   std::vector supported_ops = toco::GetPotentiallySupportedOps();
   PyObject* list = PyList_New(supported_ops.size());
   for (size_t i = 0; i < supported_ops.size(); ++i) {
-    const string& op = supported_ops[i];
+    const std::string& op = supported_ops[i];
     PyObject* op_dict = PyDict_New();
     PyDict_SetItemString(op_dict, "op", PyUnicode_FromString(op.c_str()));
     PyList_SetItem(list, i, op_dict);

From 4747be646ba3a7cfe979afeb9fcaafdca5f1c800 Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Wed, 17 Jun 2020 10:46:23 -0700
Subject: [PATCH 0409/1390] [tf.data service] Increase the number of
 client-side uncompress threads.

This is necessary to prevent uncompression from becoming the bottleneck.

The change required updating the unit tests because now the `distribute` transformation may prefetch up to 16 elements.

PiperOrigin-RevId: 316919714
Change-Id: I4e0c0b2985792a2a2a0f216de2143a645076b1c8
---
 .../data/experimental/ops/data_service_ops.py |  5 +++-
 .../kernel_tests/data_service_ops_test.py     | 27 +++++++++++--------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/tensorflow/python/data/experimental/ops/data_service_ops.py b/tensorflow/python/data/experimental/ops/data_service_ops.py
index 39790d843ba..01ec155a89c 100644
--- a/tensorflow/python/data/experimental/ops/data_service_ops.py
+++ b/tensorflow/python/data/experimental/ops/data_service_ops.py
@@ -240,8 +240,11 @@ def _distribute(processing_mode,
         task_refresh_interval_hint_ms=task_refresh_interval_hint_ms)
     # TODO(b/157105111): Make this an autotuned parallel map when we have a way
     # to limit memory usage.
+    # The value 16 is chosen based on experience with pipelines that require
+    # more than 8 parallel calls to prevent this stage from being a bottleneck.
     dataset = dataset.map(
-        lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec))
+        lambda x: compression_ops.uncompress(x, output_spec=uncompressed_spec),
+        num_parallel_calls=16)
 
     # Disable autosharding for shared jobs.
     if job_name:
diff --git a/tensorflow/python/data/kernel_tests/data_service_ops_test.py b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
index d316009ce0c..2356a866d6e 100644
--- a/tensorflow/python/data/kernel_tests/data_service_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
@@ -201,13 +201,18 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
     self._new_worker = server_lib.WorkerServer(
         port=port, master_address=self._master._address, protocol=PROTOCOL)
 
-    # The dataset starts over now that we read from the new worker.
-    for i in range(num_elements):
+    # There may have been some elements prefetched from the first worker
+    # before it was stopped.
+    while True:
+      val = next(iterator).numpy()
+      if val == 0:
+        break
+
+    # The dataset starts over now that we read from the new worker.
+    # TODO(b/157086991): Iterate until end of sequence when we support
+    # detecting lost workers.
+    for i in range(1, num_elements // 2):
       val = next(iterator).numpy()
-      if val == midpoint and i != midpoint:
-        # There may have been one last element prefetched from the first worker
-        # before it was stopped.
-        val = next(iterator).numpy()
       self.assertEqual(i, val)
 
   @combinations.generate(test_base.eager_only_combinations())
@@ -248,7 +253,7 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
 
   @combinations.generate(test_base.eager_only_combinations())
   def testSharedJobName(self):
-    num_elements = 10
+    num_elements = 100
     master_address = self.create_cluster(1)
     ds = dataset_ops.Dataset.range(num_elements)
     ds1 = _make_distributed_dataset(ds, master_address, job_name="job_name")
@@ -256,7 +261,7 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
     iter1 = iter(ds1)
     iter2 = iter(ds2)
     results = []
-    for _ in range(3):
+    for _ in range(num_elements // 5):
       results.append(next(iter1).numpy())
       results.append(next(iter2).numpy())
     for elem in iter1:
@@ -291,7 +296,7 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
 
   @combinations.generate(test_base.eager_only_combinations())
   def testSharedJobNameRepeat(self):
-    num_elements = 10
+    num_elements = 100
     num_repetitions = 3
     master_address = self.create_cluster(1)
     ds = dataset_ops.Dataset.range(num_elements)
@@ -302,9 +307,9 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
     results = []
     iter1 = iter(ds1)
     iter2 = iter(ds2)
-    for _ in range(((num_elements * num_repetitions) // 2) - 1):
+    for _ in range((num_elements * num_repetitions) // 5):
       results.append(next(iter1).numpy())
-    for _ in range(((num_elements * num_repetitions) // 2) - 1):
+    for _ in range((num_elements * num_repetitions) // 5):
       results.append(next(iter2).numpy())
     for elem in iter1:
       results.append(elem.numpy())

From 2ba59dab2c22a592cb47660ecdb12463e457139c Mon Sep 17 00:00:00 2001
From: Thomas O'Malley 
Date: Wed, 17 Jun 2020 10:56:38 -0700
Subject: [PATCH 0410/1390] Simplify Layer.add_udpate in v2 and update
 version_selector to use v1 inside a tf.compat.v1.wrap_function.

No longer track unused Layer.updates in v2.

PiperOrigin-RevId: 316921838
Change-Id: I4698a0c925528594f402f824705d66b8a1ae7b72
---
 tensorflow/python/keras/engine/base_layer.py  | 53 +++----------------
 .../python/keras/engine/sequential_test.py    | 33 +++++-------
 .../python/keras/layers/normalization_test.py | 19 ++-----
 .../python/keras/layers/wrappers_test.py      |  3 --
 .../python/keras/utils/version_utils.py       | 33 ++++++++----
 5 files changed, 50 insertions(+), 91 deletions(-)

diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index a0ee25417c0..5ddce951491 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -1733,54 +1733,15 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
       inputs: Deprecated, will be automatically inferred.
     """
     call_context = base_layer_utils.call_context()
-
-    if (ds_context.has_strategy() and
-        ds_context.in_cross_replica_context() and
-        # When saving the model, the distribution strategy context should be
-        # ignored, following the default path for adding updates.
-        not call_context.saving):
-      # Updates don't need to be run in a cross-replica context.
+    # No need to run updates during Functional API construction.
+    if call_context.in_keras_graph:
       return
 
-    updates = generic_utils.to_list(updates)
-
-    # All updates can be run immediately in Eager or in a tf.function.
-    if base_layer_utils.is_in_eager_or_tf_function():
-      if not call_context.frozen:
-        for update in updates:
-          if callable(update):
-            update()
-      return
-
-    def process_update(x):
-      """Standardize update ops.
-
-      Arguments:
-        x: Tensor, op, or callable.
-
-      Returns:
-        An update op.
-      """
-      if callable(x):
-        update = lambda: process_update(x())
-        if not ops.executing_eagerly_outside_functions():
-          # In V1 mode, call the callable right away and process. This is needed
-          # for TPU strategy.
-          return update()
-      elif isinstance(x, ops.Operation):
-        update = x
-      elif hasattr(x, 'op'):
-        update = x.op
-      else:
-        update = ops.convert_to_tensor_v2(x)
-      return update
-
-    updates = [process_update(x) for x in updates]
-    # Non-callable Updates are run automatically inside `call` in V2, so
-    # they do not need to be tracked later.
-    if ops.executing_eagerly_outside_functions() and call_context.in_call:
-      updates = [u for u in updates if callable(u)]
-    self._updates.extend(updates)
+    # Callable updates are disabled by setting `trainable=False`.
+    if not call_context.frozen:
+      for update in nest.flatten(updates):
+        if callable(update):
+          update()
 
   def set_weights(self, weights):
     """Sets the weights of the layer, from Numpy arrays.
diff --git a/tensorflow/python/keras/engine/sequential_test.py b/tensorflow/python/keras/engine/sequential_test.py
index 9589d24fc57..773ce003656 100644
--- a/tensorflow/python/keras/engine/sequential_test.py
+++ b/tensorflow/python/keras/engine/sequential_test.py
@@ -231,33 +231,28 @@ class TestSequential(keras_parameterized.TestCase):
     inner_model.trainable = True
     self.assertEqual(len(model.trainable_weights), 4)
 
+  @keras_parameterized.run_all_keras_modes
   def test_sequential_update_disabling(self):
     val_a = np.random.random((10, 4))
     val_out = np.random.random((10, 4))
 
-    with self.cached_session():
-      model = keras.models.Sequential()
-      model.add(keras.layers.BatchNormalization(input_shape=(4,)))
-      assert model.updates
+    model = keras.models.Sequential()
+    model.add(keras.layers.BatchNormalization(input_shape=(4,)))
 
-      model.trainable = False
-      assert not model.updates
+    model.trainable = False
+    model.compile('sgd', 'mse')
 
-      model.compile('sgd', 'mse')
-      assert not model.updates
+    x1 = model.predict(val_a)
+    model.train_on_batch(val_a, val_out)
+    x2 = model.predict(val_a)
+    self.assertAllClose(x1, x2, atol=1e-7)
 
-      x1 = model.predict(val_a)
-      model.train_on_batch(val_a, val_out)
-      x2 = model.predict(val_a)
-      self.assertAllClose(x1, x2, atol=1e-7)
+    model.trainable = True
+    model.compile('sgd', 'mse')
 
-      model.trainable = True
-      model.compile('sgd', 'mse')
-      assert model.updates
-
-      model.train_on_batch(val_a, val_out)
-      x2 = model.predict(val_a)
-      assert np.abs(np.sum(x1 - x2)) > 1e-5
+    model.train_on_batch(val_a, val_out)
+    x2 = model.predict(val_a)
+    assert np.abs(np.sum(x1 - x2)) > 1e-5
 
   @keras_parameterized.run_all_keras_modes
   def test_sequential_deferred_build_serialization(self):
diff --git a/tensorflow/python/keras/layers/normalization_test.py b/tensorflow/python/keras/layers/normalization_test.py
index ef43bcf5d22..39992f7580a 100644
--- a/tensorflow/python/keras/layers/normalization_test.py
+++ b/tensorflow/python/keras/layers/normalization_test.py
@@ -325,18 +325,18 @@ class BatchNormalizationV2Test(keras_parameterized.TestCase):
       norm(inp)
 
   def test_updates_in_wrap_function(self):
-    layer = normalization.BatchNormalization()
 
     def my_func():
+      layer = normalization.BatchNormalization()
       x = array_ops.ones((10, 1))
-      return layer(x, training=True)
+      y = layer(x, training=True)
+      # Updates should be tracked in a `wrap_function`.
+      self.assertLen(layer.updates, 2)
+      return y
 
     wrapped_fn = wrap_function.wrap_function(my_func, [])
     wrapped_fn()
 
-    # Updates should be tracked in a `wrap_function`.
-    self.assertLen(layer.updates, 2)
-
   @keras_parameterized.run_all_keras_modes
   def test_basic_batchnorm_v2_none_shape_and_virtual_batch_size(self):
     # Test case for GitHub issue for 32380
@@ -392,15 +392,11 @@ class NormalizationLayersGraphModeOnlyTest(
       model.compile(gradient_descent.GradientDescentOptimizer(0.01), 'mse')
       model.train_on_batch(x, x)
 
-      self.assertLen(bn.updates, 4)
-
       # Test model-level reuse
       x3 = keras.layers.Input(shape=(10,))
       y3 = model(x3)
       new_model = keras.models.Model(x3, y3, name='new_model')
 
-      self.assertLen(new_model.updates, 6)
-      self.assertLen(model.updates, 6)
       new_model.compile(gradient_descent.GradientDescentOptimizer(0.01), 'mse')
       new_model.train_on_batch(x, x)
 
@@ -415,10 +411,7 @@ class NormalizationLayersGraphModeOnlyTest(
       model = keras.models.Model(a, b)
 
       model.trainable = False
-      assert not model.updates
-
       model.compile(gradient_descent.GradientDescentOptimizer(0.01), 'mse')
-      assert not model.updates
 
       x1 = model.predict(val_a)
       model.train_on_batch(val_a, val_out)
@@ -427,7 +420,6 @@ class NormalizationLayersGraphModeOnlyTest(
 
       model.trainable = True
       model.compile(gradient_descent.GradientDescentOptimizer(0.01), 'mse')
-      assert model.updates
 
       model.train_on_batch(val_a, val_out)
       x2 = model.predict(val_a)
@@ -435,7 +427,6 @@ class NormalizationLayersGraphModeOnlyTest(
 
       layer.trainable = False
       model.compile(gradient_descent.GradientDescentOptimizer(0.01), 'mse')
-      assert not model.updates
 
       x1 = model.predict(val_a)
       model.train_on_batch(val_a, val_out)
diff --git a/tensorflow/python/keras/layers/wrappers_test.py b/tensorflow/python/keras/layers/wrappers_test.py
index a73177fff12..5ee794dd1ef 100644
--- a/tensorflow/python/keras/layers/wrappers_test.py
+++ b/tensorflow/python/keras/layers/wrappers_test.py
@@ -234,13 +234,10 @@ class TimeDistributedTest(keras_parameterized.TestCase):
     x = keras.layers.Input(shape=(3, 2))
     layer = keras.layers.TimeDistributed(keras.layers.BatchNormalization())
     _ = layer(x)
-    self.assertEqual(len(layer.updates), 2)
     self.assertEqual(len(layer.trainable_weights), 2)
     layer.trainable = False
-    assert not layer.updates
     assert not layer.trainable_weights
     layer.trainable = True
-    assert len(layer.updates) == 2
     assert len(layer.trainable_weights) == 2
 
   def test_TimeDistributed_with_masked_embedding_and_unspecified_shape(self):
diff --git a/tensorflow/python/keras/utils/version_utils.py b/tensorflow/python/keras/utils/version_utils.py
index 551a07d2422..d3796dcbf92 100644
--- a/tensorflow/python/keras/utils/version_utils.py
+++ b/tensorflow/python/keras/utils/version_utils.py
@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
 from tensorflow.python.util import lazy_loader
 
@@ -51,8 +52,8 @@ class ModelVersionSelector(object):
   """Chooses between Keras v1 and v2 Model class."""
 
   def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
-    eager_enabled = ops.executing_eagerly_outside_functions()
-    cls = swap_class(cls, training.Model, training_v1.Model, eager_enabled)
+    use_v2 = should_use_v2()
+    cls = swap_class(cls, training.Model, training_v1.Model, use_v2)  # pylint: disable=self-cls-assignment
     return super(ModelVersionSelector, cls).__new__(cls)
 
 
@@ -60,8 +61,8 @@ class LayerVersionSelector(object):
   """Chooses between Keras v1 and v2 Layer class."""
 
   def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
-    eager_enabled = ops.executing_eagerly_outside_functions()
-    cls = swap_class(cls, base_layer.Layer, base_layer_v1.Layer, eager_enabled)
+    use_v2 = should_use_v2()
+    cls = swap_class(cls, base_layer.Layer, base_layer_v1.Layer, use_v2)  # pylint: disable=self-cls-assignment
     return super(LayerVersionSelector, cls).__new__(cls)
 
 
@@ -69,10 +70,10 @@ class TensorBoardVersionSelector(object):
   """Chooses between Keras v1 and v2 TensorBoard callback class."""
 
   def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
-    eager_enabled = ops.executing_eagerly_outside_functions()
+    use_v2 = should_use_v2()
     start_cls = cls
     cls = swap_class(start_cls, callbacks.TensorBoard, callbacks_v1.TensorBoard,
-                     eager_enabled)
+                     use_v2)
     if start_cls == callbacks_v1.TensorBoard and cls == callbacks.TensorBoard:
       # Since the v2 class is not a subclass of the v1 class, __init__ has to
       # be called manually.
@@ -80,19 +81,33 @@ class TensorBoardVersionSelector(object):
     return super(TensorBoardVersionSelector, cls).__new__(cls)
 
 
-def swap_class(cls, v2_cls, v1_cls, eager_enabled):
+def should_use_v2():
+  """Determine if v1 or v2 version should be used."""
+  if context.executing_eagerly():
+    return True
+  elif ops.executing_eagerly_outside_functions():
+    # Check for a v1 `wrap_function` FuncGraph.
+    # Code inside a `wrap_function` is treated like v1 code.
+    graph = ops.get_default_graph()
+    if (getattr(graph, "name", False) and
+        graph.name.startswith("wrapped_function")):
+      return False
+    return True
+
+
+def swap_class(cls, v2_cls, v1_cls, use_v2):
   """Swaps in v2_cls or v1_cls depending on graph mode."""
   if cls == object:
     return cls
 
   if cls in (v2_cls, v1_cls):
-    if eager_enabled:
+    if use_v2:
       return v2_cls
     return v1_cls
 
   # Recursively search superclasses to swap in the right Keras class.
   cls.__bases__ = tuple(
-      swap_class(base, v2_cls, v1_cls, eager_enabled) for base in cls.__bases__)
+      swap_class(base, v2_cls, v1_cls, use_v2) for base in cls.__bases__)
   return cls
 
 

From f8fd28e4a0e694ca808da57e10d1e3b40e4f2fb9 Mon Sep 17 00:00:00 2001
From: Lu Wang 
Date: Wed, 17 Jun 2020 10:58:32 -0700
Subject: [PATCH 0411/1390] Create the Java API doc section on tensorflow.org

PiperOrigin-RevId: 316922276
Change-Id: I1fd4458017daecfd7256313480a4b3d1602a9310
---
 tensorflow/lite/g3doc/_book.yaml        | 2 ++
 tensorflow/lite/g3doc/api_docs/index.md | 6 +++---
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/g3doc/_book.yaml b/tensorflow/lite/g3doc/_book.yaml
index 0da3d152090..abb18870003 100644
--- a/tensorflow/lite/g3doc/_book.yaml
+++ b/tensorflow/lite/g3doc/_book.yaml
@@ -198,6 +198,8 @@ upper_tabs:
       - title: "Overview"
         status: external
         path: /api_docs/python/tf/lite
+      - heading: "Android (Java)"
+      - include: /lite/api_docs/java/_toc.yaml
       - heading: "C++"
       - title: Overview
         path: /lite/api_docs/cc/
diff --git a/tensorflow/lite/g3doc/api_docs/index.md b/tensorflow/lite/g3doc/api_docs/index.md
index 533f5881eb4..5db55fb28a3 100644
--- a/tensorflow/lite/g3doc/api_docs/index.md
+++ b/tensorflow/lite/g3doc/api_docs/index.md
@@ -4,7 +4,7 @@ The API reference documentation provides detailed information for each of the
 classes and methods in the TensorFlow Lite library. Choose your preferred
 platform from the list below.
 
-*   [Python API reference](/api_docs/python/tf/lite)
-*   Android API reference (coming soon)
+*   [Python API reference](https://tensorflow.org/api_docs/python/tf/lite)
+*   [Android (Java) API reference](https://tensorflow.org/lite/api_docs/java/org/tensorflow/lite/package-summary)
 *   iOS API reference (coming soon)
-*   [C++ API reference](/lite/api_docs/cc/)
+*   [C++ API reference](https://tensorflow.org/lite/api_docs/cc)

From 2a8bbb92b74a5dae9dd7e0cc3aed9b79231a06c3 Mon Sep 17 00:00:00 2001
From: Thomas O'Malley 
Date: Wed, 17 Jun 2020 10:59:06 -0700
Subject: [PATCH 0412/1390] Reduce TensorShape.__init__ overhead by 50%.

TensorShape.__init__ is on the hotpath because a TensorShape is
created the first time EagerTensor.shape is called. The
TensorShape is created from EagerTensor._shape_tuple, which is
a tuple of ints. This change optimizes the code for this common path.

PiperOrigin-RevId: 316922384
Change-Id: I063ea393450123ea4150972e5c73647f03a29cf5
---
 .../data/kernel_tests/from_generator_test.py       |  4 ++--
 tensorflow/python/eager/benchmarks_test.py         | 14 ++++++++++++++
 tensorflow/python/framework/tensor_shape.py        | 12 +++++++++---
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/data/kernel_tests/from_generator_test.py b/tensorflow/python/data/kernel_tests/from_generator_test.py
index a08f54a7101..386108f0de7 100644
--- a/tensorflow/python/data/kernel_tests/from_generator_test.py
+++ b/tensorflow/python/data/kernel_tests/from_generator_test.py
@@ -465,8 +465,8 @@ class FromGeneratorTest(test_base.DatasetTestBase, parameterized.TestCase):
       for _ in range(10):
         yield [20]
 
-    with self.assertRaisesRegexp(
-        TypeError, r"Failed to convert '\[\[1\]\]' to a shape"):
+    with self.assertRaisesRegex(TypeError,
+                                r"Dimension value must be integer or None"):
       dataset_ops.Dataset.from_generator(
           generator, output_types=(dtypes.int64), output_shapes=[[1]])
 
diff --git a/tensorflow/python/eager/benchmarks_test.py b/tensorflow/python/eager/benchmarks_test.py
index b7c8395790a..24e86c77a14 100644
--- a/tensorflow/python/eager/benchmarks_test.py
+++ b/tensorflow/python/eager/benchmarks_test.py
@@ -50,6 +50,7 @@ from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -1441,6 +1442,19 @@ class MicroBenchmarks(benchmarks_test_base.MicroBenchmarksBase):
 
     self._run(fn, 10000)
 
+  def benchmark_tf_tensor_shape_creation_overhead(self):
+    # A `TensorShape` is created the first time `EagerTensor.shape` is
+    # called, which puts `TensorShape.__init__` on the hotpath. The
+    # `TensorShape` is created from `EagerTensor._shape_tuple`.
+
+    x = array_ops.ones((1, 1))
+    shape_tuple = x._shape_tuple()
+
+    def fn():
+      tensor_shape.TensorShape(shape_tuple)
+
+    self._run(fn, 100000)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/framework/tensor_shape.py b/tensorflow/python/framework/tensor_shape.py
index fd229b6691a..20508f37eb7 100644
--- a/tensorflow/python/framework/tensor_shape.py
+++ b/tensorflow/python/framework/tensor_shape.py
@@ -184,10 +184,14 @@ class Dimension(object):
 
   def __init__(self, value):
     """Creates a new Dimension with the given value."""
-    if value is None:
+    if isinstance(value, int):  # Most common case.
+      if value < 0:
+        raise ValueError("Dimension %d must be >= 0" % value)
+      self._value = value
+    elif value is None:
       self._value = None
     elif isinstance(value, Dimension):
-      self._value = value
+      self._value = value._value
     else:
       try:
         # int(...) compensates for the int/long dichotomy on Python 2.X.
@@ -748,7 +752,9 @@ class TensorShape(object):
     Raises:
       TypeError: If dims cannot be converted to a list of dimensions.
     """
-    if dims is None:
+    if isinstance(dims, (tuple, list)):  # Most common case.
+      self._dims = [Dimension(d) for d in dims]
+    elif dims is None:
       self._dims = None
     elif isinstance(dims, tensor_shape_pb2.TensorShapeProto):
       if dims.unknown_rank:

From c8dd07ae28a5bdc291cf43adb128fd11a6b2d6f7 Mon Sep 17 00:00:00 2001
From: Bruce Fontaine 
Date: Wed, 17 Jun 2020 11:00:38 -0700
Subject: [PATCH 0413/1390] Fix error check in TPUEmbedding to work when used
 in outside compilation.

PiperOrigin-RevId: 316922750
Change-Id: Ie6b4c83e54f3e6d90fbe38fe8da0eea84312c382
---
 tensorflow/python/tpu/tpu_embedding_v2.py     | 20 ++++++++----
 .../python/tpu/tpu_embedding_v2_test.py       | 31 ++++++++++++++++---
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index f7a383c440c..e5cfba7c587 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -1024,11 +1024,8 @@ class TPUEmbedding(tracking.AutoTrackable):
   def _raise_error_for_inputs_not_on_cpu(self, features):
     """Checks all tensors in features to see are placed on the CPU."""
 
-    # expand_composites here is important, we need to check the device of each
-    # underlying tensor.
-    for path, input_tensor in nest.flatten_with_joined_string_paths(
-        features, expand_composites=True):
-      spec = tf_device.DeviceSpec.from_string(input_tensor.device)
+    def check_device(path, device_string):
+      spec = tf_device.DeviceSpec.from_string(device_string)
       if spec.device_type == "TPU":
         raise ValueError(
             "Received input tensor {} which is on a TPU input device {}. Input "
@@ -1037,7 +1034,18 @@ class TPUEmbedding(tracking.AutoTrackable):
             "setting the 'experimental_prefetch_to_device' option of the "
             "dataset distribution function. See the documentation of the "
             "enqueue method for an example.".format(
-                path, input_tensor.device))
+                path, device_string))
+
+    # expand_composites here is important, we need to check the device of each
+    # underlying tensor.
+    for path, input_tensor in nest.flatten_with_joined_string_paths(
+        features, expand_composites=True):
+      if (input_tensor.op.type == "Identity" and
+          input_tensor.op.inputs[0].op.type == "TPUReplicatedInput"):
+        for tensor in input_tensor.op.inputs[0].op.inputs:
+          check_device(path, tensor.device)
+      else:
+        check_device(path, input_tensor.device)
 
   def enqueue(self, features, weights=None, training=True, name=None):
     """Enqueues id tensors for embedding lookup.
diff --git a/tensorflow/python/tpu/tpu_embedding_v2_test.py b/tensorflow/python/tpu/tpu_embedding_v2_test.py
index ebaf2791055..ff09085f3f1 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2_test.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2_test.py
@@ -727,10 +727,33 @@ class TPUEmbeddingTest(parameterized.TestCase, test.TestCase):
       def get_activations():
         return mid_level_api.dequeue()
 
-      sparse_features = next(sparse_iter)
-      mid_level_api.enqueue(sparse_features, training=False)
-      sparse_activations = strategy.run(get_activations)
-      return sparse_activations
+      features = next(sparse_iter)
+      mid_level_api.enqueue(features, training=False)
+      activations = strategy.run(get_activations)
+      return activations
+
+    with self.assertRaisesRegex(ValueError, 'which is on a TPU input device'):
+      test_fn()
+
+  @parameterized.parameters([True, False])
+  def test_enqueue_cpu_tensor_with_outside_compilation(self, use_mlir):
+    if use_mlir:
+      config.enable_mlir_bridge()
+
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    input_fn = self._create_dense_input_fn(strategy)
+    sparse_iter = iter(strategy.experimental_distribute_datasets_from_function(
+        input_fn))
+
+    @def_function.function
+    def test_fn():
+      def get_activations(features):
+        mid_level_api.enqueue(features, training=False)
+        return mid_level_api.dequeue()
+
+      activations = strategy.run(get_activations, args=(next(sparse_iter),))
+      return activations
 
     with self.assertRaisesRegex(ValueError, 'which is on a TPU input device'):
       test_fn()

From ab7cb8336c780ffe02ec983e8910cf62250ec3e1 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Wed, 17 Jun 2020 11:15:35 -0700
Subject: [PATCH 0414/1390] Use the separate matrix/vector scaling factor
 version of MatrixBatchVectorMultiplyAccumulate in projection too.

PiperOrigin-RevId: 316926050
Change-Id: I64e9febce8590231c78a90a0a5aec5da11996195
---
 tensorflow/lite/kernels/lstm_eval.cc | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 65f68b34251..4ac9e538317 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -874,16 +874,13 @@ inline void LstmStepHybrid(
       tensor_utils::BatchQuantizeFloats(
           output_gate_scratch, n_batch, n_cell, quantized_cell_state_ptr,
           scaling_factors, zero_points, asymmetric_quantize_inputs);
-      for (int b = 0; b < n_batch; ++b) {
-        scaling_factors_scratch[b] =
-            scaling_factors[b] * projection_weights_scale;
-      }
       tensor_utils::MatrixBatchVectorMultiplyAccumulate(
           projection_weights_ptr, n_output, n_cell, quantized_cell_state_ptr,
-          scaling_factors_scratch, n_batch, output_state_ptr,
+          projection_weights_scale, scaling_factors, n_batch, output_state_ptr,
           /*per_channel_scale=*/nullptr,
           asymmetric_quantize_inputs ? zero_points : nullptr, accum_scratch_ptr,
-          projection_weights_row_sums, compute_row_sums, context);
+          projection_weights_row_sums, compute_row_sums,
+          scaling_factors_scratch, context);
     }
     if (params->proj_clip > 0.0) {
       tensor_utils::ClipVector(output_state_ptr, n_batch * n_output,

From 5e7fc9584a983c3e53e2c845017acded9b19f180 Mon Sep 17 00:00:00 2001
From: Ayush Dubey 
Date: Wed, 17 Jun 2020 11:21:20 -0700
Subject: [PATCH 0415/1390] Add a unit test demonstrating collective ops with
 different groups of devices.

Collectives configured this way can be used to implement all-reduce for batch
norm with a subset of all available devices.

PiperOrigin-RevId: 316927203
Change-Id: Ic288d01134776efbe0e49a83fd8030f721890725
---
 tensorflow/python/ops/collective_ops_test.py | 35 ++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tensorflow/python/ops/collective_ops_test.py b/tensorflow/python/ops/collective_ops_test.py
index 47c25fcafc0..9727593a1c5 100644
--- a/tensorflow/python/ops/collective_ops_test.py
+++ b/tensorflow/python/ops/collective_ops_test.py
@@ -581,6 +581,41 @@ class CollectiveOpTest(test.TestCase):
       results = sess.run(run_ops)
       self.assertEqual(results, [3., 3., 3., 3.])
 
+  @test_util.run_v2_only
+  def testMultipleGroups(self):
+    context._reset_context()
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    context.ensure_initialized()
+    num_elements = 4
+
+    @def_function.function
+    def run_all_reduce(group_size, group_key):
+      instance_key = group_key
+      input_value = [group_key for i in range(num_elements)]
+      collectives = []
+      for device_idx in range(group_size):
+        with ops.device('/CPU:{}'.format(device_idx)):
+          input_tensor = constant_op.constant(input_value)
+          collectives.append(collective_ops.all_reduce(
+              input_tensor, group_size, group_key, instance_key, merge_op='Add',
+              final_op='Id'))
+      return collectives
+
+    def run_and_assert(group_size, group_key):
+      for reduced_tensor in run_all_reduce(group_size, group_key):
+        self.assertAllEqual(
+            [group_key * group_size for i in range(num_elements)],
+            reduced_tensor.numpy())
+
+    run_and_assert(group_size=2, group_key=1)
+    run_and_assert(group_size=3, group_key=2)
+
 
 if __name__ == '__main__':
   test.main()

From c6a3ab159f0d19958198e315d2f4f0943725b840 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Wed, 17 Jun 2020 11:25:14 -0700
Subject: [PATCH 0416/1390] Simplify broadcast plus compare

PiperOrigin-RevId: 316927924
Change-Id: If7f3ce209cbff3720c19a60d3e713167e1c6e8c6
---
 .../xla/service/algebraic_simplifier.cc       | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
index ce2a801fccd..130661bf1cd 100755
--- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc
+++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -2815,6 +2815,28 @@ Status AlgebraicSimplifierVisitor::HandleCompare(HloInstruction* compare) {
   HloInstruction* lhs;
   HloInstruction* rhs;
   CHECK(Match(compare, m::Compare(m::Op(&lhs), m::Op(&rhs))));
+  {
+    // compare(broadcast(a) + x, broadcast(b)) ==>
+    //   compare(x, broadcast(b-a))
+    HloInstruction *x, *a, *b;
+    if (Match(compare,
+              m::Compare(
+                  m::AddAnyOrder(m::Op(&x), m::Broadcast(m::Op(&a).WithShape(
+                                                m::Shape().IsScalar()))),
+                  m::Broadcast(m::Op(&b).WithShape(m::Shape().IsScalar()))))) {
+      if (ShapeUtil::ElementIsSigned(x->shape())) {
+        HloInstruction* sub =
+            computation_->AddInstruction(HloInstruction::CreateBinary(
+                b->shape(), HloOpcode::kSubtract, b, a));
+        HloInstruction* broadcast = computation_->AddInstruction(
+            HloInstruction::CreateBroadcast(x->shape(), sub, {}));
+        HloInstruction* new_compare = computation_->AddInstruction(
+            HloInstruction::CreateCompare(compare->shape(), x, broadcast,
+                                          compare->comparison_direction()));
+        return ReplaceInstruction(compare, new_compare);
+      }
+    }
+  }
 
   if (compare->comparison_direction() == ComparisonDirection::kLt &&
       lhs->opcode() == HloOpcode::kIota && IsAll(rhs, 0)) {

From ef55a40b374d7310e4ce3149d86395d403403d0d Mon Sep 17 00:00:00 2001
From: Yong Tang 
Date: Wed, 17 Jun 2020 18:38:24 +0000
Subject: [PATCH 0417/1390] Move tf.equal shape inference test to
 math_ops_test.cc

also added additional shape inference test cases

Signed-off-by: Yong Tang 
---
 tensorflow/core/ops/math_ops_test.cc           | 18 ++++++++++++++++++
 .../python/autograph/operators/logical_test.py | 14 --------------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/ops/math_ops_test.cc b/tensorflow/core/ops/math_ops_test.cc
index a2837d88bde..2b65f88042c 100644
--- a/tensorflow/core/ops/math_ops_test.cc
+++ b/tensorflow/core/ops/math_ops_test.cc
@@ -604,4 +604,22 @@ TEST(MathOpsTest, SobolSample) {
 
   INFER_OK(op, "[];[];[]", "[?,?]");
 }
+
+TEST(MathOpsTest, EqualOp) {
+  ShapeInferenceTestOp op("Equal");
+  AddNodeAttr("incompatible_shape_error", true, &op.node_def);
+
+  INFER_OK(op, "?;?", "?");
+  INFER_OK(op, "[1,2];?", "?");
+  INFER_OK(op, "?;[1,2]", "?");
+
+  INFER_OK(op, "[1,2,3];[1]", "[d0_0,d0_1,d0_2]");
+  INFER_OK(op, "[?,2,1];[1,3]", "[d0_0,d0_1,d1_1]");
+  INFER_OK(op, "[1,?,3];[3,1]", "[d0_0,d1_0,d0_2]");
+  INFER_OK(op, "[1,2,3];[2,1,3]", "[d1_0,d0_1,d0_2]");
+
+  // Note: Test case for GitHub issue 40471
+  INFER_OK(op, "[?,10,1];[?,1,4]", "[?,d0_1,d1_2]");
+  INFER_OK(op, "[10,?,1];[1,?,4]", "[d0_0,?,d1_2]");
+}
 }  // end namespace tensorflow
diff --git a/tensorflow/python/autograph/operators/logical_test.py b/tensorflow/python/autograph/operators/logical_test.py
index 0eab302a825..e22f39932d1 100644
--- a/tensorflow/python/autograph/operators/logical_test.py
+++ b/tensorflow/python/autograph/operators/logical_test.py
@@ -19,9 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.autograph.operators import logical
-from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
 
@@ -85,18 +83,6 @@ class LogicalOperatorsTest(test.TestCase):
       t = logical.not_(self._tf_false())
       self.assertEqual(self.evaluate(t), True)
 
-  # Test case for GitHub issue 40471
-  def test_equal_output_shapes(self):
-
-    @def_function.function(input_signature=[
-        tensor_spec.TensorSpec([None, 10, 1]),
-        tensor_spec.TensorSpec([None, 1, 4])])
-    def f(x, y):
-      z = x == y
-      return z
-
-    self.assertAllEqual(f.get_concrete_function().output_shapes, [None, 10, 4])
-
 
 if __name__ == '__main__':
   test.main()

From 653131dd38e9bbde2b9163d756ca4d9cfa69e1a5 Mon Sep 17 00:00:00 2001
From: Pavithra Vijay 
Date: Wed, 17 Jun 2020 11:32:39 -0700
Subject: [PATCH 0418/1390] Remove automatic control dep wrapping from layers
 in v2.

PiperOrigin-RevId: 316929712
Change-Id: Ic1a7d125776eeb0c7654e321dd6f2351c8656a16
---
 tensorflow/python/keras/engine/base_layer.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/tensorflow/python/keras/engine/base_layer.py b/tensorflow/python/keras/engine/base_layer.py
index 5ddce951491..fbec5382a08 100644
--- a/tensorflow/python/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/engine/base_layer.py
@@ -40,7 +40,6 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import execute
 from tensorflow.python.eager import function
 from tensorflow.python.eager import monitoring
-from tensorflow.python.framework import auto_control_deps
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
@@ -1105,17 +1104,7 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
 
           try:
             with ops.enable_auto_cast_variables(self._compute_dtype_object):
-              # Add auto_control_deps in V2 when they are not already added by
-              # a `tf.function`.
-              if (ops.executing_eagerly_outside_functions() and
-                  not base_layer_utils.is_in_eager_or_tf_function()):
-                with auto_control_deps.AutomaticControlDependencies() as acd:
-                  outputs = call_fn(cast_inputs, *args, **kwargs)
-                  # Wrap Tensors in `outputs` in `tf.identity` to avoid
-                  # circular dependencies.
-                  outputs = base_layer_utils.mark_as_return(outputs, acd)
-              else:
-                outputs = call_fn(cast_inputs, *args, **kwargs)
+              outputs = call_fn(cast_inputs, *args, **kwargs)
 
           except errors.OperatorNotAllowedInGraphError as e:
             raise TypeError('You are attempting to use Python control '

From f24487e6190eb4637f5e4988737cbe17e91a231a Mon Sep 17 00:00:00 2001
From: Brian Zhao 
Date: Wed, 17 Jun 2020 11:38:22 -0700
Subject: [PATCH 0419/1390] Disabling test until bug fix lands.

PiperOrigin-RevId: 316930923
Change-Id: I8aa57e5627fa649474c469bd5c92a713c4d9bd75
---
 tensorflow/c/experimental/saved_model/core/ops/BUILD | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/c/experimental/saved_model/core/ops/BUILD b/tensorflow/c/experimental/saved_model/core/ops/BUILD
index b42e93c3716..aa909c692ca 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/ops/BUILD
@@ -77,6 +77,9 @@ tf_cc_test(
     srcs = [
         "variable_ops_test.cc",
     ],
+    tags = [
+        "no_windows",  # TODO(b/159210739): Remove this tag after fixing the bug.
+    ],
     deps = [
         ":owned_eager_context",
         ":owned_tensor",

From 8d5171bad7374206cb5389ec1df83ed0189fdf6d Mon Sep 17 00:00:00 2001
From: Saurabh Saxena 
Date: Wed, 17 Jun 2020 12:35:43 -0700
Subject: [PATCH 0420/1390] Split Abstract interfaces into Abstract and
 ImmediateExecution interfaces. The Abstract interfaces are shared with
 tracing mode. Introduce an AbstractFunction which handles the conversion
 between MLIR function and FunctionDef and the runtime can query whichever
 representation is suitable. Right now this only supports GetFunctionDef but
 an API for fetching the MLIR function directly will be added in future
 changes.

PiperOrigin-RevId: 316942774
Change-Id: I1abebbe853b98dd0048bab9fc092252f4caf3d1b
---
 tensorflow/c/c_api_experimental.cc            |  2 +-
 tensorflow/c/eager/BUILD                      | 98 +++++++++++++++----
 tensorflow/c/eager/abstract_context.h         | 69 +++++++++++++
 tensorflow/c/eager/abstract_function.h        | 46 +++++++++
 ...ation_interface.h => abstract_operation.h} | 47 ++++-----
 tensorflow/c/eager/abstract_tensor_handle.h   | 45 +++++++++
 tensorflow/c/eager/c_api.cc                   | 21 ++--
 tensorflow/c/eager/c_api_experimental.cc      |  2 +-
 ...erface.h => immediate_execution_context.h} | 36 +++----
 .../c/eager/immediate_execution_operation.h   | 53 ++++++++++
 ....h => immediate_execution_tensor_handle.h} | 22 ++---
 tensorflow/c/eager/tfe_context_internal.h     |  4 +-
 tensorflow/c/eager/tfe_op_internal.h          |  6 +-
 .../c/eager/tfe_tensorhandle_internal.h       |  6 +-
 .../c/experimental/saved_model/core/BUILD     |  4 +-
 .../saved_model/core/concrete_function.cc     |  4 +-
 .../saved_model/core/concrete_function.h      | 10 +-
 .../c/experimental/saved_model/core/ops/BUILD | 11 ++-
 .../core/ops/owned_eager_context.h            | 10 +-
 .../saved_model/core/ops/owned_eager_op.h     | 10 +-
 .../core/ops/owned_tensor_handle.h            |  6 +-
 .../saved_model/core/ops/variable_ops.cc      | 33 ++++---
 .../saved_model/core/ops/variable_ops.h       | 20 ++--
 .../c/experimental/saved_model/internal/BUILD |  4 +-
 .../saved_model/internal/tensorhandle_list.cc |  2 +-
 .../internal/tensorhandle_list_type.h         |  4 +-
 tensorflow/core/common_runtime/eager/BUILD    | 18 ++--
 .../core/common_runtime/eager/context.cc      |  2 -
 .../core/common_runtime/eager/context.h       | 16 +--
 tensorflow/core/common_runtime/eager/core.cc  | 23 +++--
 .../common_runtime/eager/eager_operation.cc   | 16 +--
 .../common_runtime/eager/eager_operation.h    | 18 ++--
 .../common_runtime/eager/tensor_handle.cc     |  2 +-
 .../core/common_runtime/eager/tensor_handle.h | 12 +--
 tensorflow/python/eager/pywrap_tfe_src.cc     |  4 +-
 35 files changed, 489 insertions(+), 197 deletions(-)
 create mode 100644 tensorflow/c/eager/abstract_context.h
 create mode 100644 tensorflow/c/eager/abstract_function.h
 rename tensorflow/c/eager/{operation_interface.h => abstract_operation.h} (80%)
 create mode 100644 tensorflow/c/eager/abstract_tensor_handle.h
 rename tensorflow/c/eager/{context_interface.h => immediate_execution_context.h} (78%)
 create mode 100644 tensorflow/c/eager/immediate_execution_operation.h
 rename tensorflow/c/eager/{tensor_handle_interface.h => immediate_execution_tensor_handle.h} (74%)

diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc
index e9e6d470c68..831c6a0ad40 100644
--- a/tensorflow/c/c_api_experimental.cc
+++ b/tensorflow/c/c_api_experimental.cc
@@ -624,7 +624,7 @@ void TFE_InferShapes(TFE_Op* tfe_op, TF_ShapeAndTypeList* input_shapes,
 
   const int num_inputs = input_shapes->num_items;
   NodeDef node_def;
-  tensorflow::AbstractOperationInterface* op = tensorflow::unwrap(tfe_op);
+  tensorflow::ImmediateExecutionOperation* op = tensorflow::unwrap(tfe_op);
   node_def.set_name(op->Name());
   node_def.set_op(op->Name());
   for (int i = 0; i < num_inputs; ++i) {
diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD
index 9d3c79e0ae7..5f7ab4a1f59 100644
--- a/tensorflow/c/eager/BUILD
+++ b/tensorflow/c/eager/BUILD
@@ -38,9 +38,10 @@ tf_cuda_library(
             "//tensorflow/core:portable_tensorflow_lib_lite",
         ],
         "//conditions:default": [
-            ":context_interface",
-            ":operation_interface",
-            ":tensor_handle_interface",
+            ":immediate_execution_context",
+            ":immediate_execution_operation",
+            ":immediate_execution_tensor_handle",
+            ":abstract_tensor_handle",
             ":tfe_context_internal",
             ":tfe_cancellation_manager_internal",
             ":tfe_executor_internal",
@@ -101,13 +102,17 @@ tf_cuda_library(
 filegroup(
     name = "pywrap_required_hdrs",
     srcs = [
+        "abstract_context.h",
+        "abstract_function.h",
+        "abstract_operation.h",
+        "abstract_tensor_handle.h",
         "c_api_experimental.h",
         "c_api_internal.h",
         "c_api_unified_experimental.h",
-        "context_interface.h",
         "dlpack.h",
-        "operation_interface.h",
-        "tensor_handle_interface.h",
+        "immediate_execution_context.h",
+        "immediate_execution_operation.h",
+        "immediate_execution_tensor_handle.h",
         "tfe_cancellation_manager_internal.h",
         "tfe_executor_internal.h",
         "tfe_monitoring_internal.h",
@@ -163,12 +168,22 @@ cc_library(
 )
 
 cc_library(
-    name = "tensor_handle_interface",
-    hdrs = ["tensor_handle_interface.h"],
+    name = "abstract_tensor_handle",
+    hdrs = ["abstract_tensor_handle.h"],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [],
+)
+
+cc_library(
+    name = "immediate_execution_tensor_handle",
+    hdrs = ["immediate_execution_tensor_handle.h"],
     visibility = [
         "//tensorflow:internal",
     ],
     deps = [
+        ":abstract_tensor_handle",
         "//tensorflow/c:tensor_interface",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
@@ -177,13 +192,13 @@ cc_library(
 )
 
 cc_library(
-    name = "operation_interface",
-    hdrs = ["operation_interface.h"],
+    name = "abstract_operation",
+    hdrs = ["abstract_operation.h"],
     visibility = [
         "//tensorflow:internal",
     ],
     deps = [
-        ":tensor_handle_interface",
+        ":abstract_tensor_handle",
         "//tensorflow/c:tensor_interface",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
@@ -193,14 +208,58 @@ cc_library(
 )
 
 cc_library(
-    name = "context_interface",
-    hdrs = ["context_interface.h"],
+    name = "immediate_execution_operation",
+    hdrs = ["immediate_execution_operation.h"],
     visibility = [
         "//tensorflow:internal",
     ],
     deps = [
-        ":operation_interface",
-        ":tensor_handle_interface",
+        ":abstract_operation",
+        ":abstract_tensor_handle",
+        ":immediate_execution_tensor_handle",
+        "//tensorflow/c:tensor_interface",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "abstract_context",
+    hdrs = ["abstract_context.h"],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":abstract_function",
+        ":abstract_operation",
+    ],
+)
+
+cc_library(
+    name = "abstract_function",
+    hdrs = ["abstract_function.h"],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        "//tensorflow/c:c_api",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
+    ],
+)
+
+cc_library(
+    name = "immediate_execution_context",
+    hdrs = ["immediate_execution_context.h"],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":abstract_context",
+        ":immediate_execution_operation",
+        ":immediate_execution_tensor_handle",
         "//tensorflow/c:tensor_interface",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
@@ -217,7 +276,7 @@ cc_library(
         "//tensorflow:internal",
     ],
     deps = [
-        ":context_interface",
+        ":immediate_execution_context",
         "//tensorflow/c:conversion_macros",
     ],
 )
@@ -277,7 +336,7 @@ cc_library(
         "//tensorflow:internal",
     ],
     deps = [
-        ":operation_interface",
+        ":immediate_execution_operation",
         "//tensorflow/c:conversion_macros",
     ],
 )
@@ -300,7 +359,7 @@ cc_library(
         "//tensorflow:internal",
     ],
     deps = [
-        ":tensor_handle_interface",
+        ":immediate_execution_tensor_handle",
         "//tensorflow/c:conversion_macros",
     ],
 )
@@ -480,6 +539,9 @@ tf_cuda_library(
             ":tfe_context_internal",
             ":tfe_op_internal",
             ":tfe_tensorhandle_internal",
+            ":abstract_operation",
+            ":abstract_context",
+            ":abstract_tensor_handle",
             "//tensorflow/c:c_api",
             "//tensorflow/c:c_api_internal",
             "//tensorflow/core:core_cpu",
diff --git a/tensorflow/c/eager/abstract_context.h b/tensorflow/c/eager/abstract_context.h
new file mode 100644
index 00000000000..59c726349ac
--- /dev/null
+++ b/tensorflow/c/eager/abstract_context.h
@@ -0,0 +1,69 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_
+
+#include 
+
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+
+namespace tensorflow {
+
+// Abstract interface to a context.
+//
+// This serves as a factory for creating `AbstractOperation`s and for
+// registering traced functions.
+// Operations creation within a context can only be executed in that context
+// (for now at least).
+// Implementations of the context may contain some state e.g. an execution
+// environment, a traced representation etc.
+class AbstractContext {
+ protected:
+  enum AbstractContextKind { kTracing, kImmediateExecution };
+  explicit AbstractContext(AbstractContextKind kind) : kind_(kind) {}
+  virtual ~AbstractContext() {}
+
+ public:
+  AbstractContextKind getKind() const { return kind_; }
+
+  // Release any underlying resources, including the interface object.
+  //
+  // WARNING: The destructor of this class is marked as protected to disallow
+  // clients from directly destroying this object since it may manage it's own
+  // lifetime through ref counting. Thus clients MUST call Release() in order to
+  // destroy an instance of this class.
+  virtual void Release() = 0;
+
+  // Creates an operation builder and ties it to this context.
+  // The returned object can be used for setting operation's attributes,
+  // adding inputs and finally executing (immediately or lazily as in tracing)
+  // it in this context.
+  virtual AbstractOperation* CreateOperation() = 0;
+
+  // Registers a function with this context, after this the function is
+  // available to be called/referenced by its name in this context.
+  virtual Status RegisterFunction(AbstractFunction*) = 0;
+  // Remove a function. 'func' argument is the name of a previously added
+  // FunctionDef. The name is in fdef.signature.name.
+  virtual Status RemoveFunction(const string& func) = 0;
+
+ private:
+  const AbstractContextKind kind_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_CONTEXT_H_
diff --git a/tensorflow/c/eager/abstract_function.h b/tensorflow/c/eager/abstract_function.h
new file mode 100644
index 00000000000..e322b31f2b4
--- /dev/null
+++ b/tensorflow/c/eager/abstract_function.h
@@ -0,0 +1,46 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_
+
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+// A traced function: this hides the complexity of converting the serialized
+// representation between various supported formats e.g. FunctionDef and Mlir
+// function.
+class AbstractFunction {
+ protected:
+  enum AbstractFunctionKind { kGraphFunc, kMlirFunc };
+  explicit AbstractFunction(AbstractFunctionKind kind) : kind_(kind) {}
+
+ public:
+  // Returns which subclass is this instance of.
+  AbstractFunctionKind getKind() const { return kind_; }
+  virtual ~AbstractFunction() = default;
+
+  // Returns the AbstractFunction as a FunctionDef.
+  virtual Status GetFunctionDef(FunctionDef**) = 0;
+
+ private:
+  const AbstractFunctionKind kind_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_FUNCTION_H_
diff --git a/tensorflow/c/eager/operation_interface.h b/tensorflow/c/eager/abstract_operation.h
similarity index 80%
rename from tensorflow/c/eager/operation_interface.h
rename to tensorflow/c/eager/abstract_operation.h
index 844ba6c14bd..da4b6ecb75e 100644
--- a/tensorflow/c/eager/operation_interface.h
+++ b/tensorflow/c/eager/abstract_operation.h
@@ -12,24 +12,29 @@ 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.
 ==============================================================================*/
-#ifndef TENSORFLOW_C_EAGER_OPERATION_INTERFACE_H_
-#define TENSORFLOW_C_EAGER_OPERATION_INTERFACE_H_
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_
 
 #include "absl/types/span.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
 #include "tensorflow/c/tensor_interface.h"
-#include "tensorflow/core/framework/device_attributes.pb.h"
-#include "tensorflow/core/framework/op_def.pb.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/status.h"
 
-struct TFE_Op;
-
 namespace tensorflow {
 
 // Abstract interface to an operation.
-class AbstractOperationInterface {
+// This interface allows building and executing an operation in either
+// tracing or immediate execution mode.
+class AbstractOperation {
+ protected:
+  enum AbstractOperationKind { kTracing, kImmediateExecution };
+  explicit AbstractOperation(AbstractOperationKind kind) : kind_(kind) {}
+  virtual ~AbstractOperation() {}
+
  public:
+  AbstractOperationKind getKind() const { return kind_; }
+
   // Release any underlying resources, including the interface object.
   //
   // WARNING: The destructor of this class is marked as protected to disallow
@@ -38,7 +43,6 @@ class AbstractOperationInterface {
   // clients MUST call Release() in order to destroy an instance of this class.
   virtual void Release() = 0;
 
-  virtual void Clear() = 0;
   virtual Status Reset(const char* op, const char* raw_device_name) = 0;
 
   virtual const string& Name() const = 0;
@@ -66,12 +70,10 @@ class AbstractOperationInterface {
   // existing and given constraints will be performed.
   virtual Status SetDeviceName(const char* name) = 0;
 
-  virtual Status AddInput(AbstractTensorHandleInterface* input) = 0;
-  virtual Status AddInputList(
-      absl::Span inputs) = 0;
-  virtual Status Execute(absl::Span retvals,
+  virtual Status AddInput(AbstractTensorHandle* input) = 0;
+  virtual Status AddInputList(absl::Span inputs) = 0;
+  virtual Status Execute(absl::Span retvals,
                          int* num_retvals) = 0;
-  virtual const tensorflow::OpDef* OpDef() const = 0;
 
   virtual Status SetAttrString(const char* attr_name, const char* data,
                                size_t length) = 0;
@@ -82,7 +84,7 @@ class AbstractOperationInterface {
   virtual Status SetAttrShape(const char* attr_name, const int64_t* dims,
                               const int num_dims) = 0;
   virtual Status SetAttrFunction(const char* attr_name,
-                                 const AbstractOperationInterface* value) = 0;
+                                 const AbstractOperation* value) = 0;
   virtual Status SetAttrFunctionName(const char* attr_name, const char* value,
                                      size_t length) = 0;
   virtual Status SetAttrTensor(const char* attr_name,
@@ -102,19 +104,12 @@ class AbstractOperationInterface {
   virtual Status SetAttrShapeList(const char* attr_name, const int64_t** dims,
                                   const int* num_dims, int num_values) = 0;
   virtual Status SetAttrFunctionList(
-      const char* attr_name,
-      absl::Span values) = 0;
+      const char* attr_name, absl::Span values) = 0;
 
-  virtual Status InputLength(const char* input_name, int* length) = 0;
-  virtual Status OutputLength(const char* output_name, int* length) = 0;
-
-  // Experimental
-  virtual Status SetUseXla(bool enable) = 0;
-
- protected:
-  virtual ~AbstractOperationInterface() {}
+ private:
+  const AbstractOperationKind kind_;
 };
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_C_EAGER_OPERATION_INTERFACE_H_
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_OPERATION_H_
diff --git a/tensorflow/c/eager/abstract_tensor_handle.h b/tensorflow/c/eager/abstract_tensor_handle.h
new file mode 100644
index 00000000000..14acac29bb9
--- /dev/null
+++ b/tensorflow/c/eager/abstract_tensor_handle.h
@@ -0,0 +1,45 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
+#define TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
+
+namespace tensorflow {
+
+// Abstract interface to a Tensor handle in either tracing or immediate
+// execution mode.
+class AbstractTensorHandle {
+ protected:
+  enum AbstractTensorHandleKind { kTracing, kImmediateExecution };
+  explicit AbstractTensorHandle(AbstractTensorHandleKind kind) : kind_(kind) {}
+  virtual ~AbstractTensorHandle() {}
+
+ public:
+  AbstractTensorHandleKind getKind() const { return kind_; }
+
+  // Release any underlying resources, including the interface object.
+  //
+  // WARNING: The destructor of this class is marked as protected to disallow
+  // clients from directly destroying this object since it may manage it's own
+  // lifetime through ref counting. Thus this must be allocated on the heap and
+  // clients MUST call Release() in order to destroy an instance of this class.
+  virtual void Release() = 0;
+
+ private:
+  const AbstractTensorHandleKind kind_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc
index fdc91675f8b..4be3cdd7c2d 100644
--- a/tensorflow/c/eager/c_api.cc
+++ b/tensorflow/c/eager/c_api.cc
@@ -21,6 +21,8 @@ limitations under the License.
 #include 
 #include 
 
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+
 // clang-format off
 #include "tensorflow/core/platform/platform.h"
 // clang-format on
@@ -31,8 +33,8 @@ limitations under the License.
 #include "tensorflow/c/c_api_internal.h"
 #include "tensorflow/c/eager/c_api_experimental.h"
 #include "tensorflow/c/eager/c_api_internal.h"
-#include "tensorflow/c/eager/operation_interface.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/eager/tfe_context_internal.h"
 #include "tensorflow/c/eager/tfe_op_internal.h"
 #include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
@@ -1119,7 +1121,7 @@ size_t TFE_TensorHandleDeviceMemorySize(TFE_TensorHandle* h,
 
 TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name,
                   TF_Status* status) {
-  tensorflow::AbstractOperationInterface* new_op =
+  tensorflow::ImmediateExecutionOperation* new_op =
       tensorflow::unwrap(ctx)->CreateOperation();
   status->status = new_op->Reset(op_or_function_name, nullptr);
   if (!status->status.ok()) {
@@ -1164,7 +1166,9 @@ void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* input, TF_Status* status) {
 void TFE_OpAddInputList(TFE_Op* op, TFE_TensorHandle** inputs, int num_inputs,
                         TF_Status* status) {
   status->status = tensorflow::unwrap(op)->AddInputList(
-      {tensorflow::unwrap(inputs), static_cast(num_inputs)});
+      {reinterpret_cast(
+           tensorflow::unwrap(inputs)),
+       static_cast(num_inputs)});
 }
 
 TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name,
@@ -1324,7 +1328,9 @@ void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name,
 void TFE_OpSetAttrFunctionList(TFE_Op* op, const char* attr_name,
                                const TFE_Op** value, int num_values) {
   auto s = tensorflow::unwrap(op)->SetAttrFunctionList(
-      attr_name, {tensorflow::unwrap(value), static_cast(num_values)});
+      attr_name, {reinterpret_cast(
+                      tensorflow::unwrap(value)),
+                  static_cast(num_values)});
   if (!s.ok()) {
     LOG(WARNING) << "Unable to set attribute: " << attr_name;
   }
@@ -1368,7 +1374,10 @@ TF_CAPI_EXPORT extern int TFE_OpGetOutputLength(TFE_Op* op,
 void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals,
                  TF_Status* status) {
   status->status = tensorflow::unwrap(op)->Execute(
-      absl::MakeSpan(tensorflow::unwrap(retvals), *num_retvals), num_retvals);
+      absl::MakeSpan(reinterpret_cast(
+                         tensorflow::unwrap(retvals)),
+                     *num_retvals),
+      num_retvals);
 }
 
 TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h,
diff --git a/tensorflow/c/eager/c_api_experimental.cc b/tensorflow/c/eager/c_api_experimental.cc
index 0d71b11531b..9937fd7551f 100644
--- a/tensorflow/c/eager/c_api_experimental.cc
+++ b/tensorflow/c/eager/c_api_experimental.cc
@@ -38,7 +38,7 @@ using tensorflow::string;
 void TFE_OpReset(TFE_Op* op_to_reset, const char* op_or_function_name,
                  const char* raw_device_name, TF_Status* status) {
   if (op_to_reset) {
-    tensorflow::AbstractOperationInterface* op =
+    tensorflow::ImmediateExecutionOperation* op =
         tensorflow::unwrap(op_to_reset);
     op->Clear();
     status->status = op->Reset(op_or_function_name, raw_device_name);
diff --git a/tensorflow/c/eager/context_interface.h b/tensorflow/c/eager/immediate_execution_context.h
similarity index 78%
rename from tensorflow/c/eager/context_interface.h
rename to tensorflow/c/eager/immediate_execution_context.h
index e5a770a6826..0e3fe8cd4e1 100644
--- a/tensorflow/c/eager/context_interface.h
+++ b/tensorflow/c/eager/immediate_execution_context.h
@@ -12,15 +12,16 @@ 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.
 ==============================================================================*/
-#ifndef TENSORFLOW_C_EAGER_CONTEXT_INTERFACE_H_
-#define TENSORFLOW_C_EAGER_CONTEXT_INTERFACE_H_
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_
 
 #include 
 
 #include "absl/types/optional.h"
 #include "absl/types/span.h"
-#include "tensorflow/c/eager/operation_interface.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/tensor_interface.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/numeric_types.h"
@@ -34,16 +35,9 @@ namespace tensorflow {
 //
 // A context is responsible for creating key objects such as Tensors,
 // TensorHandles & Operations.
-class AbstractContextInterface {
+class ImmediateExecutionContext : public AbstractContext {
  public:
-  // Release any underlying resources, including the interface object.
-  //
-  // WARNING: The destructor of this class is marked as protected to disallow
-  // clients from directly destroying this object since it may manage it's own
-  // lifetime through ref counting. Thus clients MUST call Release() in order to
-  // destroy an instance of this class.
-  virtual void Release() = 0;
-
+  static constexpr AbstractContextKind kKind = kImmediateExecution;
   // Optimized scalar creation functions
   virtual AbstractTensorInterface* CreateInt64Scalar(int64 value) = 0;
   virtual AbstractTensorInterface* CreateUint64Scalar(uint64 value) = 0;
@@ -74,15 +68,15 @@ class AbstractContextInterface {
                                                 void* memory_releaser_arg) = 0;
 
   // Create a handle to wrap and manage a Tensor
-  virtual AbstractTensorHandleInterface* CreateLocalHandle(
+  virtual ImmediateExecutionTensorHandle* CreateLocalHandle(
       AbstractTensorInterface* t) = 0;
   // Copy the handle to another device.
-  virtual AbstractTensorHandleInterface* CopyTensorHandleToDevice(
-      AbstractTensorHandleInterface* handle, const char* device_name,
+  virtual ImmediateExecutionTensorHandle* CopyTensorHandleToDevice(
+      ImmediateExecutionTensorHandle* handle, const char* device_name,
       Status* status) = 0;
 
   // Create an operation to perform op execution
-  virtual AbstractOperationInterface* CreateOperation() = 0;
+  ImmediateExecutionOperation* CreateOperation() override = 0;
 
   // Returns whether the runtime is backed by TFRT or the legacy TF Eager
   // Runtime. This is necessary to decouple runtime-dependent
@@ -107,14 +101,12 @@ class AbstractContextInterface {
   // be executed as an op. Return error if the function with the same name
   // already exists.
   virtual Status AddFunctionDef(const FunctionDef& fdef) = 0;
-  // Remove a function. 'func' argument is the name of a previously added
-  // FunctionDef. The name is in fdef.signature.name.
-  virtual Status RemoveFunction(const string& func) = 0;
 
  protected:
-  virtual ~AbstractContextInterface() {}
+  ImmediateExecutionContext() : AbstractContext(kKind) {}
+  ~ImmediateExecutionContext() override {}
 };
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_C_EAGER_CONTEXT_INTERFACE_H_
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_
diff --git a/tensorflow/c/eager/immediate_execution_operation.h b/tensorflow/c/eager/immediate_execution_operation.h
new file mode 100644
index 00000000000..31413b5b4b9
--- /dev/null
+++ b/tensorflow/c/eager/immediate_execution_operation.h
@@ -0,0 +1,53 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/casts.h"
+#include "tensorflow/core/platform/status.h"
+
+struct TFE_Op;
+
+namespace tensorflow {
+
+// Abstract interface to an operation.
+class ImmediateExecutionOperation : public AbstractOperation {
+ public:
+  static constexpr AbstractOperationKind kKind = kImmediateExecution;
+  virtual void Clear() = 0;
+
+  virtual const tensorflow::OpDef* OpDef() const = 0;
+
+  virtual Status InputLength(const char* input_name, int* length) = 0;
+  virtual Status OutputLength(const char* output_name, int* length) = 0;
+
+  // Experimental
+  virtual Status SetUseXla(bool enable) = 0;
+
+ protected:
+  ImmediateExecutionOperation() : AbstractOperation(kKind) {}
+  ~ImmediateExecutionOperation() override {}
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_OPERATION_H_
diff --git a/tensorflow/c/eager/tensor_handle_interface.h b/tensorflow/c/eager/immediate_execution_tensor_handle.h
similarity index 74%
rename from tensorflow/c/eager/tensor_handle_interface.h
rename to tensorflow/c/eager/immediate_execution_tensor_handle.h
index 1ca40daec41..1f5a77e54ee 100644
--- a/tensorflow/c/eager/tensor_handle_interface.h
+++ b/tensorflow/c/eager/immediate_execution_tensor_handle.h
@@ -12,9 +12,10 @@ 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.
 ==============================================================================*/
-#ifndef TENSORFLOW_C_EAGER_TENSOR_HANDLE_INTERFACE_H_
-#define TENSORFLOW_C_EAGER_TENSOR_HANDLE_INTERFACE_H_
+#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_
+#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_
 
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
 #include "tensorflow/c/tensor_interface.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/status.h"
@@ -30,15 +31,9 @@ namespace tensorflow {
 // files. The interface lists the common functionality that must be provided by
 // any concrete implementation. However, in cases where the true concrete class
 // is needed a static_cast can be applied.
-class AbstractTensorHandleInterface {
+class ImmediateExecutionTensorHandle : public AbstractTensorHandle {
  public:
-  // Release any underlying resources, including the interface object.
-  //
-  // WARNING: The destructor of this class is marked as protected to disallow
-  // clients from directly destroying this object since it may manage it's own
-  // lifetime through ref counting. Thus this must be allocated on the heap and
-  // clients MUST call Release() in order to destroy an instance of this class.
-  virtual void Release() = 0;
+  static constexpr AbstractTensorHandleKind kKind = kImmediateExecution;
 
   // Returns tensor dtype.
   virtual tensorflow::DataType DataType() const = 0;
@@ -57,12 +52,13 @@ class AbstractTensorHandleInterface {
   virtual AbstractTensorInterface* Resolve(Status* status) = 0;
 
   // Return a copy of the handle.
-  virtual AbstractTensorHandleInterface* Copy() = 0;
+  virtual ImmediateExecutionTensorHandle* Copy() = 0;
 
  protected:
-  virtual ~AbstractTensorHandleInterface() {}
+  ImmediateExecutionTensorHandle() : AbstractTensorHandle(kKind) {}
+  ~ImmediateExecutionTensorHandle() override {}
 };
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_C_EAGER_TENSOR_HANDLE_INTERFACE_H_
+#endif  // TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_TENSOR_HANDLE_H_
diff --git a/tensorflow/c/eager/tfe_context_internal.h b/tensorflow/c/eager/tfe_context_internal.h
index 1d29bee9ee3..1f2035317fa 100644
--- a/tensorflow/c/eager/tfe_context_internal.h
+++ b/tensorflow/c/eager/tfe_context_internal.h
@@ -16,7 +16,7 @@ limitations under the License.
 #define TENSORFLOW_C_EAGER_TFE_CONTEXT_INTERNAL_H_
 
 #include "tensorflow/c/conversion_macros.h"
-#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
 
 // Wraps a pointer to a context implementation.
 //
@@ -28,7 +28,7 @@ typedef struct TFE_Context TFE_Context;
 
 namespace tensorflow {
 
-DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractContextInterface, TFE_Context);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionContext, TFE_Context);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/c/eager/tfe_op_internal.h b/tensorflow/c/eager/tfe_op_internal.h
index 6ca7f741d16..3fe94d358b6 100644
--- a/tensorflow/c/eager/tfe_op_internal.h
+++ b/tensorflow/c/eager/tfe_op_internal.h
@@ -16,7 +16,7 @@ limitations under the License.
 #define TENSORFLOW_C_EAGER_TFE_OP_INTERNAL_H_
 
 #include "tensorflow/c/conversion_macros.h"
-#include "tensorflow/c/eager/operation_interface.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
 
 // Wraps a pointer to an operation implementation.
 //
@@ -28,8 +28,8 @@ typedef struct TFE_Op TFE_Op;
 
 namespace tensorflow {
 
-DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractOperationInterface, TFE_Op);
-DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractOperationInterface*, TFE_Op*);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionOperation, TFE_Op);
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionOperation*, TFE_Op*);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/c/eager/tfe_tensorhandle_internal.h b/tensorflow/c/eager/tfe_tensorhandle_internal.h
index 543e5f1d932..308e8c24e2c 100644
--- a/tensorflow/c/eager/tfe_tensorhandle_internal.h
+++ b/tensorflow/c/eager/tfe_tensorhandle_internal.h
@@ -16,7 +16,7 @@ limitations under the License.
 #define TENSORFLOW_C_EAGER_TFE_TENSORHANDLE_INTERNAL_H_
 
 #include "tensorflow/c/conversion_macros.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 
 // Wraps a pointer to a tensor handle implementation.
 //
@@ -28,9 +28,9 @@ typedef struct TFE_TensorHandle TFE_TensorHandle;
 
 namespace tensorflow {
 
-DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractTensorHandleInterface,
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionTensorHandle,
                             TFE_TensorHandle);
-DEFINE_CONVERSION_FUNCTIONS(tensorflow::AbstractTensorHandleInterface*,
+DEFINE_CONVERSION_FUNCTIONS(tensorflow::ImmediateExecutionTensorHandle*,
                             TFE_TensorHandle*);
 
 }  // namespace tensorflow
diff --git a/tensorflow/c/experimental/saved_model/core/BUILD b/tensorflow/c/experimental/saved_model/core/BUILD
index 2e817ed02e0..dbe1b6d656c 100644
--- a/tensorflow/c/experimental/saved_model/core/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/BUILD
@@ -23,8 +23,8 @@ cc_library(
     ],
     deps = [
         ":function_metadata",
-        "//tensorflow/c/eager:operation_interface",
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:immediate_execution_operation",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
         "//tensorflow/core:protos_all_cc",
     ],
 )
diff --git a/tensorflow/c/experimental/saved_model/core/concrete_function.cc b/tensorflow/c/experimental/saved_model/core/concrete_function.cc
index d5da2ca9bf4..41bae4352fc 100644
--- a/tensorflow/c/experimental/saved_model/core/concrete_function.cc
+++ b/tensorflow/c/experimental/saved_model/core/concrete_function.cc
@@ -15,12 +15,12 @@ limitations under the License.
 
 #include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
 
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
 
 namespace tensorflow {
 
-const std::vector&
+const std::vector&
 ConcreteFunction::GetCaptures() const {
   return captures_;
 }
diff --git a/tensorflow/c/experimental/saved_model/core/concrete_function.h b/tensorflow/c/experimental/saved_model/core/concrete_function.h
index 6f8a5375277..22535641ef5 100644
--- a/tensorflow/c/experimental/saved_model/core/concrete_function.h
+++ b/tensorflow/c/experimental/saved_model/core/concrete_function.h
@@ -18,8 +18,8 @@ limitations under the License.
 
 #include 
 
-#include "tensorflow/c/eager/operation_interface.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
 #include "tensorflow/core/framework/function.pb.h"
 
@@ -38,15 +38,15 @@ class ConcreteFunction {
   virtual ~ConcreteFunction() = 0;
 
   // This method returns the "Call" Op used to execute the function.
-  virtual AbstractOperationInterface* GetCallOp() = 0;
+  virtual ImmediateExecutionOperation* GetCallOp() = 0;
 
-  const std::vector& GetCaptures()
+  const std::vector& GetCaptures()
       const;
   const FunctionMetadata& GetFunctionMetadata() const;
 
  private:
   FunctionMetadata metadata_;
-  std::vector captures_;
+  std::vector captures_;
   FunctionDef* function_;
 };
 
diff --git a/tensorflow/c/experimental/saved_model/core/ops/BUILD b/tensorflow/c/experimental/saved_model/core/ops/BUILD
index aa909c692ca..8c4c41c6d75 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/ops/BUILD
@@ -20,7 +20,7 @@ cc_library(
         "owned_eager_op.h",
     ],
     deps = [
-        "//tensorflow/c/eager:operation_interface",
+        "//tensorflow/c/eager:immediate_execution_operation",
     ],
 )
 
@@ -30,7 +30,7 @@ cc_library(
         "owned_tensor_handle.h",
     ],
     deps = [
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
         "//tensorflow/core/common_runtime/eager:tensor_handle",
     ],
 )
@@ -39,7 +39,7 @@ cc_library(
     name = "owned_eager_context",
     hdrs = ["owned_eager_context.h"],
     deps = [
-        "//tensorflow/c/eager:context_interface",
+        "//tensorflow/c/eager:immediate_execution_context",
         "//tensorflow/core/common_runtime/eager:context",
     ],
 )
@@ -63,8 +63,9 @@ cc_library(
     deps = [
         ":owned_eager_op",
         ":owned_tensor_handle",
-        "//tensorflow/c/eager:context_interface",
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:immediate_execution_context",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
index 300059cd069..d944fcb51a2 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_context.h
@@ -18,14 +18,14 @@ limitations under the License.
 
 #include 
 
-#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/core/common_runtime/eager/context.h"
 
 namespace tensorflow {
 namespace internal {
 
-struct AbstractContextInterfaceDeleter {
-  void operator()(AbstractContextInterface* p) const {
+struct ImmediateExecutionContextDeleter {
+  void operator()(ImmediateExecutionContext* p) const {
     if (p != nullptr) {
       p->Release();
     }
@@ -43,8 +43,8 @@ struct EagerContextDeleter {
 }  // namespace internal
 
 using AbstractContextPtr =
-    std::unique_ptr;
+    std::unique_ptr;
 
 using EagerContextPtr =
     std::unique_ptr;
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
index c6b21578820..b3a08334a97 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h
@@ -18,13 +18,13 @@ limitations under the License.
 
 #include 
 
-#include "tensorflow/c/eager/operation_interface.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
 
 namespace tensorflow {
 namespace internal {
 
-struct AbstractOperationInterfaceDeleter {
-  void operator()(AbstractOperationInterface* p) const {
+struct ImmediateExecutionOperationDeleter {
+  void operator()(ImmediateExecutionOperation* p) const {
     if (p != nullptr) {
       p->Release();
     }
@@ -34,8 +34,8 @@ struct AbstractOperationInterfaceDeleter {
 }  // namespace internal
 
 using AbstractOpPtr =
-    std::unique_ptr;
+    std::unique_ptr;
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
index e98d6554afb..c52ebaa2479 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
+++ b/tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h
@@ -18,7 +18,7 @@ limitations under the License.
 
 #include 
 
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/core/common_runtime/eager/tensor_handle.h"
 
 namespace tensorflow {
@@ -33,7 +33,7 @@ struct TensorHandleDeleter {
 };
 
 struct AbstractTensorHandleDeleter {
-  void operator()(AbstractTensorHandleInterface* p) const {
+  void operator()(ImmediateExecutionTensorHandle* p) const {
     if (p != nullptr) {
       p->Release();
     }
@@ -46,7 +46,7 @@ using TensorHandlePtr =
     std::unique_ptr;
 
 using AbstractTensorHandlePtr =
-    std::unique_ptr;
 
 }  // namespace tensorflow
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
index a3b3ace7be9..eb06662722e 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
@@ -16,7 +16,8 @@ limitations under the License.
 #include "tensorflow/c/experimental/saved_model/core/ops/variable_ops.h"
 
 #include "absl/types/span.h"
-#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/c/experimental/saved_model/core/ops/owned_eager_op.h"
 #include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -32,7 +33,7 @@ namespace internal {
 static const char kNoSharingResourceID[] =
     "cd2c89b7-88b7-44c8-ad83-06c2a9158347";
 
-Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
+Status CreateUninitializedResourceVariable(ImmediateExecutionContext* ctx,
                                            DataType dtype, TensorShape shape,
                                            AbstractTensorHandlePtr* handle) {
   AbstractOpPtr varhandle_op = AbstractOpPtr(ctx->CreateOperation());
@@ -50,17 +51,20 @@ Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
   TF_RETURN_IF_ERROR(varhandle_op->SetAttrString(
       "shared_name", kNoSharingResourceID, strlen(kNoSharingResourceID)));
 
-  AbstractTensorHandleInterface* var_handle = nullptr;
+  AbstractTensorHandle* var_handle = nullptr;
   int num_retvals = 1;
   TF_RETURN_IF_ERROR(varhandle_op->Execute(
       absl::MakeSpan(&var_handle, num_retvals), &num_retvals));
-  handle->reset(var_handle);
+  if (var_handle->getKind() != ImmediateExecutionTensorHandle::kKind) {
+    return errors::Internal("Unexpected tensor handle kind.");
+  }
+  handle->reset(reinterpret_cast(var_handle));
   return Status();
 }
 
-Status AssignVariable(AbstractContextInterface* ctx,
-                      AbstractTensorHandleInterface* variable_handle,
-                      DataType dtype, AbstractTensorHandleInterface* value) {
+Status AssignVariable(ImmediateExecutionContext* ctx,
+                      ImmediateExecutionTensorHandle* variable_handle,
+                      DataType dtype, ImmediateExecutionTensorHandle* value) {
   AbstractOpPtr assign_op(ctx->CreateOperation());
   TF_RETURN_IF_ERROR(assign_op->Reset("AssignVariableOp", nullptr));
   TF_RETURN_IF_ERROR(assign_op->SetAttrType("dtype", dtype));
@@ -72,24 +76,27 @@ Status AssignVariable(AbstractContextInterface* ctx,
   return Status();
 }
 
-Status ReadVariable(AbstractContextInterface* ctx,
-                    AbstractTensorHandleInterface* variable_handle,
+Status ReadVariable(ImmediateExecutionContext* ctx,
+                    ImmediateExecutionTensorHandle* variable_handle,
                     DataType dtype, AbstractTensorHandlePtr* output) {
   AbstractOpPtr read_op = AbstractOpPtr(ctx->CreateOperation());
   TF_RETURN_IF_ERROR(read_op->Reset("ReadVariableOp", nullptr));
   TF_RETURN_IF_ERROR(read_op->SetAttrType("dtype", dtype));
   TF_RETURN_IF_ERROR(read_op->AddInput(variable_handle));
 
-  AbstractTensorHandleInterface* value = nullptr;
+  AbstractTensorHandle* value = nullptr;
   int num_retvals = 1;
   TF_RETURN_IF_ERROR(
       read_op->Execute(absl::MakeSpan(&value, num_retvals), &num_retvals));
-  output->reset(value);
+  if (value->getKind() != ImmediateExecutionTensorHandle::kKind) {
+    return errors::Internal("Unexpected tensor handle kind.");
+  }
+  output->reset(reinterpret_cast(value));
   return Status();
 }
 
-Status DestroyResource(AbstractContextInterface* ctx,
-                       AbstractTensorHandleInterface* handle) {
+Status DestroyResource(ImmediateExecutionContext* ctx,
+                       ImmediateExecutionTensorHandle* handle) {
   AbstractOpPtr destroy_op = AbstractOpPtr(ctx->CreateOperation());
   TF_RETURN_IF_ERROR(destroy_op->Reset("DestroyResourceOp", nullptr));
   TF_RETURN_IF_ERROR(destroy_op->SetAttrBool("ignore_lookup_error", true));
diff --git a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
index 8a410328b9e..038b2c3d62a 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
+++ b/tensorflow/c/experimental/saved_model/core/ops/variable_ops.h
@@ -16,8 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H
 #define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H
 
-#include "tensorflow/c/eager/context_interface.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/experimental/saved_model/core/ops/owned_tensor_handle.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
@@ -30,7 +30,7 @@ namespace internal {
 // TensorHandle associated with the variable. This is equivalent to creating an
 // unitialized TF2 tf.Variable.
 // https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L1867-L1872
-Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
+Status CreateUninitializedResourceVariable(ImmediateExecutionContext* ctx,
                                            DataType dtype, TensorShape shape,
                                            AbstractTensorHandlePtr* handle);
 
@@ -39,22 +39,22 @@ Status CreateUninitializedResourceVariable(AbstractContextInterface* ctx,
 // underlying variable for `variable_handle`. Note that it is illegal to assign
 // a variable to a Tensor with a different dtype than what the variable was
 // created with.
-Status AssignVariable(AbstractContextInterface* ctx,
-                      AbstractTensorHandleInterface* variable_handle,
-                      DataType dtype, AbstractTensorHandleInterface* value);
+Status AssignVariable(ImmediateExecutionContext* ctx,
+                      ImmediateExecutionTensorHandle* variable_handle,
+                      DataType dtype, ImmediateExecutionTensorHandle* value);
 
 // Executes a ReadVariableOp using `ctx`. This reads the underlying variable
 // value of `variable_handle` and copies the value to `output`. `dtype` must be
 // the dtype of the variable associated with `variable_handle`.
-Status ReadVariable(AbstractContextInterface* ctx,
-                    AbstractTensorHandleInterface* variable_handle,
+Status ReadVariable(ImmediateExecutionContext* ctx,
+                    ImmediateExecutionTensorHandle* variable_handle,
                     DataType dtype, AbstractTensorHandlePtr* output);
 
 // Executes DestroyResourceOp on `handle`, using `ctx`. This is equivalent to
 // the cleanup that occurs in a tf.Variable's EagerResourceDeleter:
 // https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L289-L290
-Status DestroyResource(AbstractContextInterface* ctx,
-                       AbstractTensorHandleInterface* handle);
+Status DestroyResource(ImmediateExecutionContext* ctx,
+                       ImmediateExecutionTensorHandle* handle);
 
 }  // namespace internal
 }  // namespace tensorflow
diff --git a/tensorflow/c/experimental/saved_model/internal/BUILD b/tensorflow/c/experimental/saved_model/internal/BUILD
index 72474940c16..888c284bb12 100644
--- a/tensorflow/c/experimental/saved_model/internal/BUILD
+++ b/tensorflow/c/experimental/saved_model/internal/BUILD
@@ -178,7 +178,7 @@ cc_library(
         ":tensorhandle_list_type",
         "//tensorflow/c:c_api_macros",
         "//tensorflow/c/eager:c_api",
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
         "//tensorflow/c/eager:tfe_tensorhandle_internal",
     ],
 )
@@ -190,7 +190,7 @@ cc_library(
     ],
     deps = [
         "//tensorflow/c:conversion_macros",
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
     ],
 )
 
diff --git a/tensorflow/c/experimental/saved_model/internal/tensorhandle_list.cc b/tensorflow/c/experimental/saved_model/internal/tensorhandle_list.cc
index 7d018658101..c8f00c1f7c0 100644
--- a/tensorflow/c/experimental/saved_model/internal/tensorhandle_list.cc
+++ b/tensorflow/c/experimental/saved_model/internal/tensorhandle_list.cc
@@ -17,7 +17,7 @@ limitations under the License.
 
 #include 
 
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
 #include "tensorflow/c/experimental/saved_model/internal/tensorhandle_list_type.h"
 
diff --git a/tensorflow/c/experimental/saved_model/internal/tensorhandle_list_type.h b/tensorflow/c/experimental/saved_model/internal/tensorhandle_list_type.h
index 8cbec2806a8..566417df025 100644
--- a/tensorflow/c/experimental/saved_model/internal/tensorhandle_list_type.h
+++ b/tensorflow/c/experimental/saved_model/internal/tensorhandle_list_type.h
@@ -19,7 +19,7 @@ limitations under the License.
 #include 
 
 #include "tensorflow/c/conversion_macros.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 
 // Internal structures used by the SavedModel C API. These are likely to
 // change and should not be depended on.
@@ -29,7 +29,7 @@ typedef struct TF_TensorHandleList TF_TensorHandleList;
 namespace tensorflow {
 
 DEFINE_CONVERSION_FUNCTIONS(
-    std::vector,
+    std::vector,
     TF_TensorHandleList)
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index e4f4c483209..fb69bcb7ab5 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -32,6 +32,8 @@ tf_cuda_library(
         ":tensor_handle",
         "//tensorflow/c:c_api_internal",
         "//tensorflow/c:tf_tensor_internal",
+        "//tensorflow/c/eager:abstract_function",
+        "//tensorflow/core/platform:errors",
     ],
     alwayslink = 1,
 )
@@ -74,9 +76,9 @@ tf_cuda_library(
         ":kernel_and_device",
         "@com_google_absl//absl/container:flat_hash_map",
         "//tensorflow/c:tf_tensor_internal",
-        "//tensorflow/c/eager:context_interface",
-        "//tensorflow/c/eager:tensor_handle_interface",
-        "//tensorflow/c/eager:operation_interface",
+        "//tensorflow/c/eager:immediate_execution_context",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/c/eager:immediate_execution_operation",
         "//tensorflow/core/distributed_runtime:rendezvous_mgr_interface",
         "//tensorflow/core/distributed_runtime:worker_env",
     ] + select({
@@ -137,8 +139,10 @@ tf_cuda_library(
         "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
         "//tensorflow/c:tf_tensor_internal",
-        "//tensorflow/c/eager:operation_interface",
-        "//tensorflow/c/eager:tensor_handle_interface",
+        "//tensorflow/c/eager:immediate_execution_operation",
+        "//tensorflow/c/eager:abstract_operation",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/c/eager:abstract_tensor_handle",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:lib",
@@ -211,7 +215,7 @@ tf_cuda_library(
             "@com_google_absl//absl/strings",
             "@com_google_absl//absl/types:variant",
             "//tensorflow/c:tf_tensor_internal",
-            "//tensorflow/c/eager:tensor_handle_interface",
+            "//tensorflow/c/eager:immediate_execution_tensor_handle",
             "//tensorflow/core:core_cpu_lib",
             "//tensorflow/core:framework",
             "//tensorflow/core:framework_internal",
@@ -496,6 +500,8 @@ cc_library(
         "//tensorflow/c:tf_tensor_internal",
         "//tensorflow/compiler/jit:common",
         "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/c/eager:abstract_function",
+        "//tensorflow/core/platform:errors",
     ] + select({
         "//tensorflow:android": [
             "//tensorflow/core:portable_tensorflow_lib_lite",
diff --git a/tensorflow/core/common_runtime/eager/context.cc b/tensorflow/core/common_runtime/eager/context.cc
index 970c2bcbb89..6dc0a3a8200 100644
--- a/tensorflow/core/common_runtime/eager/context.cc
+++ b/tensorflow/core/common_runtime/eager/context.cc
@@ -30,8 +30,6 @@ limitations under the License.
 
 #include "tensorflow/c/tf_tensor.h"
 #include "tensorflow/c/tf_tensor_internal.h"
-#include "tensorflow/c/eager/operation_interface.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
 #include "tensorflow/core/common_runtime/collective_executor_mgr.h"
 #include "tensorflow/core/common_runtime/collective_param_resolver_local.h"
 #include "tensorflow/core/common_runtime/colocation_graph.h"
diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h
index cb6d09f8f1d..141327c08cb 100644
--- a/tensorflow/core/common_runtime/eager/context.h
+++ b/tensorflow/core/common_runtime/eager/context.h
@@ -33,7 +33,7 @@ limitations under the License.
 
 #include "absl/types/optional.h"
 #include "absl/container/flat_hash_map.h"
-#include "tensorflow/c/eager/context_interface.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/core/common_runtime/composite_device.h"
 #include "tensorflow/core/common_runtime/device_factory.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
@@ -135,7 +135,7 @@ class CustomDevice {
 // TensorHandles may be placed either on custom or physical devices.
 using VariantDevice = absl::variant;
 
-class EagerContext : public AbstractContextInterface, public core::RefCounted {
+class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
  public:
   static constexpr uint64 kInvalidContextId = 0;
 
@@ -178,12 +178,14 @@ class EagerContext : public AbstractContextInterface, public core::RefCounted {
                                         MemoryReleaser memory_releaser,
                                         void* memory_releaser_arg) override;
 
-  AbstractTensorHandleInterface* CreateLocalHandle(
+  ImmediateExecutionTensorHandle* CreateLocalHandle(
       AbstractTensorInterface* t) override;
-  AbstractTensorHandleInterface* CopyTensorHandleToDevice(
-      AbstractTensorHandleInterface* handle, const char* device_name,
+  ImmediateExecutionTensorHandle* CopyTensorHandleToDevice(
+      ImmediateExecutionTensorHandle* handle, const char* device_name,
       Status* status) override;
-  AbstractOperationInterface* CreateOperation() override;
+  ImmediateExecutionOperation* CreateOperation() override;
+
+  Status RegisterFunction(AbstractFunction* f) override;
 
   bool UsesTFRT() override;
 
@@ -716,7 +718,7 @@ class EagerContext : public AbstractContextInterface, public core::RefCounted {
   std::function resource_deallocator_ = nullptr;
 };
 
-inline EagerContext* ContextFromInterface(AbstractContextInterface* context) {
+inline EagerContext* ContextFromInterface(ImmediateExecutionContext* context) {
   return down_cast(context);
 }
 
diff --git a/tensorflow/core/common_runtime/eager/core.cc b/tensorflow/core/common_runtime/eager/core.cc
index e342f6ae6cd..3d37250a4fe 100644
--- a/tensorflow/core/common_runtime/eager/core.cc
+++ b/tensorflow/core/common_runtime/eager/core.cc
@@ -13,11 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include "tensorflow/c/c_api_internal.h"
+#include "tensorflow/c/eager/abstract_function.h"
 #include "tensorflow/c/tf_tensor_internal.h"
 #include "tensorflow/core/common_runtime/eager/context.h"
 #include "tensorflow/core/common_runtime/eager/eager_operation.h"
 #include "tensorflow/core/common_runtime/eager/execute.h"
 #include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace {
 
@@ -112,8 +114,8 @@ AbstractTensorInterface* TensorHandle::Resolve(Status* status) {
   }
 }
 
-AbstractTensorHandleInterface* EagerContext::CopyTensorHandleToDevice(
-    AbstractTensorHandleInterface* handle, const char* device_name,
+ImmediateExecutionTensorHandle* EagerContext::CopyTensorHandleToDevice(
+    ImmediateExecutionTensorHandle* handle, const char* device_name,
     Status* status) {
   TensorHandle* input = TensorHandleFromInterface(handle);
   TensorHandle* result = nullptr;
@@ -158,7 +160,7 @@ AbstractTensorHandleInterface* EagerContext::CopyTensorHandleToDevice(
 // here to a circular BUILD dep issue. If we move this to context.cc, then we
 // will have the circular dependency of:
 //   context -> tensor_handle -> remote_tensor_handle_data -> context
-AbstractTensorHandleInterface* EagerContext::CreateLocalHandle(
+ImmediateExecutionTensorHandle* EagerContext::CreateLocalHandle(
     AbstractTensorInterface* t) {
   Tensor tensor = TensorFromInterface(t);
   return TensorHandle::CreateLocalHandle(std::move(tensor), /*d=*/HostCPU(),
@@ -168,14 +170,23 @@ AbstractTensorHandleInterface* EagerContext::CreateLocalHandle(
 // TODO(b/152902651): We have to keep this function here since EagerOperation
 // depends on EagerContext. Thus, the context build target can't depend on
 // EagerOperation.
-AbstractOperationInterface* EagerContext::CreateOperation() {
+ImmediateExecutionOperation* EagerContext::CreateOperation() {
   return new EagerOperation(this);
 }
 
+Status EagerContext::RegisterFunction(AbstractFunction* f) {
+  FunctionDef* fdef;
+  TF_RETURN_IF_ERROR(f->GetFunctionDef(&fdef));
+  if (!fdef) {
+    return errors::InvalidArgument("GetFunctionDef returned nullptr.");
+  }
+  return AddFunctionDef(*fdef);
+}
+
 // TODO(b/152902651): Once we move many execute.cc functions into
 // eager_operation.cc we can avoid a circular dependency between them.
-Status EagerOperation::Execute(
-    absl::Span retvals, int* num_retvals) {
+Status EagerOperation::Execute(absl::Span retvals,
+                               int* num_retvals) {
   return EagerExecute(
       this, reinterpret_cast(retvals.data()),
       num_retvals);
diff --git a/tensorflow/core/common_runtime/eager/eager_operation.cc b/tensorflow/core/common_runtime/eager/eager_operation.cc
index 090bfef46bd..073095e64d1 100644
--- a/tensorflow/core/common_runtime/eager/eager_operation.cc
+++ b/tensorflow/core/common_runtime/eager/eager_operation.cc
@@ -15,7 +15,9 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/eager_operation.h"
 
 #include "absl/types/span.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/tf_tensor_internal.h"
 #include "tensorflow/core/common_runtime/eager/attr_builder.h"
 #include "tensorflow/core/common_runtime/input_colocation_exemption_registry.h"
@@ -91,8 +93,8 @@ Status EagerOperation::SetAttrShape(const char* attr_name, const int64_t* dims,
   return Status::OK();
 }
 
-Status EagerOperation::SetAttrFunction(
-    const char* attr_name, const AbstractOperationInterface* value) {
+Status EagerOperation::SetAttrFunction(const char* attr_name,
+                                       const AbstractOperation* value) {
   AttrValue attr_value;
   NameAttrList* func = attr_value.mutable_func();
   func->set_name(value->Name());
@@ -194,8 +196,7 @@ Status EagerOperation::SetAttrShapeList(const char* attr_name,
 }
 
 Status EagerOperation::SetAttrFunctionList(
-    const char* attr_name,
-    absl::Span values) {
+    const char* attr_name, absl::Span values) {
   size_t num_values = values.size();
   std::unique_ptr funcs(new NameAttrList[num_values]);
   for (int i = 0; i < num_values; i++) {
@@ -253,14 +254,13 @@ Status EagerOperation::OutputLength(const char* output_name, int* length) {
   return Status::OK();
 }
 
-Status EagerOperation::AddInput(AbstractTensorHandleInterface* input) {
+Status EagerOperation::AddInput(AbstractTensorHandle* input) {
   TensorHandle* h = TensorHandleFromInterface(input);
   AddTensorHandle(h);
   return MaybeInferSingleInputAttrs(h);
 }
 
-Status EagerOperation::AddInputList(
-    absl::Span inputs) {
+Status EagerOperation::AddInputList(absl::Span inputs) {
   for (auto& input : inputs) {
     TensorHandle* h = TensorHandleFromInterface(input);
     AddTensorHandle(h);
diff --git a/tensorflow/core/common_runtime/eager/eager_operation.h b/tensorflow/core/common_runtime/eager/eager_operation.h
index 14268ef2630..963aed25733 100644
--- a/tensorflow/core/common_runtime/eager/eager_operation.h
+++ b/tensorflow/core/common_runtime/eager/eager_operation.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "absl/types/variant.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
 #include "tensorflow/core/common_runtime/eager/attr_builder.h"
 #include "tensorflow/core/common_runtime/eager/context.h"
 #include "tensorflow/core/common_runtime/eager/eager_executor.h"
@@ -31,7 +32,7 @@ limitations under the License.
 
 namespace tensorflow {
 
-class EagerOperation : public AbstractOperationInterface {
+class EagerOperation : public ImmediateExecutionOperation {
  public:
   explicit EagerOperation(tensorflow::EagerContext* ctx) : ctx_(*ctx) {}
   ~EagerOperation() override {
@@ -56,7 +57,7 @@ class EagerOperation : public AbstractOperationInterface {
   }
 
   // Replaces the previous device name with the given one (see
-  // AbstractOperationInterface::SetDeviceName for more details).
+  // AbstractOperation::SetDeviceName for more details).
   //
   // This also resets the internal device pointer, unless the given name refers
   // to a known custom device, in which case the internal device pointer is
@@ -76,10 +77,9 @@ class EagerOperation : public AbstractOperationInterface {
 
   Status SetAttrValue(const char* attr_name, const AttrValue& value);
 
-  Status AddInput(AbstractTensorHandleInterface* input) override;
-  Status AddInputList(
-      absl::Span inputs) override;
-  Status Execute(absl::Span retvals,
+  Status AddInput(AbstractTensorHandle* input) override;
+  Status AddInputList(absl::Span inputs) override;
+  Status Execute(absl::Span retvals,
                  int* num_retvals) override;
   const tensorflow::OpDef* OpDef() const override { return op_def_; };
 
@@ -92,7 +92,7 @@ class EagerOperation : public AbstractOperationInterface {
   Status SetAttrShape(const char* attr_name, const int64_t* dims,
                       const int num_dims) override;
   Status SetAttrFunction(const char* attr_name,
-                         const AbstractOperationInterface* value) override;
+                         const AbstractOperation* value) override;
   Status SetAttrFunctionName(const char* attr_name, const char* data,
                              size_t length) override;
   Status SetAttrTensor(const char* attr_name,
@@ -111,7 +111,7 @@ class EagerOperation : public AbstractOperationInterface {
                           const int* num_dims, int num_values) override;
   Status SetAttrFunctionList(
       const char* attr_name,
-      absl::Span values) override;
+      absl::Span values) override;
 
   Status InputLength(const char* input_name, int* length) override;
   Status OutputLength(const char* output_name, int* length) override;
@@ -235,7 +235,7 @@ inline void EagerOperation::UpdateInput(int i, TensorHandle* h) {
 }
 
 inline EagerOperation* OperationFromInterface(
-    AbstractOperationInterface* operation) {
+    ImmediateExecutionOperation* operation) {
   return down_cast(operation);
 }
 
diff --git a/tensorflow/core/common_runtime/eager/tensor_handle.cc b/tensorflow/core/common_runtime/eager/tensor_handle.cc
index 9b82c556cd0..9e607c97683 100644
--- a/tensorflow/core/common_runtime/eager/tensor_handle.cc
+++ b/tensorflow/core/common_runtime/eager/tensor_handle.cc
@@ -1071,7 +1071,7 @@ const char* TensorHandle::BackingDeviceName(Status* status) const {
   }
 }
 
-tensorflow::AbstractTensorHandleInterface* TensorHandle::Copy() {
+tensorflow::ImmediateExecutionTensorHandle* TensorHandle::Copy() {
   Ref();
   return this;
 }
diff --git a/tensorflow/core/common_runtime/eager/tensor_handle.h b/tensorflow/core/common_runtime/eager/tensor_handle.h
index 5e7638ae03c..a14df475e0f 100644
--- a/tensorflow/core/common_runtime/eager/tensor_handle.h
+++ b/tensorflow/core/common_runtime/eager/tensor_handle.h
@@ -31,7 +31,7 @@ limitations under the License.
 // clang-format on
 
 #include "absl/types/variant.h"
-#include "tensorflow/c/eager/tensor_handle_interface.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/core/common_runtime/device.h"
 #include "tensorflow/core/common_runtime/eager/eager_executor.h"
 #include "tensorflow/core/common_runtime/eager/tensor_handle_data.h"
@@ -53,7 +53,7 @@ class EagerContext;
 // Associates a Tensor and a Device, used in the eager runtime. Internal version
 // of the TFE_TensorHandle struct and the python EagerTensor class
 // (unrelated to python TensorHandle).
-class TensorHandle : public AbstractTensorHandleInterface,
+class TensorHandle : public ImmediateExecutionTensorHandle,
                      public core::RefCounted {
   // TensorHandle for dtype != DT_RESOURCE
   TensorHandle(tensorflow::Tensor&& t, Device* d, Device* op_device,
@@ -121,7 +121,7 @@ class TensorHandle : public AbstractTensorHandleInterface,
   const char* BackingDeviceName(Status* status) const override;
   AbstractTensorInterface* Resolve(Status* status) override;
 
-  AbstractTensorHandleInterface* Copy() override;
+  ImmediateExecutionTensorHandle* Copy() override;
 
   // Return the Tensor from the default device.
   Status Tensor(const tensorflow::Tensor** t) const;
@@ -372,12 +372,12 @@ const VariantDevice kVariantDeviceNull = static_cast(nullptr);
 // Returns the device backing the resource. Else, returns nullptr.
 Device* GetResourceDevice(const ResourceHandle& handle, EagerContext* ctx);
 
-class TensorHandleInterface : public AbstractTensorHandleInterface {
+class TensorHandleInterface : public ImmediateExecutionTensorHandle {
  public:
 };
 
-inline TensorHandle* TensorHandleFromInterface(
-    AbstractTensorHandleInterface* handle) {
+template 
+inline TensorHandle* TensorHandleFromInterface(T* handle) {
   return down_cast(handle);
 }
 
diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc
index 639f623bd1a..b9ff474caab 100644
--- a/tensorflow/python/eager/pywrap_tfe_src.cc
+++ b/tensorflow/python/eager/pywrap_tfe_src.cc
@@ -2008,7 +2008,7 @@ bool ListContainsNone(PyObject* list) {
 
 static PyTapeTensor TapeTensorFromTensor(PyObject* tensor) {
   if (EagerTensor_CheckExact(tensor)) {
-    tensorflow::AbstractTensorHandleInterface* handle =
+    tensorflow::ImmediateExecutionTensorHandle* handle =
         tensorflow::unwrap(EagerTensor_Handle(tensor));
     tensorflow::int64 id = PyEagerTensor_ID(tensor);
     tensorflow::DataType dtype =
@@ -3869,7 +3869,7 @@ tensorflow::Status TFE_Py_EncodeTensor(PyObject* arg,
                                        bool include_tensor_ranks_only,
                                        EncodeResult* result) {
   if (EagerTensor_CheckExact(arg)) {
-    tensorflow::AbstractTensorHandleInterface* handle =
+    tensorflow::ImmediateExecutionTensorHandle* handle =
         tensorflow::unwrap(EagerTensor_Handle(arg));
 
     absl::StrAppend(&result->str, kDType,

From cb61c42938010f3ff33b55ef4b41651c22fb7652 Mon Sep 17 00:00:00 2001
From: Matt Conley 
Date: Mon, 15 Jun 2020 11:22:30 -0700
Subject: [PATCH 0421/1390] [TFTRT] Add Dynamic Shape Tests for ConvertSquare

Co-authored-by: Tamas Feher 

- Modify ConvertSquare tests to use newer TFTRT testing API
- Add INT32 as a supported dtype for TFTRT ConvertSquare
---
 .../tf2tensorrt/convert/convert_nodes.cc      |  5 ++
 .../tf2tensorrt/convert/convert_nodes_test.cc | 64 +++++++------------
 2 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
index 20ee5ffd8f8..28b27959afc 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
@@ -4424,8 +4424,13 @@ Status ConvertSquare(OpConverterParams* params) {
   const auto& inputs = params->inputs;
   const auto& node_def = params->node_def;
   TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}}));
+#if IS_TRT_VERSION_GE(6, 0, 1, 0)
+  TF_RETURN_IF_ERROR(AllowDataTypes(
+      *params, {DataType::DT_FLOAT, DataType::DT_HALF, DataType::DT_INT32}));
+#else
   TF_RETURN_IF_ERROR(
       AllowDataTypes(*params, {DataType::DT_FLOAT, DataType::DT_HALF}));
+#endif
   if (params->validation_only) return Status::OK();
 
   // Constant 2 with same rank as input
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
index 450831910f6..1192b563e57 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
@@ -2754,58 +2754,40 @@ TEST_F(OpConverterTest, ConvertQuantize) {
   }
 }
 
-template 
-void TestConvertSquare(OpConverterTest* test) {
-  test->Reset();
-  typedef typename EnumToDataType::Type CType;
-
-  Scope s = Scope::NewRootScope();
-  auto input = ops::Placeholder(s.WithOpName("input"), dtype);
-  auto square = ops::Square(s.WithOpName("my_square"), input);
-  NodeDef node_def = square.operation.node()->def();
-
-  test->AddTestTensor("input", {1, 20}, /*batch_size=*/1,
-                      TfDataTypeToTrt(dtype));
-  test->RunValidationAndConversion(node_def);
-  TRT_TensorOrWeights output;
-  TF_EXPECT_OK(test->GetTensorOrWeights("my_square", &output));
-  ASSERT_TRUE(output.is_tensor());
-  ExpectTrtDimsEqualsArray({1, 20}, output.tensor()->getDimensions());
-
-  const int num_inputs = 20;
-  std::vector inputs(num_inputs);
-  std::vector expected_outputs(num_inputs);
-  for (int i = 0; i < num_inputs; ++i) {
-    const CType value = CType(i - 9);
-    inputs[i] = value;
-    expected_outputs[i] = value * value;
-  }
-  const DataVec input_data{{"input", test->AsTensor(inputs)}};
-  // Engine outputs are converted to FP16 automatically if we set FP16 mode in
-  // the builder.
-  DataVec output_data{{"my_square", test->ConstructTensor(num_inputs)}};
-  TF_EXPECT_OK(test->BuildAndRun(input_data, &output_data));
-  ExpectArrayNear(expected_outputs, GetSpanForData(output_data[0]));
-}
-
-TEST_F(OpConverterTest, ConvertSquare) {
+TEST_P(OpConverterTest2, ConvertSquare) {
   {
     // Input is weights, should fail.
     Reset();
     Scope s = Scope::NewRootScope();
-    auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
+    auto input = ops::Placeholder(s.WithOpName("input"), tf_type);
     auto square = ops::Square(s.WithOpName("my_square"), input);
     NodeDef node_def = square.operation.node()->def();
-    AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, -5, 6});
+    AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, -5, 6}, tf_type);
     RunValidationAndConversion(
         node_def, error::UNIMPLEMENTED,
         "The input \"x\" for Square must be a tensor, at my_square");
   }
 
-  // OK. Note that kINT32 is not supported by IElementWiseLayer, so we don't
-  // test DT_INT32 type here.
-  TestConvertSquare(this);
-  TestConvertSquare(this);
+  Reset();
+
+  Scope s = Scope::NewRootScope();
+  auto input = ops::Placeholder(s.WithOpName("input"), tf_type);
+  auto square = ops::Square(s.WithOpName("my_square"), input);
+  NodeDef node_def = square.operation.node()->def();
+
+  const int num_inputs = 20;
+  std::vector inputs(num_inputs);
+  std::vector expected_outputs(num_inputs);
+
+  for (int i = 0; i < num_inputs; ++i) {
+    const float value = (i - 9);
+    inputs[i] = value;
+    expected_outputs[i] = value * value;
+  }
+  AddTestTensor("input", {1, 1, 20}, tf_type, inputs);
+
+  TestOpConverter("my_square", node_def, {1, 1, 20}, Status::OK(), Status::OK(),
+                  ArrayFloatNear(expected_outputs, 0));
 }
 
 #if IS_TRT_VERSION_GE(5, 1, 0, 0)

From d5ca984c5314d2e683eca87539d03ee4e35e3ee6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 13:05:59 -0700
Subject: [PATCH 0422/1390] Allow the compiler to vectorize the loop.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

```
name                                                                                old time/op             new time/op             delta
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x112x112x64_f2x2x64_s2x2_SAME               18.6ms ± 5%             18.5ms ±13%     ~           (p=0.912 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x128_f2x2x128_s2x2_SAME               12.7ms ±12%             12.7ms ±17%     ~           (p=0.684 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x256_f2x2x256_s2x2_SAME               12.2ms ± 8%             11.2ms ± 4%   -8.21%         (p=0.001 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x112x112x64_f2x2x64_s2x2_VALID              18.7ms ±20%             18.6ms ±23%     ~            (p=0.278 n=9+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x128_f2x2x128_s2x2_VALID              12.5ms ±15%             11.4ms ± 2%   -8.98%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x256_f2x2x256_s2x2_VALID              11.6ms ± 8%             11.1ms ± 2%   -4.22%          (p=0.011 n=9+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x64_f1x1x64_s1x1_SAME                 4.57ms ± 3%             4.34ms ± 1%   -5.04%          (p=0.000 n=8+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x64_f1x1x256_s1x1_SAME                12.0ms ± 4%             11.5ms ± 2%   -4.32%          (p=0.000 n=8+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x256_f1x1x64_s1x1_SAME                20.0ms ±31%             20.6ms ±17%     ~           (p=0.912 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x56x56x64_f3x3x64_s1x1_SAME                 36.5ms ±21%             32.0ms ± 1%  -12.30%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x128_f1x1x128_s1x1_SAME               3.71ms ±17%             3.33ms ± 1%  -10.47%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x128_f1x1x512_s1x1_SAME               11.8ms ±16%             10.5ms ± 1%  -11.37%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x512_f1x1x128_s1x1_SAME               13.1ms ±13%             11.4ms ± 2%  -13.36%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x28x28x512_f3x3x128_s1x1_SAME                142ms ±12%              124ms ± 1%  -13.22%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x14x14x256_f1x1x256_s1x1_SAME               3.51ms ±14%             3.18ms ±20%   -9.43%        (p=0.009 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x14x14x256_f1x1x1024_s1x1_SAME              14.0ms ±18%             12.0ms ± 4%  -13.80%         (p=0.012 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x14x14x1024_f1x1x256_s1x1_SAME              12.8ms ±18%             11.1ms ± 2%  -13.57%         (p=0.001 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in8x14x14x256_f3x3x256_s1x1_SAME               23.0ms ±18%             19.9ms ± 4%  -13.38%         (p=0.004 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x112x112x64_f2x2x64_s2x2_SAME              45.3ms ± 9%             40.5ms ± 4%  -10.74%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x128_f2x2x128_s2x2_SAME              33.2ms ±13%             28.8ms ± 2%  -13.11%         (p=0.001 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x256_f2x2x256_s2x2_SAME              31.5ms ±15%             26.7ms ± 2%  -15.13%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x112x112x64_f2x2x64_s2x2_VALID             45.5ms ± 8%             41.3ms ± 9%   -9.31%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x128_f2x2x128_s2x2_VALID             33.2ms ±12%             28.8ms ± 2%  -13.38%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x256_f2x2x256_s2x2_VALID             31.6ms ±14%             26.7ms ± 1%  -15.53%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x64_f1x1x64_s1x1_SAME                11.1ms ±15%              9.4ms ± 3%  -15.29%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x64_f1x1x256_s1x1_SAME               27.3ms ±13%             23.5ms ± 2%  -13.97%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x256_f1x1x64_s1x1_SAME               46.2ms ± 8%             40.9ms ± 5%  -11.54%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x56x56x64_f3x3x64_s1x1_SAME                81.5ms ±17%             64.2ms ± 1%  -21.17%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x128_f1x1x128_s1x1_SAME              8.30ms ±16%             6.80ms ± 1%  -18.02%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x128_f1x1x512_s1x1_SAME              27.7ms ±14%             23.3ms ± 2%  -15.87%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x512_f1x1x128_s1x1_SAME              34.9ms ±14%             28.8ms ± 2%  -17.41%        (p=0.000 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x28x28x512_f3x3x128_s1x1_SAME               300ms ±12%              250ms ± 1%  -16.80%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x14x14x256_f1x1x256_s1x1_SAME              7.91ms ±15%             7.17ms ±21%   -9.36%        (p=0.029 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x14x14x256_f1x1x1024_s1x1_SAME             35.6ms ±15%             33.4ms ±13%   -6.21%        (p=0.023 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x14x14x1024_f1x1x256_s1x1_SAME             28.0ms ± 3%             26.7ms ± 1%   -4.81%          (p=0.000 n=8+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in16x14x14x256_f3x3x256_s1x1_SAME              41.4ms ± 3%             39.6ms ± 4%   -4.43%          (p=0.001 n=9+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x112x112x64_f2x2x64_s2x2_SAME              90.0ms ± 3%             85.5ms ±11%   -5.00%         (p=0.010 n=9+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x128_f2x2x128_s2x2_SAME              69.0ms ± 1%             65.8ms ± 3%   -4.68%          (p=0.000 n=9+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x256_f2x2x256_s2x2_SAME              57.8ms ± 2%             55.8ms ± 2%   -3.54%          (p=0.000 n=9+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x112x112x64_f2x2x64_s2x2_VALID             89.3ms ± 2%             86.1ms ± 5%   -3.57%         (p=0.006 n=9+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x128_f2x2x128_s2x2_VALID             68.7ms ± 2%             65.7ms ± 3%   -4.31%          (p=0.001 n=8+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x256_f2x2x256_s2x2_VALID             57.1ms ± 1%             55.6ms ± 1%   -2.58%         (p=0.000 n=8+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x64_f1x1x64_s1x1_SAME                23.6ms ± 9%             21.9ms ± 1%   -7.27%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x64_f1x1x256_s1x1_SAME               51.0ms ± 3%             48.4ms ± 5%   -5.13%         (p=0.001 n=8+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x256_f1x1x64_s1x1_SAME               90.2ms ± 3%             85.4ms ± 5%   -5.29%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x56x56x64_f3x3x64_s1x1_SAME                 143ms ± 7%              133ms ± 3%   -6.48%          (p=0.000 n=9+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x128_f1x1x128_s1x1_SAME              14.5ms ± 4%             14.0ms ± 2%   -3.82%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x128_f1x1x512_s1x1_SAME              53.8ms ± 3%             51.5ms ± 2%   -4.43%         (p=0.000 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x512_f1x1x128_s1x1_SAME              69.8ms ± 5%             66.9ms ± 8%   -4.15%         (p=0.010 n=10+9)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x28x28x512_f3x3x128_s1x1_SAME               533ms ± 3%              508ms ± 0%   -4.65%         (p=0.000 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x14x14x256_f1x1x256_s1x1_SAME              14.0ms ± 6%             14.9ms ±21%     ~           (p=0.481 n=10+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x14x14x256_f1x1x1024_s1x1_SAME             67.2ms ± 3%             71.0ms ±22%     ~            (p=0.278 n=9+10)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x14x14x1024_f1x1x256_s1x1_SAME             57.5ms ± 4%             56.6ms ± 5%     ~            (p=0.055 n=10+8)
BM_Conv2DBackpropInput_fp32_NHWC_cpu_in32x14x14x256_f3x3x256_s1x1_SAME              78.1ms ± 5%             75.5ms ± 3%   -3.32%          (p=0.011 n=9+8)
```

PiperOrigin-RevId: 316949219
Change-Id: I1a1fb517a5c28d489da9762b650577b61bf4e0de
---
 tensorflow/core/kernels/conv_grad_input_ops.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/conv_grad_input_ops.cc b/tensorflow/core/kernels/conv_grad_input_ops.cc
index fd2f569a8b8..2dd63d1f4d0 100644
--- a/tensorflow/core/kernels/conv_grad_input_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_input_ops.cc
@@ -76,7 +76,7 @@ template 
 void Col2im(const T* col_data, const int depth, const int height,
             const int width, const int filter_h, const int filter_w,
             const int pad_t, const int pad_l, const int pad_b, const int pad_r,
-            const int stride_h, const int stride_w, T* im_data) {
+            const int stride_h, const int stride_w, T* __restrict im_data) {
   int height_col = (height + pad_t + pad_b - filter_h) / stride_h + 1;
   int width_col = (width + pad_l + pad_r - filter_w) / stride_w + 1;
   int h_pad = -pad_t;
@@ -87,7 +87,6 @@ void Col2im(const T* col_data, const int depth, const int height,
       for (int ih = h_pad; ih < h_pad + filter_h; ++ih) {
         for (int iw = w_pad; iw < w_pad + filter_w; ++iw) {
           if (ih >= 0 && ih < height && iw >= 0 && iw < width) {
-            // TODO(andydavis) Vectorize this loop (if compiler does not).
             for (int i = 0; i < depth; ++i) {
               im_patch_data[i] += col_data[i];
             }

From 9426d35abc2ea3a5e790a928641749c5783e4b65 Mon Sep 17 00:00:00 2001
From: Lukasz Kaiser 
Date: Wed, 17 Jun 2020 13:07:01 -0700
Subject: [PATCH 0423/1390] Add the vdot op to tf-numpy.

PiperOrigin-RevId: 316949446
Change-Id: I9b0718ec108486096e032729481ec9129863b429
---
 tensorflow/python/ops/numpy_ops/np_math_ops.py      | 10 ++++++++++
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py |  6 ++++++
 2 files changed, 16 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index b32f78bee5a..abfd9087ffd 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -309,6 +309,16 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):  # pylint: disable=mis
   return _bin_op(f, a, b)
 
 
+@np_utils.np_doc_only(np.vdot)
+def vdot(a, b):  # pylint: disable=missing-docstring
+  a, b = np_array_ops._promote_dtype(a, b)
+  a = np_array_ops.reshape(a, [-1])
+  b = np_array_ops.reshape(b, [-1])
+  if a.dtype == np_dtypes.complex128 or a.dtype == np_dtypes.complex64:
+    a = conj(a)
+  return dot(a, b)
+
+
 @np_utils.np_doc(np.power)
 def power(x1, x2):
   return _bin_op(math_ops.pow, x1, x2)
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
index a3dfbb6c871..cb5326bcded 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -124,6 +124,12 @@ class MathTest(test.TestCase, parameterized.TestCase):
       np_math_ops.matmul(
           np_array_ops.ones([2, 3], np.int32), np_array_ops.ones([], np.int32))
 
+  def testVDot(self):
+    operands = [([[1, 2], [3, 4]], [[3, 4], [6, 7]]),
+                ([[1, 2], [3, 4]], [3, 4, 6, 7])]
+    return self._testBinaryOp(
+        np_math_ops.vdot, np.vdot, 'vdot', operands=operands)
+
   def _testUnaryOp(self, math_fun, np_fun, name):
 
     def run_test(a):

From cca9b615b264651f5c70d604bb29146ddfbadede Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Wed, 17 Jun 2020 13:12:33 -0700
Subject: [PATCH 0424/1390] [tfdbg2] Fork local_cli_wrapper_test for keras
 related tests.

PiperOrigin-RevId: 316950499
Change-Id: I428273592694426f72427e6236c68bdfb4e95eba
---
 tensorflow/python/debug/BUILD                 |  2 -
 .../debug/wrappers/local_cli_wrapper_test.py  | 39 +------------------
 2 files changed, 1 insertion(+), 40 deletions(-)

diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD
index 1ef0504ecb8..bb3bbbd87c4 100644
--- a/tensorflow/python/debug/BUILD
+++ b/tensorflow/python/debug/BUILD
@@ -788,7 +788,6 @@ cuda_py_test(
         "//tensorflow/python:training",
         "//tensorflow/python:variables",
         "//tensorflow/python/data/ops:dataset_ops",
-        "//tensorflow/python/keras",
         "//third_party/py/numpy",
     ],
 )
@@ -1400,7 +1399,6 @@ py_test(
         "//tensorflow/python:state_ops",
         "//tensorflow/python:training",
         "//tensorflow/python:variables",
-        "//tensorflow/python/keras",
         "//third_party/py/numpy",
     ],
 )
diff --git a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
index 30bb99387b2..ab33a4af030 100644
--- a/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
+++ b/tensorflow/python/debug/wrappers/local_cli_wrapper_test.py
@@ -22,10 +22,10 @@ import tempfile
 
 import numpy as np
 
-from tensorflow.python.debug.cli import cli_config
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session
+from tensorflow.python.debug.cli import cli_config
 from tensorflow.python.debug.cli import cli_shared
 from tensorflow.python.debug.cli import debugger_cli_common
 from tensorflow.python.debug.cli import ui_factory
@@ -36,9 +36,6 @@ from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.lib.io import file_io
-from tensorflow.python.keras import backend
-from tensorflow.python.keras.engine import sequential
-from tensorflow.python.keras.layers import core
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
@@ -832,40 +829,6 @@ class LocalCLIDebugWrapperSessionTest(test_util.TensorFlowTestCase):
     run_output = wrapped_sess.run([])
     self.assertEqual([], run_output)
 
-  def testDebuggingKerasFitWithSkippedRunsWorks(self):
-    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
-        [["run"], ["run"], ["run", "-t", "10"]], self.sess)
-
-    backend.set_session(wrapped_sess)
-
-    model = sequential.Sequential()
-    model.add(core.Dense(4, input_shape=[2], activation="relu"))
-    model.add(core.Dense(1))
-    model.compile(loss="mse", optimizer="sgd")
-
-    x = np.zeros([8, 2])
-    y = np.zeros([8, 1])
-    model.fit(x, y, epochs=2)
-
-    self.assertEqual(2, len(wrapped_sess.observers["debug_dumps"]))
-
-  def testDebuggingKerasFitWithProfilingWorks(self):
-    wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
-        [["run", "-p"]] * 10, self.sess)
-
-    backend.set_session(wrapped_sess)
-
-    model = sequential.Sequential()
-    model.add(core.Dense(4, input_shape=[2], activation="relu"))
-    model.add(core.Dense(1))
-    model.compile(loss="mse", optimizer="sgd")
-
-    x = np.zeros([8, 2])
-    y = np.zeros([8, 1])
-    model.fit(x, y, epochs=2)
-
-    self.assertEqual(0, len(wrapped_sess.observers["debug_dumps"]))
-
   def testRunsWithEmptyNestedFetchWorks(self):
     wrapped_sess = LocalCLIDebuggerWrapperSessionForTest(
         [["run"]], self.sess, dump_root="")

From f8195170f88259829ab11538bd5d4b068dfa44a8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 13:23:53 -0700
Subject: [PATCH 0425/1390] Enable a disabled testcase, since it is fixed.

PiperOrigin-RevId: 316952692
Change-Id: I70d295634012685097b2324b156b2552218d62d7
---
 .../keras/layers/preprocessing/text_vectorization_test.py       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py b/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
index 992f47efc85..e7f61e94724 100644
--- a/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/text_vectorization_test.py
@@ -1503,7 +1503,7 @@ class TextVectorizationSavingTest(
     loaded_model = keras.models.load_model(output_path)
     self.assertAllEqual(loaded_model.predict(input_array), expected_output)
 
-  def DISABLE_test_saving_with_tfidf(self):
+  def test_saving_with_tfidf(self):
     vocab_data = ["earth", "wind", "and", "fire"]
     tfidf_data = [.5, .25, .2, .125]
     input_array = np.array([["earth", "wind", "and", "earth"],

From 56e71dd0e77507fc5200a650de5c736271df7f8a Mon Sep 17 00:00:00 2001
From: Pavithra Vijay 
Date: Wed, 17 Jun 2020 13:42:03 -0700
Subject: [PATCH 0426/1390] Removing identity backtracking from entropy losses.

PiperOrigin-RevId: 316956157
Change-Id: I91130052e29e69ae131fe8aad0bbd1d4d42b00f1
---
 tensorflow/python/keras/backend.py            | 65 ++++++++-----------
 .../keras/tests/add_loss_correctness_test.py  | 14 ++++
 2 files changed, 42 insertions(+), 37 deletions(-)

diff --git a/tensorflow/python/keras/backend.py b/tensorflow/python/keras/backend.py
index 391c695b18f..9330425272f 100644
--- a/tensorflow/python/keras/backend.py
+++ b/tensorflow/python/keras/backend.py
@@ -4637,12 +4637,6 @@ def softsign(x):
   return nn.softsign(x)
 
 
-def _backtrack_identity(tensor):
-  while tensor.op.type == 'Identity':
-    tensor = tensor.op.inputs[0]
-  return tensor
-
-
 @keras_export('keras.backend.categorical_crossentropy')
 @dispatch.add_dispatch_support
 def categorical_crossentropy(target, output, from_logits=False, axis=-1):
@@ -4695,17 +4689,16 @@ def categorical_crossentropy(target, output, from_logits=False, axis=-1):
     return nn.softmax_cross_entropy_with_logits_v2(
         labels=target, logits=output, axis=axis)
 
-  if not isinstance(output, (ops.EagerTensor, variables_module.Variable)):
-    output = _backtrack_identity(output)
-    if output.op.type == 'Softmax':
-      # When softmax activation function is used for output operation, we
-      # use logits from the softmax function directly to compute loss in order
-      # to prevent collapsing zero when training.
-      # See b/117284466
-      assert len(output.op.inputs) == 1
-      output = output.op.inputs[0]
-      return nn.softmax_cross_entropy_with_logits_v2(
-          labels=target, logits=output, axis=axis)
+  if (not isinstance(output, (ops.EagerTensor, variables_module.Variable)) and
+      output.op.type == 'Softmax'):
+    # When softmax activation function is used for output operation, we
+    # use logits from the softmax function directly to compute loss in order
+    # to prevent collapsing zero when training.
+    # See b/117284466
+    assert len(output.op.inputs) == 1
+    output = output.op.inputs[0]
+    return nn.softmax_cross_entropy_with_logits_v2(
+        labels=target, logits=output, axis=axis)
 
   # scale preds so that the class probas of each sample sum to 1
   output = output / math_ops.reduce_sum(output, axis, True)
@@ -4740,17 +4733,16 @@ def sparse_categorical_crossentropy(target, output, from_logits=False, axis=-1):
   target = ops.convert_to_tensor_v2(target)
   output = ops.convert_to_tensor_v2(output)
 
-  if not from_logits and not isinstance(
-      output, (ops.EagerTensor, variables_module.Variable)):
-    output = _backtrack_identity(output)
-    if output.op.type == 'Softmax':
-      # When softmax activation function is used for output operation, we
-      # use logits from the softmax function directly to compute loss in order
-      # to prevent collapsing zero when training.
-      # See b/117284466
-      assert len(output.op.inputs) == 1
-      output = output.op.inputs[0]
-      from_logits = True
+  if (not from_logits and
+      not isinstance(output, (ops.EagerTensor, variables_module.Variable)) and
+      output.op.type == 'Softmax'):
+    # When softmax activation function is used for output operation, we
+    # use logits from the softmax function directly to compute loss in order
+    # to prevent collapsing zero when training.
+    # See b/117284466
+    assert len(output.op.inputs) == 1
+    output = output.op.inputs[0]
+    from_logits = True
 
   if not from_logits:
     epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
@@ -4821,15 +4813,14 @@ def binary_crossentropy(target, output, from_logits=False):
   if from_logits:
     return nn.sigmoid_cross_entropy_with_logits(labels=target, logits=output)
 
-  if not isinstance(output, (ops.EagerTensor, variables_module.Variable)):
-    output = _backtrack_identity(output)
-    if output.op.type == 'Sigmoid':
-      # When sigmoid activation function is used for output operation, we
-      # use logits from the sigmoid function directly to compute loss in order
-      # to prevent collapsing zero when training.
-      assert len(output.op.inputs) == 1
-      output = output.op.inputs[0]
-      return nn.sigmoid_cross_entropy_with_logits(labels=target, logits=output)
+  if (not isinstance(output, (ops.EagerTensor, variables_module.Variable)) and
+      output.op.type == 'Sigmoid'):
+    # When sigmoid activation function is used for output operation, we
+    # use logits from the sigmoid function directly to compute loss in order
+    # to prevent collapsing zero when training.
+    assert len(output.op.inputs) == 1
+    output = output.op.inputs[0]
+    return nn.sigmoid_cross_entropy_with_logits(labels=target, logits=output)
 
   epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
   output = clip_ops.clip_by_value(output, epsilon_, 1. - epsilon_)
diff --git a/tensorflow/python/keras/tests/add_loss_correctness_test.py b/tensorflow/python/keras/tests/add_loss_correctness_test.py
index a19eec75ffb..f99b285489d 100644
--- a/tensorflow/python/keras/tests/add_loss_correctness_test.py
+++ b/tensorflow/python/keras/tests/add_loss_correctness_test.py
@@ -34,6 +34,7 @@ from tensorflow.python.keras import testing_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training.rmsprop import RMSPropOptimizer
 
 MAE = losses.MeanAbsoluteError
@@ -450,6 +451,19 @@ class TestAddLossCorrectness(keras_parameterized.TestCase):
           'Expected a symbolic Tensors or a callable for the loss value'):
         model.add_loss(model.weights[0])
 
+  @keras_parameterized.run_all_keras_modes
+  def test_add_entropy_loss_on_functional_model(self):
+    inputs = Input(shape=(1,))
+    targets = Input(shape=(1,))
+    outputs = testing_utils.Bias()(inputs)
+    model = Model([inputs, targets], outputs)
+    model.add_loss(losses.binary_crossentropy(targets, outputs))
+    model.compile('sgd', run_eagerly=testing_utils.should_run_eagerly())
+    with test.mock.patch.object(logging, 'warning') as mock_log:
+      model.fit([self.x, self.y], batch_size=3, epochs=5)
+      self.assertNotIn('Gradients do not exist for variables',
+                       str(mock_log.call_args))
+
 
 if __name__ == '__main__':
   test.main()

From 8e7be6f71aaa97f8c2ef806da8945a9d11f99830 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen 
Date: Wed, 17 Jun 2020 13:56:34 -0700
Subject: [PATCH 0427/1390] Populate side-effect decorators for HashTable ops.

PiperOrigin-RevId: 316958908
Change-Id: I8e75fe538bb3dfc4b6d0742f63072fe87f133a4d
---
 tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td | 6 +++---
 tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td       | 3 +++
 tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h   | 4 ++++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index d403462e6a6..7ed5a215ab8 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -3862,7 +3862,7 @@ table will be immutable.
   );
 
   let results = (outs
-    TF_ResourceTensor:$table_handle
+    Res:$table_handle
   );
 }
 
@@ -4814,7 +4814,7 @@ table. It must also be of the same type as the table values.
   }];
 
   let arguments = (ins
-    TF_ResourceTensor:$table_handle,
+    Arg:$table_handle,
     TF_Tensor:$keys,
     TF_Tensor:$default_value
   );
@@ -4838,7 +4838,7 @@ The tensor `values` must be of the type of the table values.
   }];
 
   let arguments = (ins
-    TF_ResourceTensor:$table_handle,
+    Arg:$table_handle,
     TF_Tensor:$keys,
     TF_Tensor:$values
   );
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
index 17424b54fc2..aac03061718 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td
@@ -91,14 +91,17 @@ class TF_ResourceBase :
 def TF_VariableResource : TF_ResourceBase<"Variable">;
 def TF_StackResource : TF_ResourceBase<"Stack">;
 def TF_TensorArrayResource : TF_ResourceBase<"TensorArray">;
+def TF_TableResource : TF_ResourceBase<"Table">;
 
 def TF_VariableRead : MemRead;
 def TF_StackRead : MemRead;
 def TF_TensorArrayRead : MemRead;
+def TF_TableRead : MemRead;
 
 def TF_VariableWrite : MemWrite;
 def TF_StackWrite : MemWrite;
 def TF_TensorArrayWrite : MemWrite;
+def TF_TableWrite: MemWrite;
 
 //===----------------------------------------------------------------------===//
 // TensorFlow op definitions
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
index 9be61b1db39..ab1d7935bad 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h
@@ -35,6 +35,10 @@ struct TensorArray : ::mlir::SideEffects::Resource::Base {
   StringRef getName() final { return "TensorArray"; }
 };
 
+struct Table : ::mlir::SideEffects::Resource::Base {
+  StringRef getName() final { return "Table"; }
+};
+
 }  // namespace ResourceEffects
 }  // namespace TF
 }  // namespace mlir

From cbf8f57413d8a14c32515e88292012324a0990f0 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen 
Date: Wed, 17 Jun 2020 14:16:51 -0700
Subject: [PATCH 0428/1390] Import initialization graph in SignatureDef
 SavedModels as an MLIR function in TF saved model dialect.

PiperOrigin-RevId: 316963067
Change-Id: I6c49ba4cfa3832f136bec36700aea695ad116737
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |   4 +-
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  25 +++
 .../mlir/tensorflow/ir/tf_saved_model_ops.td  |  24 +++
 .../tests/tf_saved_model/common_v1.py         |   1 +
 .../tests/tf_saved_model/hash_table_v1.py     |  92 +++++++++++
 .../tensorflow/tests/tf_saved_model_ops.mlir  |   5 +
 .../tests/tf_saved_model_ops_invalid.mlir     |  33 ++++
 .../mlir/tensorflow/translate/import_model.cc | 149 +++++++++++++-----
 8 files changed, 290 insertions(+), 43 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 904ccb7e820..17ed0e36a28 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -661,7 +661,9 @@ cc_library(
         ":tensorflow_types",
         ":translate_utils",
         "//tensorflow/cc/saved_model:bundle_v2",
+        "//tensorflow/cc/saved_model:constants",
         "//tensorflow/cc/saved_model:loader_lite",
+        "//tensorflow/cc/saved_model:loader_util",
         "//tensorflow/compiler/jit:shape_inference_helpers",
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/tf2xla:functionalize_control_flow",
@@ -673,6 +675,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/utils:transitive_fanin",
+        "//tensorflow/core/platform:protobuf_internal",
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
@@ -682,7 +685,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:StandardOps",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index 140a778770c..6af70158e14 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -76,6 +76,23 @@ static LogicalResult Verify(GlobalTensorOp global_tensor) {
   return success();
 }
 
+static LogicalResult Verify(SessionInitializerOp session_initializer) {
+  mlir::SymbolTable symbol_table(
+      session_initializer.getParentOfType());
+
+  auto init_func_op =
+      symbol_table.lookup(session_initializer.initializer());
+  if (!init_func_op)
+    return session_initializer.emitOpError()
+           << "the initializer function does not exist";
+
+  if (!init_func_op.getType().getResults().empty())
+    return session_initializer.emitOpError()
+           << "the initializer function should have no output";
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc.inc"
 
@@ -220,6 +237,14 @@ static LogicalResult VerifySavedModelModule(
       }
     }
   }
+
+  auto session_initializers = module.getOps();
+  if (std::distance(session_initializers.begin(), session_initializers.end()) >
+      1) {
+    return (*++session_initializers.begin()).emitError()
+           << "there must be no more than one session_initializer op";
+  }
+
   SymbolTable symbol_table(module);
   auto symbol_uses = SymbolTable::getSymbolUses(&module.getBodyRegion());
   if (!symbol_uses.hasValue()) {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
index 4431a160edf..497f4d90cb9 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
@@ -128,4 +128,28 @@ def TfSavedModel_GlobalTensorOp : TfSavedModel_Op<"global_tensor"> {
   let verifier = [{ return Verify(*this); }];
 }
 
+def TfSavedModel_SessionInitializerOp: TfSavedModel_Op<"session_initializer"> {
+  let summary = "Initializes TensorFlow session state.";
+  let description = [{
+    Represents a session initializer function initializes TensorFlow session
+    state. It is used to initialize resources in the saved model before calling
+    any exported functions. There must be no more than one session initializer
+    in a saved model.
+
+    The `initializer` represents the initialization function. The function have
+    no output and this function should be only called once.
+
+    This is used, for example, to initialize hash tables stored in resources and
+    accessed by resource name (rather than as resource handles or bound inputs
+    which is how `global_tensor`s are referenced)
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$initializer
+  );
+
+
+  let verifier = [{ return Verify(*this); }];
+}
+
 #endif // SAVED_MODEL_DIALECT
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
index 7171f63bb05..51ccbeb1fbd 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
@@ -84,6 +84,7 @@ def do_test(signature_def_map, show_debug_info=False):
     builder.add_meta_graph_and_variables(
         sess, [tf.saved_model.tag_constants.SERVING],
         signature_def_map,
+        main_op=tf.tables_initializer(),
         strip_default_attrs=True)
     builder.save()
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
new file mode 100644
index 00000000000..64847434b82
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
@@ -0,0 +1,92 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# RUN: %p/hash_table_v1 | FileCheck %s
+
+# pylint: disable=missing-docstring,line-too-long
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v1 as tf
+from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
+
+# Verify that the tf.versions attribute exists. It is difficult to enforce
+# contents, since the version numbers change over time. The conversion logic
+# itself is verified in the common graphdef converter, so here just assert
+# it is being invoked.
+# CHECK: module
+# CHECK-SAME: tf.versions
+# CHECK-SAME: bad_consumers
+# CHECK-SAME: min_consumer
+# CHECK-SAME: producer
+
+# CHECK: "tf_saved_model.session_initializer"() {initializer = [[init:@.*]]} : () -> ()
+# CHECK: "tf_saved_model.global_tensor"()
+
+# CHECK:      func {{@[a-zA-Z_0-9]+}}(
+# CHECK-SAME: [[ARG0:%.*]]: tensor
+# CHECK-SAME: [[ARG1:%.*]]: tensor ()
+
   // Representation for constants: (immutable) global tensor.
   // CHECK: tf_saved_model.global_tensor
   "tf_saved_model.global_tensor"() {
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index c055c6c9f56..544600cf6b8 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -258,3 +258,36 @@ module attributes {tf_saved_model.semantics} {
   // expected-error@+1 {{'type' attribute for immutable 'tf_saved_model.global_tensor' should have a static shape}}
   "tf_saved_model.global_tensor"() { sym_name = "v", type = tensor, value = dense<1.> : tensor<1xf32> } : () -> ()
 }
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function does not exist}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function should have no output}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  // expected-error@+1 {{there must be no more than one session_initializer op}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 820d0ce31fb..3cff4217215 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -60,6 +60,8 @@ limitations under the License.
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
@@ -99,6 +101,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/protobuf_internal.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
@@ -116,6 +119,7 @@ using mlir::NamedAttrList;
 using mlir::TensorType;
 using mlir::TF::VarHandleOp;
 using mlir::tf_saved_model::GlobalTensorOp;
+using mlir::tf_saved_model::SessionInitializerOp;
 using stream_executor::port::StatusOr;
 
 namespace {
@@ -2955,6 +2959,13 @@ void SortSavedModelModule(mlir::ModuleOp module) {
     named_global_tensor.global_tensor.getOperation()->moveBefore(
         &module.getBody()->front());
   }
+
+  auto initializers = module.getOps();
+  if (!initializers.empty()) {
+    (*initializers.begin())
+        .getOperation()
+        ->moveBefore(&module.getBody()->front());
+  }
 }
 
 Status CreateSavedModelIR(
@@ -3241,17 +3252,29 @@ class SavedModelSignatureDefImporter {
                                  absl::Span exported_names,
                                  mlir::MLIRContext* context)
       : bundle_(bundle),
+        flib_def_(OpRegistry::Global(), graph_def().library()),
+        debug_info_(),
         exported_names_(exported_names),
-        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {}
+        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {
+    // debug_info might not be loaded with loader_lite.
+    if (bundle_.debug_info != nullptr) debug_info_ = *bundle_.debug_info;
+  }
 
   // Converts the SavedModel to the SavedModel dialect. Creates an MLIR function
   // for each signature.
   StatusOr ConvertSignatures();
-  Status ConvertSignature(const GraphDef& graphdef,
-                          const std::string& sig_def_key,
-                          const SignatureDef& signature_def,
-                          const GraphDebugInfo& debug_info,
-                          const FunctionLibraryDefinition& flib_def);
+  Status ConvertSignature(const std::string& sig_def_key,
+                          const SignatureDef& signature_def);
+
+  // Converts the initialization graph in the SavedModel to an MLIR function.
+  Status ConvertInitializer();
+
+  // Converts a graph with feeds and fetches to an MLIR function.
+  StatusOr ConvertGraph(
+      const std::string& name,
+      const std::vector>& inputs,
+      const std::vector>& outputs,
+      const std::vector control_outputs);
 
   // Creates GlobalTensorOp for each variable and moves each VarHandle op to
   // the enclosing function's arguments.
@@ -3273,18 +3296,62 @@ class SavedModelSignatureDefImporter {
   GraphImportConfig::InputArrays ParseInputArrays(
       const std::vector>& inputs);
 
+  const GraphDef& graph_def() const {
+    return bundle_.meta_graph_def.graph_def();
+  }
+  const FunctionLibraryDefinition& flib_def() const { return flib_def_; }
+  const GraphDebugInfo& debug_info() const { return debug_info_; }
+
   const SavedModelBundle& bundle_;
+  FunctionLibraryDefinition flib_def_;
+  GraphDebugInfo debug_info_;
   absl::Span exported_names_;
   mlir::OwningModuleRef module_;
 };
 
+Status SavedModelSignatureDefImporter::ConvertInitializer() {
+  std::vector asset_file_defs;
+  TF_RETURN_IF_ERROR(
+      internal::GetAssetFileDefs(bundle_.meta_graph_def, &asset_file_defs));
+
+  if (!asset_file_defs.empty())
+    return errors::Unimplemented(
+        absl::StrCat("Assets are not supported in signaturedef importer"));
+
+  std::string init_node_name;
+  TF_RETURN_IF_ERROR(
+      internal::GetInitOp("", bundle_.meta_graph_def, &init_node_name));
+
+  if (init_node_name.empty()) return Status::OK();
+
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(init_node_name, {}, {}, {init_node_name}));
+
+  mlir::SymbolTable symbol_table(*sub_module);
+
+  auto init_func_op = symbol_table.lookup(init_node_name);
+
+  init_func_op.removeAttr("tf.entry_function");
+
+  mlir::OpBuilder builder(module_->getBodyRegion());
+
+  builder.create(
+      module_->getLoc(), builder.getSymbolRefAttr(init_func_op.getName()));
+
+  // Move the converted functions to top level MLIR module.
+  auto* block = module_->getBody();
+  auto* sub_block = sub_module->getBody();
+  block->getOperations().splice(
+      mlir::Block::iterator(block->getTerminator()), sub_block->getOperations(),
+      sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
+
+  return Status::OK();
+}
+
 StatusOr
 SavedModelSignatureDefImporter::ConvertSignatures() {
   const auto& signatures = bundle_.GetSignatures();
-  const auto& graphdef = bundle_.meta_graph_def.graph_def();
-  PopulateTfVersions(module_.get(), graphdef.versions());
-
-  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graphdef.library());
+  PopulateTfVersions(module_.get(), graph_def().versions());
 
   // debug_info might not be loaded with loader_lite.
   GraphDebugInfo debug_info;
@@ -3307,9 +3374,10 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
       continue;
     }
 
-    TF_RETURN_IF_ERROR(ConvertSignature(graphdef, sig_def_key, signature_def,
-                                        debug_info, flib_def));
+    TF_RETURN_IF_ERROR(ConvertSignature(sig_def_key, signature_def));
   }
+
+  TF_RETURN_IF_ERROR(ConvertInitializer());
   TF_RETURN_IF_ERROR(LiftVariables());
 
   mlir::OpBuilder builder(module_->getBodyRegion());
@@ -3320,10 +3388,32 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
   return std::move(module_);
 }
 
+StatusOr SavedModelSignatureDefImporter::ConvertGraph(
+    const std::string& name,
+    const std::vector>& inputs,
+    const std::vector>& outputs,
+    const std::vector control_outputs) {
+  GraphImportConfig specs;
+  specs.prune_unused_nodes = true;
+  specs.inputs = ParseInputArrays(inputs);
+  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
+  specs.control_outputs = control_outputs;
+
+  // Convert sub-graphdef to sub-graph.
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.add_default_attributes = true;
+  Graph graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, graph_def(), &graph));
+
+  // Convert sub-graph to MLIR module.true
+  return GraphDefImporter::Convert(module_->getContext(), graph, debug_info(),
+                                   flib_def(), specs, name);
+}
+
 Status SavedModelSignatureDefImporter::ConvertSignature(
-    const GraphDef& graphdef, const std::string& sig_def_key,
-    const SignatureDef& signature_def, const GraphDebugInfo& debug_info,
-    const FunctionLibraryDefinition& flib_def) {
+    const std::string& sig_def_key, const SignatureDef& signature_def) {
   // Create local vectors for the input and output and sort them to be
   // deterministic. We don't want anyone to really depend on the order, client
   // should lookup argument/result mapping by attribute name.
@@ -3339,34 +3429,9 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
     return lhs.first.size() < rhs.first.size() || lhs.first > rhs.first;
   });
 
-  GraphImportConfig specs;
-  specs.prune_unused_nodes = true;
-  specs.inputs = ParseInputArrays(inputs);
-  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
-
-  // Remove unused nodes and create sub-graphdef.
-  GraphDef sub_graph_def;
-  TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(
-      graphdef, &sub_graph_def,
-      /*terminal_nodes=*/{specs.outputs.begin(), specs.outputs.end()}));
-
-  // Set the function library definitions in the pruned graphdef.
-  *sub_graph_def.mutable_library() = flib_def.ToProto();
-
-  // Convert sub-graphdef to sub-graph.
-  GraphConstructorOptions options;
-  options.allow_internal_ops = true;
-  options.add_default_attributes = true;
-  Graph sub_graph(OpRegistry::Global());
-
-  TF_RETURN_IF_ERROR(
-      ConvertGraphDefToGraph(options, sub_graph_def, &sub_graph));
-
   // Convert sub-graph to MLIR module.
-  TF_ASSIGN_OR_RETURN(
-      auto sub_module,
-      GraphDefImporter::Convert(module_->getContext(), sub_graph, debug_info,
-                                flib_def, specs, sig_def_key));
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(sig_def_key, inputs, outputs, {}));
   mlir::OpBuilder builder(sub_module->getBodyRegion());
 
   // Find the FuncOp which corresponds to current SignatureDef.

From 5b5ab0034ae0f12731943177b275e35cb9f92bb1 Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Wed, 17 Jun 2020 23:26:29 +0200
Subject: [PATCH 0429/1390] Fix distributed autocast variable assign

---
 .../experimental/autocast_variable.py         | 104 +++++++++---------
 .../experimental/autocast_variable_test.py    |  30 +++--
 2 files changed, 77 insertions(+), 57 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
index 7d0abe30581..ca6420f0c0b 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
@@ -188,61 +188,88 @@ class AutoCastVariable(variables.Variable, core.Tensor):
   def constraint(self):
     return self._variable.constraint
 
+  def _apply_assign_update(
+      self, update_fn, value, use_locking=None, name=None, read_value=True):
+    if not read_value:
+      return update_fn(value, use_locking, name, read_value)
+
+    if context.executing_eagerly() or ops.inside_function():
+      assign_op = update_fn(value, use_locking, name, False)
+      with ops.control_dependencies([assign_op]):
+        return self
+
+    # Fallback to wrapping the returned variable in graph mode if possible
+    assign_var = update_fn(value, use_locking, name, read_value)
+    if resource_variable_ops.is_resource_variable(assign_var):
+      return create_autocast_variable(assign_var)
+    return assign_var
+
+  def _apply_update(self, update_fn, *args, **kwargs):
+    update_var = update_fn(*args, **kwargs)
+    if context.executing_eagerly() or ops.inside_function():
+      with ops.control_dependencies([update_var]):
+        return self
+
+    # Fallback to wrapping the returned variable in graph mode if possible
+    if resource_variable_ops.is_resource_variable(update_var):
+      return create_autocast_variable(update_var)
+    return update_var
+
   def assign(self, value, use_locking=None, name=None, read_value=True):
-    assign_op = self._variable.assign(value, use_locking, name, read_value)
-    return _maybe_wrap(assign_op, wrap=read_value)
+    return self._apply_assign_update(
+        self._variable.assign, value, use_locking, name, read_value)
 
   def assign_add(self, delta, use_locking=None, name=None, read_value=True):
-    assign_op = self._variable.assign_add(delta, use_locking, name, read_value)
-    return _maybe_wrap(assign_op, wrap=read_value)
+    return self._apply_assign_update(
+        self._variable.assign_add, delta, use_locking, name, read_value)
 
   def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
-    assign_op = self._variable.assign_sub(delta, use_locking, name, read_value)
-    return _maybe_wrap(assign_op, wrap=read_value)
+    return self._apply_assign_update(
+        self._variable.assign_sub, delta, use_locking, name, read_value)
 
   def scatter_sub(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_sub(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_sub, sparse_delta, use_locking, name)
 
   def scatter_add(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_add(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_add, sparse_delta, use_locking, name)
 
   def scatter_max(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_max(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_max, sparse_delta, use_locking, name)
 
   def scatter_min(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_min(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_min, sparse_delta, use_locking, name)
 
   def scatter_mul(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_mul(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_mul, sparse_delta, use_locking, name)
 
   def scatter_div(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_div(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_div, sparse_delta, use_locking, name)
 
   def scatter_update(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.scatter_update(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_update, sparse_delta, use_locking, name)
 
   def batch_scatter_update(self, sparse_delta, use_locking=False, name=None):
-    var = self._variable.batch_scatter_update(sparse_delta, use_locking, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.batch_scatter_update, sparse_delta, use_locking, name)
 
   def scatter_nd_sub(self, indices, updates, name=None):
-    var = self._variable.scatter_nd_sub(indices, updates, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_nd_sub, indices, updates, name)
 
   def scatter_nd_add(self, indices, updates, name=None):
-    var = self._variable.scatter_nd_add(indices, updates, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_nd_add, indices, updates, name)
 
   def scatter_nd_update(self, indices, updates, name=None):
-    var = self._variable.scatter_nd_update(indices, updates, name)
-    return _maybe_wrap(var)
+    return self._apply_update(
+        self._variable.scatter_nd_update, indices, updates, name)
 
   def load(self, value, session=None):
     return self._variable.load(value, session)
@@ -462,24 +489,3 @@ def create_autocast_variable(variable):
       # pylint: enable=missing-format-attribute
 
   return AutoCastDistributedVariable(variable)
-
-
-def _maybe_wrap(variable, wrap=True):
-  """Creates an AutoCastVariable that wraps another variable if applicable.
-
-  This function is used to wrap the return value of AutoCastVariable.assign.
-  Unfortunately MirroredVariable.assign will (incorrectly) return a Mirrored
-  value instead of a MirroredVariable. So we cannot properly wrap it in an
-  AutoCastVariable. We return the original variable in that case.
-
-  Args:
-    variable: A tf.Variable or op.
-    wrap: A boolean to define whether to wrap the variable in an
-      AutoCastVariable or not.
-
-  Returns:
-    An AutoCastVariable if wrap is True and variable is a resource variable.
-  """
-  if wrap and resource_variable_ops.is_resource_variable(variable):
-    return create_autocast_variable(variable)
-  return variable
diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index c45015b644e..940bd07c813 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -304,8 +304,8 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertAllClose(3., self.evaluate(x.assign_sub(3.)))
 
         # Assign multiple times
-        # This currently only works if no strategy is used
-        if not ds_context.has_strategy():
+        # This currently doesn't work in graph mode
+        if context.executing_eagerly() or ops.inside_function():
           assign = x.assign(1.)
           self.assertAllClose(1., self.evaluate(assign))
           self.assertAllClose(0., self.evaluate(assign.assign(0.)))
@@ -343,6 +343,20 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         # assign still expect float32 value even if in float16 scope
         run_and_check()
 
+  @combinations.generate(maybe_distribute)
+  def test_assign_tf_function(self, distribution):
+    with distribution.scope():
+      x = get_var(0., dtypes.float32)
+      x = autocast_variable.create_autocast_variable(x)
+
+      @def_function.function
+      def run_assign():
+        return x.assign(1.).assign_add(3.).assign_add(3.).assign_sub(2.)
+
+      with ops.get_default_graph()._enable_auto_casting_variables(
+          dtypes.float16):
+        self.assertAllClose(5., self.evaluate(run_assign()))
+
   @combinations.generate(maybe_distribute)
   def test_assign_stays_in_true_dtype(self, distribution):
     with distribution.scope():
@@ -357,18 +371,18 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
           dtypes.float16):
         # Variable should be increased, despite it appearing to be the same
         # float16 value.
-        self.assertEqual(1. + small_val,
-                         self.evaluate(x.assign(1. + small_tensor)))
+        self.evaluate(x.assign(1. + small_tensor))
+        self.assertEqual(1. + small_val, self.evaluate(x._variable))
         self.assertEqual(1., self.evaluate(x.value()))
-      self.assertEqual(1. + small_val, self.evaluate(x.value()))
+      self.assertEqual(1. + small_val, self.evaluate(x))
 
       self.evaluate(x.assign(1.))
       with ops.get_default_graph()._enable_auto_casting_variables(
           dtypes.float16):
-        self.assertEqual(1. + small_val,
-                         self.evaluate(x.assign_add(small_tensor)))
+        self.evaluate(x.assign_add(small_tensor))
+        self.assertEqual(1. + small_val, self.evaluate(x._variable))
         self.assertEqual(1., self.evaluate(x.value()))
-      self.assertEqual(1. + small_val, self.evaluate(x.value()))
+      self.assertEqual(1. + small_val, self.evaluate(x))
 
   @combinations.generate(maybe_distribute)
   def test_checkpoint(self, distribution):

From 92d68dd5f10ce8ef40852f5bb207258ce4126edf Mon Sep 17 00:00:00 2001
From: Francois Chollet 
Date: Wed, 17 Jun 2020 14:30:33 -0700
Subject: [PATCH 0430/1390] Avoid deprecation warning related to the use of
 `collections.Sequence`. It will stop working in Python 3.8.

PiperOrigin-RevId: 316965801
Change-Id: Ia44313b1920653a0dd0a94d404ac914b08239c43
---
 tensorflow/python/keras/engine/data_adapter.py | 18 ------------------
 tensorflow/python/keras/layers/recurrent.py    |  9 ++++++---
 2 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/tensorflow/python/keras/engine/data_adapter.py b/tensorflow/python/keras/engine/data_adapter.py
index 469355dd722..29a99137982 100644
--- a/tensorflow/python/keras/engine/data_adapter.py
+++ b/tensorflow/python/keras/engine/data_adapter.py
@@ -19,7 +19,6 @@ from __future__ import division
 from __future__ import print_function
 
 import abc
-import collections
 import contextlib
 import functools
 import itertools
@@ -57,7 +56,6 @@ try:
   from scipy import sparse as scipy_sparse  # pylint: disable=g-import-not-at-top
 except ImportError:
   scipy_sparse = None
-
 try:
   import pandas as pd  # pylint: disable=g-import-not-at-top
 except ImportError:
@@ -786,7 +784,6 @@ class GeneratorDataAdapter(DataAdapter):
     # Since we have to know the dtype of the python generator when we build the
     # dataset, we have to look at a batch to infer the structure.
     peek, x = self._peek_and_restore(x)
-    assert_not_namedtuple(peek)
     peek = self._standardize_batch(peek)
     peek = _process_tensorlike(peek)
 
@@ -1070,21 +1067,6 @@ def broadcast_sample_weight_modes(target_structure, sample_weight_modes):
   return sample_weight_modes
 
 
-def assert_not_namedtuple(x):
-  if (isinstance(x, tuple) and
-      # TODO(b/144192902): Use a namedtuple checking utility.
-      hasattr(x, "_fields") and
-      isinstance(x._fields, collections.Sequence) and
-      all(isinstance(f, six.string_types) for f in x._fields)):
-    raise ValueError(
-        "Received namedtuple ({}) with fields `{}` as input. namedtuples "
-        "cannot, in general, be unambiguously resolved into `x`, `y`, "
-        "and `sample_weight`. For this reason Keras has elected not to "
-        "support them. If you would like the value to be unpacked, "
-        "please explicitly convert it to a tuple before passing it to "
-        "Keras.".format(x.__class__, x._fields))
-
-
 class DataHandler(object):
   """Handles iterating over epoch-level `tf.data.Iterator` objects."""
 
diff --git a/tensorflow/python/keras/layers/recurrent.py b/tensorflow/python/keras/layers/recurrent.py
index 0ce17c6101e..78a4a33a533 100644
--- a/tensorflow/python/keras/layers/recurrent.py
+++ b/tensorflow/python/keras/layers/recurrent.py
@@ -19,8 +19,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import collections
-
 import numpy as np
 
 from tensorflow.python.distribute import distribution_strategy_context as ds_context
@@ -49,6 +47,11 @@ from tensorflow.python.util import nest
 from tensorflow.python.util.tf_export import keras_export
 from tensorflow.tools.docs import doc_controls
 
+try:
+  from collections import abc as collections_abc  # pylint: disable=g-import-not-at-top
+except ImportError:  # For Python 2
+  import collections as collections_abc  # pylint: disable=g-import-not-at-top
+
 
 RECURRENT_DROPOUT_WARNING_MSG = (
     'RNN `implementation=2` is not supported when `recurrent_dropout` is set. '
@@ -828,7 +831,7 @@ class RNN(Layer):
     # input shape: `(samples, time (padded with zeros), input_dim)`
     # note that the .build() method of subclasses MUST define
     # self.input_spec and self.state_spec with complete input shapes.
-    if (isinstance(inputs, collections.Sequence)
+    if (isinstance(inputs, collections_abc.Sequence)
         and not isinstance(inputs, tuple)):
       # get initial_state from full input spec
       # as they could be copied to multiple GPU.

From b920ae92628e58362a3ef5fd55b3ef6e1a2ff617 Mon Sep 17 00:00:00 2001
From: Hanhan Wang 
Date: Wed, 17 Jun 2020 14:31:35 -0700
Subject: [PATCH 0431/1390] Internal change

PiperOrigin-RevId: 316965990
Change-Id: Ic5ee6b8b3fc04c7d5ca1548f8b638ba9c12a32d6
---
 third_party/mlir/BUILD      | 1 -
 third_party/mlir/test.BUILD | 1 -
 2 files changed, 2 deletions(-)

diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index cb0b2f9dc8e..476b8566265 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -3286,7 +3286,6 @@ cc_library(
         ":LinalgTransforms",
         ":Pass",
         ":StandardOps",
-        ":StandardOpsTransforms",
         ":Support",
         ":Transforms",
         ":VectorToLLVM",
diff --git a/third_party/mlir/test.BUILD b/third_party/mlir/test.BUILD
index 14c2ba7778e..23287ce28d6 100644
--- a/third_party/mlir/test.BUILD
+++ b/third_party/mlir/test.BUILD
@@ -166,7 +166,6 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:StandardOps",
-        "@llvm-project//mlir:StandardOpsTransforms",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TargetNVVMIR",
         "@llvm-project//mlir:TargetROCDLIR",

From 69e36409c0eb96379f8a4e8d5219f64382aaf75e Mon Sep 17 00:00:00 2001
From: Karim Nosir 
Date: Wed, 17 Jun 2020 14:32:39 -0700
Subject: [PATCH 0432/1390] Update documentation for FC op in hexagon delegate.
 FC now supports relu activation function

PiperOrigin-RevId: 316966229
Change-Id: If5a42ecca8aa9b6e94474e75614153f50ca8ae3b
---
 tensorflow/lite/delegates/hexagon/README.md                  | 2 +-
 tensorflow/lite/delegates/hexagon/builders/matmul_builder.cc | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/lite/delegates/hexagon/README.md b/tensorflow/lite/delegates/hexagon/README.md
index 226f3a61077..deff36e80b8 100644
--- a/tensorflow/lite/delegates/hexagon/README.md
+++ b/tensorflow/lite/delegates/hexagon/README.md
@@ -74,7 +74,7 @@ are verified in `IsNodeSupportedByHexagon`:
       - depth_multiplier == 1
       - dilation only supported when stride == 1
       - Otherwise, stride height/width <= 3
-* FullyConnected (without any activation)
+* FullyConnected
 * Hardswish
 * L2Normalization (without any activation)
 * Logistic (aka Sigmoid)
diff --git a/tensorflow/lite/delegates/hexagon/builders/matmul_builder.cc b/tensorflow/lite/delegates/hexagon/builders/matmul_builder.cc
index 0757ea6180e..6189294c3a1 100644
--- a/tensorflow/lite/delegates/hexagon/builders/matmul_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/matmul_builder.cc
@@ -151,7 +151,6 @@ TfLiteStatus AddFullyConnectedHelper(const TfLiteIntArray* inputs,
 // Data (8-bit), Weights (const, 8-bit) => MatMul => MatMul out (int32)
 // MatMul out (int32), Bias (int32) => QuantizedBiasAdd => BiasAdd out (int32)
 // BiasAdd out (int32) => Requantize_32to8 => Output (8-bit)
-// TODO(b/129276536): Add activation support.
 TfLiteStatus MatMulWithConstWeightsOpBuilder::PopulateSubGraph(
     const TfLiteIntArray* inputs, const TfLiteIntArray* outputs,
     TfLiteContext* context) {

From b6ff68822a59578f942e4fb8076757da8db278ae Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Wed, 17 Jun 2020 14:32:53 -0700
Subject: [PATCH 0433/1390] Manually define `GpuAtomic{Max,Min}` for `long
 long`.

CUDA only defines `atomic{Max,Min}` for `long long` for compute capability 3.5 or higher. Hence, to prevent compile errors due to failed instantiation when compiling on `sm30` (or other similar ones), we manually define the corresponding wrappers.

PiperOrigin-RevId: 316966290
Change-Id: I813e331309d12bdc1c06f7a74cb45c20c1833a41
---
 tensorflow/core/util/gpu_device_functions.h | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h
index de55b5c33c4..d4e09a7fc98 100644
--- a/tensorflow/core/util/gpu_device_functions.h
+++ b/tensorflow/core/util/gpu_device_functions.h
@@ -820,6 +820,11 @@ __device__ inline tensorflow::uint64 GpuAtomicMax(tensorflow::uint64* ptr,
   return detail::GpuAtomicCasHelper(
       ptr, [value](tensorflow::uint64 a) { return max(a, value); });
 }
+
+__device__ inline int64 GpuAtomicMax(int64* ptr, int64 value) {
+  return detail::GpuAtomicCasHelper(ptr,
+                                    [value](int64 a) { return max(a, value); });
+}
 #endif
 CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicMax, CudaAtomicMax);
 
@@ -885,6 +890,11 @@ __device__ inline tensorflow::uint64 GpuAtomicMin(tensorflow::uint64* ptr,
   return detail::GpuAtomicCasHelper(
       ptr, [value](tensorflow::uint64 a) { return min(a, value); });
 }
+
+__device__ inline int64 GpuAtomicMin(int64* ptr, int64 value) {
+  return detail::GpuAtomicCasHelper(ptr,
+                                    [value](int64 a) { return min(a, value); });
+}
 #endif
 CREATE_CUDA_DEVICE_FUNCTION_ALIAS(GpuAtomicMin, CudaAtomicMin);
 

From ed2b3d6e1ec7999cd44e84c62d3c99a4b8f047b8 Mon Sep 17 00:00:00 2001
From: Jaesung Chung 
Date: Wed, 17 Jun 2020 14:47:39 -0700
Subject: [PATCH 0434/1390] Fix msan test in tf_saved_model_lift_variables test

Now, all the tenors, generated from lift_variables_test.h, has initialized
values.

PiperOrigin-RevId: 316969327
Change-Id: If7f68aacf0c931430804dfaa1d73c0ea4be70c75
---
 .../transforms/lift_variables_test_pass.h        | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
index faecdf04368..0e6d844bed3 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/lift_variables_test_pass.h
@@ -96,15 +96,19 @@ class FakeSession : public tensorflow::Session {
     for (const std::string& output_name : output_names) {
       Tensor output;
       if (output_name == "dense/bias") {
-        outputs->push_back(
-            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({50})));
+        Tensor t = Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({50}));
+        t.flat().setZero();
+        outputs->push_back(t);
       } else if (output_name == "dense/kernel") {
-        outputs->push_back(
-            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({100, 50})));
+        Tensor t =
+            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({100, 50}));
+        t.flat().setZero();
+        outputs->push_back(t);
       } else {
         // Create a scalar float tensor.
-        outputs->push_back(
-            Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({})));
+        Tensor t = Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({}));
+        t.flat()(0) = 1.0f;
+        outputs->push_back(t);
       }
     }
     return Status::OK();

From 7a928592460e51c453ebaaecf965ef9ec963890b Mon Sep 17 00:00:00 2001
From: Ken Franko 
Date: Wed, 17 Jun 2020 15:06:41 -0700
Subject: [PATCH 0435/1390] Add SaveOptions/CheckpointOptions to
 keras.Models.save_weights and keras_call_backs.ModelCheckpoint.

PiperOrigin-RevId: 316973333
Change-Id: I43f5b59ece4b862db41ab0e99f3c8df0a0d3b901
---
 tensorflow/python/keras/callbacks.py          | 30 +++++++++++++---
 tensorflow/python/keras/callbacks_test.py     | 34 +++++++++++++++++++
 tensorflow/python/keras/engine/training.py    | 10 ++++--
 .../golden/v1/tensorflow.keras.-model.pbtxt   |  2 +-
 .../v1/tensorflow.keras.-sequential.pbtxt     |  2 +-
 ...ow.keras.callbacks.-model-checkpoint.pbtxt |  2 +-
 ...low.keras.experimental.-linear-model.pbtxt |  2 +-
 ....keras.experimental.-wide-deep-model.pbtxt |  2 +-
 .../v1/tensorflow.keras.models.-model.pbtxt   |  2 +-
 .../tensorflow.keras.models.-sequential.pbtxt |  2 +-
 .../golden/v2/tensorflow.keras.-model.pbtxt   |  2 +-
 .../v2/tensorflow.keras.-sequential.pbtxt     |  2 +-
 ...ow.keras.callbacks.-model-checkpoint.pbtxt |  2 +-
 ...low.keras.experimental.-linear-model.pbtxt |  2 +-
 ....keras.experimental.-wide-deep-model.pbtxt |  2 +-
 .../v2/tensorflow.keras.models.-model.pbtxt   |  2 +-
 .../tensorflow.keras.models.-sequential.pbtxt |  2 +-
 17 files changed, 82 insertions(+), 20 deletions(-)

diff --git a/tensorflow/python/keras/callbacks.py b/tensorflow/python/keras/callbacks.py
index 1bca5419774..1fae5abd84b 100644
--- a/tensorflow/python/keras/callbacks.py
+++ b/tensorflow/python/keras/callbacks.py
@@ -54,7 +54,9 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import summary_ops_v2
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.profiler import profiler_v2 as profiler
+from tensorflow.python.saved_model import save_options as save_options_lib
 from tensorflow.python.training import checkpoint_management
+from tensorflow.python.training.saving import checkpoint_options as checkpoint_options_lib
 from tensorflow.python.util import nest
 from tensorflow.python.util.compat import collections_abc
 from tensorflow.python.util.tf_export import keras_export
@@ -1115,6 +1117,9 @@ class ModelCheckpoint(Callback):
         epochs, the monitored metric may potentially be less reliable (it
         could reflect as little as 1 batch, since the metrics get reset every
         epoch). Defaults to `'epoch'`.
+      options: Optional `tf.train.CheckpointOptions` object if
+        `save_weights_only` is true or optional `tf.saved_model.SavedOptions`
+        object if `save_weights_only` is false.
       **kwargs: Additional arguments for backwards compatibility. Possible key
         is `period`.
   """
@@ -1127,6 +1132,7 @@ class ModelCheckpoint(Callback):
                save_weights_only=False,
                mode='auto',
                save_freq='epoch',
+               options=None,
                **kwargs):
     super(ModelCheckpoint, self).__init__()
     self._supports_tf_logs = True
@@ -1140,6 +1146,20 @@ class ModelCheckpoint(Callback):
     self._batches_seen_since_last_saving = 0
     self._last_batch_seen = 0
 
+    if save_weights_only:
+      if options is None or isinstance(
+          options, checkpoint_options_lib.CheckpointOptions):
+        self._options = options or checkpoint_options_lib.CheckpointOptions()
+      else:
+        raise TypeError('If save_weights_only is True, then `options` must be'
+                        'either None or a tf.train.CheckpointOptions')
+    else:
+      if options is None or isinstance(options, save_options_lib.SaveOptions):
+        self._options = options or save_options_lib.SaveOptions()
+      else:
+        raise TypeError('If save_weights_only is False, then `options` must be'
+                        'either None or a tf.saved_model.SaveOptions')
+
     # Deprecated field `load_weights_on_restart` is for loading the checkpoint
     # file from `filepath` at the start of `model.fit()`
     # TODO(rchao): Remove the arg during next breaking release.
@@ -1269,9 +1289,10 @@ class ModelCheckpoint(Callback):
                                                self.best, current, filepath))
               self.best = current
               if self.save_weights_only:
-                self.model.save_weights(filepath, overwrite=True)
+                self.model.save_weights(
+                    filepath, overwrite=True, options=self._options)
               else:
-                self.model.save(filepath, overwrite=True)
+                self.model.save(filepath, overwrite=True, options=self._options)
             else:
               if self.verbose > 0:
                 print('\nEpoch %05d: %s did not improve from %0.5f' %
@@ -1280,9 +1301,10 @@ class ModelCheckpoint(Callback):
           if self.verbose > 0:
             print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
           if self.save_weights_only:
-            self.model.save_weights(filepath, overwrite=True)
+            self.model.save_weights(
+                filepath, overwrite=True, options=self._options)
           else:
-            self.model.save(filepath, overwrite=True)
+            self.model.save(filepath, overwrite=True, options=self._options)
 
         self._maybe_remove_file()
       except IOError as e:
diff --git a/tensorflow/python/keras/callbacks_test.py b/tensorflow/python/keras/callbacks_test.py
index 28f85304688..d180e85a1d9 100644
--- a/tensorflow/python/keras/callbacks_test.py
+++ b/tensorflow/python/keras/callbacks_test.py
@@ -49,9 +49,11 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import summary_ops_v2
 from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import save_options as save_options_lib
 from tensorflow.python.summary import summary_iterator
 from tensorflow.python.training import adam
 from tensorflow.python.training import checkpoint_management
+from tensorflow.python.training.saving import checkpoint_options as checkpoint_options_lib
 
 try:
   import h5py  # pylint:disable=g-import-not-at-top
@@ -666,6 +668,38 @@ class KerasCallbacksTest(keras_parameterized.TestCase):
         mode=mode,
         save_freq=3)
 
+    # Case 9: `ModelCheckpoint` with valid and invalid `options` argument.
+    with self.assertRaisesRegexp(TypeError, 'tf.train.CheckpointOptions'):
+      keras.callbacks.ModelCheckpoint(
+          filepath,
+          monitor=monitor,
+          save_best_only=save_best_only,
+          save_weights_only=True,
+          mode=mode,
+          options=save_options_lib.SaveOptions())
+    with self.assertRaisesRegexp(TypeError, 'tf.saved_model.SaveOptions'):
+      keras.callbacks.ModelCheckpoint(
+          filepath,
+          monitor=monitor,
+          save_best_only=save_best_only,
+          save_weights_only=False,
+          mode=mode,
+          options=checkpoint_options_lib.CheckpointOptions())
+    keras.callbacks.ModelCheckpoint(
+        filepath,
+        monitor=monitor,
+        save_best_only=save_best_only,
+        save_weights_only=True,
+        mode=mode,
+        options=checkpoint_options_lib.CheckpointOptions())
+    keras.callbacks.ModelCheckpoint(
+        filepath,
+        monitor=monitor,
+        save_best_only=save_best_only,
+        save_weights_only=False,
+        mode=mode,
+        options=save_options_lib.SaveOptions())
+
   def _get_dummy_resource_for_model_checkpoint_testing(self):
 
     def get_input_datasets():
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index 5567e1733a7..ccd184a8bc4 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -1979,7 +1979,11 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
     save.save_model(self, filepath, overwrite, include_optimizer, save_format,
                     signatures, options)
 
-  def save_weights(self, filepath, overwrite=True, save_format=None):
+  def save_weights(self,
+                   filepath,
+                   overwrite=True,
+                   save_format=None,
+                   options=None):
     """Saves all layer weights.
 
     Either saves in HDF5 or in TensorFlow format based on the `save_format`
@@ -2032,6 +2036,8 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
         save_format: Either 'tf' or 'h5'. A `filepath` ending in '.h5' or
             '.keras' will default to HDF5 if `save_format` is `None`. Otherwise
             `None` defaults to 'tf'.
+        options: Optional `tf.train.CheckpointOptions` object that specifies
+            options for saving weights.
 
     Raises:
         ImportError: If h5py is not available when attempting to save in HDF5
@@ -2093,7 +2099,7 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
              'the TensorFlow format the optimizer\'s state will not be '
              'saved.\n\nConsider using a TensorFlow optimizer from `tf.train`.')
             % (optimizer,))
-      self._trackable_saver.save(filepath, session=session)
+      self._trackable_saver.save(filepath, session=session, options=options)
       # Record this checkpoint so it's visible from tf.train.latest_checkpoint.
       checkpoint_management.update_checkpoint_state_internal(
           save_dir=os.path.dirname(filepath),
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
index b62814e81cb..6318e577087 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.-model.pbtxt
@@ -302,7 +302,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
index 7485a0b3c62..9b7b7736746 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.-sequential.pbtxt
@@ -320,7 +320,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.callbacks.-model-checkpoint.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
index 5fb646e1c63..e6cc7aee5a8 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
@@ -5,7 +5,7 @@ tf_class {
   is_instance: ""
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'filepath\', \'monitor\', \'verbose\', \'save_best_only\', \'save_weights_only\', \'mode\', \'save_freq\'], varargs=None, keywords=kwargs, defaults=[\'val_loss\', \'0\', \'False\', \'False\', \'auto\', \'epoch\'], "
+    argspec: "args=[\'self\', \'filepath\', \'monitor\', \'verbose\', \'save_best_only\', \'save_weights_only\', \'mode\', \'save_freq\', \'options\'], varargs=None, keywords=kwargs, defaults=[\'val_loss\', \'0\', \'False\', \'False\', \'auto\', \'epoch\', \'None\'], "
   }
   member_method {
     name: "on_batch_begin"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
index bf980e5d116..976eb49d4c8 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-linear-model.pbtxt
@@ -303,7 +303,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
index c214a5c3419..500aa28eae7 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.experimental.-wide-deep-model.pbtxt
@@ -303,7 +303,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
index 86868c9d17f..ad0edc64606 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-model.pbtxt
@@ -302,7 +302,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
index 05aa19a915a..b38c669df0f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.models.-sequential.pbtxt
@@ -320,7 +320,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
index b62814e81cb..6318e577087 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.-model.pbtxt
@@ -302,7 +302,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
index 7485a0b3c62..9b7b7736746 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.-sequential.pbtxt
@@ -320,7 +320,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.callbacks.-model-checkpoint.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
index 5fb646e1c63..e6cc7aee5a8 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.callbacks.-model-checkpoint.pbtxt
@@ -5,7 +5,7 @@ tf_class {
   is_instance: ""
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'filepath\', \'monitor\', \'verbose\', \'save_best_only\', \'save_weights_only\', \'mode\', \'save_freq\'], varargs=None, keywords=kwargs, defaults=[\'val_loss\', \'0\', \'False\', \'False\', \'auto\', \'epoch\'], "
+    argspec: "args=[\'self\', \'filepath\', \'monitor\', \'verbose\', \'save_best_only\', \'save_weights_only\', \'mode\', \'save_freq\', \'options\'], varargs=None, keywords=kwargs, defaults=[\'val_loss\', \'0\', \'False\', \'False\', \'auto\', \'epoch\', \'None\'], "
   }
   member_method {
     name: "on_batch_begin"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
index bf980e5d116..976eb49d4c8 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-linear-model.pbtxt
@@ -303,7 +303,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
index c214a5c3419..500aa28eae7 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.experimental.-wide-deep-model.pbtxt
@@ -303,7 +303,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
index 86868c9d17f..ad0edc64606 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-model.pbtxt
@@ -302,7 +302,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
index 05aa19a915a..b38c669df0f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.models.-sequential.pbtxt
@@ -320,7 +320,7 @@ tf_class {
   }
   member_method {
     name: "save_weights"
-    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\'], varargs=None, keywords=None, defaults=[\'True\', \'None\'], "
+    argspec: "args=[\'self\', \'filepath\', \'overwrite\', \'save_format\', \'options\'], varargs=None, keywords=None, defaults=[\'True\', \'None\', \'None\'], "
   }
   member_method {
     name: "set_weights"

From b186ba0334e98ebbee5eacf7fc1c66897eaa19bc Mon Sep 17 00:00:00 2001
From: Li Lao 
Date: Wed, 17 Jun 2020 15:21:23 -0700
Subject: [PATCH 0436/1390] Add TraceMeProducer/Consumer for
 SharedBatchScheduler.

PiperOrigin-RevId: 316976310
Change-Id: I1cd2a03390aedd7e8e85b2826ab3aadd096bafdd
---
 tensorflow/core/kernels/batching_util/BUILD   |  2 ++
 .../kernels/batching_util/batch_scheduler.h   | 21 ++++++++++++++++++-
 .../batching_util/shared_batch_scheduler.h    | 18 +++++++++++-----
 .../core/profiler/lib/connected_traceme.h     |  1 +
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/batching_util/BUILD b/tensorflow/core/kernels/batching_util/BUILD
index 370c96c6e7f..803eb2e9048 100644
--- a/tensorflow/core/kernels/batching_util/BUILD
+++ b/tensorflow/core/kernels/batching_util/BUILD
@@ -70,6 +70,7 @@ cc_library(
         ":batch_scheduler_hdrs",
         ":periodic_function_dynamic",
         "//tensorflow/core:framework_headers_lib",
+        "//tensorflow/core/profiler/lib:connected_traceme",
         "//tensorflow/core/profiler/lib:traceme",
     ],
 )
@@ -81,6 +82,7 @@ cc_library(
         ":batch_scheduler",
         ":periodic_function_dynamic",
         "//tensorflow/core:lib",
+        "//tensorflow/core/profiler/lib:connected_traceme",
         "//tensorflow/core/profiler/lib:traceme",
     ],
     alwayslink = 1,
diff --git a/tensorflow/core/kernels/batching_util/batch_scheduler.h b/tensorflow/core/kernels/batching_util/batch_scheduler.h
index e418f8acbb1..d0e1d20bed4 100644
--- a/tensorflow/core/kernels/batching_util/batch_scheduler.h
+++ b/tensorflow/core/kernels/batching_util/batch_scheduler.h
@@ -77,7 +77,8 @@ class BatchTask {
 template 
 class Batch {
  public:
-  Batch() = default;
+  Batch();
+  explicit Batch(uint64 traceme_context_id);
   virtual ~Batch();  // Blocks until the batch is closed.
 
   // Appends 'task' to the batch. After calling AddTask(), the newly-added task
@@ -113,6 +114,9 @@ class Batch {
   // Marks the batch as closed. Dies if called more than once.
   void Close();
 
+  // Returns the TraceMe context id of this batch.
+  uint64 traceme_context_id() const;
+
  private:
   mutable mutex mu_;
 
@@ -125,6 +129,9 @@ class Batch {
   // Whether the batch has been closed.
   Notification closed_;
 
+  // The TracMe context id.
+  const uint64 traceme_context_id_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(Batch);
 };
 
@@ -187,6 +194,13 @@ class BatchScheduler {
 //////////
 // Implementation details follow. API users need not read.
 
+template 
+Batch::Batch() : Batch(0) {}
+
+template 
+Batch::Batch(uint64 traceme_context_id)
+    : traceme_context_id_(traceme_context_id) {}
+
 template 
 Batch::~Batch() {
   WaitUntilClosed();
@@ -275,6 +289,11 @@ void Batch::Close() {
   closed_.Notify();
 }
 
+template 
+uint64 Batch::traceme_context_id() const {
+  return traceme_context_id_;
+}
+
 }  // namespace serving
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h b/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
index 66bdff933d8..e47e069eff5 100644
--- a/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
+++ b/tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/profiler/lib/connected_traceme.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
@@ -311,6 +312,9 @@ class Queue {
   // The enqueued batches. See the invariants in the class comments above.
   std::deque>> batches_ TF_GUARDED_BY(mu_);
 
+  // The counter of the TraceMe context ids.
+  uint64 traceme_context_id_counter_ TF_GUARDED_BY(mu_) = 0;
+
   // The time at which the first task was added to the open (back-most) batch
   // in 'batches_'. Valid iff that batch contains at least one task.
   uint64 open_batch_start_time_micros_ TF_GUARDED_BY(mu_);
@@ -529,8 +533,6 @@ Queue::~Queue() {
 
 template 
 Status Queue::Schedule(std::unique_ptr* task) {
-  profiler::TraceMe trace_me(
-      [task] { return strings::StrCat("Schedule:", (*task)->size()); });
   if ((*task)->size() > options_.max_batch_size) {
     return errors::InvalidArgument("Task size ", (*task)->size(),
                                    " is larger than maximum batch size ",
@@ -554,6 +556,10 @@ Status Queue::Schedule(std::unique_ptr* task) {
     if (batches_.back()->empty()) {
       open_batch_start_time_micros_ = env_->NowMicros();
     }
+    profiler::TraceMeProducer trace_me(
+        [&] { return strings::StrCat("Schedule:", (*task)->size()); },
+        profiler::ContextType::kSharedBatchScheduler,
+        batches_.back()->traceme_context_id());
     batches_.back()->AddTask(std::move(*task));
 
     if (!schedulable_batch_) {
@@ -621,8 +627,10 @@ std::unique_ptr> Queue::ScheduleBatch() {
 
 template 
 void Queue::ProcessBatch(std::unique_ptr> batch) {
-  profiler::TraceMe trace_me(
-      [&batch] { return strings::StrCat("ProcessBatch:", batch->size()); });
+  profiler::TraceMeConsumer trace_me(
+      [&batch] { return strings::StrCat("ProcessBatch:", batch->size()); },
+      profiler::ContextType::kSharedBatchScheduler,
+      batch->traceme_context_id());
   process_batch_callback_(std::move(batch));
 
   {
@@ -665,7 +673,7 @@ bool Queue::IsEmptyInternal() const {
 template 
 void Queue::StartNewBatch() {
   batches_.back()->Close();
-  batches_.emplace_back(new Batch);
+  batches_.emplace_back(new Batch(++traceme_context_id_counter_));
 }
 
 template 
diff --git a/tensorflow/core/profiler/lib/connected_traceme.h b/tensorflow/core/profiler/lib/connected_traceme.h
index cbc610af407..ed8b4ac1ad2 100644
--- a/tensorflow/core/profiler/lib/connected_traceme.h
+++ b/tensorflow/core/profiler/lib/connected_traceme.h
@@ -28,6 +28,7 @@ namespace profiler {
 enum class ContextType : int {
   kGeneric,
   kTfExecutor,
+  kSharedBatchScheduler,
 };
 
 /*

From 19c51afbf1fd2f724cabc42313ce84e31a4defcc Mon Sep 17 00:00:00 2001
From: Jaesung Chung 
Date: Wed, 17 Jun 2020 15:42:14 -0700
Subject: [PATCH 0437/1390] Do not emit op errors on TF -> TFL legalization

Applying legalization patterns will emit unwanted, transient errors
when the replaced TFLite ops do not meet the sanity checks. In order to
ignore the transient errors, the following lines override a diagnostic
handler with an no-op handler only while this pass runs.

PiperOrigin-RevId: 316980278
Change-Id: Idef14e13f36ff0ee3c4bb1a401f92ba217042dbb
---
 .../compiler/mlir/lite/converter_gen.cc       | 18 +++++-----------
 .../mlir/lite/ir/tfl_op_interfaces.td         |  3 +--
 .../compiler/mlir/lite/tests/legalize-tf.mlir |  7 +++++++
 .../mlir/lite/transforms/legalize_tf.cc       | 21 ++++++++++++++++---
 .../mlir/lite/transforms/runtime_verify.cc    |  3 +--
 5 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/converter_gen.cc b/tensorflow/compiler/mlir/lite/converter_gen.cc
index 6df569a8031..edead2037a3 100644
--- a/tensorflow/compiler/mlir/lite/converter_gen.cc
+++ b/tensorflow/compiler/mlir/lite/converter_gen.cc
@@ -446,7 +446,7 @@ static void GenOperandResultVerifier(raw_ostream &os,
     auto desc =
         definit->getDef()->getValueAsString("tflRuntimeTypeDescription");
 
-    // Emit a loop to check all the dynamic values in the pack.
+    // Emit a loop to check all operands.
     os << formatv("    for (Value v : top.getODS{0}{1}s({2})) {{\n",
                   // Capitalize the first letter to match the function name
                   valueKind.substr(0, 1).upper(), valueKind.substr(1),
@@ -455,14 +455,10 @@ static void GenOperandResultVerifier(raw_ostream &os,
     os << "      (void)v;\n"
        << "      if (!("
        << tgfmt(pred.getCondition(), &fctx.withSelf("v.getType()")) << ")) {\n"
-       << "        if (failure_on_operand_type_mismatch) {\n"
        << formatv(
-              "        return op->emitOpError(\"{0} #\") << index "
+              "      return op->emitOpError(\"{0} #\") << index "
               "<< \" must be {1}, but got \" << v.getType();\n",
               valueKind, desc)
-       << "        } else {\n"
-       << "          return ::mlir::LogicalResult::Failure;\n"
-       << "        }\n"
        << "      }\n"  // if
        << "      ++index;\n"
        << "    }\n";  // for
@@ -487,8 +483,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, RecordKeeper &records) {
 
     mlir::tblgen::FmtContext verify_ctx;
     os << "::mlir::LogicalResult " << op.getCppClassName()
-       << "::VerifyTflRuntimeConstraints(::mlir::Operation *op, bool "
-          "failure_on_operand_type_mismatch) {\n";
+       << "::VerifyTflRuntimeConstraints(::mlir::Operation *op) {\n";
     os << "  auto top = cast<" << op.getCppClassName() << ">(op); (void)top;\n";
     verify_ctx.withOp("top");
 
@@ -529,11 +524,8 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, RecordKeeper &records) {
 
       mlir::tblgen::Pred pred(dyn_cast(val->getValue()));
       os << tgfmt(
-          "  if (!($0)) {\n    "
-          "    if (failure_on_operand_type_mismatch) {\n"
-          "      return top.emitOpError(\"failed to verify that $1\");\n"
-          "    } else {\n"
-          "      return ::mlir::LogicalResult::Failure;\n  }\n  }\n",
+          "  if (!($0))\n"
+          "    return top.emitOpError(\"failed to verify that $1\");\n",
           &verify_ctx, tgfmt(pred.getCondition(), &verify_ctx), desc);
     }
     os << "  return top.verify();\n}\n";
diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_op_interfaces.td b/tensorflow/compiler/mlir/lite/ir/tfl_op_interfaces.td
index 23101113a6f..a79d79b5970 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_op_interfaces.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_op_interfaces.td
@@ -94,8 +94,7 @@ def TFL_RuntimeVerification : OpInterface<"TflRuntimeVerifyOpInterface"> {
   let methods = [
     StaticInterfaceMethod<
       [{Returns whether the op's operands/results are supported by runtime.}],
-      "LogicalResult", "VerifyTflRuntimeConstraints",
-      (ins "Operation*":$op, "bool":$failure_on_operand_type_mismatch)
+      "LogicalResult", "VerifyTflRuntimeConstraints", (ins "Operation*":$op)
     >,
   ];
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir
index 1ae789f5468..5756fa6dec2 100644
--- a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir
@@ -990,6 +990,13 @@ func @batch_to_space_nd(%arg0: tensor<4x2x2x3xf32>, %arg1: tensor<2xi32>, %arg2:
   // CHECK: "tfl.batch_to_space_nd"(%arg0, %arg1, %arg2) : (tensor<4x2x2x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor
 }
 
+func @batch_to_space_nd_unsupported(%arg0: tensor, %arg1: tensor<3xi32>, %arg2: tensor<3x2xi32>) -> tensor {
+  %0 = "tf.BatchToSpaceND"(%arg0, %arg1, %arg2) : (tensor, tensor<3xi32>, tensor<3x2xi32>) -> tensor
+  return %0 : tensor
+  // CHECK-LABEL: batch_to_space_nd_unsupported
+  // CHECK: "tf.BatchToSpaceND"
+}
+
 func @space_to_batch_nd(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<2xi32>, %arg2: tensor<2x2xi32>) -> tensor {
   %0 = "tf.SpaceToBatchND"(%arg0, %arg1, %arg2) : (tensor<1x4x4x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor
   return %0 : tensor
diff --git a/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc b/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
index 46ed134d7ee..1328a2baf5d 100644
--- a/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc
@@ -28,9 +28,11 @@ limitations under the License.
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/Threading.h"
 #include "mlir/Dialect/Quant/FakeQuantSupport.h"  // from @llvm-project
 #include "mlir/Dialect/Quant/UniformSupport.h"  // from @llvm-project
 #include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
@@ -767,13 +769,26 @@ void LegalizeTF::runOnFunction() {
             [](Operation* op) {
               auto tfl_op = dyn_cast_or_null(op);
               if (!tfl_op) return false;
-              return succeeded(tfl_op.VerifyTflRuntimeConstraints(
-                  tfl_op.getOperation(),
-                  /*failure_on_operand_type_mismatch=*/false));
+              return succeeded(tfl_op.VerifyTflRuntimeConstraints(op));
             }));
   } else {
     target.addLegalDialect();
   }
+
+  // Ignore transient errors by registering an no-op handler.
+  // Applying legalization patterns will emit unwanted, transient errors when
+  // the replaced TFLite ops do not meet the sanity checks. In order to ignore
+  // the transient errors, the following lines override a diagnostic handler
+  // with an no-op handler only while this pass runs.
+  uint64_t current_thread_id = llvm::get_threadid();
+  ScopedDiagnosticHandler scoped_diag_handler(
+      context, [¤t_thread_id](Diagnostic&) -> LogicalResult {
+        // Consume only errors that are coming from the same thread in order not
+        // to ignore errors from other passes that are running. Things running
+        // in the pass manager can be multi-threaded.
+        return success(current_thread_id == llvm::get_threadid());
+      });
+
   // Keep trying to convert.
   // TODO(karimnosseir): This is similar to what apply greedy patterns does.
   // Look if there is a function that tries until it converge.
diff --git a/tensorflow/compiler/mlir/lite/transforms/runtime_verify.cc b/tensorflow/compiler/mlir/lite/transforms/runtime_verify.cc
index 3268329b1c1..cc2e691180e 100644
--- a/tensorflow/compiler/mlir/lite/transforms/runtime_verify.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/runtime_verify.cc
@@ -34,8 +34,7 @@ class RuntimeVerifyPass
 
 void RuntimeVerifyPass::runOnFunction() {
   getFunction().walk([&](TflRuntimeVerifyOpInterface op) {
-    if (failed(op.VerifyTflRuntimeConstraints(
-            op.getOperation(), /*failure_on_operand_type_mismatch=*/true)))
+    if (failed(op.VerifyTflRuntimeConstraints(op.getOperation())))
       signalPassFailure();
   });
 }

From db3640aaf1e1b520fa59cc134f822cbb7a2ed270 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 15:43:59 -0700
Subject: [PATCH 0438/1390] tf.numpy Add a README.md

PiperOrigin-RevId: 316980572
Change-Id: I441a3d9dc0d36b663a47721f2335b4178a71db65
---
 tensorflow/python/ops/numpy_ops/README.md | 94 +++++++++++++++++++++++
 1 file changed, 94 insertions(+)
 create mode 100644 tensorflow/python/ops/numpy_ops/README.md

diff --git a/tensorflow/python/ops/numpy_ops/README.md b/tensorflow/python/ops/numpy_ops/README.md
new file mode 100644
index 00000000000..3dc37423d26
--- /dev/null
+++ b/tensorflow/python/ops/numpy_ops/README.md
@@ -0,0 +1,94 @@
+# tf.experimental.numpy
+
+This module provides a subset of numpy API, built on top of TensorFlow
+operations. APIs are based on numpy 1.16 version.
+
+The set of supported APIs may be expanded over time. Also future releases may
+change the baseline version of numpy API being supported. A list of some
+systematic differences with numpy are listed later in the "Differences with
+Numpy" section.
+
+## Types
+
+The module provide an `ndarray` class which wraps an immutable `tf.Tensor`.
+Additional functions are provided which accept array-like objects. Here
+array-like objects includes `ndarrays` as defined by this module, as well as
+`tf.Tensor`, in addition to types accepted by `numpy`.
+
+A subset of `numpy` dtypes are supported, along with `tf.bfloat16`.
+Additionally, support is provided for selecting the default float type
+(`np.float32` vs `np.float64`) given that some applications may prefer lower
+precision.
+
+## Device Support
+
+Given that `ndarray` and functions wrap TensorFlow constructs, the code will
+have GPU and TPU support on par with TensorFlow. Also the code can be wrapped
+with `tf.function` and XLA compiled. Device placement can be controlled by using
+`with tf.device` scopes.
+
+## Graph and Eager Modes
+
+Eager mode execution should typically match numpy semantics of executing
+op-by-op. However the same code can be executed in graph mode, by putting it
+inside a `tf.function`. This can change behavior of certain operations since
+symbolic execution may not have information that is computed during runtime.
+
+Some differences are:
+
+  * Shapes can be incomplete or unknown. This means that `ndarray.shape`,
+    `ndarray.size` and `ndarray.ndim` can return `ndarray` objects instead of
+    returning integer (or tuple of integer) values.
+  * Python control flow based on `ndarray` values may not work and may have to
+    be rewritten to use `tf.cond` or `tf.while_loop`. Note that autograph
+    conversion as part of `tf.function` should still work.
+  * `__len__`, `__iter__` and `__index__` properties of `ndarray` may similarly
+    not work in graph mode.
+
+## Mutation and Variables
+
+`ndarrays` currently wrap immutable `tf.Tensor`. Also currently mutation
+operations like slice assigns are not supported. This may change in the future.
+
+There is currently no explict construct on par with `tf.Variable`. However one
+can directly construct a `tf.Variable` and use that with the numpy APIs in this
+module. See section on Interop.
+
+## Interop
+
+The numpy API calls can be interleaved with TensorFlow calls without incurring
+Tensor data copies. This is true even if the `ndarray` or `tf.Tensor` is placed
+on a non-CPU device.
+
+Additionally, one could put these calls in a `with tf.GradientTape()` context to
+compute gradients through the numpy API calls. Similarly, code vectorization can
+be done using `tf.vectorized_map()`.
+
+In general, the expected behavior should be on par with that of code involving
+`tf.Tensor` and running stateless TensorFlow functions on them.
+
+## Array Interface
+
+The `ndarray` class implements the `__array__` interface. This should allow
+these objects to be passed into contexts that expect a `numpy` or array-like
+object (e.g. matplotlib).
+
+
+## Differences with Numpy
+
+Here is a non-exhaustive list of differences:
+
+  * Not all dtypes are currently supported. e.g. `np.float96`, `np.float128`.
+    `np.object`, `np.str`, `np.recarray` types are not supported.
+  * `ndarray` storage is in C order only. Fortran order, views, stride_tricks
+    are not supported.
+  * Only a subset of functions and modules are supported. This set would be
+    expanded over time. For supported functions, some arguments or argument
+    values may not be supported. This differences are listed in the function
+    comments.
+  * Buffer mutation is currently not supported. `ndarrays` wrap immutable
+    tensors. This means that output buffer arguments (e..g `out` in ufuncs) are
+    not supported
+  * full `ufunc` support is not provided.
+  * Numpy C API is not supported. Numpy's Cython and Swig integration are not
+    supported.

From eb0d46affb1eccfb4de1463ef9c70fdedc227d22 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 15:44:12 -0700
Subject: [PATCH 0439/1390] Import initialization graph in SignatureDef
 SavedModels as an MLIR function in TF saved model dialect.

PiperOrigin-RevId: 316980623
Change-Id: I3dafe9303ad65192ca0c137943be3cd1ea0e74e7
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |   4 +-
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  25 ---
 .../mlir/tensorflow/ir/tf_saved_model_ops.td  |  24 ---
 .../tests/tf_saved_model/common_v1.py         |   1 -
 .../tests/tf_saved_model/hash_table_v1.py     |  92 -----------
 .../tensorflow/tests/tf_saved_model_ops.mlir  |   5 -
 .../tests/tf_saved_model_ops_invalid.mlir     |  33 ----
 .../mlir/tensorflow/translate/import_model.cc | 149 +++++-------------
 8 files changed, 43 insertions(+), 290 deletions(-)
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 17ed0e36a28..904ccb7e820 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -661,9 +661,7 @@ cc_library(
         ":tensorflow_types",
         ":translate_utils",
         "//tensorflow/cc/saved_model:bundle_v2",
-        "//tensorflow/cc/saved_model:constants",
         "//tensorflow/cc/saved_model:loader_lite",
-        "//tensorflow/cc/saved_model:loader_util",
         "//tensorflow/compiler/jit:shape_inference_helpers",
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/tf2xla:functionalize_control_flow",
@@ -675,7 +673,6 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/utils:transitive_fanin",
-        "//tensorflow/core/platform:protobuf_internal",
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
@@ -685,6 +682,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:StandardOps",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index 6af70158e14..140a778770c 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -76,23 +76,6 @@ static LogicalResult Verify(GlobalTensorOp global_tensor) {
   return success();
 }
 
-static LogicalResult Verify(SessionInitializerOp session_initializer) {
-  mlir::SymbolTable symbol_table(
-      session_initializer.getParentOfType());
-
-  auto init_func_op =
-      symbol_table.lookup(session_initializer.initializer());
-  if (!init_func_op)
-    return session_initializer.emitOpError()
-           << "the initializer function does not exist";
-
-  if (!init_func_op.getType().getResults().empty())
-    return session_initializer.emitOpError()
-           << "the initializer function should have no output";
-
-  return success();
-}
-
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc.inc"
 
@@ -237,14 +220,6 @@ static LogicalResult VerifySavedModelModule(
       }
     }
   }
-
-  auto session_initializers = module.getOps();
-  if (std::distance(session_initializers.begin(), session_initializers.end()) >
-      1) {
-    return (*++session_initializers.begin()).emitError()
-           << "there must be no more than one session_initializer op";
-  }
-
   SymbolTable symbol_table(module);
   auto symbol_uses = SymbolTable::getSymbolUses(&module.getBodyRegion());
   if (!symbol_uses.hasValue()) {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
index 497f4d90cb9..4431a160edf 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
@@ -128,28 +128,4 @@ def TfSavedModel_GlobalTensorOp : TfSavedModel_Op<"global_tensor"> {
   let verifier = [{ return Verify(*this); }];
 }
 
-def TfSavedModel_SessionInitializerOp: TfSavedModel_Op<"session_initializer"> {
-  let summary = "Initializes TensorFlow session state.";
-  let description = [{
-    Represents a session initializer function initializes TensorFlow session
-    state. It is used to initialize resources in the saved model before calling
-    any exported functions. There must be no more than one session initializer
-    in a saved model.
-
-    The `initializer` represents the initialization function. The function have
-    no output and this function should be only called once.
-
-    This is used, for example, to initialize hash tables stored in resources and
-    accessed by resource name (rather than as resource handles or bound inputs
-    which is how `global_tensor`s are referenced)
-  }];
-
-  let arguments = (ins
-    FlatSymbolRefAttr:$initializer
-  );
-
-
-  let verifier = [{ return Verify(*this); }];
-}
-
 #endif // SAVED_MODEL_DIALECT
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
index 51ccbeb1fbd..7171f63bb05 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
@@ -84,7 +84,6 @@ def do_test(signature_def_map, show_debug_info=False):
     builder.add_meta_graph_and_variables(
         sess, [tf.saved_model.tag_constants.SERVING],
         signature_def_map,
-        main_op=tf.tables_initializer(),
         strip_default_attrs=True)
     builder.save()
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
deleted file mode 100644
index 64847434b82..00000000000
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-
-# RUN: %p/hash_table_v1 | FileCheck %s
-
-# pylint: disable=missing-docstring,line-too-long
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow.compat.v1 as tf
-from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
-
-# Verify that the tf.versions attribute exists. It is difficult to enforce
-# contents, since the version numbers change over time. The conversion logic
-# itself is verified in the common graphdef converter, so here just assert
-# it is being invoked.
-# CHECK: module
-# CHECK-SAME: tf.versions
-# CHECK-SAME: bad_consumers
-# CHECK-SAME: min_consumer
-# CHECK-SAME: producer
-
-# CHECK: "tf_saved_model.session_initializer"() {initializer = [[init:@.*]]} : () -> ()
-# CHECK: "tf_saved_model.global_tensor"()
-
-# CHECK:      func {{@[a-zA-Z_0-9]+}}(
-# CHECK-SAME: [[ARG0:%.*]]: tensor
-# CHECK-SAME: [[ARG1:%.*]]: tensor ()
-
   // Representation for constants: (immutable) global tensor.
   // CHECK: tf_saved_model.global_tensor
   "tf_saved_model.global_tensor"() {
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index 544600cf6b8..c055c6c9f56 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -258,36 +258,3 @@ module attributes {tf_saved_model.semantics} {
   // expected-error@+1 {{'type' attribute for immutable 'tf_saved_model.global_tensor' should have a static shape}}
   "tf_saved_model.global_tensor"() { sym_name = "v", type = tensor, value = dense<1.> : tensor<1xf32> } : () -> ()
 }
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // expected-error@+1 {{the initializer function does not exist}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // expected-error@+1 {{the initializer function should have no output}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> {
-    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
-    return %0 : tensor<1xf32>
-  }
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  // expected-error@+1 {{there must be no more than one session_initializer op}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> {
-    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
-    return %0 : tensor<1xf32>
-  }
-}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 3cff4217215..820d0ce31fb 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -60,8 +60,6 @@ limitations under the License.
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
-#include "tensorflow/cc/saved_model/constants.h"
-#include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
@@ -101,7 +99,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/platform/protobuf_internal.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
@@ -119,7 +116,6 @@ using mlir::NamedAttrList;
 using mlir::TensorType;
 using mlir::TF::VarHandleOp;
 using mlir::tf_saved_model::GlobalTensorOp;
-using mlir::tf_saved_model::SessionInitializerOp;
 using stream_executor::port::StatusOr;
 
 namespace {
@@ -2959,13 +2955,6 @@ void SortSavedModelModule(mlir::ModuleOp module) {
     named_global_tensor.global_tensor.getOperation()->moveBefore(
         &module.getBody()->front());
   }
-
-  auto initializers = module.getOps();
-  if (!initializers.empty()) {
-    (*initializers.begin())
-        .getOperation()
-        ->moveBefore(&module.getBody()->front());
-  }
 }
 
 Status CreateSavedModelIR(
@@ -3252,29 +3241,17 @@ class SavedModelSignatureDefImporter {
                                  absl::Span exported_names,
                                  mlir::MLIRContext* context)
       : bundle_(bundle),
-        flib_def_(OpRegistry::Global(), graph_def().library()),
-        debug_info_(),
         exported_names_(exported_names),
-        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {
-    // debug_info might not be loaded with loader_lite.
-    if (bundle_.debug_info != nullptr) debug_info_ = *bundle_.debug_info;
-  }
+        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {}
 
   // Converts the SavedModel to the SavedModel dialect. Creates an MLIR function
   // for each signature.
   StatusOr ConvertSignatures();
-  Status ConvertSignature(const std::string& sig_def_key,
-                          const SignatureDef& signature_def);
-
-  // Converts the initialization graph in the SavedModel to an MLIR function.
-  Status ConvertInitializer();
-
-  // Converts a graph with feeds and fetches to an MLIR function.
-  StatusOr ConvertGraph(
-      const std::string& name,
-      const std::vector>& inputs,
-      const std::vector>& outputs,
-      const std::vector control_outputs);
+  Status ConvertSignature(const GraphDef& graphdef,
+                          const std::string& sig_def_key,
+                          const SignatureDef& signature_def,
+                          const GraphDebugInfo& debug_info,
+                          const FunctionLibraryDefinition& flib_def);
 
   // Creates GlobalTensorOp for each variable and moves each VarHandle op to
   // the enclosing function's arguments.
@@ -3296,62 +3273,18 @@ class SavedModelSignatureDefImporter {
   GraphImportConfig::InputArrays ParseInputArrays(
       const std::vector>& inputs);
 
-  const GraphDef& graph_def() const {
-    return bundle_.meta_graph_def.graph_def();
-  }
-  const FunctionLibraryDefinition& flib_def() const { return flib_def_; }
-  const GraphDebugInfo& debug_info() const { return debug_info_; }
-
   const SavedModelBundle& bundle_;
-  FunctionLibraryDefinition flib_def_;
-  GraphDebugInfo debug_info_;
   absl::Span exported_names_;
   mlir::OwningModuleRef module_;
 };
 
-Status SavedModelSignatureDefImporter::ConvertInitializer() {
-  std::vector asset_file_defs;
-  TF_RETURN_IF_ERROR(
-      internal::GetAssetFileDefs(bundle_.meta_graph_def, &asset_file_defs));
-
-  if (!asset_file_defs.empty())
-    return errors::Unimplemented(
-        absl::StrCat("Assets are not supported in signaturedef importer"));
-
-  std::string init_node_name;
-  TF_RETURN_IF_ERROR(
-      internal::GetInitOp("", bundle_.meta_graph_def, &init_node_name));
-
-  if (init_node_name.empty()) return Status::OK();
-
-  TF_ASSIGN_OR_RETURN(auto sub_module,
-                      ConvertGraph(init_node_name, {}, {}, {init_node_name}));
-
-  mlir::SymbolTable symbol_table(*sub_module);
-
-  auto init_func_op = symbol_table.lookup(init_node_name);
-
-  init_func_op.removeAttr("tf.entry_function");
-
-  mlir::OpBuilder builder(module_->getBodyRegion());
-
-  builder.create(
-      module_->getLoc(), builder.getSymbolRefAttr(init_func_op.getName()));
-
-  // Move the converted functions to top level MLIR module.
-  auto* block = module_->getBody();
-  auto* sub_block = sub_module->getBody();
-  block->getOperations().splice(
-      mlir::Block::iterator(block->getTerminator()), sub_block->getOperations(),
-      sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
-
-  return Status::OK();
-}
-
 StatusOr
 SavedModelSignatureDefImporter::ConvertSignatures() {
   const auto& signatures = bundle_.GetSignatures();
-  PopulateTfVersions(module_.get(), graph_def().versions());
+  const auto& graphdef = bundle_.meta_graph_def.graph_def();
+  PopulateTfVersions(module_.get(), graphdef.versions());
+
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graphdef.library());
 
   // debug_info might not be loaded with loader_lite.
   GraphDebugInfo debug_info;
@@ -3374,10 +3307,9 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
       continue;
     }
 
-    TF_RETURN_IF_ERROR(ConvertSignature(sig_def_key, signature_def));
+    TF_RETURN_IF_ERROR(ConvertSignature(graphdef, sig_def_key, signature_def,
+                                        debug_info, flib_def));
   }
-
-  TF_RETURN_IF_ERROR(ConvertInitializer());
   TF_RETURN_IF_ERROR(LiftVariables());
 
   mlir::OpBuilder builder(module_->getBodyRegion());
@@ -3388,32 +3320,10 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
   return std::move(module_);
 }
 
-StatusOr SavedModelSignatureDefImporter::ConvertGraph(
-    const std::string& name,
-    const std::vector>& inputs,
-    const std::vector>& outputs,
-    const std::vector control_outputs) {
-  GraphImportConfig specs;
-  specs.prune_unused_nodes = true;
-  specs.inputs = ParseInputArrays(inputs);
-  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
-  specs.control_outputs = control_outputs;
-
-  // Convert sub-graphdef to sub-graph.
-  GraphConstructorOptions options;
-  options.allow_internal_ops = true;
-  options.add_default_attributes = true;
-  Graph graph(OpRegistry::Global());
-
-  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, graph_def(), &graph));
-
-  // Convert sub-graph to MLIR module.true
-  return GraphDefImporter::Convert(module_->getContext(), graph, debug_info(),
-                                   flib_def(), specs, name);
-}
-
 Status SavedModelSignatureDefImporter::ConvertSignature(
-    const std::string& sig_def_key, const SignatureDef& signature_def) {
+    const GraphDef& graphdef, const std::string& sig_def_key,
+    const SignatureDef& signature_def, const GraphDebugInfo& debug_info,
+    const FunctionLibraryDefinition& flib_def) {
   // Create local vectors for the input and output and sort them to be
   // deterministic. We don't want anyone to really depend on the order, client
   // should lookup argument/result mapping by attribute name.
@@ -3429,9 +3339,34 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
     return lhs.first.size() < rhs.first.size() || lhs.first > rhs.first;
   });
 
+  GraphImportConfig specs;
+  specs.prune_unused_nodes = true;
+  specs.inputs = ParseInputArrays(inputs);
+  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
+
+  // Remove unused nodes and create sub-graphdef.
+  GraphDef sub_graph_def;
+  TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(
+      graphdef, &sub_graph_def,
+      /*terminal_nodes=*/{specs.outputs.begin(), specs.outputs.end()}));
+
+  // Set the function library definitions in the pruned graphdef.
+  *sub_graph_def.mutable_library() = flib_def.ToProto();
+
+  // Convert sub-graphdef to sub-graph.
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.add_default_attributes = true;
+  Graph sub_graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(
+      ConvertGraphDefToGraph(options, sub_graph_def, &sub_graph));
+
   // Convert sub-graph to MLIR module.
-  TF_ASSIGN_OR_RETURN(auto sub_module,
-                      ConvertGraph(sig_def_key, inputs, outputs, {}));
+  TF_ASSIGN_OR_RETURN(
+      auto sub_module,
+      GraphDefImporter::Convert(module_->getContext(), sub_graph, debug_info,
+                                flib_def, specs, sig_def_key));
   mlir::OpBuilder builder(sub_module->getBodyRegion());
 
   // Find the FuncOp which corresponds to current SignatureDef.

From 02be9b5ef8b3faac54feea8ee99ff14c62d40704 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Wed, 17 Jun 2020 15:46:30 -0700
Subject: [PATCH 0440/1390] Remove erroneous end-of-namespace comment at the
 end of a function.

PiperOrigin-RevId: 316981028
Change-Id: I06c1063ba6206fe301f8fd6b89e3ac435ec56ba6
---
 tensorflow/lite/kernels/lstm_eval.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 4ac9e538317..2e7f300f9a9 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -893,7 +893,7 @@ inline void LstmStepHybrid(
     std::copy_n(output_state_ptr + b * n_output, n_output,
                 output_ptr + b * output_batch_leading_dim);
   }
-}  // namespace
+}
 
 // Fully quantized lstm kernel for 16 bit gate matmul output.
 //

From 09bdeb13ebb20f12c76e3d0bd852f42b2f63ffef Mon Sep 17 00:00:00 2001
From: Andy Ly 
Date: Wed, 17 Jun 2020 15:46:55 -0700
Subject: [PATCH 0441/1390] Fix updating results of TPU cluster with
 parallel_execute results when TPU cluster results are not perfectly
 forwarded.

After TPUExtractHeadTailOutsideCompilation some results left of the TPU cluster may not be used by an optional tf_device.replicate op if there is data parallelism. Instead, all results should be remapped if they are used outside of parallel_execute.

PiperOrigin-RevId: 316981114
Change-Id: I5529074857e06cfe26a7141c262a6229fe848be6
---
 .../tpu_extract_outside_compilation.mlir      | 22 +++++++++++++++++-
 .../tpu_extract_outside_compilation.cc        | 23 ++++++++-----------
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu_extract_outside_compilation.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu_extract_outside_compilation.mlir
index 6bb8e99d796..d88489f5da0 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tpu_extract_outside_compilation.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu_extract_outside_compilation.mlir
@@ -262,7 +262,6 @@ func @single_outside_compiled_input_output_single_outside_compilation(%arg0: ten
   return %1 : tensor
 }
 
-
 // Tests extraction of a single outside compiled cluster with multiple input/output.
 
 // CHECK-LABEL: func @multiple_outside_compiled_input_output_single_outside_compilation
@@ -439,3 +438,24 @@ func @multiple_outside_compiled_inputs_single_outside_compilation(%arg0: tensor<
 
   return %1 : tensor
 }
+
+// Tests only directly used results of tpu cluster are remapped with
+// parallel_execute.
+
+// CHECK-LABEL: func @remapped_results
+func @remapped_results(%arg0: tensor) -> tensor {
+  %0 = "tf.A"(%arg0) : (tensor) -> tensor
+  // CHECK: %[[REPLICATE:[0-9]*]]:2 = tf_device.replicate
+  // CHECK:   %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:2 = "tf_device.parallel_execute"
+  // CHECK: tf_device.return %[[PARALLEL_EXECUTE_OUTPUT]]#1 : tensor
+  %1:2 = tf_device.replicate([%0, %arg0] as %ri_0: tensor) {n = 2 : i32} {
+    %2:2 = "tf_device.cluster"() ( {
+      %3 = "tf.A"() : () -> (tensor)
+      %4 = "tf.B"(%3) {_xla_outside_compilation = "cluster1"} : (tensor) -> (tensor)
+      %5:2 = "tf.C"(%4) : (tensor) -> (tensor, tensor)
+      tf_device.return %5#0, %5#1 : tensor, tensor
+    }) {cluster_attr = "cluster_attr"} : () -> (tensor, tensor)
+    tf_device.return %2#1 : tensor
+  }
+  return %1 : tensor
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
index 54600faca4b..a2a19108326 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
@@ -108,18 +108,6 @@ tf_device::LaunchOp CreateLaunchOpForOutsideCluster(
   return launch_op;
 }
 
-// Propagates the return from `parallel_execute_op` to parent replicate
-// op if it exists.
-void PropagateParallelExecuteReturnToReplicate(
-    tf_device::ParallelExecuteOp parallel_execute_op) {
-  // Update the return for the parallel_execute op parent.
-  auto replicate = llvm::dyn_cast_or_null(
-      parallel_execute_op.getParentOp());
-  if (replicate)
-    replicate.GetBody().getTerminator()->setOperands(
-        parallel_execute_op.execute_outputs());
-}
-
 // Extracts all externally provided operands of `cluster_ops`.
 llvm::SmallSetVector GetExternalOperands(
     llvm::ArrayRef cluster_ops) {
@@ -305,7 +293,16 @@ void CreateParallelExecuteFromOutsideClusters(
   tpu_cluster.getOperation()->moveBefore(
       parallel_execute_tpu_block.getTerminator());
 
-  PropagateParallelExecuteReturnToReplicate(parallel_execute_op);
+  // Remap cluster results with parallel_execute results if user is outside of
+  // parallel_execute.
+  for (auto result :
+       llvm::zip(tpu_cluster.getResults(), parallel_execute_op.getResults())) {
+    Value tpu_cluster_result = std::get<0>(result);
+    Value parallel_execute_result = std::get<1>(result);
+    for (auto& use : llvm::make_early_inc_range(tpu_cluster_result.getUses()))
+      if (!parallel_execute_op.getOperation()->isProperAncestor(use.getOwner()))
+        use.set(parallel_execute_result);
+  }
 }
 
 void TPUExtractOutsideCompilation::runOnFunction() {

From 259ce6bd35c487ca39b375f16dc91d3636a9e838 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 15:47:01 -0700
Subject: [PATCH 0442/1390] Move the ScatterExpander pass after the
 DynamicPadder pass on cpu. DynamicPadder will pad the indices of
 scatter(operand 1) if indices is of dynamic size. If ScatterExpander pass is
 run first, scatter will be rewrite to a sequence of other ops and this pad
 will not be performed and give incorrect answer.

PiperOrigin-RevId: 316981132
Change-Id: I2662d783e73a0467ebd4f1af21f7ac6cf93faef5
---
 tensorflow/compiler/xla/service/cpu/cpu_compiler.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index b2416ac2799..31b9fe1c920 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -277,12 +277,12 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn(
   pipeline.AddPass(
       cost_model,
       /*convert_batch_groups_only=*/false);
-  pipeline.AddPass();
   pipeline.AddPass(
       /*rewrite_training_op=*/true,
       /*rewrite_inference_op=*/true,
       /*rewrite_grad_op=*/true);
   pipeline.AddPass();
+  pipeline.AddPass();
   pipeline.AddPass();
   pipeline.AddPass(target_machine_features);
   {

From 2053b238d0d717aef4dfcf4be03daa43d7eeede2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 15:50:24 -0700
Subject: [PATCH 0443/1390] add two internal stats.

PiperOrigin-RevId: 316981769
Change-Id: Ia53c128b82fa750942acb5441ae6c8dc31478b07
---
 tensorflow/core/profiler/utils/xplane_schema.cc | 5 ++++-
 tensorflow/core/profiler/utils/xplane_schema.h  | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index 197dab75d3b..be53dcbdc01 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -166,6 +166,8 @@ const StatTypeMap& GetStatTypeMap() {
       {"is_eager", kIsEager},
       {"tf_function_call", kTfFunctionCall},
       {"tracing_count", kTfFunctionTracingCount},
+      {"flops", kFlops},
+      {"bytes_accessed", kBytesAccessed},
       // Performance counter related.
       {"Raw Value", kRawValue},
       {"Scaled Value", kScaledValue},
@@ -227,7 +229,8 @@ bool IsInternalStat(absl::optional stat_type) {
       StatType::kKernelDetails, StatType::kLevel0,
       StatType::kProducerType,  StatType::kProducerId,
       StatType::kConsumerType,  StatType::kConsumerId,
-      StatType::kIsRoot,        StatType::kIsAsync};
+      StatType::kIsRoot,        StatType::kIsAsync,
+      StatType::kFlops,         StatType::kBytesAccessed};
   return stat_type.has_value() && kInternalStats->contains(*stat_type);
 }
 
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 8b999dc6f9f..a31814cef06 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -153,6 +153,8 @@ enum StatType {
   kIsEager,
   kTfFunctionCall,
   kTfFunctionTracingCount,
+  kFlops,
+  kBytesAccessed,
   // Performance counter related.
   kRawValue,
   kScaledValue,

From 97afd248868c7f28c197abde87cf610b550bdef9 Mon Sep 17 00:00:00 2001
From: Kaixi Hou 
Date: Wed, 17 Jun 2020 10:00:23 -0700
Subject: [PATCH 0444/1390] Relu grad GPU uses 8 float16 element vector

---
 tensorflow/core/kernels/relu_op_gpu.cu.cc | 89 +++++++++++++++++++++--
 1 file changed, 82 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/kernels/relu_op_gpu.cu.cc b/tensorflow/core/kernels/relu_op_gpu.cu.cc
index 27fd5f64249..983cc127863 100644
--- a/tensorflow/core/kernels/relu_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/relu_op_gpu.cu.cc
@@ -35,6 +35,7 @@ namespace tensorflow {
 
 typedef Eigen::GpuDevice GPUDevice;
 
+static constexpr int VectorSizeElements = 8;
 namespace functor {
 
 // This kernel computes ReluGrad by processing one half2, two fp16, at a time.
@@ -93,6 +94,66 @@ __global__ void ReluGradHalfKernel(const Eigen::half* __restrict__ gradient,
   }
 }
 
+__global__ void ReluGradHalfKernelVector(
+    const Eigen::half* __restrict__ gradient,
+    const Eigen::half* __restrict__ feature,
+    Eigen::half* __restrict__ backprop, int32 count) {
+  int32 half8_count = count / VectorSizeElements;
+  int32 index = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (index < half8_count) {
+    // Cast to xx_h8 for vector load and store.
+    float4 gradient_h8 = reinterpret_cast(gradient)[index];
+    float4 feature_h8 = reinterpret_cast(feature)[index];
+    float4* p_backprop_h8 = reinterpret_cast(backprop) + index;
+
+    half2 *gradient_h2 = reinterpret_cast(&gradient_h8);
+    half2 *feature_h2 = reinterpret_cast(&feature_h8);
+    float4 backprop_h8;
+    half2* p_backprop_h2 = reinterpret_cast(&backprop_h8);
+
+    // Fast path, when half2 primitives are available.
+#if __CUDA_ARCH__ >= 530
+    const half2 kZeroH2 = __float2half2_rn(0.f);
+#endif
+    for (int i = 0; i < VectorSizeElements / 2; i++) {
+#if __CUDA_ARCH__ >= 530
+      // mask = (feature > 0)
+      half2 mask_h2 = __hgt2(feature_h2[i], kZeroH2);
+      // backprop = mask * gradient
+      half2 backprop_h2 = __hmul2(mask_h2, gradient_h2[i]);
+#else
+      // Fall back: convert half2 to float2 for processing.
+      float2 feature_f2 = __half22float2(feature_h2[i]);
+      float2 gradient_f2 = __half22float2(gradient_h2[i]);
+      float2 backprop_f2 =
+          make_float2((feature_f2.x > 0.0f) ? float(gradient_f2.x) : 0.0f,
+                      (feature_f2.y > 0.0f) ? float(gradient_f2.y) : 0.0f);
+      // Convert back to half2.
+      half2 backprop_h2 = __float22half2_rn(backprop_f2);
+#endif
+      p_backprop_h2[i] = backprop_h2;
+    }
+    // Write back the result.
+    *p_backprop_h8 = backprop_h8;
+  }
+
+  int remaining_count = (count % VectorSizeElements);
+
+  if (index < remaining_count) {
+    // Use first threads to process the remaining elements.
+    Eigen::half grad_h = gradient[half8_count * VectorSizeElements + index];
+    Eigen::half feature_h = feature[half8_count * VectorSizeElements + index];
+
+    float grad_f = static_cast(grad_h);
+    float feature_f = static_cast(feature_h);
+    float backprop_f = (feature_f > 0) ? grad_f : 0;
+
+    Eigen::half backprop_h(backprop_f);
+    backprop[half8_count * VectorSizeElements + index] = backprop_h;
+  }
+}
+
 template 
 struct ReluGrad {
   // Computes ReluGrad backprop.
@@ -108,15 +169,29 @@ struct ReluGrad {
     // NOTE: When the activation is exactly zero, we do not propagate the
     // associated gradient value. This allows the output of the Relu to be used,
     // as well as its input.
+    auto gradient_ptr = reinterpret_cast(gradient.data());
+    auto feature_ptr = reinterpret_cast(feature.data());
+    auto backprop_ptr = reinterpret_cast(backprop.data());
+    bool aligned = gradient_ptr % 16 == 0 && feature_ptr % 16 == 0 &&
+                   backprop_ptr % 16 == 0;
     int32 count = gradient.size();
-    if (count == 0) return;
-    int32 half2_count = Eigen::divup(count, 2);
     constexpr int32 kThreadInBlock = 512;
-    GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize(
-        half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock);
-    TF_CHECK_OK(GpuLaunchKernel(
-        ReluGradHalfKernel, config.block_count, config.thread_per_block, 0,
-        d.stream(), gradient.data(), feature.data(), backprop.data(), count));
+    if (count == 0) return;
+    if (aligned) {
+      int32 half8_count = Eigen::divup(count, VectorSizeElements);
+      int32 kBlock = Eigen::divup(half8_count, kThreadInBlock);
+      TF_CHECK_OK(GpuLaunchKernel(
+          ReluGradHalfKernelVector, kBlock, kThreadInBlock,
+          0, d.stream(), gradient.data(), feature.data(), backprop.data(),
+          count));
+    } else {
+      int32 half2_count = Eigen::divup(count, 2);
+      GpuLaunchConfig config = GetGpuLaunchConfigFixedBlockSize(
+          half2_count, d, ReluGradHalfKernel, 0, kThreadInBlock);
+      TF_CHECK_OK(GpuLaunchKernel(
+          ReluGradHalfKernel, config.block_count, config.thread_per_block, 0,
+          d.stream(), gradient.data(), feature.data(), backprop.data(), count));
+    }
   }
 };
 

From eb07cc9d8f370d1c99c317ade64f9dd210b61c61 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Wed, 17 Jun 2020 23:20:11 +0000
Subject: [PATCH 0445/1390] Map dtypes to classes and add whitelist op set

---
 tensorflow/python/framework/python_op_gen.cc | 30 ++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index ca0c5d9ef1a..217033f1c31 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -45,6 +45,36 @@ const int kRightMargin = 78;
 
 constexpr char kEagerFallbackSuffix[] = "_eager_fallback";
 
+std::unordered_map dtypes_map {
+      {"_dtypes.float16", "_dtypes.Float16"},
+      {"_dtypes.float32", "_dtypes.Float32"},
+      {"_dtypes.float64", "_dtypes.Float64"},
+      {"_dtypes.bfloat16", "_dtypes.BFloat16"},
+      {"_dtypes.complex64", "_dtypes.Complex64"},
+      {"_dtypes.complex128", "_dtypes.Complex128"},
+      {"_dtypes.int8", "_dtypes.Int8"},
+      {"_dtypes.uint8", "_dtypes.UInt8"},
+      {"_dtypes.uint16", "_dtypes.UInt16"},
+      {"_dtypes.uint32", "_dtypes.UInt32"},
+      {"_dtypes.uint64", "_dtypes.UInt64"},
+      {"_dtypes.int16", "_dtypes.Int16"},
+      {"_dtypes.int32", "_dtypes.Int32"},
+      {"_dtypes.int64", "_dtypes.Int64"},
+      {"_dtypes.bool", "_dtypes.Bool"},
+      {"_dtypes.string", "_dtypes.String"},
+      {"_dtypes.qint8", "_dtypes.QInt8"},
+      {"_dtypes.quint8", "_dtypes.QUInt8"},
+      {"_dtypes.qint16", "_dtypes.QInt16"},
+      {"_dtypes.quint16", "_dtypes.QUInt16"},
+      {"_dtypes.qint32", "_dtypes.QInt32"},
+      {"_dtypes.resource", "_dtypes.Resource"},
+      {"_dtypes.variant", "_dtypes.Variant"}
+};
+
+// Add op name to this set to add type annotations
+std::unordered_set type_annotate_ops {
+};
+
 string AttrVarName(const string& attr_name,
                    std::unordered_map* attr_expressions) {
   const string var = strings::StrCat("_attr_", attr_name);

From 06075a3baa5f57a8e7c84ac814d7350e4582e017 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Wed, 17 Jun 2020 23:34:23 +0000
Subject: [PATCH 0446/1390] Add function to generate TypeVars for each op

---
 tensorflow/python/framework/python_op_gen.cc | 53 ++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 217033f1c31..e8e670de57c 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -178,6 +178,8 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
 
   void AddRawOpExport(const string& parameters);
 
+  void GenerateTypeVars();
+
   void AddAttrForArg(const string& attr, int arg_index) {
     gtl::InsertIfNotPresent(&inferred_attrs_, attr,
                             op_def_.input_arg(arg_index).name());
@@ -397,6 +399,53 @@ string GenEagerPythonOp::Code() {
   return prelude_ + result_;
 }
 
+// Generate TypeVars using attrs
+void GenEagerPythonOp::GenerateTypeVars() {
+  bool added_typevar = false;
+  for (int i = 0; i allowed_types;
+      bool has_dtype_half = false;
+      for (int t : attr.allowed_values().list().type()) {
+        if (t == 19) { // DT_HALF = 19;
+          has_dtype_half = true;
+          break;
+        }
+        DataType dtype = static_cast(t);
+        const string py_dtype = python_op_gen_internal::DataTypeToPython(dtype, "_dtypes.");
+        if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
+          allowed_types.emplace_back(dtypes_map[py_dtype]);
+        }
+      }
+
+      // Do not create a type variable that includes the dtype half
+      if (has_dtype_half) continue;
+
+      // If all dtypes are allowed, add them all
+      if (allowed_types.empty()) {
+        for (std::pair map_dtype : dtypes_map) {
+          allowed_types.emplace_back(map_dtype.second);
+        }
+      }
+
+      std::sort(allowed_types.begin(), allowed_types.end());
+
+      string typevar_dtypes;
+      for (std::vector::iterator it = allowed_types.begin(); it != allowed_types.end(); ++it) {
+        if (!typevar_dtypes.empty()) strings::StrAppend(&typevar_dtypes, ", ");
+        strings::StrAppend(&typevar_dtypes, *it);
+      }
+
+      const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
+      strings::StrAppend(&result_, type_var_name, " = TypeVar(\"", type_var_name, "\", ", typevar_dtypes,")\n");
+      added_typevar = true;
+    }
+  }
+
+  if(added_typevar) strings::StrAppend(&result_, "\n");
+}
+
 void GenEagerPythonOp::HandleGraphMode(
     const string& function_setup, const std::vector& output_sizes) {
   strings::StrAppend(&result_, "  # Add nodes to the TensorFlow graph.\n");
@@ -720,6 +769,9 @@ void GenEagerPythonOp::AddEagerFunctionTeardown(
 bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
     const string& parameters, const std::vector& output_sizes,
     const string& eager_not_allowed_error) {
+  if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
+    GenerateTypeVars();
+  }
   if (api_def_.visibility() == ApiDef::VISIBLE) {
     strings::StrAppend(&result_, "@_dispatch.add_dispatch_list\n");
   }
@@ -1047,6 +1099,7 @@ from tensorflow.python.util.deprecation import deprecated_endpoints
 from tensorflow.python.util import dispatch as _dispatch
 from tensorflow.python.util.tf_export import tf_export
 
+from typing import TypeVar
 )");
 
   for (const auto& op_def : ops.op()) {

From c87ccd156017b2cad24740b375a78f322f3951ef Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Wed, 17 Jun 2020 23:51:42 +0000
Subject: [PATCH 0447/1390] Add type annotations to op parameters

---
 tensorflow/python/framework/python_op_gen.cc | 69 +++++++++++++++++++-
 1 file changed, 66 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index e8e670de57c..777603c502d 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -179,6 +179,7 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
   void AddRawOpExport(const string& parameters);
 
   void GenerateTypeVars();
+  string GetTypeAnnotatedParams();
 
   void AddAttrForArg(const string& attr, int arg_index) {
     gtl::InsertIfNotPresent(&inferred_attrs_, attr,
@@ -343,10 +344,15 @@ string GenEagerPythonOp::Code() {
   }
 
   string parameters;
-  for (const auto& param : params_no_default_) {
-    if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
-    strings::StrAppend(¶meters, param.GetRenameTo());
+  if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
+    strings::StrAppend(¶meters, GetTypeAnnotatedParams());
+  } else {
+    for (const auto& param : params_no_default_) {
+      if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
+      strings::StrAppend(¶meters, param.GetRenameTo());
+    }
   }
+
   string parameters_with_defaults = parameters;
   for (const auto& param_and_default : params_with_default_) {
     if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
@@ -399,6 +405,63 @@ string GenEagerPythonOp::Code() {
   return prelude_ + result_;
 }
 
+string GenEagerPythonOp::GetTypeAnnotatedParams() {
+  // holds mappings from param name to its type annotation
+  std::unordered_map param_type_map;
+  for (int i = 0; i
Date: Wed, 17 Jun 2020 23:59:45 +0000
Subject: [PATCH 0448/1390] Generate type annotations for op return values

---
 tensorflow/python/framework/python_op_gen.cc | 54 ++++++++++++++++++--
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 777603c502d..8c4d0f5b753 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -180,7 +180,7 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
 
   void GenerateTypeVars();
   string GetTypeAnnotatedParams();
-
+  void AddReturnTypeAnnotation();
   void AddAttrForArg(const string& attr, int arg_index) {
     gtl::InsertIfNotPresent(&inferred_attrs_, attr,
                             op_def_.input_arg(arg_index).name());
@@ -413,7 +413,7 @@ string GenEagerPythonOp::GetTypeAnnotatedParams() {
     if (attr.type() == "type") {
       bool has_dtype_half = false;
       for (int t : attr.allowed_values().list().type()) {
-        if (t == 19) { // DT_HALF = 19;
+        if (t == 19) { // DT_HALF = 19
           has_dtype_half = true;
           break;
         }
@@ -471,7 +471,7 @@ void GenEagerPythonOp::GenerateTypeVars() {
       std::vector allowed_types;
       bool has_dtype_half = false;
       for (int t : attr.allowed_values().list().type()) {
-        if (t == 19) { // DT_HALF = 19;
+        if (t == 19) { // DT_HALF = 19
           has_dtype_half = true;
           break;
         }
@@ -509,6 +509,51 @@ void GenEagerPythonOp::GenerateTypeVars() {
   if(added_typevar) strings::StrAppend(&result_, "\n");
 }
 
+void GenEagerPythonOp::AddReturnTypeAnnotation() {
+  string return_type = "";
+  if (op_def_.output_arg_size() == 1) {
+    const auto& arg = op_def_.output_arg(0);
+    // If the "type" field is set, the return Tensor has a single DataType
+    if (arg.type() != 0) {
+      const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
+      if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
+        strings::StrAppend(&return_type, "_ops.Tensor[", dtypes_map[py_dtype], "]");
+      }
+    }
+    else {
+      for (int i = 0; i allowed_types;
+          for (int t : attr.allowed_values().list().type()) {
+            // Do not add type annotations when return type can be half
+            if (t == 19) return; // DT_HALF = 19
+            DataType dtype = static_cast(t);
+            const string py_dtype = python_op_gen_internal::DataTypeToPython(dtype, "_dtypes.");
+            allowed_types.emplace_back(py_dtype);
+          }
+
+          std::sort(allowed_types.begin(), allowed_types.end());
+
+          string typevar_dtypes;
+          for (std::vector::iterator it = allowed_types.begin(); it != allowed_types.end(); ++it) {
+            if (!typevar_dtypes.empty()) strings::StrAppend(&typevar_dtypes, ", ");
+            strings::StrAppend(&typevar_dtypes, *it);
+          }
+
+          const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
+          strings::StrAppend(&return_type, "_ops.Tensor[", type_var_name, "]");
+        }
+      }
+    }
+
+    if (!return_type.empty()) {
+      result_.erase(result_.length() - 2);
+      strings::StrAppend(&result_, " -> ", return_type, ":\n");
+    }
+  }
+}
+
 void GenEagerPythonOp::HandleGraphMode(
     const string& function_setup, const std::vector& output_sizes) {
   strings::StrAppend(&result_, "  # Add nodes to the TensorFlow graph.\n");
@@ -841,6 +886,9 @@ bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
 
   AddExport();
   AddDefLine(function_name_, parameters);
+  if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
+    AddReturnTypeAnnotation();
+  }
   AddDocStringDescription();
   AddDocStringArgs();
   AddDocStringInputs();

From bae911298b8b79ffd24464adc8fc123566a79c24 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:16:59 -0700
Subject: [PATCH 0449/1390] Additional check for Intel specialization.

PiperOrigin-RevId: 316997503
Change-Id: Ibc97cd77c5e15bb021546aac754e2953b050f0f3
---
 tensorflow/lite/delegates/gpu/cl/cl_device.cc | 27 +++++++++++++++++++
 tensorflow/lite/delegates/gpu/cl/cl_device.h  |  1 +
 .../delegates/gpu/cl/kernels/conv_powervr.cc  |  2 +-
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.cc b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
index 64e07428515..271fbce61ce 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_device.cc
+++ b/tensorflow/lite/delegates/gpu/cl/cl_device.cc
@@ -501,6 +501,33 @@ bool CLDevice::IsCL20OrHigher() const {
          info_.cl_version != OpenCLVersion::CL_1_2;
 }
 
+bool CLDevice::SupportsSubGroupWithSize(int sub_group_size) const {
+  if (IsIntel()) {
+    if (SupportsExtension("cl_intel_required_subgroup_size")) {
+      size_t sub_groups_count;
+      cl_int error =
+          clGetDeviceInfo(id_, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/, 0,
+                          nullptr, &sub_groups_count);
+      if (error != CL_SUCCESS) {
+        return false;
+      }
+      std::vector sub_group_sizes(sub_groups_count);
+      error = clGetDeviceInfo(id_, 0x4108 /*CL_DEVICE_SUB_GROUP_SIZES_INTEL*/,
+                              sizeof(size_t) * sub_groups_count,
+                              sub_group_sizes.data(), nullptr);
+      if (error != CL_SUCCESS) {
+        return false;
+      }
+      for (int i = 0; i < sub_groups_count; ++i) {
+        if (sub_group_sizes[i] == sub_group_size) {
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+
 bool CLDevice::IsAdreno() const { return info_.vendor == Vendor::QUALCOMM; }
 
 bool CLDevice::IsAdreno3xx() const {
diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.h b/tensorflow/lite/delegates/gpu/cl/cl_device.h
index ae6a1d11af6..68abcf3e202 100644
--- a/tensorflow/lite/delegates/gpu/cl/cl_device.h
+++ b/tensorflow/lite/delegates/gpu/cl/cl_device.h
@@ -179,6 +179,7 @@ class CLDevice {
   bool SupportsFP32RTN() const;
   bool SupportsFP16RTN() const;
   bool IsCL20OrHigher() const;
+  bool SupportsSubGroupWithSize(int sub_group_size) const;
   bool IsAdreno() const;
   bool IsAdreno3xx() const;
   bool IsAdreno4xx() const;
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
index bd694e7cc4f..c20cfdbeaa3 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_powervr.cc
@@ -874,7 +874,7 @@ ConvPowerVR::ConvParams ConvPowerVR::GuessBestParams(
     if (definition.precision != CalculationsPrecision::F32_F16 &&
         device.SupportsExtension("cl_khr_subgroups") &&
         device.SupportsExtension("cl_intel_required_subgroup_size") &&
-        device.IsCL20OrHigher()) {
+        device.IsCL20OrHigher() && device.SupportsSubGroupWithSize(16)) {
       conv_params.weights_upload_type =
           WeightsUploadType::PRIVATE_MEM_SIMD16_BROADCAST;
     } else {

From 1f05cc5973c1182c35cd4b11ee8cf0f01ed707dc Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:18:45 -0700
Subject: [PATCH 0450/1390] Better FullyConnected/ConvTransposed selection for
 Intel.

PiperOrigin-RevId: 316997745
Change-Id: I28befdd528917c3846ff6ae79b0f8427389dfc39
---
 .../delegates/gpu/cl/kernels/convolution_transposed_3x3.cc     | 2 +-
 .../delegates/gpu/cl/kernels/convolution_transposed_4x4.cc     | 2 +-
 .../gpu/cl/selectors/convolution_transposed_selector.cc        | 1 +
 .../delegates/gpu/cl/selectors/fully_connected_selector.cc     | 3 +++
 4 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_3x3.cc b/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_3x3.cc
index 4a68eda1d95..9b028721d2d 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_3x3.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_3x3.cc
@@ -269,7 +269,7 @@ ConvolutionTransposed3x3::ConvolutionTransposed3x3(
       work_group_launch_order_(2, 0, 1) {
   if (device.IsPowerVR()) {
     weights_upload_type_ = WeightsUploadType::LOCAL_MEM_ASYNC;
-  } else if (device.IsNvidia()) {
+  } else if (device.IsNvidia() || device.IsIntel()) {
     weights_upload_type_ = WeightsUploadType::LOCAL_MEM_BY_THREADS;
   } else if (device.IsAMD()) {
     weights_upload_type_ = WeightsUploadType::CONSTANT_MEM;
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_4x4.cc b/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_4x4.cc
index 0f7f90989e8..209b675087e 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_4x4.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/convolution_transposed_4x4.cc
@@ -270,7 +270,7 @@ ConvolutionTransposed4x4::ConvolutionTransposed4x4(
     : GPUOperation(definition) {
   if (device.IsPowerVR()) {
     weights_upload_type_ = WeightsUploadType::LOCAL_MEM_ASYNC;
-  } else if (device.IsNvidia()) {
+  } else if (device.IsNvidia() || device.IsIntel()) {
     weights_upload_type_ = WeightsUploadType::LOCAL_MEM_BY_THREADS;
   } else if (device.IsAMD()) {
     weights_upload_type_ = WeightsUploadType::CONSTANT_MEM;
diff --git a/tensorflow/lite/delegates/gpu/cl/selectors/convolution_transposed_selector.cc b/tensorflow/lite/delegates/gpu/cl/selectors/convolution_transposed_selector.cc
index 12e99b57aa7..5fdfdca073e 100644
--- a/tensorflow/lite/delegates/gpu/cl/selectors/convolution_transposed_selector.cc
+++ b/tensorflow/lite/delegates/gpu/cl/selectors/convolution_transposed_selector.cc
@@ -112,6 +112,7 @@ absl::Status SelectConvolutionTransposed(
     case Vendor::POWERVR:
     case Vendor::NVIDIA:
     case Vendor::AMD:
+    case Vendor::INTEL:
       return SelectConvolutionTransposedPowerVR(attr, creation_context, op_def,
                                                 ptr);
     case Vendor::MALI:
diff --git a/tensorflow/lite/delegates/gpu/cl/selectors/fully_connected_selector.cc b/tensorflow/lite/delegates/gpu/cl/selectors/fully_connected_selector.cc
index 12a1d726368..eacbea8b586 100644
--- a/tensorflow/lite/delegates/gpu/cl/selectors/fully_connected_selector.cc
+++ b/tensorflow/lite/delegates/gpu/cl/selectors/fully_connected_selector.cc
@@ -109,6 +109,9 @@ absl::Status SelectFullyConnected(const FullyConnectedAttributes& attr,
       return SelectFullyConnectedAdreno(attr, creation_context, op_def,
                                         batch_size, ptr);
     case Vendor::POWERVR:
+    case Vendor::AMD:
+    case Vendor::NVIDIA:
+    case Vendor::INTEL:
       return SelectFullyConnectedPowerVR(attr, creation_context, op_def,
                                          batch_size, ptr);
     case Vendor::MALI:

From 2a6d9a1e811f2c0d459839854597f8ca6b5c8750 Mon Sep 17 00:00:00 2001
From: Jose Baiocchi 
Date: Wed, 17 Jun 2020 17:18:58 -0700
Subject: [PATCH 0451/1390] Add test for op_metrics_db_utils IdleTimeRatio

PiperOrigin-RevId: 316997793
Change-Id: I0b89e697d661b2ea1fb95e7a6d51800b407617d0
---
 tensorflow/core/profiler/utils/BUILD          | 11 +++++
 .../utils/op_metrics_db_utils_test.cc         | 46 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 tensorflow/core/profiler/utils/op_metrics_db_utils_test.cc

diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index ece58802661..0262c5659b7 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -77,6 +77,17 @@ cc_library(
     ],
 )
 
+tf_cc_test(
+    name = "op_metrics_db_utils_test",
+    srcs = ["op_metrics_db_utils_test.cc"],
+    deps = [
+        ":op_metrics_db_utils",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/profiler/protobuf:op_metrics_proto_cc",
+    ],
+)
+
 cc_library(
     name = "op_utils",
     srcs = ["op_utils.cc"],
diff --git a/tensorflow/core/profiler/utils/op_metrics_db_utils_test.cc b/tensorflow/core/profiler/utils/op_metrics_db_utils_test.cc
new file mode 100644
index 00000000000..12c68426b2e
--- /dev/null
+++ b/tensorflow/core/profiler/utils/op_metrics_db_utils_test.cc
@@ -0,0 +1,46 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/profiler/utils/op_metrics_db_utils.h"
+
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/profiler/protobuf/op_metrics.pb.h"
+
+namespace tensorflow {
+namespace profiler {
+namespace {
+
+constexpr double kMaxError = 1E-10;
+
+TEST(OpMetricsDbTest, IdleTimeRatio) {
+  OpMetricsDb metrics_db_0;
+  metrics_db_0.set_total_time_ps(100000000);
+  metrics_db_0.set_total_op_time_ps(60000000);
+  EXPECT_NEAR(0.4, IdleTimeRatio(metrics_db_0), kMaxError);
+
+  OpMetricsDb metrics_db_1;
+  metrics_db_1.set_total_time_ps(200000000);
+  metrics_db_1.set_total_op_time_ps(150000000);
+  EXPECT_NEAR(0.25, IdleTimeRatio(metrics_db_1), kMaxError);
+
+  OpMetricsDb metrics_db_2;
+  metrics_db_1.set_total_time_ps(0);
+  metrics_db_1.set_total_op_time_ps(0);
+  EXPECT_NEAR(1.0, IdleTimeRatio(metrics_db_2), kMaxError);
+}
+
+}  // namespace
+}  // namespace profiler
+}  // namespace tensorflow

From 73cf8263c771f1813416af9310469c7041e7637e Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:21:24 -0700
Subject: [PATCH 0452/1390] Pooling converted to new style. Merged 2D and 3D
 versions into one.

PiperOrigin-RevId: 316998142
Change-Id: I92c020476f085e6160a02282c1edafabdd72ca30
---
 .../lite/delegates/gpu/cl/kernels/pooling.cc  | 671 +++++++-----------
 .../lite/delegates/gpu/cl/kernels/pooling.h   |  40 +-
 .../lite/delegates/gpu/cl/kernels/util.cc     |   9 -
 .../lite/delegates/gpu/cl/kernels/util.h      |   5 -
 .../lite/delegates/gpu/cl/tensor_type.cc      |  36 +-
 .../lite/delegates/gpu/cl/tensor_type.h       |  10 +
 6 files changed, 313 insertions(+), 458 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc
index e292f2dad7d..922d484c57d 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.cc
@@ -25,366 +25,307 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetAveragePoolingKernelCode(
-    const OperationDef& op_def, bool stride_correction, const CLDevice& device,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+std::string GetAveragePoolingKernelCode(const OperationDef& op_def,
+                                        bool stride_correction,
+                                        const CLDevice& device,
+                                        Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  src_desc->SetTextureAddressMode(GetFastestZeroMode(device));
+  if (op_def.IsBatchSupported()) {
+    src_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    dst_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
+  if (op_def.dst_tensors[0].HasAxis(Axis::WIDTH)) {
+    args->AddInt("kernel_size_x");
+    args->AddInt("padding_x");
+    args->AddInt("stride_x");
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
+    args->AddInt("kernel_size_y");
+    args->AddInt("padding_y");
+    args->AddInt("stride_y");
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    args->AddInt("kernel_size_z");
+    args->AddInt("padding_z");
+    args->AddInt("stride_z");
+  }
 
-  const auto address_mode = GetFastestZeroMode(device);
+  std::map axis_to_src_coord = {
+      {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
 
-  std::string c = GetCommonDefines(op_def.precision);
+  std::map axis_to_dst_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector src_coords;
+  std::vector dst_coords;
+  for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_dst_coord[axis]);
+    }
+    if (op_def.src_tensors[0].HasAxis(axis)) {
+      src_coords.push_back(axis_to_src_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
+  }
 
   const bool manual_clamp =
       op_def.src_tensors[0].storage_type == TensorStorageType::BUFFER ||
       op_def.src_tensors[0].storage_type == TensorStorageType::IMAGE_BUFFER;
 
+  std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
-  c += "    int2 kernel_size,          \n";
-  c += "    int2 padding,              \n";
-  c += "    int2 stride                \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = get_global_id(1);\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = get_global_id(1);\n";
+  }
   c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
   c += "  float4 r = (float4)(0.0f);\n";
   c += "  float window_size = 0.0;\n";
   if (stride_correction) {
     c += "  int xs = " +
-         GetXStrideCorrected("X", "src_size.w", "stride.x", "padding.x") +
+         GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
+                             "args.padding_x") +
          ";\n";
   } else {
-    c += "  int xs = X * stride.x + padding.x;\n";
+    c += "  int xs = X * args.stride_x + args.padding_x;\n";
   }
-  c += "  int ys = Y * stride.y + padding.y;\n";
-  c += "  for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
+  c += "  int ys = Y * args.stride_y + args.padding_y;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int ds = D * args.stride_z + args.padding_z;\n";
+    c += "  for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+    c += "    int d_c = ds + kz;\n";
+    c += "    if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
+  }
+  c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
   c += "    int y_c = ys + ky;\n";
-  c += "    bool outside_y = y_c < 0 || y_c >= src_size.y;\n";
-  c += "    for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
+  c += "    bool outside_y = y_c < 0 || y_c >= args.src_tensor.Height();\n";
+  c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
   if (op_def.IsBatchSupported()) {
-    c += "      int x_c = xs + kx * src_size.w;\n";
+    c += "      int x_c = xs + kx * args.src_tensor.Batch();\n";
   } else {
     c += "      int x_c = xs + kx;\n";
   }
-  c += "      bool outside = outside_y || x_c < 0 || x_c >= src_size.x;\n";
+  c += "      bool outside = outside_y || x_c < 0 || x_c >= "
+       "args.src_tensor.Width();\n";
   if (manual_clamp) {
-    c += "     r += !outside ? " +
-         src_tensor.ReadAsFloatWHS("x_c", "y_c", "Z") + " : (float4)(0.0f);\n";
+    c += "     r += !outside ? args.src_tensor.Read(" + src_coord +
+         ") : "
+         "(float4)(0.0f);\n";
   } else {
-    c += "      r += " +
-         src_tensor.ReadAsFloatWHS("x_c", "y_c", "Z", address_mode) + ";\n";
+    c += "      r += args.src_tensor.Read(" + src_coord + ");\n";
   }
   c += "        window_size += !outside ? 1.0 : 0.0;\n";
   c += "    }\n";
   c += "  }\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  }  // Depth\n";
+  }
   // If window_size==0, window covered nothing. This situation is a sign of
   // incorrectly constructed operation. NaNs are expected as output.
   c += "  FLT4 result = TO_FLT4(r / window_size);\n";
-  const LinkingContext context{"result", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("result", "X", "Y", "Z");
+  c += "  args.dst_tensor.Write(result, " + dst_coord + ");\n";
   c += "}\n";
 
   return c;
 }
 
-std::string GetAveragePooling3DKernelCode(
-    const OperationDef& op_def, bool stride_correction, const CLDevice& device,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHDSPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
-
-  const auto address_mode = GetFastestZeroMode(device);
-
-  std::string c = GetCommonDefines(op_def.precision);
-
-  c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
+std::string GetMaxPoolingKernelCode(const OperationDef& op_def,
+                                    bool stride_correction, bool output_indices,
+                                    Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
   if (op_def.IsBatchSupported()) {
-    c += "    int batch_size,          \n";
+    src_desc->SetStateVar("BatchedWidth", "true");
   }
-  c += "    int4 kernel_size,          \n";
-  c += "    int4 padding,              \n";
-  c += "    int4 stride                \n";
-  c += ") {\n";
-  c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int linear_id_z = get_global_id(2);\n";
-  c += "  int S = linear_id_z % dst_size.w;\n";
-  c += "  int Z = linear_id_z / dst_size.w;\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
-  c += "  float4 r = (float4)(0.0f);\n";
-  c += "  float window_size = 0.0;\n";
-  if (stride_correction) {
-    c += "  int xs = " +
-         GetXStrideCorrected("X", "batch_size", "stride.x", "padding.x") +
-         ";\n";
-  } else {
-    c += "  int xs = X * stride.x + padding.x;\n";
-  }
-  c += "  int ys = Y * stride.y + padding.y;\n";
-  c += "  int zs = Z * stride.z + padding.z;\n";
-  c += "  for (int kz = 0; kz < kernel_size.z; ++kz) {\n";
-  c += "    int z_c = zs + kz;\n";
-  c += "    if (z_c < 0 || z_c >= src_size.z) continue;\n";
-  c += "    for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-  c += "      int y_c = ys + ky;\n";
-  c += "      if (y_c < 0 || y_c >= src_size.y) continue;\n";
-  c += "      for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
   if (op_def.IsBatchSupported()) {
-    c += "        int x_c = xs + kx * batch_size;\n";
-  } else {
-    c += "        int x_c = xs + kx;\n";
+    dst_desc->SetStateVar("BatchedWidth", "true");
   }
-  c += "        if(x_c < 0 || x_c >= src_size.x) continue;\n";
-  c += "        r += " +
-       src_tensor.ReadAsFloatWHDS("x_c", "y_c", "z_c", "S", address_mode) +
-       ";\n";
-  c += "        window_size += 1.0;\n";
-  c += "      }\n";
-  c += "    }\n";
-  c += "  }\n";
-  // If window_size==0, window covered nothing. This situation is a sign of
-  // incorrectly constructed operation. NaNs are expected as output.
-  c += "  FLT4 result = TO_FLT4(r / window_size);\n";
-  const LinkingContext context{"result", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHDS("result", "X", "Y", "Z", "S");
-  c += "}\n";
-
-  return c;
-}
-
-std::string GetMaxPoolingKernelCode(
-    const OperationDef& op_def, bool stride_correction,
-    const std::vector& linked_operations,
-    bool output_indices) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
-  const auto dst_ind_def =
-      output_indices ? op_def.dst_tensors[1] : op_def.dst_tensors[0];
-  TensorCodeGenerator indices_tensor(
-      "dst_indices", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      dst_ind_def);
-
-  std::string c = GetCommonDefines(op_def.precision);
-
-  c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
   if (output_indices) {
-    c += indices_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
+    auto dst_ind_desc =
+        absl::make_unique(op_def.dst_tensors[1]);
+    if (op_def.IsBatchSupported()) {
+      dst_ind_desc->SetStateVar("BatchedWidth", "true");
+    }
+    args->AddObjectRef("dst_indices", AccessType::WRITE,
+                       std::move(dst_ind_desc));
   }
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
-  c += "    int2 kernel_size,          \n";
-  c += "    int2 padding,              \n";
-  c += "    int2 stride                \n";
-  c += ") {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::WIDTH)) {
+    args->AddInt("kernel_size_x");
+    args->AddInt("padding_x");
+    args->AddInt("stride_x");
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
+    args->AddInt("kernel_size_y");
+    args->AddInt("padding_y");
+    args->AddInt("stride_y");
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    args->AddInt("kernel_size_z");
+    args->AddInt("padding_z");
+    args->AddInt("stride_z");
+  }
+
+  std::map axis_to_src_coord = {
+      {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::map axis_to_dst_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector src_coords;
+  std::vector dst_coords;
+  for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_dst_coord[axis]);
+    }
+    if (op_def.src_tensors[0].HasAxis(axis)) {
+      src_coords.push_back(axis_to_src_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
+  }
+
+  std::string c = GetCommonDefines(op_def.precision);
+  c += "__kernel void main_function(\n";
+  c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = get_global_id(1);\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = get_global_id(1);\n";
+  }
   c += "  int Z = get_global_id(2);\n";
-  c +=
-      "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return; \n";
-  c += "  FLT4 maximum = (FLT4)(-10000.0f);\n";
-  if (output_indices) {
-    c += "  FLT4 indexes = (FLT4)(0.0f);\n";
-    c += "  FLT index_counter = (FLT)(0.1f);\n";
-  }
-  if (stride_correction) {
-    c += "  int xs = " +
-         GetXStrideCorrected("X", "src_size.w", "stride.x", "padding.x") +
-         ";\n";
-  } else {
-    c += "  int xs = X * stride.x + padding.x;\n";
-  }
-  c += "  int ys = Y * stride.y + padding.y;\n";
-  c += "  for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-  c += "    int y_c = ys + ky;\n";
-  c += "    bool outside_y = y_c < 0 || y_c >= src_size.y;\n";
-  c += "    for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
-  if (op_def.IsBatchSupported()) {
-    c += "      int x_c = xs + kx * src_size.w;\n";
-  } else {
-    c += "      int x_c = xs + kx;\n";
-  }
-  c += "      bool outside_x = x_c < 0 || x_c >= src_size.x;\n";
-  c += "      if (!outside_x && !outside_y) {\n";
-  c += "        FLT4 src = " + src_tensor.ReadWHS("x_c", "y_c", "Z") + ";\n";
-  if (output_indices) {
-    c += "        if (src.x > maximum.x) {\n";
-    c += "          indexes.x = index_counter;\n";
-    c += "          maximum.x = src.x;\n";
-    c += "        }\n";
-    c += "        if (src.y > maximum.y) {\n";
-    c += "          indexes.y = index_counter;\n";
-    c += "          maximum.y = src.y;\n";
-    c += "        }\n";
-    c += "        if (src.z > maximum.z) {\n";
-    c += "          indexes.z = index_counter;\n";
-    c += "          maximum.z = src.z;\n";
-    c += "        }\n";
-    c += "        if (src.w > maximum.w) {\n";
-    c += "          indexes.w = index_counter;\n";
-    c += "          maximum.w = src.w;\n";
-    c += "        }\n";
-    c += "        index_counter += (FLT)(1.0f);\n";
-  } else {
-    c += "        maximum = max(src, maximum);\n";
-  }
-  c += "      }\n";
-  c += "    }\n";
-  c += "  }\n";
-  const LinkingContext context{"maximum", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("maximum", "X", "Y", "Z");
-  if (output_indices) {
-    c += "  " + indices_tensor.WriteWHS("indexes", "X", "Y", "Z");
-  }
-  c += "}\n";
-
-  return c;
-}
-
-std::string GetMaxPooling3DKernelCode(
-    const OperationDef& op_def, bool stride_correction,
-    const std::vector& linked_operations,
-    bool output_indices) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHDSPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
-  const auto dst_ind_def =
-      output_indices ? op_def.dst_tensors[1] : op_def.dst_tensors[0];
-  TensorCodeGenerator indices_tensor(
-      "dst_indices",
-      WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      dst_ind_def);
-
-  std::string c = GetCommonDefines(op_def.precision);
-
-  c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ);
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  if (output_indices) {
-    c += indices_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  }
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size,             \n";
-  if (op_def.IsBatchSupported()) {
-    c += "    int batch_size,          \n";
-  }
-  c += "    int4 kernel_size,          \n";
-  c += "    int4 padding,              \n";
-  c += "    int4 stride                \n";
-  c += ") {\n";
-  c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int linear_id_z = get_global_id(2);\n";
-  c += "  int S = linear_id_z % dst_size.w;\n";
-  c += "  int Z = linear_id_z / dst_size.w;\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
   c += "  FLT4 maximum = (FLT4)(-10000.0f);\n";
   if (output_indices) {
     c += "  FLT4 indexes = (FLT4)(0.0f);\n";
   }
   if (stride_correction) {
     c += "  int xs = " +
-         GetXStrideCorrected("X", "batch_size", "stride.x", "padding.x") +
+         GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
+                             "args.padding_x") +
          ";\n";
   } else {
-    c += "  int xs = X * stride.x + padding.x;\n";
+    c += "  int xs = X * args.stride_x + args.padding_x;\n";
   }
-  c += "  int ys = Y * stride.y + padding.y;\n";
-  c += "  int zs = Z * stride.z + padding.z;\n";
-  c += "  for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
+  c += "  int ys = Y * args.stride_y + args.padding_y;\n";
+  c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
   c += "    int y_c = ys + ky;\n";
-  c += "    if (y_c < 0 || y_c >= src_size.y) continue;\n";
-  c += "    for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
+  c += "    if (y_c < 0 || y_c >= args.src_tensor.Height()) continue;\n";
+  c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
   if (op_def.IsBatchSupported()) {
-    c += "      int x_c = xs + kx * batch_size;\n";
+    c += "      int x_c = xs + kx * args.src_tensor.Batch();\n";
   } else {
     c += "      int x_c = xs + kx;\n";
   }
-  c += "      if (x_c < 0 || x_c >= src_size.x) continue;\n";
-  c += "      for (int kz = 0; kz < kernel_size.z; ++kz) {\n";
-  c += "        int z_c = zs + kz;\n";
-  c += "        if (z_c < 0 || z_c >= src_size.z) continue;\n";
-  c += "        FLT4 src = " + src_tensor.ReadWHDS("x_c", "y_c", "z_c", "S") +
-       ";\n";
-  if (output_indices) {
-    c += "        FLT index_counter = (FLT)((ky * kernel_size.x + kx) * "
-         "kernel_size.z + kz) + (FLT)(0.1f);\n";
-    c += "        if (src.x > maximum.x) {\n";
-    c += "          indexes.x = index_counter;\n";
-    c += "          maximum.x = src.x;\n";
-    c += "        }\n";
-    c += "        if (src.y > maximum.y) {\n";
-    c += "          indexes.y = index_counter;\n";
-    c += "          maximum.y = src.y;\n";
-    c += "        }\n";
-    c += "        if (src.z > maximum.z) {\n";
-    c += "          indexes.z = index_counter;\n";
-    c += "          maximum.z = src.z;\n";
-    c += "        }\n";
-    c += "        if (src.w > maximum.w) {\n";
-    c += "          indexes.w = index_counter;\n";
-    c += "          maximum.w = src.w;\n";
-    c += "        }\n";
-  } else {
-    c += "        maximum = max(src, maximum);\n";
+  c += "      if (x_c < 0 || x_c >= args.src_tensor.Width()) continue;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "    int ds = D * args.stride_z + args.padding_z;\n";
+    c += "    for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+    c += "    int d_c = ds + kz;\n";
+    c += "      if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
+  }
+  c += "      FLT4 src = args.src_tensor.Read(" + src_coord + ");\n";
+  if (output_indices) {
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      c += "      FLT index_counter = (FLT)((ky * args.kernel_size_x + kx) * "
+           "args.kernel_size_z + kz) + (FLT)(0.1f);\n";
+    } else {
+      c += "      FLT index_counter = (FLT)(ky * args.kernel_size_x + kx) + "
+           "(FLT)(0.1f);\n";
+    }
+    c += "      if (src.x > maximum.x) {\n";
+    c += "        indexes.x = index_counter;\n";
+    c += "        maximum.x = src.x;\n";
+    c += "      }\n";
+    c += "      if (src.y > maximum.y) {\n";
+    c += "        indexes.y = index_counter;\n";
+    c += "        maximum.y = src.y;\n";
+    c += "      }\n";
+    c += "      if (src.z > maximum.z) {\n";
+    c += "        indexes.z = index_counter;\n";
+    c += "        maximum.z = src.z;\n";
+    c += "      }\n";
+    c += "      if (src.w > maximum.w) {\n";
+    c += "        indexes.w = index_counter;\n";
+    c += "        maximum.w = src.w;\n";
+    c += "      }\n";
+  } else {
+    c += "      maximum = max(src, maximum);\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "    }  // Depth\n";
   }
-  c += "      };\n";
   c += "    }\n";
   c += "  }\n";
-  const LinkingContext context{"maximum", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHDS("maximum", "X", "Y", "Z", "S");
+  c += "  args.dst_tensor.Write(maximum, " + dst_coord + ");\n";
   if (output_indices) {
-    c += "  " + indices_tensor.WriteWHDS("indexes", "X", "Y", "Z", "S");
+    c += "  args.dst_indices.Write(indexes, " + dst_coord + ");\n";
   }
   c += "}\n";
+
   return c;
 }
-
 }  // namespace
 
 Pooling::Pooling(const OperationDef& definition,
                  const Pooling2DAttributes& attr)
     : GPUOperation(definition),
-      stride_(attr.strides.w, attr.strides.h),
-      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h),
-      kernel_size_(attr.kernel.w, attr.kernel.h),
+      stride_(attr.strides.w, attr.strides.h, 0, 0),
+      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h, 0, 0),
+      kernel_size_(attr.kernel.w, attr.kernel.h, 0, 0),
+      type_(attr.type),
+      output_indices_(attr.output_indices) {}
+
+Pooling::Pooling(const OperationDef& definition,
+                 const Pooling3DAttributes& attr)
+    : GPUOperation(definition),
+      stride_(attr.strides.w, attr.strides.h, attr.strides.d, 0),
+      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h,
+               -attr.padding.prepended.d, 0),
+      kernel_size_(attr.kernel.w, attr.kernel.h, attr.kernel.d, 0),
       type_(attr.type),
       output_indices_(attr.output_indices) {}
 
@@ -419,44 +360,56 @@ absl::Status Pooling::Compile(const CreationContext& creation_context) {
   switch (type_) {
     case PoolingType::AVERAGE:
       code = GetAveragePoolingKernelCode(definition_, stride_correction,
-                                         *creation_context.device,
-                                         linked_operations_);
+                                         *creation_context.device, &args_);
       break;
     case PoolingType::MAX:
       code = GetMaxPoolingKernelCode(definition_, stride_correction,
-                                     linked_operations_, output_indices_);
+                                     output_indices_, &args_);
       break;
     default:
       return absl::InvalidArgumentError(
           "You should create another kernel with this params");
       break;
   }
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status Pooling::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  if (output_indices_) {
-    RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[1]->GetMemoryPtrForWriting()));
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  if (definition_.dst_tensors[0].HasAxis(Axis::WIDTH)) {
+    RETURN_IF_ERROR(args_.SetInt("stride_x", stride_.x));
+    RETURN_IF_ERROR(args_.SetInt("padding_x", padding_.x * src_[0]->Batch()));
+    RETURN_IF_ERROR(args_.SetInt("kernel_size_x", kernel_size_.x));
   }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(kernel_size_));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(padding_.x * src_[0]->Batch(), padding_.y)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(stride_));
-
-  return absl::OkStatus();
+  if (definition_.dst_tensors[0].HasAxis(Axis::HEIGHT)) {
+    RETURN_IF_ERROR(args_.SetInt("stride_y", stride_.y));
+    RETURN_IF_ERROR(args_.SetInt("padding_y", padding_.y));
+    RETURN_IF_ERROR(args_.SetInt("kernel_size_y", kernel_size_.y));
+  }
+  if (definition_.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    RETURN_IF_ERROR(args_.SetInt("stride_z", stride_.z));
+    RETURN_IF_ERROR(args_.SetInt("padding_z", padding_.z));
+    RETURN_IF_ERROR(args_.SetInt("kernel_size_z", kernel_size_.z));
+  }
+  if (output_indices_) {
+    RETURN_IF_ERROR(args_.SetObjectRef("dst_indices", dst_[1]));
+  }
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 Pooling::GetGridSize() const {
   const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
-  const int grid_y = dst_[0]->Height();
+  const int grid_y = dst_[0]->Height() * dst_[0]->Depth();
   const int grid_z = dst_[0]->Slices();
   return int3(grid_x, grid_y, grid_z);
 }
@@ -476,107 +429,9 @@ Pooling CreatePooling(const OperationDef& definition,
   return Pooling(definition, attr);
 }
 
-Pooling3D::Pooling3D(const OperationDef& definition,
-                     const Pooling3DAttributes& attr)
-    : GPUOperation(definition),
-      stride_(attr.strides.w, attr.strides.h, attr.strides.d),
-      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h,
-               -attr.padding.prepended.d),
-      kernel_size_(attr.kernel.w, attr.kernel.h, attr.kernel.d),
-      type_(attr.type),
-      output_indices_(attr.output_indices) {}
-
-Pooling3D::Pooling3D(Pooling3D&& kernel)
-    : GPUOperation(std::move(kernel)),
-      stride_(kernel.stride_),
-      padding_(kernel.padding_),
-      kernel_size_(kernel.kernel_size_),
-      type_(kernel.type_),
-      output_indices_(kernel.output_indices_),
-      kernel_(std::move(kernel.kernel_)),
-      work_group_size_(kernel.work_group_size_) {}
-
-Pooling3D& Pooling3D::operator=(Pooling3D&& kernel) {
-  if (this != &kernel) {
-    std::swap(stride_, kernel.stride_);
-    std::swap(padding_, kernel.padding_);
-    std::swap(kernel_size_, kernel.kernel_size_);
-    std::swap(type_, kernel.type_);
-    std::swap(output_indices_, kernel.output_indices_);
-    kernel_ = std::move(kernel.kernel_);
-    std::swap(work_group_size_, kernel.work_group_size_);
-    GPUOperation::operator=(std::move(kernel));
-  }
-  return *this;
-}
-
-absl::Status Pooling3D::Compile(const CreationContext& creation_context) {
-  std::string code;
-  const bool stride_correction =
-      definition_.IsBatchSupported() && stride_.x != 1;
-  switch (type_) {
-    case PoolingType::AVERAGE:
-      code = GetAveragePooling3DKernelCode(definition_, stride_correction,
-                                           *creation_context.device,
-                                           linked_operations_);
-      break;
-    case PoolingType::MAX:
-      code = GetMaxPooling3DKernelCode(definition_, stride_correction,
-                                       linked_operations_, output_indices_);
-      break;
-    default:
-      return absl::InvalidArgumentError(
-          "You should create another kernel with this params");
-      break;
-  }
-  return creation_context.cache->GetOrCreateCLKernel(
-      code, "main_function", *creation_context.context,
-      *creation_context.device, &kernel_);
-}
-
-absl::Status Pooling3D::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  if (output_indices_) {
-    RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[1]->GetMemoryPtrForWriting()));
-  }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHDS()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHDS()));
-  if (definition_.IsBatchSupported()) {
-    RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->Batch()));
-  }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(
-      int4(kernel_size_.x, kernel_size_.y, kernel_size_.z, 1)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(
-      int4(padding_.x * src_[0]->Batch(), padding_.y, padding_.z, 1)));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int4(stride_.x, stride_.y, stride_.z, 1)));
-
-  return absl::OkStatus();
-}
-
-int3 Pooling3D::GetGridSize() const {
-  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
-  const int grid_y = dst_[0]->Height();
-  const int grid_z = dst_[0]->Slices() * dst_[0]->Depth();
-  return int3(grid_x, grid_y, grid_z);
-}
-
-absl::Status Pooling3D::Tune(const TuningParameters& params) {
-  RETURN_IF_ERROR(BindArguments());
-  return GetBestWorkGroup(params, kernel_, GetGridSize(), &work_group_size_);
-}
-
-absl::Status Pooling3D::AddToQueue(CLCommandQueue* queue) {
-  RETURN_IF_ERROR(BindArguments());
-  return queue->DispatchImplicit(kernel_, GetGridSize(), work_group_size_);
-}
-
-Pooling3D CreatePooling3D(const OperationDef& definition,
-                          const Pooling3DAttributes& attr) {
-  return Pooling3D(definition, attr);
+Pooling CreatePooling(const OperationDef& definition,
+                      const Pooling3DAttributes& attr) {
+  return Pooling(definition, attr);
 }
 
 }  // namespace cl
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.h b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.h
index 09d2d5260f7..20719c90ae3 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/pooling.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/pooling.h
@@ -30,6 +30,7 @@ namespace cl {
 class Pooling : public GPUOperation {
  public:
   Pooling(const OperationDef& definition, const Pooling2DAttributes& attr);
+  Pooling(const OperationDef& definition, const Pooling3DAttributes& attr);
   absl::Status AddToQueue(CLCommandQueue* queue) override;
   absl::Status Tune(const TuningParameters& params) override;
 
@@ -45,9 +46,9 @@ class Pooling : public GPUOperation {
   absl::Status BindArguments();
   int3 GetGridSize() const;
 
-  int2 stride_;
-  int2 padding_;
-  int2 kernel_size_;
+  int4 stride_;
+  int4 padding_;
+  int4 kernel_size_;
 
   PoolingType type_;
   bool output_indices_;
@@ -59,37 +60,8 @@ class Pooling : public GPUOperation {
 Pooling CreatePooling(const OperationDef& definition,
                       const Pooling2DAttributes& attr);
 
-class Pooling3D : public GPUOperation {
- public:
-  Pooling3D(const OperationDef& definition, const Pooling3DAttributes& attr);
-  absl::Status AddToQueue(CLCommandQueue* queue) override;
-  absl::Status Tune(const TuningParameters& params) override;
-
-  absl::Status Compile(const CreationContext& creation_context) override;
-
-  // Move only
-  Pooling3D(Pooling3D&& kernel);
-  Pooling3D& operator=(Pooling3D&& kernel);
-  Pooling3D(const Pooling3D&) = delete;
-  Pooling3D& operator=(const Pooling3D&) = delete;
-
- private:
-  absl::Status BindArguments();
-  int3 GetGridSize() const;
-
-  int3 stride_;
-  int3 padding_;
-  int3 kernel_size_;
-
-  PoolingType type_;
-  bool output_indices_;
-
-  CLKernel kernel_;
-  int3 work_group_size_ = int3(8, 4, 1);
-};
-
-Pooling3D CreatePooling3D(const OperationDef& definition,
-                          const Pooling3DAttributes& attr);
+Pooling CreatePooling(const OperationDef& definition,
+                      const Pooling3DAttributes& attr);
 
 }  // namespace cl
 }  // namespace gpu
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
index 3161a73a18f..8cfff18b4ee 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
@@ -60,15 +60,6 @@ std::string GetImageModifier(AccessType access) {
   }
 }
 
-std::string TextureAddressModeToString(TextureAddressMode address_mode) {
-  switch (address_mode) {
-    case TextureAddressMode::DONT_CARE:
-      return "smp_none";
-    case TextureAddressMode::ZERO:
-      return "smp_zero";
-  }
-}
-
 }  // namespace
 
 std::string GetCommonDefines(CalculationsPrecision precision) {
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/util.h b/tensorflow/lite/delegates/gpu/cl/kernels/util.h
index 836a95f7407..3a51d064b40 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/util.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/util.h
@@ -36,11 +36,6 @@ namespace cl {
 
 std::string GetCommonDefines(CalculationsPrecision precision);
 
-enum class TextureAddressMode {
-  DONT_CARE,  // translated to CLK_ADDRESS_NONE
-  ZERO,       // translated to CLK_ADDRESS_CLAMP
-};
-
 struct WHSPoint {
   std::string w_name;
   std::string h_name;
diff --git a/tensorflow/lite/delegates/gpu/cl/tensor_type.cc b/tensorflow/lite/delegates/gpu/cl/tensor_type.cc
index 3b6c686a99a..8e048675697 100644
--- a/tensorflow/lite/delegates/gpu/cl/tensor_type.cc
+++ b/tensorflow/lite/delegates/gpu/cl/tensor_type.cc
@@ -45,6 +45,15 @@ std::string GetWriteImageFromDataType(DataType data_type) {
 
 }  // namespace
 
+std::string TextureAddressModeToString(TextureAddressMode address_mode) {
+  switch (address_mode) {
+    case TextureAddressMode::DONT_CARE:
+      return "smp_none";
+    case TextureAddressMode::ZERO:
+      return "smp_zero";
+  }
+}
+
 std::string ToString(TensorStorageType type) {
   switch (type) {
     case TensorStorageType::UNKNOWN:
@@ -271,8 +280,10 @@ std::string TensorDescriptor::Read(DataType read_as_type,
     case TensorStorageType::TEXTURE_3D:
     case TensorStorageType::SINGLE_TEXTURE_2D:
     case TensorStorageType::TEXTURE_ARRAY:
-      return absl::StrCat(read_as, "(", image_type, ", smp_none, ",
-                          global_address, ")");
+      return absl::StrCat(
+          read_as, "(", image_type,
+          ", " + TextureAddressModeToString(ModeFromState()) + ", ",
+          global_address, ")");
     case TensorStorageType::IMAGE_BUFFER:
       return absl::StrCat(read_as, "(image_buffer, ", global_address, ")");
     case TensorStorageType::UNKNOWN:
@@ -500,6 +511,14 @@ bool TensorDescriptor::HasAxis(Axis axis) const {
   return false;
 }
 
+void TensorDescriptor::SetTextureAddressMode(TextureAddressMode mode) {
+  if (mode == TextureAddressMode::ZERO) {
+    state_vars_["TextureMode"] = "ZERO";
+  } else {
+    state_vars_["TextureMode"] = "DONT_CARE";
+  }
+}
+
 bool TensorDescriptor::ParseCoordsFromArgs(const std::vector& args,
                                            int offset, std::string* xc,
                                            std::string* yc, std::string* zc,
@@ -549,6 +568,19 @@ bool TensorDescriptor::IsBatchedWidth() const {
   return it != state_vars_.end() && it->second == "true";
 }
 
+TextureAddressMode TensorDescriptor::ModeFromState() const {
+  auto it = state_vars_.find("TextureMode");
+  if (it != state_vars_.end()) {
+    if (it->second == "ZERO") {
+      return TextureAddressMode::ZERO;
+    } else {
+      return TextureAddressMode::DONT_CARE;
+    }
+  } else {
+    return TextureAddressMode::DONT_CARE;
+  }
+}
+
 }  // namespace cl
 }  // namespace gpu
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/gpu/cl/tensor_type.h b/tensorflow/lite/delegates/gpu/cl/tensor_type.h
index 58ebfc51ec4..2d4ae0c7335 100644
--- a/tensorflow/lite/delegates/gpu/cl/tensor_type.h
+++ b/tensorflow/lite/delegates/gpu/cl/tensor_type.h
@@ -27,6 +27,13 @@ namespace tflite {
 namespace gpu {
 namespace cl {
 
+enum class TextureAddressMode {
+  DONT_CARE,  // translated to CLK_ADDRESS_NONE
+  ZERO,       // translated to CLK_ADDRESS_CLAMP
+};
+
+std::string TextureAddressModeToString(TextureAddressMode address_mode);
+
 enum class TensorStorageType {
   UNKNOWN,
   BUFFER,
@@ -71,6 +78,7 @@ struct TensorDescriptor : public GPUObjectDescriptor {
   GPUResources GetGPUResources(AccessType access_type) const override;
 
   bool HasAxis(Axis axis) const;
+  void SetTextureAddressMode(TextureAddressMode mode);
 
   absl::Status GetLinkingContextFromWriteSelector(
       const std::vector& args, std::string* value_name,
@@ -106,6 +114,8 @@ struct TensorDescriptor : public GPUObjectDescriptor {
 
   bool IsBatchedWidth() const;
 
+  TextureAddressMode ModeFromState() const;
+
   absl::Status GetDataTypeFromTemplateArgs(const std::string& template_arg,
                                            DataType* result) const;
 

From b1933d67e5507bc1433c6de1700bcb0b8f3a6aec Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:22:03 -0700
Subject: [PATCH 0453/1390] Added support of different layouts for src/dst. For
 example src - HWC, dst - BHWC, or vice versa.

PiperOrigin-RevId: 316998239
Change-Id: I89b07923020f185c356bb0b63926bbe81be55cb5
---
 .../lite/delegates/gpu/cl/kernels/strided_slice.cc     | 10 +++++-----
 tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc  |  9 +++++----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc b/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
index 2cf65f24447..d0c4e432f3a 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/strided_slice.cc
@@ -42,11 +42,12 @@ std::string GetStridedSliceCode(const OperationDef& op_def, bool alignedx4,
   args->AddInt("stride_z");
   args->AddInt("stride_b");
 
-  const std::string dst_batch = op_def.IsBatchSupported() ? "B" : "";
+  const std::string batch_id =
+      op_def.dst_tensors[0].HasAxis(Axis::BATCH) ? "B" : "0";
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
   c += "$0) {\n";
-  if (op_def.IsBatchSupported()) {
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
     c += "  int linear_id = get_global_id(0);\n";
     c += "  int X = linear_id / args.dst_tensor.Batch();\n";
     c += "  int B = linear_id % args.dst_tensor.Batch();\n";
@@ -62,11 +63,10 @@ std::string GetStridedSliceCode(const OperationDef& op_def, bool alignedx4,
   c += "  } \n";
   c += "  int s_x = X * args.stride_x + args.offset_x;\n";
   c += "  int s_y = Y * args.stride_y + args.offset_y;\n";
-  if (op_def.IsBatchSupported()) {
-    c += "  int s_b = B * args.stride_b + args.offset_b;\n";
+  if (op_def.src_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int s_b = " + batch_id + " * args.stride_b + args.offset_b;\n";
     c += "  args.src_tensor.SetBatchRef(s_b);\n";
   }
-  const std::string src_batch = op_def.IsBatchSupported() ? "s_b" : "";
   if (alignedx4) {
     c += "  int s_z = Z + args.offset_z;\n";
     c += "  FLT4 result = args.src_tensor.Read(s_x, s_y, s_z);\n";
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc b/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
index e12c44566b7..cacfd52542d 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/transpose.cc
@@ -36,11 +36,12 @@ std::string GetTransposeCode(
       "dst_tensor", AccessType::WRITE,
       absl::make_unique(op_def.dst_tensors[0]));
 
-  const std::string batch_id = op_def.IsBatchSupported() ? "B" : "";
+  const std::string batch_id =
+      op_def.dst_tensors[0].HasAxis(Axis::BATCH) ? "B" : "0";
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
   c += "$0) {\n";
-  if (op_def.IsBatchSupported()) {
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
     c += "  int linear_id = get_global_id(0);\n";
     c += "  int X = linear_id / args.dst_tensor.Batch();\n";
     c += "  int B = linear_id % args.dst_tensor.Batch();\n";
@@ -65,7 +66,7 @@ std::string GetTransposeCode(
   remap[attr.perm.w] = 2;
   remap[attr.perm.c] = 3;
   if (attr.perm.c == 3) {  // optimized reading when no channels permutation
-    const std::string bhw[] = {"B", "Y", "X"};
+    const std::string bhw[] = {batch_id, "Y", "X"};
     if (op_def.src_tensors[0].HasAxis(Axis::BATCH)) {
       c += "  args.src_tensor.SetBatchRef(" + bhw[remap[0]] + ");\n";
     }
@@ -80,7 +81,7 @@ std::string GetTransposeCode(
     c += "  for (int i = 0; i < 4; ++i) {\n";
     c += "    int dst_channel = Z * 4 + i;\n";
     c += "    if (dst_channel < args.dst_tensor.Channels()) {\n";
-    const std::string bhwc[] = {"B", "Y", "X", "dst_channel"};
+    const std::string bhwc[] = {batch_id, "Y", "X", "dst_channel"};
     if (op_def.src_tensors[0].HasAxis(Axis::BATCH)) {
       c += "      args.src_tensor.SetBatchRef(" + bhwc[remap[0]] + ");\n";
     }

From 274a0f944e50922759fe8c263f74beeb24884137 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:29:06 -0700
Subject: [PATCH 0454/1390] ConcatXY converted to new style. Added support of
 concatenation in Batch and Depth axis.

PiperOrigin-RevId: 316999295
Change-Id: I94f2168f2861790b3a30c79b2b3476aa44c55748
---
 .../delegates/gpu/cl/kernels/concat_xy.cc     | 160 +++++++++++-------
 .../gpu/cl/selectors/simple_selectors.cc      |   6 +-
 2 files changed, 99 insertions(+), 67 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/concat_xy.cc b/tensorflow/lite/delegates/gpu/cl/kernels/concat_xy.cc
index ef7915afba5..5476cc22965 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/concat_xy.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/concat_xy.cc
@@ -15,7 +15,9 @@ limitations under the License.
 
 #include "tensorflow/lite/delegates/gpu/cl/kernels/concat_xy.h"
 
+#include 
 #include 
+#include 
 
 #include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
 #include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
@@ -27,51 +29,93 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetConcatKernelCode(
-    const OperationDef& op_def, int tensors_count,
-    const std::vector& linked_operations) {
-  std::vector srcs(tensors_count);
-  for (int i = 0; i < tensors_count; ++i) {
-    const std::string tensor_name = "src_data_" + std::to_string(i);
-    const std::string width = "src_size_" + std::to_string(i) + ".x";
-    const std::string height = "src_size_" + std::to_string(i) + ".y";
-    srcs[i] =
-        TensorCodeGenerator(tensor_name, WHSPoint{width, height, "dst_size.z"},
-                            op_def.src_tensors[i]);
+std::string GetConcatKernelCode(const OperationDef& op_def,
+                                const ConcatAttributes& attr, Arguments* args) {
+  std::vector tensor_names(op_def.src_tensors.size());
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    tensor_names[i] = "src_tensor_" + std::to_string(i);
+    args->AddObjectRef(
+        tensor_names[i], AccessType::READ,
+        absl::make_unique(op_def.src_tensors[0]));
+  }
+  args->AddObjectRef(
+      "dst_tensor", AccessType::WRITE,
+      absl::make_unique(op_def.dst_tensors[0]));
+
+  std::map axis_to_selector = {
+      {Axis::WIDTH, "Width"}, {Axis::HEIGHT, "Height"},
+      {Axis::DEPTH, "Depth"}, {Axis::CHANNELS, "Channels"},
+      {Axis::BATCH, "Batch"},
+  };
+  std::map axis_to_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "S"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector src_coords;
+  std::vector dst_coords;
+  for (auto axis :
+       {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS, Axis::BATCH}) {
+    if (op_def.src_tensors[0].HasAxis(axis) && axis != Axis::BATCH) {
+      if (axis == attr.axis) {
+        src_coords.push_back("coord");
+      } else {
+        src_coords.push_back(axis_to_coord[axis]);
+      }
+    }
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
   }
-  TensorCodeGenerator dst("dst_data",
-                          WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-                          op_def.dst_tensors[0]);
 
   std::string c = GetCommonDefines(op_def.precision);
-
   c += "__kernel void main_function(\n";
-  for (const auto& src : srcs) {
-    c += src.GetDeclaration(AccessType::READ) + ",\n";
+  c += "$0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id_0 = get_global_id(0);\n";
+    c += "  int X = linear_id_0 / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id_0 % args.dst_tensor.Batch();\n";
+  } else {
+    c += "  int X = get_global_id(0);\n";
   }
-  c += dst.GetDeclaration(AccessType::WRITE);
-  c += GetArgsDeclaration(linked_operations);
-  for (int i = 0; i < tensors_count; ++i) {
-    const std::string uniform_name = "src_size_" + std::to_string(i);
-    c += "    int4 " + uniform_name + ",\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = get_global_id(1);\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = get_global_id(1);\n";
   }
-  c += "    int4 dst_size  \n";
-  c += ") {\n";
-  c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int Z = get_global_id(2);\n";
-  c += "  if (Z >= dst_size.z) return;\n";
-  for (int i = 0; i < tensors_count; ++i) {
-    const std::string size_name = "src_size_" + std::to_string(i);
-    c += "  if (X < " + size_name + ".x && Y < " + size_name + ".y) { \n";
-    c += "    FLT4 result = " + srcs[i].ReadWHS("X", "Y", "Z") + ";\n";
-    c += "    int dst_x = X + " + size_name + ".z;\n";
-    c += "    int dst_y = Y + " + size_name + ".w;\n";
-    const LinkingContext context{"result", "dst_x", "dst_y", "Z"};
-    c += PostProcess(linked_operations, context);
-    c += "    " + dst.WriteWHS("result", "dst_x", "dst_y", "Z");
+  c += "  int S = get_global_id(2);\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  FLT4 result = (FLT4)(0.0f);\n";
+  c += "  int coord = " + axis_to_coord[attr.axis] + ";\n";
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    const std::string field =
+        "args." + tensor_names[i] + "." + axis_to_selector[attr.axis] + "()";
+    c += "  if (coord >= 0 && coord < " + field + ") { \n";
+    if (op_def.src_tensors[i].HasAxis(Axis::BATCH)) {
+      if (attr.axis == Axis::BATCH) {
+        c += "  args." + tensor_names[i] + ".SetBatchRef(coord);\n";
+      } else {
+        c += "  args." + tensor_names[i] + ".SetBatchRef(B);\n";
+      }
+    }
+    c += "    result = args." + tensor_names[i] + ".Read(" + src_coord + ");\n";
     c += "  } \n";
+    c += "  coord -= " + field + ";\n";
   }
+  c += "  args.dst_tensor.Write(result, " + dst_coord + ");\n";
   c += "}\n";
   return c;
 }
@@ -97,46 +141,32 @@ ConcatXY& ConcatXY::operator=(ConcatXY&& operation) {
 }
 
 absl::Status ConcatXY::Compile(const CreationContext& creation_context) {
-  const auto code =
-      GetConcatKernelCode(definition_, tensors_count_, linked_operations_);
+  std::string code = GetConcatKernelCode(definition_, attr_, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status ConcatXY::BindArguments() {
-  kernel_.ResetBindingCounter();
-  for (int i = 0; i < tensors_count_; ++i) {
-    RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[i]->GetMemoryPtr()));
-  }
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  int x_offset = 0;
-  int y_offset = 0;
-  for (int i = 0; i < tensors_count_; ++i) {
-    const int width = src_[i]->Width() * src_[i]->Batch();
-    const int height = src_[i]->Height();
+  for (int i = 0; i < definition_.src_tensors.size(); ++i) {
     RETURN_IF_ERROR(
-        kernel_.SetBytesAuto(int4(width, height, x_offset, y_offset)));
-    x_offset += attr_.axis == Axis::WIDTH ? width : 0;
-    y_offset += attr_.axis == Axis::HEIGHT ? height : 0;
+        args_.SetObjectRef("src_tensor_" + std::to_string(i), src_[i]));
   }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHSB()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 ConcatXY::GetGridSize() const {
-  int max_src_width = 0;
-  int max_src_height = 0;
-  for (int i = 0; i < tensors_count_; ++i) {
-    max_src_width = std::max(max_src_width, src_[i]->Width());
-    max_src_height = std::max(max_src_height, src_[i]->Height());
-  }
-
-  const int grid_x = max_src_width * dst_[0]->Batch();
-  const int grid_y = max_src_height;
+  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = dst_[0]->Height() * dst_[0]->Depth();
   const int grid_z = dst_[0]->Slices();
-
   return int3(grid_x, grid_y, grid_z);
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/selectors/simple_selectors.cc b/tensorflow/lite/delegates/gpu/cl/selectors/simple_selectors.cc
index 5fc04d12822..d9bd70cc6b3 100644
--- a/tensorflow/lite/delegates/gpu/cl/selectors/simple_selectors.cc
+++ b/tensorflow/lite/delegates/gpu/cl/selectors/simple_selectors.cc
@@ -105,8 +105,10 @@ absl::Status SelectConcat(const ConcatAttributes& attr,
       *ptr = absl::make_unique(std::move(operation));
       return absl::OkStatus();
     }
-    case Axis::WIDTH:
-    case Axis::HEIGHT: {
+    case Axis::BATCH:
+    case Axis::DEPTH:
+    case Axis::HEIGHT:
+    case Axis::WIDTH: {
       ConcatXY operation = CreateConcatXY(op_def, attr, channels.size());
       *ptr = absl::make_unique(std::move(operation));
       return absl::OkStatus();

From 70e2387ecc190d415cde234c28ce5f5ae6b90bea Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Wed, 17 Jun 2020 17:29:11 -0700
Subject: [PATCH 0455/1390] Merged Reshape code generation for HWC/BHWC
 layouts. Added support of different layouts for src/dst. For example src -
 HWC, dst - BHWC, or vice versa.

PiperOrigin-RevId: 316999310
Change-Id: I20bd9a12afba8bdcb832565f09350440349041bd
---
 .../lite/delegates/gpu/cl/kernels/reshape.cc  | 79 ++++++-------------
 .../delegates/gpu/cl/kernels/reshapex4.cc     | 69 ++++++----------
 2 files changed, 47 insertions(+), 101 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc b/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
index a99fff0a1da..5abfad60c1b 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/reshape.cc
@@ -25,56 +25,6 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetReshapeBatchedCode(const OperationDef& op_def, Arguments* args) {
-  args->AddObjectRef(
-      "src_tensor", AccessType::READ,
-      absl::make_unique(op_def.src_tensors[0]));
-  args->AddObjectRef(
-      "dst_tensor", AccessType::WRITE,
-      absl::make_unique(op_def.dst_tensors[0]));
-
-  std::string c = GetCommonDefines(op_def.precision);
-  c += "__kernel void main_function(\n";
-  c += "$0) {\n";
-  c += "  int linear_id = get_global_id(0);\n";
-  c += "  int X = linear_id / args.dst_tensor.Batch();\n";
-  c += "  int B = linear_id % args.dst_tensor.Batch();\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
-       "Z >= args.dst_tensor.Slices()) { \n";
-  c += "    return; \n";
-  c += "  } \n";
-  c += "  FLT temps[4];\n";
-  c += "  temps[0] = (FLT)(0.0f);\n";
-  c += "  temps[1] = (FLT)(0.0f);\n";
-  c += "  temps[2] = (FLT)(0.0f);\n";
-  c += "  temps[3] = (FLT)(0.0f);\n";
-  c += "  int base = ((B * args.dst_tensor.Height() + Y) * "
-       "args.dst_tensor.Width() + X) * args.dst_tensor.Channels() + Z * 4;\n";
-  c += "  for (int i = 0; i < 4; ++i) {\n";
-  c += "    int dst_channel = Z * 4 + i;\n";
-  c += "    if (dst_channel < args.dst_tensor.Channels()) {;\n";
-  c += "      int p = base + i;\n";
-  c += "      int src_c = p % args.src_tensor.Channels();\n";
-  c += "      p = p / args.src_tensor.Channels();\n";
-  c += "      int src_x = p % args.src_tensor.Width();\n";
-  c += "      p = p / args.src_tensor.Width();\n";
-  c += "      int src_y = p % args.src_tensor.Height();\n";
-  c += "      int src_b = p / args.src_tensor.Height();\n";
-  c += "      int src_z = src_c / 4;\n";
-  c += "      int src_sub_ch = src_c % 4;\n";
-  c += "      FLT4 t = args.src_tensor.Read(src_x, src_y, src_z, src_b);\n";
-  c += "      FLT t_ar[4] = {t.x, t.y, t.z, t.w};\n";
-  c += "      temps[i] = t_ar[src_sub_ch];\n";
-  c += "    }\n";
-  c += "  }\n";
-  c += "  FLT4 result = (FLT4)(temps[0], temps[1], temps[2], temps[3]);\n";
-  c += "  args.dst_tensor.Write(result, X, Y, Z, B);\n";
-  c += "}\n";
-  return c;
-}
-
 std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   args->AddObjectRef(
       "src_tensor", AccessType::READ,
@@ -86,7 +36,14 @@ std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
   c += "$0) {\n";
-  c += "  int X = get_global_id(0);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = get_global_id(0);\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = get_global_id(0);\n";
+  }
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
   c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
@@ -98,8 +55,13 @@ std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   c += "  temps[1] = (FLT)(0.0f);\n";
   c += "  temps[2] = (FLT)(0.0f);\n";
   c += "  temps[3] = (FLT)(0.0f);\n";
-  c += "  int base = (Y * args.dst_tensor.Width() + X) * "
-       "args.dst_tensor.Channels() + Z * 4;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int base = B;\n";
+  } else {
+    c += "  int base = 0;\n";
+  }
+  c += "  base = ((base * args.dst_tensor.Height() + Y) * "
+       "args.dst_tensor.Width() + X) * args.dst_tensor.Channels() + Z * 4;\n";
   c += "  for (int i = 0; i < 4; ++i) {\n";
   c += "    int dst_channel = Z * 4 + i;\n";
   c += "    if (dst_channel < args.dst_tensor.Channels()) {;\n";
@@ -107,7 +69,12 @@ std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   c += "      int src_c = p % args.src_tensor.Channels();\n";
   c += "      p = p / args.src_tensor.Channels();\n";
   c += "      int src_x = p % args.src_tensor.Width();\n";
-  c += "      int src_y = p / args.src_tensor.Width();\n";
+  c += "      p = p / args.src_tensor.Width();\n";
+  c += "      int src_y = p % args.src_tensor.Height();\n";
+  if (op_def.src_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int src_b = p / args.src_tensor.Height();\n";
+    c += "  args.src_tensor.SetBatchRef(src_b);\n";
+  }
   c += "      int src_z = src_c / 4;\n";
   c += "      int src_sub_ch = src_c % 4;\n";
   c += "      FLT4 t = args.src_tensor.Read(src_x, src_y, src_z);\n";
@@ -137,9 +104,7 @@ Reshape& Reshape::operator=(Reshape&& operation) {
 }
 
 absl::Status Reshape::Compile(const CreationContext& creation_context) {
-  std::string code = definition_.IsBatchSupported()
-                         ? GetReshapeBatchedCode(definition_, &args_)
-                         : GetReshapeCode(definition_, &args_);
+  std::string code = GetReshapeCode(definition_, &args_);
   std::string element_wise_code;
   RETURN_IF_ERROR(
       MergeOperations(linked_operations_, &args_, &element_wise_code));
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc b/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
index 0847fce5836..3edbe637aa2 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/reshapex4.cc
@@ -25,40 +25,6 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GetReshapeBatchedCode(const OperationDef& op_def, Arguments* args) {
-  args->AddObjectRef(
-      "src_tensor", AccessType::READ,
-      absl::make_unique(op_def.src_tensors[0]));
-  args->AddObjectRef(
-      "dst_tensor", AccessType::WRITE,
-      absl::make_unique(op_def.dst_tensors[0]));
-
-  std::string c = GetCommonDefines(op_def.precision);
-  c += "__kernel void main_function(\n";
-  c += "$0) {\n";
-  c += "  int linear_id = get_global_id(0);\n";
-  c += "  int X = linear_id / args.dst_tensor.Batch();\n";
-  c += "  int B = linear_id % args.dst_tensor.Batch();\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
-       "Z >= args.dst_tensor.Slices()) { \n";
-  c += "    return; \n";
-  c += "  } \n";
-  c += "  int dst_bhwc4 = ((B * args.dst_tensor.Height() + Y) * "
-       "args.dst_tensor.Width() + X) * args.dst_tensor.Slices() + Z;\n";
-  c += "  int src_z = dst_bhwc4 % args.src_tensor.Slices();\n";
-  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Slices();\n";
-  c += "  int src_x = dst_bhwc4 % args.src_tensor.Width();\n";
-  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Width();\n";
-  c += "  int src_y = dst_bhwc4 % args.src_tensor.Height();\n";
-  c += "  int src_b = dst_bhwc4 / args.src_tensor.Height();\n";
-  c += "  FLT4 result = args.src_tensor.Read(src_x, src_y, src_z, src_b);\n";
-  c += "  args.dst_tensor.Write(result, X, Y, Z, B);\n";
-  c += "}\n";
-  return c;
-}
-
 std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   args->AddObjectRef(
       "src_tensor", AccessType::READ,
@@ -70,19 +36,36 @@ std::string GetReshapeCode(const OperationDef& op_def, Arguments* args) {
   std::string c = GetCommonDefines(op_def.precision);
   c += "__kernel void main_function(\n";
   c += "$0) {\n";
-  c += "  int X = get_global_id(0);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = get_global_id(0);\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = get_global_id(0);\n";
+  }
   c += "  int Y = get_global_id(1);\n";
   c += "  int Z = get_global_id(2);\n";
   c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
        "Z >= args.dst_tensor.Slices()) { \n";
   c += "    return; \n";
   c += "  } \n";
-  c += "  int dst_hwc4 = (Y * args.dst_tensor.Width() + X) * "
-       "args.dst_tensor.Slices() + Z;\n";
-  c += "  int src_z = dst_hwc4 % args.src_tensor.Slices();\n";
-  c += "  dst_hwc4 = dst_hwc4 / args.src_tensor.Slices();\n";
-  c += "  int src_x = dst_hwc4 % args.src_tensor.Width();\n";
-  c += "  int src_y = dst_hwc4 / args.src_tensor.Width();\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int dst_bhwc4 = B;\n";
+  } else {
+    c += "  int dst_bhwc4 = 0;\n";
+  }
+  c += "  dst_bhwc4 = ((dst_bhwc4 * args.dst_tensor.Height() + Y) * "
+       "args.dst_tensor.Width() + X) * args.dst_tensor.Slices() + Z;\n";
+  c += "  int src_z = dst_bhwc4 % args.src_tensor.Slices();\n";
+  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Slices();\n";
+  c += "  int src_x = dst_bhwc4 % args.src_tensor.Width();\n";
+  c += "  dst_bhwc4 = dst_bhwc4 / args.src_tensor.Width();\n";
+  c += "  int src_y = dst_bhwc4 % args.src_tensor.Height();\n";
+  if (op_def.src_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int src_b = dst_bhwc4 / args.src_tensor.Height();\n";
+    c += "  args.src_tensor.SetBatchRef(src_b);\n";
+  }
   c += "  FLT4 result = args.src_tensor.Read(src_x, src_y, src_z);\n";
   c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
   c += "}\n";
@@ -105,9 +88,7 @@ Reshapex4& Reshapex4::operator=(Reshapex4&& operation) {
 }
 
 absl::Status Reshapex4::Compile(const CreationContext& creation_context) {
-  std::string code = definition_.IsBatchSupported()
-                         ? GetReshapeBatchedCode(definition_, &args_)
-                         : GetReshapeCode(definition_, &args_);
+  std::string code = GetReshapeCode(definition_, &args_);
   std::string element_wise_code;
   RETURN_IF_ERROR(
       MergeOperations(linked_operations_, &args_, &element_wise_code));

From adf0573f9d874b568c9680a6d3cf310e555dcd62 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 17:37:35 -0700
Subject: [PATCH 0456/1390] Qualify uses of std::string

PiperOrigin-RevId: 317000789
Change-Id: I6f847b235496d1dc8f8b3380e21ce890566a9a88
---
 .../tests/fuse_binary_into_following_affine_test.cc      | 5 +++--
 .../tests/fuse_binary_into_preceding_affine_test.cc      | 5 +++--
 .../toco/graph_transformations/tests/lstm_utils_test.cc  | 4 ++--
 .../tests/resolve_constant_concatenation_test.cc         | 8 ++++----
 .../graph_transformations/tests/unpack_quantize_test.cc  | 9 +++++----
 5 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
index 2cba6824cfb..d6cf31f4211 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
@@ -43,14 +43,15 @@ class FuseBinaryIntoFollowingAffineTest : public ::testing::Test {
 
   void SetUp() override { model_.reset(new Model); }
 
-  void CreateArray(const string& name, const std::vector& shape) {
+  void CreateArray(const std::string& name, const std::vector& shape) {
     Array& array = model_->GetOrCreateArray(name);
     array.data_type = ArrayDataType::kFloat;
     Shape* array_shape = array.mutable_shape();
     *(array_shape->mutable_dims()) = shape;
   }
 
-  void CreateConstantArray(const string& name, const std::vector& shape,
+  void CreateConstantArray(const std::string& name,
+                           const std::vector& shape,
                            const std::vector& data) {
     CreateArray(name, shape);
     Array& array = model_->GetOrCreateArray(name);
diff --git a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
index b5c321c1a26..6c3dc7dc761 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
@@ -43,14 +43,15 @@ class FuseBinaryIntoPrecedingAffineTest : public ::testing::Test {
 
   void SetUp() override { model_.reset(new Model); }
 
-  void CreateArray(const string& name, const std::vector& shape) {
+  void CreateArray(const std::string& name, const std::vector& shape) {
     Array& array = model_->GetOrCreateArray(name);
     array.data_type = ArrayDataType::kFloat;
     Shape* array_shape = array.mutable_shape();
     *(array_shape->mutable_dims()) = shape;
   }
 
-  void CreateConstantArray(const string& name, const std::vector& shape,
+  void CreateConstantArray(const std::string& name,
+                           const std::vector& shape,
                            const std::vector& data) {
     CreateArray(name, shape);
     Array& array = model_->GetOrCreateArray(name);
diff --git a/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc b/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
index bdb27e8af2e..204e197e186 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
@@ -46,12 +46,12 @@ class CopyArrayDataTest : public ::testing::Test {
                       int src_dim_1, int src_dim_2,
                       std::initializer_list dst_data, int dst_dim_1,
                       int dst_dim_2) {
-    string src_array = "src_array";
+    std::string src_array = "src_array";
     src_buffer_ = CreateFloatArrayBuffer(
         model, &src_array,
         src_dim_2 == 1 ? Shape({src_dim_1}) : Shape({src_dim_1, src_dim_2}));
     PopulateBuffer(src_buffer_, src_data);
-    string dst_array = "dst_array";
+    std::string dst_array = "dst_array";
     dst_buffer_ = CreateFloatArrayBuffer(
         model, &dst_array,
         dst_dim_2 == 1 ? Shape({dst_dim_1}) : Shape({dst_dim_1, dst_dim_2}));
diff --git a/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc b/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
index bfed38ce7aa..5b0566fe074 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
@@ -107,10 +107,10 @@ class ResolveConstantConcatenationTest : public ::testing::Test {
   // together with 4 arrays as its inputs.
   // It receives the dimension of concatenation as input.
   void PrepareModel(Model* model, int axis) {
-    const string output_name("concat_op_output");
+    const std::string output_name("concat_op_output");
     model->flags.add_output_arrays(output_name);
-    std::vector concat_input_names = {"array0", "array1", "array2",
-                                              "array3"};
+    std::vector concat_input_names = {"array0", "array1", "array2",
+                                                   "array3"};
 
     const int kDim = 3;
     const int kElementPerDim = 2;
@@ -122,7 +122,7 @@ class ResolveConstantConcatenationTest : public ::testing::Test {
         {20., 21., 22., 23., 24., 25., 26., 27.},
         {30., 31., 32., 33., 34., 35., 36., 37.}};
     int cnt = 0;
-    for (const string& concat_input_name : concat_input_names) {
+    for (const std::string& concat_input_name : concat_input_names) {
       Array& in_array = model->GetOrCreateArray(concat_input_name);
       in_array.data_type = ArrayDataType::kFloat;
 
diff --git a/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc b/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
index 2dc3fb35b0f..3cc4e725463 100755
--- a/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
@@ -40,10 +40,11 @@ class UnpackQuantizeTest : public ::testing::Test {
   // 1. calculate min and max of the input.
   // 2. insert dequantization nodes after quantized outputs of Unpack operation.
   void PrepareModel(Model* model, int axis) {
-    std::vector unpack_output_names = {"unpack_out0", "unpack_out1"};
+    std::vector unpack_output_names = {"unpack_out0",
+                                                    "unpack_out1"};
     model->flags.add_output_arrays(unpack_output_names[0]);
     model->flags.add_output_arrays(unpack_output_names[1]);
-    const string unpack_input_name("unpack_op_input");
+    const std::string unpack_input_name("unpack_op_input");
 
     const int kDim = 2;
     const int kElementPerDim = 2;
@@ -75,7 +76,7 @@ class UnpackQuantizeTest : public ::testing::Test {
     // Configuring the necessary outputs. The outputs also happen to be in
     // kFloat. This is because during quantization transformation data types for
     // these arrays are going to be forced to be kUint8.
-    for (const string& unpack_output_name : unpack_output_names) {
+    for (const std::string& unpack_output_name : unpack_output_names) {
       Array& out_array = model->GetOrCreateArray(unpack_output_name);
       out_array.GetOrCreateMinMax();
       out_array.data_type = ArrayDataType::kFloat;
@@ -109,7 +110,7 @@ TEST_F(UnpackQuantizeTest, CheckUnpackPreservesQuantizationParameters) {
                   ->Run(&model, /*op_index=*/0, &modified)
                   .ok());
 
-  const string output_name = model.flags.output_arrays(0);
+  const std::string output_name = model.flags.output_arrays(0);
 
   // Quantization transformation inserts NODE_NAME_DEQUANTIZE operations,
   // effectively making them the new outputs of the array. Old outputs of the

From 56db128697fd0e9d3f689d472d53e33bf4ea1bb8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 17:56:36 -0700
Subject: [PATCH 0457/1390] Qualify uses of std::string

PiperOrigin-RevId: 317003622
Change-Id: Iae6a9a287ffd3b97dee8b9993c443db322936592
---
 .../lite/toco/logging/conversion_log_util.cc  | 67 ++++++++++---------
 .../lite/toco/logging/conversion_log_util.h   | 22 +++---
 .../toco/logging/conversion_log_util_test.cc  | 60 ++++++++---------
 3 files changed, 77 insertions(+), 72 deletions(-)

diff --git a/tensorflow/lite/toco/logging/conversion_log_util.cc b/tensorflow/lite/toco/logging/conversion_log_util.cc
index c23c305c750..55afa1370b3 100644
--- a/tensorflow/lite/toco/logging/conversion_log_util.cc
+++ b/tensorflow/lite/toco/logging/conversion_log_util.cc
@@ -34,8 +34,8 @@ namespace toco {
 
 namespace {
 
-string TryGetOperatorName(const Operator& op) {
-  string op_name;
+std::string TryGetOperatorName(const Operator& op) {
+  std::string op_name;
   if (!op.tensorflow_node_def.empty()) {
     // Parse op name from serialized NodeDef.
     tensorflow::NodeDef node_def;
@@ -63,8 +63,8 @@ string TryGetOperatorName(const Operator& op) {
   return op_name;
 }
 
-string GetOSVersion() {
-  string os_info;
+std::string GetOSVersion() {
+  std::string os_info;
 #ifdef __linux__
   utsname info;
   if (uname(&info)) {
@@ -72,12 +72,13 @@ string GetOSVersion() {
     LOG(ERROR) << "Cannot get OS info.";
     return "";
   }
-  os_info = string(info.sysname) + ";OSVer=" + string(info.release) + ";";
+  os_info =
+      std::string(info.sysname) + ";OSVer=" + std::string(info.release) + ";";
 #endif
   return os_info;
 }
 
-string ShapeToStringNoSpace(const Shape& shape) {
+std::string ShapeToStringNoSpace(const Shape& shape) {
   if (shape.dimensions_count() == 0) {
     return "[]";
   }
@@ -85,13 +86,13 @@ string ShapeToStringNoSpace(const Shape& shape) {
   return absl::StrCat("[", absl::StrJoin(shape.dims(), ","), "]");
 }
 
-string GetOperatorSignature(
+std::string GetOperatorSignature(
     const Model& model, const Operator& op,
     const std::map>&
         op_types_map) {
   // The signature of an op has the following schema:
   // INPUT:SHAPE::TYPE::OUTPUT:SHAPE::TYPE::NAME:VERSION:
-  string op_signature;
+  std::string op_signature;
   constexpr char delimiter[] = "::";
 
   // Get input shapes and types.
@@ -137,8 +138,8 @@ string GetOperatorSignature(
 
 }  // namespace
 
-std::vector GetOperatorNames(const Model& model) {
-  std::vector op_names;
+std::vector GetOperatorNames(const Model& model) {
+  std::vector op_names;
   for (const auto& op : model.operators) {
     op_names.push_back(TryGetOperatorName(*op));
   }
@@ -146,9 +147,9 @@ std::vector GetOperatorNames(const Model& model) {
 }
 
 void CountOperatorsByType(const Model& model,
-                          std::map* built_in_ops,
-                          std::map* custom_ops,
-                          std::map* select_ops) {
+                          std::map* built_in_ops,
+                          std::map* custom_ops,
+                          std::map* select_ops) {
   for (const auto& op : model.operators) {
     OperatorSignature op_signature = {op.get(), &model};
     const auto ops_by_type =
@@ -156,7 +157,7 @@ void CountOperatorsByType(const Model& model,
     tflite::details::OperatorKey op_key(op_signature, ops_by_type,
                                         true /*enable_select_tf_ops*/);
 
-    const string op_name = TryGetOperatorName(*op);
+    const std::string op_name = TryGetOperatorName(*op);
     if (op_key.is_custom_op()) {
       (*custom_ops)[op_name]++;
     } else if (op_key.is_flex_op()) {
@@ -168,8 +169,9 @@ void CountOperatorsByType(const Model& model,
 }
 
 void GetInputAndOutputTypes(
-    const Model& model, TFLITE_PROTO_NS::RepeatedPtrField* input_types,
-    TFLITE_PROTO_NS::RepeatedPtrField* output_types) {
+    const Model& model,
+    TFLITE_PROTO_NS::RepeatedPtrField* input_types,
+    TFLITE_PROTO_NS::RepeatedPtrField* output_types) {
   for (const auto& input_array : model.flags.input_arrays()) {
     const Array& array = model.GetArray(input_array.name());
     input_types->Add(ArrayDataTypeName(array.data_type));
@@ -180,15 +182,16 @@ void GetInputAndOutputTypes(
   }
 }
 
-string GetTfLiteVersion() { return TFLITE_VERSION_STRING; }
+std::string GetTfLiteVersion() { return TFLITE_VERSION_STRING; }
 
-string GetCachedOSVersion() {
-  static string* version = new string(GetOSVersion());
+std::string GetCachedOSVersion() {
+  static std::string* version = new std::string(GetOSVersion());
   return *version;
 }
 
-void GetOpSignatures(const Model& model,
-                     TFLITE_PROTO_NS::RepeatedPtrField* op_signatures) {
+void GetOpSignatures(
+    const Model& model,
+    TFLITE_PROTO_NS::RepeatedPtrField* op_signatures) {
   const auto& op_types_map =
       tflite::BuildOperatorByTypeMap(true /*enable_select_tf_ops*/);
   for (const auto& op : model.operators) {
@@ -196,7 +199,7 @@ void GetOpSignatures(const Model& model,
   }
 }
 
-string GetModelHash(const Model& model) {
+std::string GetModelHash(const Model& model) {
   // TODO(b/123519920): Implement the hash function for Model.
   // Need to consider different implementations for public/private models.
   return "";
@@ -204,18 +207,18 @@ string GetModelHash(const Model& model) {
 
 // This function scans through the error message string, extracts the part about
 // missing ops and prunes away all other information in the error info.
-string SanitizeErrorMessage(const string& error_message) {
-  const string s1 = "Ops that can be supported by the flex runtime";
-  const string s2 = "Ops that need custom implementation";
-  string pruned_message;
+std::string SanitizeErrorMessage(const std::string& error_message) {
+  const std::string s1 = "Ops that can be supported by the flex runtime";
+  const std::string s2 = "Ops that need custom implementation";
+  std::string pruned_message;
   size_t pos = error_message.find(s1);
-  if (pos != string::npos) {
+  if (pos != std::string::npos) {
     // Find the terminate point for flex op list.
     auto end = error_message.find(".", pos);
     pruned_message.append(error_message.substr(pos, end - pos + 1));
   }
   pos = error_message.find(s2);
-  if (pos != string::npos) {
+  if (pos != std::string::npos) {
     // Find the terminate point for custom op list.
     auto end = error_message.find(".", pos);
     pruned_message.append(error_message.substr(pos, end - pos + 1));
@@ -225,18 +228,18 @@ string SanitizeErrorMessage(const string& error_message) {
 
 void PopulateConversionLog(const Model& model, TocoConversionLog* log) {
   // Get the list of ops after conversion.
-  const std::vector op_names = GetOperatorNames(model);
+  const std::vector op_names = GetOperatorNames(model);
   for (const auto& op_name : op_names) {
     log->add_op_list(op_name);
   }
 
   // Get op signatures.
-  TFLITE_PROTO_NS::RepeatedPtrField op_signatures;
+  TFLITE_PROTO_NS::RepeatedPtrField op_signatures;
   GetOpSignatures(model, &op_signatures);
   log->mutable_op_signatures()->CopyFrom(op_signatures);
 
   // Get op counts by category: custom, built-in or select.
-  std::map custom_ops, select_ops, built_in_ops;
+  std::map custom_ops, select_ops, built_in_ops;
   CountOperatorsByType(model, &built_in_ops, &custom_ops, &select_ops);
   log->mutable_custom_ops()->insert(custom_ops.cbegin(), custom_ops.cend());
   log->mutable_built_in_ops()->insert(built_in_ops.cbegin(),
@@ -244,7 +247,7 @@ void PopulateConversionLog(const Model& model, TocoConversionLog* log) {
   log->mutable_select_ops()->insert(select_ops.cbegin(), select_ops.cend());
 
   // Get the model's input and output types.
-  TFLITE_PROTO_NS::RepeatedPtrField input_types, output_types;
+  TFLITE_PROTO_NS::RepeatedPtrField input_types, output_types;
   GetInputAndOutputTypes(model, &input_types, &output_types);
   log->mutable_input_tensor_types()->CopyFrom(input_types);
   log->mutable_output_tensor_types()->CopyFrom(output_types);
diff --git a/tensorflow/lite/toco/logging/conversion_log_util.h b/tensorflow/lite/toco/logging/conversion_log_util.h
index 2237615adbb..c21ec0792cc 100644
--- a/tensorflow/lite/toco/logging/conversion_log_util.h
+++ b/tensorflow/lite/toco/logging/conversion_log_util.h
@@ -25,37 +25,39 @@ namespace toco {
 
 // This function scans through the error message string, extracts the part about
 // missing ops and prunes away all other information in the error info.
-string SanitizeErrorMessage(const string& error_message);
+std::string SanitizeErrorMessage(const std::string& error_message);
 
 // Populates the TocoConversionLog proto after analyzing the model.
 void PopulateConversionLog(const Model& model, TocoConversionLog* log);
 
 // Returns the names of the operators in the model.
-std::vector GetOperatorNames(const Model& model);
+std::vector GetOperatorNames(const Model& model);
 
 // Counts the number of different types of operators in the model:
 // Built-in ops, custom ops and select ops.
 // Each map is mapping from the name of the operator (such as 'Conv') to its
 // total number of occurrences in the model.
 void CountOperatorsByType(const Model& model,
-                          std::map* built_in_ops,
-                          std::map* custom_ops,
-                          std::map* select_ops);
+                          std::map* built_in_ops,
+                          std::map* custom_ops,
+                          std::map* select_ops);
 
 // Gets the input and output types of the model. The input and output is
 // specified by model.flags.input_arrays and model.flags.output_arrays.
 void GetInputAndOutputTypes(
-    const Model& model, TFLITE_PROTO_NS::RepeatedPtrField* input_types,
-    TFLITE_PROTO_NS::RepeatedPtrField* output_types);
+    const Model& model,
+    TFLITE_PROTO_NS::RepeatedPtrField* input_types,
+    TFLITE_PROTO_NS::RepeatedPtrField* output_types);
 
 // Calculates signatures for all the ops in the model. An op signature is
 // defined by its input/output shapes and types, op name and its version.
-void GetOpSignatures(const Model& model,
-                     TFLITE_PROTO_NS::RepeatedPtrField* op_signatures);
+void GetOpSignatures(
+    const Model& model,
+    TFLITE_PROTO_NS::RepeatedPtrField* op_signatures);
 
 // TODO(b/123519920): Implement this.
 // Calculates a unique hash for the model.
-string GetModelHash(const Model& model);
+std::string GetModelHash(const Model& model);
 
 }  // namespace toco
 
diff --git a/tensorflow/lite/toco/logging/conversion_log_util_test.cc b/tensorflow/lite/toco/logging/conversion_log_util_test.cc
index c4960715f25..17111eca6d0 100644
--- a/tensorflow/lite/toco/logging/conversion_log_util_test.cc
+++ b/tensorflow/lite/toco/logging/conversion_log_util_test.cc
@@ -58,9 +58,9 @@ TEST(ConversionLogUtilTest, TestCountOperatorsByType) {
   Model model;
   // 1st Conv operator.
   std::unique_ptr conv1(new ConvOperator());
-  const string conv1_input_name = "conv_input1";
-  const string conv1_filter_name = "conv_filter1";
-  const string conv1_output_name = "conv_output1";
+  const std::string conv1_input_name = "conv_input1";
+  const std::string conv1_filter_name = "conv_filter1";
+  const std::string conv1_output_name = "conv_output1";
   conv1->inputs.push_back(conv1_input_name);
   conv1->inputs.push_back(conv1_filter_name);
   conv1->outputs.push_back(conv1_output_name);
@@ -71,9 +71,9 @@ TEST(ConversionLogUtilTest, TestCountOperatorsByType) {
 
   // 2nd Conv operator.
   std::unique_ptr conv2(new ConvOperator());
-  const string conv2_input_name = "conv_input2";
-  const string conv2_filter_name = "conv_filter2";
-  const string conv2_output_name = "conv_output2";
+  const std::string conv2_input_name = "conv_input2";
+  const std::string conv2_filter_name = "conv_filter2";
+  const std::string conv2_output_name = "conv_output2";
   conv2->inputs.push_back(conv2_input_name);
   conv2->inputs.push_back(conv2_filter_name);
   conv2->outputs.push_back(conv2_output_name);
@@ -83,7 +83,7 @@ TEST(ConversionLogUtilTest, TestCountOperatorsByType) {
 
   // Mean operator.
   std::unique_ptr mean(new MeanOperator());
-  const string mean_input_name = "mean_input";
+  const std::string mean_input_name = "mean_input";
   mean->inputs.push_back(mean_input_name);
   array_map[mean_input_name] = std::unique_ptr(new Array);
 
@@ -111,26 +111,26 @@ TEST(ConversionLogUtilTest, TestCountOperatorsByType) {
   model.operators.push_back(std::move(elu_grad));
   model.operators.push_back(std::move(my_custom_op));
 
-  std::map built_in_ops, select_ops, custom_ops;
+  std::map built_in_ops, select_ops, custom_ops;
   CountOperatorsByType(model, &built_in_ops, &custom_ops, &select_ops);
 
   EXPECT_THAT(built_in_ops,
-              UnorderedElementsAre(std::pair("Conv", 2),
-                                   std::pair("Mean", 1)));
+              UnorderedElementsAre(std::pair("Conv", 2),
+                                   std::pair("Mean", 1)));
   EXPECT_THAT(select_ops,
-              UnorderedElementsAre(std::pair("AvgPool3D", 1),
-                                   std::pair("EluGrad", 1)));
-  EXPECT_THAT(custom_ops, UnorderedElementsAre(
-                              std::pair("MyAwesomeCustomOp", 1)));
+              UnorderedElementsAre(std::pair("AvgPool3D", 1),
+                                   std::pair("EluGrad", 1)));
+  EXPECT_THAT(custom_ops, UnorderedElementsAre(std::pair(
+                              "MyAwesomeCustomOp", 1)));
 }
 
 TEST(ConversionLogUtilTest, TestGetInputAndOutputTypes) {
   Model model;
   auto& array_map = model.GetMutableArrayMap();
-  const string input1 = "conv_input";
-  const string input2 = "conv_filter";
-  const string input3 = "feature";
-  const string output = "softmax";
+  const std::string input1 = "conv_input";
+  const std::string input2 = "conv_filter";
+  const std::string input3 = "feature";
+  const std::string output = "softmax";
   array_map[input1] = std::unique_ptr(new Array);
   array_map[input1]->data_type = ArrayDataType::kFloat;
   array_map[input2] = std::unique_ptr(new Array);
@@ -149,7 +149,7 @@ TEST(ConversionLogUtilTest, TestGetInputAndOutputTypes) {
   *model.flags.add_input_arrays() = input_arrays[2];
   model.flags.add_output_arrays(output);
 
-  TFLITE_PROTO_NS::RepeatedPtrField input_types, output_types;
+  TFLITE_PROTO_NS::RepeatedPtrField input_types, output_types;
   GetInputAndOutputTypes(model, &input_types, &output_types);
 
   EXPECT_THAT(input_types, ElementsAre("float", "float", "int16"));
@@ -161,9 +161,9 @@ TEST(ConversionLogUtilTest, TestGetOpSignatures) {
   auto& array_map = model.GetMutableArrayMap();
 
   std::unique_ptr conv(new ConvOperator());
-  const string conv_input_name = "conv_input";
-  const string conv_filter_name = "conv_filter";
-  const string conv_output_name = "conv_output";
+  const std::string conv_input_name = "conv_input";
+  const std::string conv_filter_name = "conv_filter";
+  const std::string conv_output_name = "conv_output";
   conv->inputs.push_back(conv_input_name);
   conv->inputs.push_back(conv_filter_name);
   conv->outputs.push_back(conv_output_name);
@@ -177,15 +177,15 @@ TEST(ConversionLogUtilTest, TestGetOpSignatures) {
   array_map[conv_output_name]->data_type = ArrayDataType::kFloat;
   array_map[conv_output_name]->copy_shape({4, 4, 2});
 
-  const string mean_input_name = "mean_input";
-  const string mean_output_name = "mean_output";
+  const std::string mean_input_name = "mean_input";
+  const std::string mean_output_name = "mean_output";
   std::unique_ptr mean(new MeanOperator());
   mean->inputs.push_back(mean_input_name);
   mean->outputs.push_back(mean_output_name);
   array_map[mean_input_name] = std::unique_ptr(new Array);
   array_map[mean_output_name] = std::unique_ptr(new Array);
 
-  const string avg_pool_3d_output_name = "avg_pool_output";
+  const std::string avg_pool_3d_output_name = "avg_pool_output";
   auto avg_pool_3d = absl::make_unique();
   avg_pool_3d->tensorflow_op = "AvgPool3D";
   tensorflow::NodeDef node_def;
@@ -197,7 +197,7 @@ TEST(ConversionLogUtilTest, TestGetOpSignatures) {
   array_map[avg_pool_3d_output_name]->data_type = ArrayDataType::kInt32;
   array_map[avg_pool_3d_output_name]->copy_shape({2, 2});
 
-  const string custom_op_output_name = "custom_op_output";
+  const std::string custom_op_output_name = "custom_op_output";
   auto my_custom_op = absl::make_unique();
   my_custom_op->tensorflow_op = "MyAwesomeCustomOp";
   my_custom_op->inputs.push_back(avg_pool_3d_output_name);
@@ -211,7 +211,7 @@ TEST(ConversionLogUtilTest, TestGetOpSignatures) {
   model.operators.push_back(std::move(avg_pool_3d));
   model.operators.push_back(std::move(my_custom_op));
 
-  TFLITE_PROTO_NS::RepeatedPtrField op_signatures;
+  TFLITE_PROTO_NS::RepeatedPtrField op_signatures;
   GetOpSignatures(model, &op_signatures);
   EXPECT_THAT(op_signatures,
               UnorderedElementsAre(
@@ -225,14 +225,14 @@ TEST(ConversionLogUtilTest, TestGetOpSignatures) {
 }
 
 TEST(ConversionLogUtilTest, TestSanitizeErrorMessage) {
-  const string error =
+  const std::string error =
       "error: failed while converting: 'main': Ops that can be supported by "
       "the flex runtime (enabled via setting the -emit-select-tf-ops flag): "
       "ResizeNearestNeighbor,ResizeNearestNeighbor. Ops that need custom "
       "implementation (enabled via setting the -emit-custom-ops flag): "
       "CombinedNonMaxSuppression.\nTraceback (most recent call last): File "
       "/usr/local/bin/toco_from_protos, line 8, in ";
-  const string pruned_error =
+  const std::string pruned_error =
       "Ops that can be supported by "
       "the flex runtime (enabled via setting the -emit-select-tf-ops flag): "
       "ResizeNearestNeighbor,ResizeNearestNeighbor.Ops that need custom "
@@ -242,7 +242,7 @@ TEST(ConversionLogUtilTest, TestSanitizeErrorMessage) {
 }
 
 TEST(ConversionLogUtilTest, TestSanitizeErrorMessageNoMatching) {
-  const string error =
+  const std::string error =
       "error: failed while converting: 'main': Traceback (most recent call "
       "last): File "
       "/usr/local/bin/toco_from_protos, line 8, in ";

From c31f2ca4a29d4469f29c57735d74cbd2748c0e03 Mon Sep 17 00:00:00 2001
From: Sam Holt 
Date: Tue, 16 Jun 2020 16:50:16 +0100
Subject: [PATCH 0458/1390] fix: convolutional padding argument valid and same
 explaination

---
 .../python/keras/layers/convolutional.py      | 55 +++++++++++++++----
 .../keras/layers/convolutional_recurrent.py   |  6 ++
 tensorflow/python/keras/layers/local.py       |  2 +
 tensorflow/python/keras/layers/pooling.py     | 22 ++++++--
 .../keras/legacy_tf_layers/convolutional.py   | 42 ++++++++++++++
 tensorflow/python/keras/utils/conv_utils.py   | 12 ++++
 tensorflow/python/ops/nn_ops.py               |  3 +
 7 files changed, 128 insertions(+), 14 deletions(-)

diff --git a/tensorflow/python/keras/layers/convolutional.py b/tensorflow/python/keras/layers/convolutional.py
index 51f4e3b320a..471d94570a5 100644
--- a/tensorflow/python/keras/layers/convolutional.py
+++ b/tensorflow/python/keras/layers/convolutional.py
@@ -72,6 +72,10 @@ class Conv(Layer):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"`,  `"same"`, or `"causal"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input. `"causal"` results in causal 
+      (dilated) convolutions, e.g. `output[t]` does not depend on `input[t+1:]`.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -418,6 +422,9 @@ class Conv1D(Conv):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"`, `"causal"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
       `"causal"` results in causal (dilated) convolutions, e.g. `output[t]`
       does not depend on `input[t+1:]`. Useful when modeling temporal data
       where the model should not violate the temporal order.
@@ -571,6 +578,9 @@ class Conv2D(Conv):
       specify the same value for all spatial dimensions. Specifying any stride
       value != 1 is incompatible with specifying any `dilation_rate` value != 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs. `channels_last` corresponds
       to inputs with shape `(batch_size, height, width, channels)` while
@@ -712,6 +722,9 @@ class Conv3D(Conv):
       specify the same value for all spatial dimensions. Specifying any stride
       value != 1 is incompatible with specifying any `dilation_rate` value != 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs. `channels_last` corresponds
       to inputs with shape `batch_shape + (spatial_dim1, spatial_dim2,
@@ -833,6 +846,9 @@ class Conv1DTranspose(Conv1D):
       time dimension. Specifying a stride value != 1 is incompatible with
       specifying a `dilation_rate` value != 1. Defaults to 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     output_padding: An integer specifying the amount of padding along
       the time dimension of the output tensor.
       The amount of output padding must be lower than the stride.
@@ -1083,6 +1099,9 @@ class Conv2DTranspose(Conv2D):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     output_padding: An integer or tuple/list of 2 integers,
       specifying the amount of padding along the height and width
       of the output tensor.
@@ -1371,19 +1390,22 @@ class Conv3DTranspose(Conv3D):
 
   Arguments:
     filters: Integer, the dimensionality of the output space
-        (i.e. the number of output filters in the convolution).
+      (i.e. the number of output filters in the convolution).
     kernel_size: An integer or tuple/list of 3 integers, specifying the
-        depth, height and width of the 3D convolution window.
-        Can be a single integer to specify the same value for
-        all spatial dimensions.
+      depth, height and width of the 3D convolution window.
+      Can be a single integer to specify the same value for
+      all spatial dimensions.
     strides: An integer or tuple/list of 3 integers,
-        specifying the strides of the convolution along the depth, height
-          and width.
-        Can be a single integer to specify the same value for
-        all spatial dimensions.
-        Specifying any stride value != 1 is incompatible with specifying
-        any `dilation_rate` value != 1.
+      specifying the strides of the convolution along the depth, height
+        and width.
+      Can be a single integer to specify the same value for
+      all spatial dimensions.
+      Specifying any stride value != 1 is incompatible with specifying
+      any `dilation_rate` value != 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     output_padding: An integer or tuple/list of 3 integers,
       specifying the amount of padding along the depth, height, and
       width.
@@ -1681,6 +1703,9 @@ class SeparableConv(Conv):
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1885,6 +1910,10 @@ class SeparableConv1D(SeparableConv):
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"`, `"same"`, or `"causal"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input. `"causal"` results in causal 
+      (dilated) convolutions, e.g. `output[t]` does not depend on `input[t+1:]`.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -2070,6 +2099,9 @@ class SeparableConv2D(SeparableConv):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -2230,6 +2262,9 @@ class DepthwiseConv2D(Conv2D):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: one of `'valid'` or `'same'` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     depth_multiplier: The number of depthwise convolution output channels
       for each input channel.
       The total number of depthwise convolution output
diff --git a/tensorflow/python/keras/layers/convolutional_recurrent.py b/tensorflow/python/keras/layers/convolutional_recurrent.py
index 19831429b73..54196f8725c 100644
--- a/tensorflow/python/keras/layers/convolutional_recurrent.py
+++ b/tensorflow/python/keras/layers/convolutional_recurrent.py
@@ -434,6 +434,9 @@ class ConvLSTM2DCell(DropoutRNNCellMixin, Layer):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       It defaults to the `image_data_format` value found in your
@@ -710,6 +713,9 @@ class ConvLSTM2D(ConvRNN2D):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
diff --git a/tensorflow/python/keras/layers/local.py b/tensorflow/python/keras/layers/local.py
index 3e9c0f9c0a3..c33c88f3a3d 100644
--- a/tensorflow/python/keras/layers/local.py
+++ b/tensorflow/python/keras/layers/local.py
@@ -67,6 +67,7 @@ class LocallyConnected1D(Layer):
           any `dilation_rate` value != 1.
       padding: Currently only supports `"valid"` (case-insensitive).
           `"same"` may be supported in the future.
+          `"valid"` means no padding.
       data_format: A string,
           one of `channels_last` (default) or `channels_first`.
           The ordering of the dimensions in the inputs.
@@ -375,6 +376,7 @@ class LocallyConnected2D(Layer):
           all spatial dimensions.
       padding: Currently only support `"valid"` (case-insensitive).
           `"same"` will be supported in future.
+          `"valid"` means no padding.
       data_format: A string,
           one of `channels_last` (default) or `channels_first`.
           The ordering of the dimensions in the inputs.
diff --git a/tensorflow/python/keras/layers/pooling.py b/tensorflow/python/keras/layers/pooling.py
index ff7d157acad..51dc5131a8a 100644
--- a/tensorflow/python/keras/layers/pooling.py
+++ b/tensorflow/python/keras/layers/pooling.py
@@ -164,8 +164,9 @@ class MaxPooling1D(Pooling1D):
       for each pooling step.
       If None, it will default to `pool_size`.
     padding: One of `"valid"` or `"same"` (case-insensitive).
-      "valid" adds no padding.  "same" adds padding such that if the stride
-      is 1, the output shape is the same as the input shape.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -209,6 +210,9 @@ class AveragePooling1D(Pooling1D):
       E.g. 2 will halve the input.
       If None, it will default to `pool_size`.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -419,8 +423,9 @@ class MaxPooling2D(Pooling2D):
       Strides values.  Specifies how far the pooling window moves
       for each pooling step. If None, it will default to `pool_size`.
     padding: One of `"valid"` or `"same"` (case-insensitive).
-      "valid" adds no zero padding.  "same" adds padding such that if the stride
-      is 1, the output shape is the same as input shape.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -475,6 +480,9 @@ class AveragePooling2D(Pooling2D):
       Strides values.
       If None, it will default to `pool_size`.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -617,6 +625,9 @@ class MaxPooling3D(Pooling3D):
       `(2, 2, 2)` will halve the size of the 3D input in each dimension.
     strides: tuple of 3 integers, or None. Strides values.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
@@ -667,6 +678,9 @@ class AveragePooling3D(Pooling3D):
       `(2, 2, 2)` will halve the size of the 3D input in each dimension.
     strides: tuple of 3 integers, or None. Strides values.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string,
       one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
diff --git a/tensorflow/python/keras/legacy_tf_layers/convolutional.py b/tensorflow/python/keras/legacy_tf_layers/convolutional.py
index 4c91251a0e7..4fd53531fd1 100644
--- a/tensorflow/python/keras/legacy_tf_layers/convolutional.py
+++ b/tensorflow/python/keras/legacy_tf_layers/convolutional.py
@@ -46,6 +46,9 @@ class Conv1D(keras_layers.Conv1D, base.Layer):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -157,6 +160,9 @@ def conv1d(inputs,
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -242,6 +248,9 @@ class Conv2D(keras_layers.Conv2D, base.Layer):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -360,6 +369,9 @@ def conv2d(inputs,
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -449,6 +461,9 @@ class Conv3D(keras_layers.Conv3D, base.Layer):
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -568,6 +583,9 @@ def conv3d(inputs,
       Specifying any stride value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -652,6 +670,9 @@ class SeparableConv1D(keras_layers.SeparableConv1D, base.Layer):
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -761,6 +782,9 @@ class SeparableConv2D(keras_layers.SeparableConv2D, base.Layer):
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -897,6 +921,9 @@ def separable_conv1d(inputs,
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1019,6 +1046,9 @@ def separable_conv2d(inputs,
       Specifying any `stride` value != 1 is incompatible with specifying
       any `dilation_rate` value != 1.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1117,6 +1147,9 @@ class Conv2DTranspose(keras_layers.Conv2DTranspose, base.Layer):
       of the convolution. Can be a single integer to specify the same value for
       all spatial dimensions.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1223,6 +1256,9 @@ def conv2d_transpose(inputs,
       of the convolution. Can be a single integer to specify the same value for
       all spatial dimensions.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1295,6 +1331,9 @@ class Conv3DTranspose(keras_layers.Conv3DTranspose, base.Layer):
       Can be a single integer to specify the same value for all spatial
       dimensions.
     padding: One of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
@@ -1396,6 +1435,9 @@ def conv3d_transpose(inputs,
       of the convolution. Can be a single integer to specify the same value for
       all spatial dimensions.
     padding: one of `"valid"` or `"same"` (case-insensitive).
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     data_format: A string, one of `channels_last` (default) or `channels_first`.
       The ordering of the dimensions in the inputs.
       `channels_last` corresponds to inputs with shape
diff --git a/tensorflow/python/keras/utils/conv_utils.py b/tensorflow/python/keras/utils/conv_utils.py
index f38fdc18252..e8ee866d958 100644
--- a/tensorflow/python/keras/utils/conv_utils.py
+++ b/tensorflow/python/keras/utils/conv_utils.py
@@ -264,6 +264,9 @@ def conv_kernel_mask(input_shape, kernel_shape, strides, padding):
       receptive field.
     strides: tuple of size N, strides along each spatial dimension.
     padding: type of padding, string `"same"` or `"valid"`.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
 
   Returns:
     A boolean 2N-D `np.ndarray` of shape
@@ -338,6 +341,9 @@ def conv_kernel_idxs(input_shape, kernel_shape, strides, padding, filters_in,
       receptive field.
     strides: tuple of size N, strides along each spatial dimension.
     padding: type of padding, string `"same"` or `"valid"`.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     filters_in: `int`, number if filters in the input to the layer.
     filters_out: `int', number if filters in the output of the layer.
     data_format: string, "channels_first" or "channels_last".
@@ -430,6 +436,9 @@ def conv_connected_inputs(input_shape, kernel_shape, output_position, strides,
       in the output of the convolution.
     strides: tuple of size N, strides along each spatial dimension.
     padding: type of padding, string `"same"` or `"valid"`.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
 
   Returns:
     N ranges `[[p_in_left1, ..., p_in_right1], ...,
@@ -468,6 +477,9 @@ def conv_output_shape(input_shape, kernel_shape, strides, padding):
       receptive field.
     strides: tuple of size N, strides along each spatial dimension.
     padding: type of padding, string `"same"` or `"valid"`.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
 
   Returns:
     tuple of size N: `(d_out1, ..., d_outN)`, spatial shape of the output.
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 1318f575737..5a9a63637f6 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -940,6 +940,9 @@ def convolution(
     filter: An (N+2)-D `Tensor` with the same type as `input` and shape
       `spatial_filter_shape + [in_channels, out_channels]`.
     padding: A string, either `"VALID"` or `"SAME"`. The padding algorithm.
+      `"valid"` means no padding. `"same"` results in padding evenly to 
+      the left/right or up/down of the input such that output has the same 
+      height/width dimension as the input.
     strides: Optional.  Sequence of N ints >= 1.  Specifies the output stride.
       Defaults to [1]*N.  If any value of strides is > 1, then all values of
       dilation_rate must be 1.

From c0ba8a09a7edf2b617d10024f7e0ec9abe21d73e Mon Sep 17 00:00:00 2001
From: Xinyi Wang 
Date: Wed, 17 Jun 2020 18:24:26 -0700
Subject: [PATCH 0459/1390] Export and Document DistributedDataset and
 DistributedIterator APIs

PiperOrigin-RevId: 317007583
Change-Id: I7d7c4615a12a19fb4fd151a0457f176ffe2cd765
---
 .../python/distribute/distribute_lib.py       | 120 ++++---
 tensorflow/python/distribute/input_lib.py     | 324 +++++++++++++++++-
 tensorflow/python/distribute/values.py        |  17 +-
 ...flow.distribute.-distributed-dataset.pbtxt |  16 +
 ...low.distribute.-distributed-iterator.pbtxt |  16 +
 .../api/golden/v2/tensorflow.distribute.pbtxt |   8 +
 6 files changed, 444 insertions(+), 57 deletions(-)
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-dataset.pbtxt
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt

diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index a6dc35507e9..b6a89463426 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -684,7 +684,8 @@ class StrategyBase(object):
         instead.
       * Use `tf.distribute.Strategy.run` to run a function
         once per replica, taking values that may be "per-replica" (e.g.
-        from a distributed dataset) and returning "per-replica" values.
+        from a `tf.distribute.DistributedDataset` object) and returning
+        "per-replica" values.
         This function is executed in "replica context", which means each
         operation is performed separately on each replica.
       * Finally use a method (such as `tf.distribute.Strategy.reduce`) to
@@ -720,7 +721,8 @@ class StrategyBase(object):
   distributed-specific behavior.
 
   You can use the `reduce` API to aggregate results across replicas and use
-  this as a return value from one iteration over the distributed dataset. Or
+  this as a return value from one iteration over a
+  `tf.distribute.DistributedDataset`. Or
   you can use `tf.keras.metrics` (such as loss, accuracy, etc.) to
   accumulate metrics across steps in a given epoch.
 
@@ -859,12 +861,12 @@ class StrategyBase(object):
     return self.run(fn, args=args)
 
   def experimental_distribute_dataset(self, dataset, options=None):
-    """Distributes a tf.data.Dataset instance provided via `dataset`.
+    """Creates `tf.distribute.DistributedDataset` from `tf.data.Dataset`.
 
-    The returned distributed dataset can be iterated over similar to how
-    regular datasets can.
-    NOTE: Currently, the user cannot add any more transformations to a
-    distributed dataset.
+    The returned `tf.distribute.DistributedDataset` can be iterated over
+    similar to how regular datasets can.
+    NOTE: The user cannot add any more transformations to a
+    `tf.distribute.DistributedDataset`.
 
     The following is an example:
 
@@ -878,48 +880,53 @@ class StrategyBase(object):
     # Distribute that dataset
     dist_dataset = strategy.experimental_distribute_dataset(dataset)
 
-    # Iterate over the distributed dataset
+    # Iterate over the `tf.distribute.DistributedDataset`
     for x in dist_dataset:
       # process dataset elements
       strategy.run(replica_fn, args=(x,))
     ```
 
-    In the code snippet above, the dataset `dist_dataset` is batched by
-    GLOBAL_BATCH_SIZE, and we iterate through it using `for x in dist_dataset`,
-    where x is one batch of data of GLOBAL_BATCH_SIZE containing N batches of
-    data of per-replica batch size, corresponding to N replicas.
-    `tf.distribute.Strategy.run` will take care of feeding
-    the right per-replica batch to the right `replica_fn` execution on each
+    In the code snippet above, the `tf.distribute.DistributedDataset`
+    `dist_dataset` is batched by `GLOBAL_BATCH_SIZE`, and we iterate through it
+    using `for x in dist_dataset`. `x` a `tf.distribute.DistributedValues`
+    containing data for all replicas, which aggregates to a batch of
+    `GLOBAL_BATCH_SIZE`. `tf.distribute.Strategy.run` will take care of feeding
+    the right per-replica data in `x` to the right `replica_fn` executed on each
     replica.
 
-    In a multi-worker setting, we will first attempt to distribute the dataset
-    by attempting to detect whether the dataset is being created out of
-    ReaderDatasets (e.g. TFRecordDataset, TextLineDataset, etc.) and if so,
-    attempting to shard the input files. Note that there has to be at least one
-    input file per worker. If you have less than one input file per worker, we
-    suggest that you should disable distributing your dataset using the method
-    below.
+    What's under the hood of this method, when we say the `tf.data.Dataset`
+    instance - `dataset` - gets distributed? It depends on how you set the
+    `tf.data.experimental.AutoShardPolicy` through
+    `tf.data.experimental.DistributeOptions`. By default, it is set to
+    `tf.data.experimental.AutoShardPolicy.AUTO`. In a multi-worker setting, we
+    will first attempt to distribute `dataset` by detecting whether `dataset` is
+    being created out of reader datasets (e.g. `tf.data.TFRecordDataset`,
+    `tf.data.TextLineDataset`, etc.) and if so, try to shard the input files.
+    Note that there has to be at least one input file per worker. If you have
+    less than one input file per worker, we suggest that you disable dataset
+    sharding across workers, by setting the
+    `tf.data.experimental.DistributeOptions.auto_shard_policy` to be
+    `tf.data.experimental.AutoShardPolicy.OFF`.
 
-    If that attempt is unsuccessful (e.g. the dataset is created from a
-    Dataset.range), we will shard the dataset evenly at the end by appending a
-    `.shard` operation to the end of the processing pipeline. This will cause
-    the entire preprocessing pipeline for all the data to be run on every
-    worker, and each worker will do redundant work. We will print a warning
-    if this method of sharding is selected.
+    If the attempt to shard by file is unsuccessful (i.e. the dataset is not
+    read from files), we will shard the dataset evenly at the end by
+    appending a `.shard` operation to the end of the processing pipeline. This
+    will cause the entire preprocessing pipeline for all the data to be run on
+    every worker, and each worker will do redundant work. We will print a
+    warning if this route is selected.
 
-    You can disable dataset sharding across workers using the
-    `auto_shard_policy` option in `tf.data.experimental.DistributeOptions`.
-
-    Within each worker, we will also split the data among all the worker
-    devices (if more than one a present), and this will happen even if
-    multi-worker sharding is disabled using the method above.
+    As mentioned before, within each worker, we will also split the data among
+    all the worker devices (if more than one a present). This will happen
+    even if multi-worker sharding is disabled.
 
     If the above batch splitting and dataset sharding logic is undesirable,
-    please use `experimental_distribute_datasets_from_function` instead, which
-    does not do any automatic splitting or sharding.
+    please use
+    `tf.distribute.Strategy.experimental_distribute_datasets_from_function`
+    instead, which does not do any automatic splitting or sharding.
 
-    You can also use the `element_spec` property of the distributed dataset
-    returned by this API to query the `tf.TypeSpec` of the elements returned
+    You can also use the `element_spec` property of the
+    `tf.distribute.DistributedDataset` instance returned by this API to query
+    the `tf.TypeSpec` of the elements returned
     by the iterator. This can be used to set the `input_signature` property
     of a `tf.function`.
 
@@ -938,12 +945,21 @@ class StrategyBase(object):
       # train model with inputs
       return
 
-    # Iterate over the distributed dataset
+    # Iterate over the `tf.distribute.DistributedDataset`
     for x in dist_dataset:
       # process dataset elements
       strategy.run(train_step, args=(x,))
     ```
 
+    Note: The order in which the data is processed by the workers when using
+    `tf.distribute.Strategy.experimental_distribute_dataset` or
+    `tf.distribute.Strategy.experimental_distribute_datasets_from_function` is
+    not guaranteed. This is typically required if you are using
+    `tf.distribute` to scale prediction. You can however insert an index for
+    each element in the batch and order outputs accordingly. Refer to [this
+    snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+    for an example of how to order outputs.
+
     Args:
       dataset: `tf.data.Dataset` that will be sharded across all replicas using
         the rules stated above.
@@ -951,8 +967,7 @@ class StrategyBase(object):
         dataset is distributed.
 
     Returns:
-      A "distributed `Dataset`", which acts like a `tf.data.Dataset` except
-      it produces "per-replica" values.
+      A `tf.distribute.DistributedDataset`.
     """
     return self._extended._experimental_distribute_dataset(dataset, options)  # pylint: disable=protected-access
 
@@ -978,10 +993,10 @@ class StrategyBase(object):
     The `dataset_fn` should take an `tf.distribute.InputContext` instance where
     information about batching and input replication can be accessed.
 
-    You can also use the `element_spec` property of the distributed dataset
-    returned by this API to query the `tf.TypeSpec` of the elements returned
-    by the iterator. This can be used to set the `input_signature` property
-    of a `tf.function`.
+    You can also use the `element_spec` property of the
+    `tf.distribute.DistributedDataset` returned by this API to query the
+    `tf.TypeSpec` of the elements returned by the iterator. This can be used to
+    set the `input_signature` property of a `tf.function`.
 
     >>> global_batch_size = 8
     >>> def dataset_fn(input_context):
@@ -1010,6 +1025,16 @@ class StrategyBase(object):
     the global batch size.  This may be computed using
     `input_context.get_per_replica_batch_size`.
 
+
+    Note: The order in which the data is processed by the workers when using
+    `tf.distribute.Strategy.experimental_distribute_dataset` or
+    `tf.distribute.Strategy.experimental_distribute_datasets_from_function` is
+    not guaranteed. This is typically required if you are using
+    `tf.distribute` to scale prediction. You can however insert an index for
+    each element in the batch and order outputs accordingly. Refer to [this
+    snippet](https://www.tensorflow.org/tutorials/distribute/input#caveats)
+    for an example of how to order outputs.
+
     Args:
       dataset_fn: A function taking a `tf.distribute.InputContext` instance and
         returning a `tf.data.Dataset`.
@@ -1017,8 +1042,7 @@ class StrategyBase(object):
         dataset is distributed.
 
     Returns:
-      A "distributed `Dataset`", which acts like a `tf.data.Dataset` except
-      it produces "per-replica" values.
+      A `tf.distribute.DistributedDataset`.
     """
     return self._extended._experimental_distribute_datasets_from_function(  # pylint: disable=protected-access
         dataset_fn, options)
@@ -1028,7 +1052,9 @@ class StrategyBase(object):
 
     Executes ops specified by `fn` on each replica. If `args` or `kwargs` have
     `tf.distribute.DistributedValues`, such as those produced by a
-    "distributed `Dataset`" or `experimental_distribute_values_from_function`
+    `tf.distribute.DistributedDataset` from
+    `tf.distribute.Strategy.experimental_distribute_dataset` or
+    `tf.distribute.Strategy.experimental_distribute_datasets_from_function`,
     when `fn` is executed on a particular replica, it will be executed with the
     component of `tf.distribute.DistributedValues` that correspond to that
     replica.
diff --git a/tensorflow/python/distribute/input_lib.py b/tensorflow/python/distribute/input_lib.py
index 85e2dac1c1d..ff468af7f87 100644
--- a/tensorflow/python/distribute/input_lib.py
+++ b/tensorflow/python/distribute/input_lib.py
@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import functools
 import sys
 
@@ -52,6 +53,8 @@ from tensorflow.python.ops.ragged import ragged_tensor
 from tensorflow.python.types import distribute as distribute_types
 from tensorflow.python.util import nest
 from tensorflow.python.util.deprecation import deprecated
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.tools.docs import doc_controls
 
 
 def get_distributed_dataset(dataset,
@@ -138,6 +141,321 @@ def get_distributed_datasets_from_function(dataset_fn,
         strategy)
 
 
+@tf_export("distribute.DistributedIterator", v1=[])
+class DistributedIteratorInterface(collections.Iterator,
+                                   distribute_types.Iterator):
+  """An iterator over `tf.distribute.DistributedDataset`.
+
+  `tf.distribute.DistributedIterator` is the primary mechanism for enumerating
+  elements of a `tf.distribute.DistributedDataset`. It supports the Python
+  Iterator protocol, which means it can be iterated over using a for-loop or by
+  fetching individual elements explicitly via `get_next()`.
+
+  You can create a `tf.distribute.DistributedIterator` by calling `iter` on
+  a `tf.distribute.DistributedDataset` or creating a python loop over a
+  `tf.distribute.DistributedDataset`.
+
+  Visit the [tutorial](https://www.tensorflow.org/tutorials/distribute/input)
+  on distributed input for more examples and caveats.
+  """
+
+  def get_next(self):
+    """Returns the next input from the iterator for all replicas.
+
+    Example use:
+
+    >>> strategy = tf.distribute.MirroredStrategy()
+    >>> dataset = tf.data.Dataset.range(100).batch(2)
+    >>> dist_dataset = strategy.experimental_distribute_dataset(dataset)
+    >>> dist_dataset_iterator = iter(dist_dataset)
+    >>> @tf.function
+    ... def one_step(input):
+    ...   return input
+    >>> step_num = 5
+    >>> for _ in range(step_num):
+    ...   strategy.run(one_step, args=(dist_dataset_iterator.get_next(),))
+    >>> strategy.experimental_local_results(dist_dataset_iterator.get_next())
+    (,)
+
+    The above example corresponds to the case where you have only one device. If
+    you have two devices, for example,
+    ```python
+    strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])
+    ```
+    Then the final line will print out:
+    ```python
+    (,
+     )
+    ```
+
+    Returns:
+      A single `tf.Tensor` or a `tf.distribute.DistributedValues` which contains
+      the next input for all replicas.
+
+    Raises:
+      `tf.errors.OutOfRangeError`: If the end of the iterator has been reached.
+    """
+    raise NotImplementedError(
+        "DistributedIterator.get_next() must be implemented in descendants.")
+
+  @property
+  def element_spec(self):
+    # pylint: disable=line-too-long
+    """The type specification of an element of `tf.distribute.DistributedIterator`.
+
+    Example usage:
+
+    >>> global_batch_size = 16
+    >>> strategy = tf.distribute.MirroredStrategy()
+    >>> dataset = tf.data.Dataset.from_tensors(([1.],[2])).repeat(100).batch(global_batch_size)
+    >>> distributed_iterator = iter(strategy.experimental_distribute_dataset(dataset))
+    >>> distributed_iterator.element_spec
+    (TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
+     TensorSpec(shape=(None, 1), dtype=tf.int32, name=None))
+
+    The above example corresponds to the case where you have only one device. If
+    you have two devices, for example,
+    ```python
+    strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])
+    ```
+    Then the final line will print out:
+    ```python
+    (PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
+                    TensorSpec(shape=(None, 1), dtype=tf.float32, name=None)),
+     PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.int32, name=None),
+                    TensorSpec(shape=(None, 1), dtype=tf.int32, name=None)))
+    ```
+
+    Returns:
+      A nested structure of `tf.TypeSpec` objects matching the structure of an
+      element of this `tf.distribute.DistributedIterator`. This returned value
+      is typically a `tf.distribute.DistributedValues` object and specifies the
+      `tf.TensorSpec` of individual components.
+    """
+    raise NotImplementedError(
+        "DistributedIterator.element_spec() must be implemented in descendants")
+
+
+@tf_export("distribute.DistributedDataset", v1=[])
+class DistributedDatasetInterface(collections.Iterable,
+                                  distribute_types.Iterable):
+  # pylint: disable=line-too-long
+  """Represents a dataset distributed among devices and machines.
+
+  A `tf.distribute.DistributedDataset` could be thought of as a "distributed"
+  dataset. When you use `tf.distribute` API to scale training to multiple
+  devices or machines, you also need to distribute the input data, which leads
+  to a `tf.distribute.DistributedDataset` instance, instead of a
+  `tf.data.Dataset` instance in the non-distributed case. In TF 2.x,
+  `tf.distribute.DistributedDataset` objects are Python iterables.
+
+  Note: `tf.distribute.DistributedDataset` instances are *not* of type
+  `tf.data.Dataset`. It only supports two usages we will mention below:
+  iteration and `element_spec`. We don't support any other APIs to transform or
+  inspect the dataset.
+
+  There are two APIs to create a `tf.distribute.DistributedDataset` object:
+  `tf.distribute.Strategy.experimental_distribute_dataset(dataset)`and
+  `tf.distribute.Strategy.experimental_distribute_datasets_from_function(dataset_fn)`.
+  *When to use which?* When you have a `tf.data.Dataset` instance, and the
+  regular batch splitting (i.e. re-batch the input `tf.data.Dataset` instance
+  with a new batch size that is equal to the global batch size divided by the
+  number of replicas in sync) and autosharding (i.e. the
+  `tf.data.experimental.AutoShardPolicy` options) work for you, use the former
+  API. Otherwise, if you are *not* using a canonical `tf.data.Dataset` instance,
+  or you would like to customize the batch splitting or sharding, you can wrap
+  these logic in a `dataset_fn` and use the latter API. Both API handles
+  prefetch to device for the user. For more details and examples, follow the
+  links to the APIs.
+
+
+  There are two main usages of a `DistributedDataset` object:
+
+  1. Iterate over it to generate the input for a single device or multiple
+  devices, which is a `tf.distribute.DistributedValues` instance. To do this,
+  you can:
+
+    * use a pythonic for-loop construct:
+
+      >>> global_batch_size = 2
+      >>> strategy = tf.distribute.MirroredStrategy()
+      >>> dataset = tf.data.Dataset.from_tensors(([1.],[1.])).repeat(4).batch(global_batch_size)
+      >>> dist_dataset = strategy.experimental_distribute_dataset(dataset)
+      >>> @tf.function
+      ... def train_step(input):
+      ...   features, labels = input
+      ...   return labels - 0.3 * features
+      >>> for x in dist_dataset:
+      ...   # train_step trains the model using the dataset elements
+      ...   loss = strategy.run(train_step, args=(x,))
+      ...   print("Loss is", loss)
+      Loss is tf.Tensor(
+      [[0.7]
+       [0.7]], shape=(2, 1), dtype=float32)
+      Loss is tf.Tensor(
+      [[0.7]
+       [0.7]], shape=(2, 1), dtype=float32)
+
+      Placing the loop inside a `tf.function` will give a performance boost.
+      However `break` and `return` are currently not supported if the loop is
+      placed inside a `tf.function`. We also don't support placing the loop
+      inside a `tf.function` when using
+      `tf.distribute.experimental.MultiWorkerMirroredStrategy` or
+      `tf.distribute.experimental.TPUStrategy` with multiple workers.
+
+    * use `__iter__` to create an explicit iterator, which is of type
+      `tf.distribute.DistributedIterator`
+
+      >>> global_batch_size = 4
+      >>> strategy = tf.distribute.MirroredStrategy()
+      >>> train_dataset = tf.data.Dataset.from_tensors(([1.],[1.])).repeat(50).batch(global_batch_size)
+      >>> train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
+      >>> @tf.function
+      ... def distributed_train_step(dataset_inputs):
+      ...   def train_step(input):
+      ...     loss = tf.constant(0.1)
+      ...     return loss
+      ...   per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
+      ...   return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,axis=None)
+      >>> EPOCHS = 2
+      >>> STEPS = 3
+      >>> for epoch in range(EPOCHS):
+      ...   total_loss = 0.0
+      ...   num_batches = 0
+      ...   dist_dataset_iterator = iter(train_dist_dataset)
+      ...   for _ in range(STEPS):
+      ...     total_loss += distributed_train_step(next(dist_dataset_iterator))
+      ...     num_batches += 1
+      ...   average_train_loss = total_loss / num_batches
+      ...   template = ("Epoch {}, Loss: {}")
+      ...   print (template.format(epoch+1, average_train_loss))
+      Epoch 1, Loss: 0.10000000894069672
+      Epoch 2, Loss: 0.10000000894069672
+
+
+    To achieve a performance improvement, you can also wrap the `strategy.run`
+    call with a `tf.range` inside a `tf.function`. This runs multiple steps in a
+    `tf.function`. Autograph will convert it to a `tf.while_loop` on the worker.
+    However, it is less flexible comparing with running a single step inside
+    `tf.function`. For example, you cannot run things eagerly or arbitrary
+    python code within the steps.
+
+
+  2. Inspect the `tf.TypeSpec` of the data generated by `DistributedDataset`.
+
+    `tf.distribute.DistributedDataset` generates
+    `tf.distribute.DistributedValues` as input to the devices. If you pass the
+    input to a `tf.function` and would like to specify the shape and type of
+    each Tensor argument to the function, you can pass a `tf.TypeSpec` object to
+    the `input_signature` argument of the `tf.function`. To get the
+    `tf.TypeSpec` of the input, you can use the `element_spec` property of the
+    `tf.distribute.DistributedDataset` or `tf.distribute.DistributedIterator`
+    object.
+
+    For example:
+
+    >>> global_batch_size = 2
+    >>> epochs = 1
+    >>> steps_per_epoch = 1
+    >>> mirrored_strategy = tf.distribute.MirroredStrategy()
+    >>> dataset = tf.data.Dataset.from_tensors(([2.])).repeat(100).batch(global_batch_size)
+    >>> dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
+    >>> @tf.function(input_signature=[dist_dataset.element_spec])
+    ... def train_step(per_replica_inputs):
+    ...   def step_fn(inputs):
+    ...     return tf.square(inputs)
+    ...   return mirrored_strategy.run(step_fn, args=(per_replica_inputs,))
+    >>> for _ in range(epochs):
+    ...   iterator = iter(dist_dataset)
+    ...   for _ in range(steps_per_epoch):
+    ...     output = train_step(next(iterator))
+    ...     print(output)
+    tf.Tensor(
+    [[4.]
+     [4.]], shape=(2, 1), dtype=float32)
+
+
+  Visit the [tutorial](https://www.tensorflow.org/tutorials/distribute/input)
+  on distributed input for more examples and caveats.
+  """
+
+  def __iter__(self):
+    """Creates an iterator for the `tf.distribute.DistributedDataset`.
+
+    The returned iterator implements the Python Iterator protocol.
+
+    Example usage:
+
+    >>> global_batch_size = 4
+    >>> strategy = tf.distribute.MirroredStrategy()
+    >>> dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4]).repeat().batch(global_batch_size)
+    >>> distributed_iterator = iter(strategy.experimental_distribute_dataset(dataset))
+    >>> print(next(distributed_iterator))
+    tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)
+
+
+    The above example corresponds to the case where you have only one device. If
+    you have two devices, for example,
+    ```python
+    strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])
+    ```
+    Then the final line will print out:
+    ```python
+    PerReplica:{
+      0: tf.Tensor([1 2], shape=(2,), dtype=int32),
+      1: tf.Tensor([3 4], shape=(2,), dtype=int32)
+    }
+    ```
+
+    Returns:
+      An `tf.distribute.DistributedIterator` instance for the given
+      `tf.distribute.DistributedDataset` object to enumerate over the
+      distributed data.
+    """
+    raise NotImplementedError("Must be implemented in descendants")
+
+  @property
+  def element_spec(self):
+    """The type specification of an element of this `tf.distribute.DistributedDataset`.
+
+    Example usage:
+
+    >>> global_batch_size = 16
+    >>> strategy = tf.distribute.MirroredStrategy()
+    >>> dataset = tf.data.Dataset.from_tensors(([1.],[2])).repeat(100).batch(global_batch_size)
+    >>> dist_dataset = strategy.experimental_distribute_dataset(dataset)
+    >>> dist_dataset.element_spec
+    (TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
+     TensorSpec(shape=(None, 1), dtype=tf.int32, name=None))
+
+    The above example corresponds to the case where you have only one device. If
+    you have two devices, for example,
+    ```python
+    strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])
+    ```
+    Then the final line will print out:
+    ```python
+    (PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
+                    TensorSpec(shape=(None, 1), dtype=tf.float32, name=None)),
+     PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.int32, name=None),
+                    TensorSpec(shape=(None, 1), dtype=tf.int32, name=None)))
+    ```
+
+    Returns:
+      A nested structure of `tf.TypeSpec` objects matching the structure of an
+      element of this `tf.distribute.DistributedDataset`. This returned value is
+      typically a `tf.distribute.DistributedValues` object and specifies the
+      `tf.TensorSpec` of individual components.
+    """
+    raise NotImplementedError(
+        "DistributedDataset.element_spec must be implemented in descendants.")
+
+  @doc_controls.do_not_generate_docs
+  def reduce(self, initial_state, reduce_func):
+    raise NotImplementedError(
+        "DistributedDataset.reduce must be implemented in descendants.")
+
+
 class InputWorkers(object):
   """A 1-to-many mapping from input worker devices to compute devices."""
 
@@ -259,9 +577,10 @@ def _get_static_shape(iterators):
     return static_shape
 
 
-class DistributedIteratorBase(distribute_types.Iterator):
+class DistributedIteratorBase(DistributedIteratorInterface):
   """Common implementation for all input iterators."""
 
+  # pylint: disable=super-init-not-called
   def __init__(self, input_workers, iterators, strategy):
     static_shape = _get_static_shape(iterators)
 
@@ -548,9 +867,10 @@ class DistributedIterator(DistributedIteratorBase,
                                    self._strategy)
 
 
-class _IterableInput(distribute_types.Iterable):
+class _IterableInput(DistributedDatasetInterface):
   """Base class for iterable inputs for distribution strategies."""
 
+  # pylint: disable=super-init-not-called
   def __init__(self, input_workers):
     assert isinstance(input_workers, InputWorkers)
     self._input_workers = input_workers
diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index c6e0eb34a7b..d0ed27c69de 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -75,20 +75,21 @@ def _on_write_update_replica(var, update_fn, value, **kwargs):
 class DistributedValues(object):
   """Base class for representing distributed values.
 
-  A subclass instance of DistributedValues is created when creating variables
-  within a distribution strategy, iterating a `tf.Dataset` or through
-  `strategy.run`.  This base class should never be instantiated
-  directly.  DistributedValues contains a value per replica.  Depending on
+  A subclass instance of `tf.distribute.DistributedValues` is created when
+  creating variables within a distribution strategy, iterating a
+  `tf.distribute.DistributedDataset` or through `tf.distribute.Strategy.run`.
+  This base class should never be instantiated directly.
+  `tf.distribute.DistributedValues` contains a value per replica. Depending on
   the subclass, the values could either be synced on update, synced on demand,
   or never synced.
 
-  DistributedValues can be reduced to obtain single value across replicas,
-  as input into `run` or the per replica values inspected
-  using `experimental_local_results`.
+  `tf.distribute.DistributedValues` can be reduced to obtain single value across
+  replicas, as input into `tf.distribute.Strategy.run` or the per-replica values
+  inspected using `tf.distribute.Strategy.experimental_local_results`.
 
   Example usage:
 
-  1. Created from Dataset:
+  1. Created from a `tf.distribute.DistributedDataset`:
 
   >>> strategy = tf.distribute.MirroredStrategy()
   >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2)
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-dataset.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-dataset.pbtxt
new file mode 100644
index 00000000000..a7b229c6c7c
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-dataset.pbtxt
@@ -0,0 +1,16 @@
+path: "tensorflow.distribute.DistributedDataset"
+tf_class {
+  is_instance: ""
+  is_instance: ""
+  member {
+    name: "element_spec"
+    mtype: ""
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "reduce"
+    argspec: "args=[\'self\', \'initial_state\', \'reduce_func\'], varargs=None, keywords=None, defaults=None"
+  }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
new file mode 100644
index 00000000000..f712d9058b9
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
@@ -0,0 +1,16 @@
+path: "tensorflow.distribute.DistributedIterator"
+tf_class {
+  is_instance: ""
+  is_instance: ""
+  member {
+    name: "element_spec"
+    mtype: ""
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "get_next"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
index 19d83909120..009cb7fe400 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.pbtxt
@@ -4,6 +4,14 @@ tf_module {
     name: "CrossDeviceOps"
     mtype: ""
   }
+  member {
+    name: "DistributedDataset"
+    mtype: ""
+  }
+  member {
+    name: "DistributedIterator"
+    mtype: ""
+  }
   member {
     name: "DistributedValues"
     mtype: ""

From a50001edf9fea03069771934e94d3b3d32ff6a19 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 17 Jun 2020 18:30:26 -0700
Subject: [PATCH 0460/1390] Internal change

PiperOrigin-RevId: 317008433
Change-Id: I5146e28d2f77c7daab4bd023f2826ec8323cff02
---
 tensorflow/tensorflow.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index f97363a919e..5da15b0a4d6 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -1817,13 +1817,13 @@ def tf_custom_op_library_additional_deps_impl():
 # tf_collected_deps will be the union of the deps of the current target
 # and the tf_collected_deps of the dependencies of this target.
 def _collect_deps_aspect_impl(target, ctx):
-    alldeps = depset()
+    direct, transitive = [], []
     if hasattr(ctx.rule.attr, "deps"):
         for dep in ctx.rule.attr.deps:
-            alldeps = depset([dep.label], transitive = [alldeps])
+            direct.append(dep.label)
             if hasattr(dep, "tf_collected_deps"):
-                alldeps = depset(transitive = [alldeps, dep.tf_collected_deps])
-    return struct(tf_collected_deps = alldeps)
+                transitive.append(dep.tf_collected_deps)
+    return struct(tf_collected_deps = depset(direct = direct, transitive = transitive))
 
 collect_deps_aspect = aspect(
     attr_aspects = ["deps"],

From ae20f08da9ce9e7336ab97cc9f77ce7a1c13ad12 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Wed, 17 Jun 2020 18:43:41 -0700
Subject: [PATCH 0461/1390] Properly support nest phi reduction in reverse
 order.

If we replaced node B with node C, then replace node A with node B, we
should redirect node A to node C instead.

PiperOrigin-RevId: 317010443
Change-Id: I165496a3d1f6571815bfd61d096e26cbba39125a
---
 .../xla/service/hlo_dataflow_analysis.cc      |  2 ++
 .../compiler/xla/service/hlo_phi_graph.cc     | 25 ++++++++++++++++++-
 .../compiler/xla/service/hlo_phi_graph.h      |  2 +-
 .../xla/service/hlo_phi_graph_test.cc         | 25 +++++++++++++++++++
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
index f19882c9347..a46d20d5808 100644
--- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
+++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc
@@ -1007,6 +1007,8 @@ void HloDataflowAnalysis::OptimizePhiValues() {
             HloValue::Id phi_id = values[0]->id();
             HloValue::Id new_id = phi_graph_.FindOptimizedValue(phi_id);
             if (new_id != phi_id) {
+              VLOG(1) << "Replacing " << values[0]->ToString() << " with "
+                      << GetValue(new_id).ToString();
               value_set->Clear();
               const HloValue& new_value = GetValue(new_id);
               value_set->AddValue(&new_value);
diff --git a/tensorflow/compiler/xla/service/hlo_phi_graph.cc b/tensorflow/compiler/xla/service/hlo_phi_graph.cc
index 9b69771dab2..a2cba3d1bff 100644
--- a/tensorflow/compiler/xla/service/hlo_phi_graph.cc
+++ b/tensorflow/compiler/xla/service/hlo_phi_graph.cc
@@ -20,10 +20,11 @@ limitations under the License.
 namespace xla {
 HloValue::Id PhiGraph::GetOptimizedId(const HloValue& value) {
   Node* node = value_id_to_node_[value.id()];
+  CHECK(!node->mark_as_dead);
   return node->value_id;
 }
 
-// Returns true if the input to a hlo value is the same as `inputs`.
+// Returns true if the inputs to a hlo value are the same as `inputs`.
 bool PhiGraph::InputsEqualTo(const HloValue& value,
                              absl::Span inputs) {
   auto iter = value_id_to_node_.find(value.id());
@@ -42,6 +43,7 @@ bool PhiGraph::InputsEqualTo(const HloValue& value,
 HloValue::Id PhiGraph::FindOptimizedValue(const HloValue::Id id) {
   auto iter = value_id_to_node_.find(id);
   CHECK(iter != value_id_to_node_.end());
+  CHECK(!iter->second->mark_as_dead);
   return iter->second->value_id;
 }
 
@@ -66,6 +68,17 @@ PhiGraph::Node* PhiGraph::CreateOrReuseNode(const HloValue& value) {
 void PhiGraph::ReplaceNodeWith(PhiGraph::Node* node, PhiGraph::Node* replace) {
   // Update users.
   CHECK(node->is_phi);
+  if (node->mark_as_dead) {
+    // The node has already been replaced with another.
+    return;
+  }
+  if (replace->mark_as_dead) {
+    // The node we are placing with has already been replaced with another node.
+    auto iter = value_id_to_node_.find(replace->value_id);
+    CHECK(iter != value_id_to_node_.end());
+    return ReplaceNodeWith(node, iter->second);
+  }
+  CHECK(!replace->mark_as_dead);
   for (Node* user : node->users) {
     absl::c_replace(user->operands, node, replace);
   }
@@ -74,6 +87,7 @@ void PhiGraph::ReplaceNodeWith(PhiGraph::Node* node, PhiGraph::Node* replace) {
   for (Node* operand : node->operands) {
     absl::c_replace(operand->users, node, replace);
   }
+
   for (HloValue::Id value_id : node_to_value_id_[node]) {
     CHECK(value_id_to_node_.contains(value_id));
     value_id_to_node_[value_id] = replace;
@@ -115,6 +129,8 @@ std::string PhiGraph::ToString() {
 }
 
 void PhiGraph::Optimize() {
+  VLOG(2) << "Optimizing phi graph:";
+  XLA_VLOG_LINES(2, ToString());
   // Set up users for each node.
   for (auto& node : node_storage_) {
     for (Node* input : node->operands) {
@@ -141,6 +157,8 @@ void PhiGraph::Optimize() {
 
       Node* node_ptr = node.get();
 
+      VLOG(2) << "Optimizing: " << node_ptr->value_id;
+
       CHECK_GE(node_ptr->operands.size(), 1);
 
       // Remove self-referencing ids from users and operands.
@@ -167,6 +185,9 @@ void PhiGraph::Optimize() {
           [&](Node* elem) { return elem == node_ptr->operands[0]; });
 
       if (all_inputs_are_same) {
+        VLOG(1) << "All inputs to node " << node_ptr->value_id
+                << " are the same, replacing it with "
+                << node_ptr->operands[0]->value_id;
         ReplaceNodeWith(node_ptr, node_ptr->operands[0]);
         changed = true;
         continue;
@@ -223,6 +244,8 @@ void PhiGraph::Optimize() {
             CHECK_EQ(node, non_phi);
             continue;
           }
+          VLOG(1) << "Replace node " << node->value_id
+                  << " in the closure with node " << non_phi->value_id;
           ReplaceNodeWith(node, non_phi);
           changed = true;
         }
diff --git a/tensorflow/compiler/xla/service/hlo_phi_graph.h b/tensorflow/compiler/xla/service/hlo_phi_graph.h
index a0eb994438e..ca0d5c5009c 100644
--- a/tensorflow/compiler/xla/service/hlo_phi_graph.h
+++ b/tensorflow/compiler/xla/service/hlo_phi_graph.h
@@ -90,7 +90,7 @@ class PhiGraph {
   // to that phi.
   absl::flat_hash_map> node_to_value_id_;
 
-  // A mapping between a HloValue and node in the phi graph.
+  // A mapping from a HloValue to node in the phi graph.
   absl::flat_hash_map value_id_to_node_;
   std::vector> node_storage_;
 };
diff --git a/tensorflow/compiler/xla/service/hlo_phi_graph_test.cc b/tensorflow/compiler/xla/service/hlo_phi_graph_test.cc
index 41f0454fe55..ee7300b160b 100644
--- a/tensorflow/compiler/xla/service/hlo_phi_graph_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_phi_graph_test.cc
@@ -82,5 +82,30 @@ TEST_F(PhiGraphTest, CircularPhi) {
   EXPECT_EQ(D.id(), phi_graph.FindOptimizedValue(C.id()));
 }
 
+TEST_F(PhiGraphTest, NestedPhiReduction) {
+  // def A = phi(B, C)
+  // def B = phi(C, E)
+  // def C = phi(A, B)
+  // def D = non-phi
+  // def E = Phi(D, D)
+  // 1. Replace E with D
+  // 2. Replace A B and C with E/D
+  PhiGraph phi_graph;
+  HloValue A = NewHloValue(true);
+  HloValue B = NewHloValue(true);
+  HloValue C = NewHloValue(true);
+  HloValue D = NewHloValue(false);
+  HloValue E = NewHloValue(true);
+  phi_graph.RegisterPhi(A, {&B, &C});
+  phi_graph.RegisterPhi(B, {&E, &C});
+  phi_graph.RegisterPhi(C, {&A, &B});
+  phi_graph.RegisterPhi(E, {&D, &D});
+  phi_graph.Optimize();
+  EXPECT_EQ(D.id(), phi_graph.FindOptimizedValue(A.id()));
+  EXPECT_EQ(D.id(), phi_graph.FindOptimizedValue(B.id()));
+  EXPECT_EQ(D.id(), phi_graph.FindOptimizedValue(C.id()));
+  EXPECT_EQ(D.id(), phi_graph.FindOptimizedValue(E.id()));
+}
+
 }  // namespace
 }  // namespace xla

From 7ae27e344cf9a501c89d785cfe4eb109a9848b47 Mon Sep 17 00:00:00 2001
From: Tiezhen WANG 
Date: Wed, 17 Jun 2020 20:04:47 -0700
Subject: [PATCH 0462/1390] TFL: selective registration: Disable linkopts not
 available on Windows.

PiperOrigin-RevId: 317019750
Change-Id: Ibce23c534a0b8fe15ada085a57aec6f87db71a9b
---
 tensorflow/lite/build_def.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/build_def.bzl b/tensorflow/lite/build_def.bzl
index 5e487395355..e6c92691b15 100644
--- a/tensorflow/lite/build_def.bzl
+++ b/tensorflow/lite/build_def.bzl
@@ -768,10 +768,10 @@ def tflite_custom_cc_library(name, models = [], srcs = [], deps = [], visibility
         name = name,
         srcs = real_srcs,
         copts = tflite_copts(),
-        linkopts = [
-            "-lm",
-            "-ldl",
-        ],
+        linkopts = select({
+            "//tensorflow:windows": [],
+            "//conditions:default": ["-lm", "-ldl"],
+        }),
         deps = depset([
             "//tensorflow/lite:framework",
             "//tensorflow/lite/kernels:builtin_ops",

From 697a0a4fe6d792b57a280e2f22fd23fe219e9d81 Mon Sep 17 00:00:00 2001
From: Thomas Raoux 
Date: Wed, 17 Jun 2020 20:43:17 -0700
Subject: [PATCH 0463/1390] Internal change

PiperOrigin-RevId: 317023916
Change-Id: I81311283fae6c28cdee2d7c87c8d062d32725bef
---
 third_party/mlir/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index 476b8566265..db75b27e78b 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -2868,9 +2868,9 @@ cc_library(
         ":SideEffects",
         ":StandardOps",
         ":Support",
-        "//third_party/vulkan_loader",
         "@llvm-project//llvm:Support",
         "@vulkan_headers",
+        "@vulkan_sdk//:sdk",
     ],
 )
 

From 4f341bb742718721563ce6dccb965c85a1fbdcf5 Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Wed, 17 Jun 2020 21:01:24 -0700
Subject: [PATCH 0464/1390] Add Bessel functions to the public api:     -
 tf.math.special.bessel_i0     - tf.math.special.bessel_i0e     -
 tf.math.special.bessel_i1     - tf.math.special.bessel_i1e     -
 tf.math.special.bessel_k0     - tf.math.special.bessel_k0e     -
 tf.math.special.bessel_k1     - tf.math.special.bessel_k1e     -
 tf.math.special.bessel_j0     - tf.math.special.bessel_j1     -
 tf.math.special.bessel_y0     - tf.math.special.bessel_y1

PiperOrigin-RevId: 317025879
Change-Id: I5c4407eda6bef0d1659b7a566979c7dbbad4ad83
---
 tensorflow/compiler/tf2xla/python/xla.py      |   5 +-
 .../api_def/base_api/api_def_BesselI0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselI0e.pbtxt  |   8 +-
 .../api_def/base_api/api_def_BesselI1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselI1e.pbtxt  |   8 +-
 .../api_def/base_api/api_def_BesselJ0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselJ1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK0e.pbtxt  |   4 +
 .../api_def/base_api/api_def_BesselK1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK1e.pbtxt  |   4 +
 .../api_def/base_api/api_def_BesselY0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselY1.pbtxt   |   4 +
 .../python_api/api_def_BesselI0e.pbtxt        |   6 -
 .../python_api/api_def_BesselI1e.pbtxt        |   6 -
 tensorflow/core/kernels/cwise_op_bessel.cc    |  29 --
 tensorflow/core/kernels/cwise_ops.h           |   6 -
 .../special_math/special_math_op_bessel.cc    |  78 +++++
 .../special_math_op_gpu_bessel.cu.cc}         |  16 +-
 .../special_math/special_math_op_misc_impl.h  |  38 +++
 tensorflow/core/ops/math_ops.cc               |   4 -
 tensorflow/core/ops/special_math_ops.cc       |  49 ++-
 .../optimization/map_vectorization_test.py    |   6 +-
 .../eager/pywrap_gradient_exclusions.cc       |   6 +-
 .../kernel_tests/cwise_ops_unary_test.py      |  20 +-
 tensorflow/python/ops/math_grad.py            | 131 ++++++++-
 .../python/ops/parallel_for/math_test.py      |   4 +-
 tensorflow/python/ops/parallel_for/pfor.py    |  14 +-
 tensorflow/python/ops/signal/window_ops.py    |   5 +-
 tensorflow/python/ops/sparse_ops.py           |   6 +-
 tensorflow/python/ops/special_math_ops.py     | 278 +++++++++++++++++-
 .../python/ops/special_math_ops_test.py       | 234 ++++++++++++++-
 .../golden/v1/tensorflow.math.special.pbtxt   |  48 +++
 .../api/golden/v1/tensorflow.raw_ops.pbtxt    |  40 +++
 .../golden/v2/tensorflow.math.special.pbtxt   |  48 +++
 .../api/golden/v2/tensorflow.raw_ops.pbtxt    |  40 +++
 36 files changed, 1028 insertions(+), 145 deletions(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselI1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselJ0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselJ1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK0e.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK1e.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselY0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselY1.pbtxt
 delete mode 100644 tensorflow/core/api_def/python_api/api_def_BesselI0e.pbtxt
 delete mode 100644 tensorflow/core/api_def/python_api/api_def_BesselI1e.pbtxt
 delete mode 100644 tensorflow/core/kernels/cwise_op_bessel.cc
 create mode 100644 tensorflow/core/kernels/special_math/special_math_op_bessel.cc
 rename tensorflow/core/kernels/{cwise_op_bessel.cu.cc => special_math/special_math_op_gpu_bessel.cu.cc} (58%)

diff --git a/tensorflow/compiler/tf2xla/python/xla.py b/tensorflow/compiler/tf2xla/python/xla.py
index c59c47e92fb..0ebca2d546f 100644
--- a/tensorflow/compiler/tf2xla/python/xla.py
+++ b/tensorflow/compiler/tf2xla/python/xla.py
@@ -37,6 +37,7 @@ from tensorflow.python.ops import gen_math_ops
 from tensorflow.python.ops import gen_random_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import special_math_ops
 
 # TODO(phawkins): provide wrappers for all XLA operators. Currently the missing
 # ops include:
@@ -103,8 +104,8 @@ sign = _unary_op(math_ops.sign)
 tanh = _unary_op(math_ops.tanh)
 
 # Bessel
-bessel_i0e = _unary_op(math_ops.bessel_i0e)
-bessel_i1e = _unary_op(math_ops.bessel_i1e)
+bessel_i0e = _unary_op(special_math_ops.bessel_i0e)
+bessel_i1e = _unary_op(special_math_ops.bessel_i1e)
 
 # Binary operators
 
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
new file mode 100644
index 00000000000..2c47960429c
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselI0"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselI0e.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselI0e.pbtxt
index 08313cebb99..7965af4916e 100644
--- a/tensorflow/core/api_def/base_api/api_def_BesselI0e.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_BesselI0e.pbtxt
@@ -1,10 +1,4 @@
 op {
   graph_op_name: "BesselI0e"
-  summary: "Computes the Bessel i0e function of `x` element-wise."
-  description: <
 {
-  StringRef getName() final { return "Table"; }
-};
-
 }  // namespace ResourceEffects
 }  // namespace TF
 }  // namespace mlir

From 64e1b489bb63c4b7ceb85bfa846ca699036561e8 Mon Sep 17 00:00:00 2001
From: Terry Heo 
Date: Wed, 17 Jun 2020 23:54:57 -0700
Subject: [PATCH 0477/1390] Enable flex delegate on tensorflow.lite.Interpreter
 Python package

Usually, flex delegate is enabled by symbol override of AcquireFlexDelegate()
function. But this approach doesn't work well with shared library.

Since pywrap_tensorflow_internal.so is available for tensorflow PIP,
I've made the following changes to enable flex delegate.
- Included flex delegate module to the pywrap_tensorflow_internal.so.
  This file already contains most TF internal logic and having TFLite flex
  delegate impacts about 72K to the output.
- Added new function of TF_AcquireFlexDelegate() in the delegate module.
- Updated logic in AcquireFlexDelegate() of interpreter_builder.cc to check
  the availability of pywrap_tensorflow_internal.so and lookup the
  TF_AcquireFlexDelegate() symbol to enable flex delegate.

Also updated python/lite_flex_test.py since flex delegate is supported with
Python API

PiperOrigin-RevId: 317044994
Change-Id: Ic5e953f4a675b3f5360a4c7d607568193103711a
---
 tensorflow/lite/delegates/flex/delegate.cc |  7 +++
 tensorflow/lite/interpreter_builder.cc     | 17 ++++++
 tensorflow/lite/python/BUILD               |  5 +-
 tensorflow/lite/python/lite_flex_test.py   | 61 +++++++++++++---------
 tensorflow/python/BUILD                    |  7 ++-
 5 files changed, 68 insertions(+), 29 deletions(-)

diff --git a/tensorflow/lite/delegates/flex/delegate.cc b/tensorflow/lite/delegates/flex/delegate.cc
index 4741bddc2f5..b8b0d4e6d01 100644
--- a/tensorflow/lite/delegates/flex/delegate.cc
+++ b/tensorflow/lite/delegates/flex/delegate.cc
@@ -136,3 +136,10 @@ TfLiteStatus FlexDelegate::CopyFromBufferHandle(
 }
 
 }  // namespace tflite
+
+// Exported C interface function which is used by AcquireFlexDelegate() at
+// interpreter_build.cc. To export the function name globally, the function name
+// must be matched with patterns in tf_version_script.lds
+extern "C" tflite::TfLiteDelegateUniquePtr TF_AcquireFlexDelegate() {
+  return tflite::AcquireFlexDelegate();
+}
diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index 43d81ef0770..d73b298e595 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -14,6 +14,9 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/interpreter_builder.h"
 
+#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
+#include 
+#endif
 #include 
 #include 
 #include 
@@ -114,6 +117,20 @@ const char* kEmptyTensorName = "";
 // For flex delegate, see also the strong override in
 // lite/delegates/flex/delegate.cc.
 TFLITE_ATTRIBUTE_WEAK Interpreter::TfLiteDelegatePtr AcquireFlexDelegate() {
+#if !defined(__ANDROID__) && !defined(__APPLE__) && !defined(_WIN32)
+  // If _pywrap_tensorflow_internal.so is available, use
+  // TF_AcquireFlexDelegate() to initialize flex delegate.
+  void* lib_tf_internal =
+      dlopen("_pywrap_tensorflow_internal.so", RTLD_NOW | RTLD_LOCAL);
+  if (lib_tf_internal) {
+    auto TF_AcquireFlexDelegate =
+        reinterpret_cast(
+            dlsym(lib_tf_internal, "TF_AcquireFlexDelegate"));
+    if (TF_AcquireFlexDelegate) {
+      return TF_AcquireFlexDelegate();
+    }
+  }
+#endif
   return Interpreter::TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {});
 }
 
diff --git a/tensorflow/lite/python/BUILD b/tensorflow/lite/python/BUILD
index d25e7d5ef8d..1b64b7d1042 100644
--- a/tensorflow/lite/python/BUILD
+++ b/tensorflow/lite/python/BUILD
@@ -193,9 +193,8 @@ py_test(
     python_version = "PY3",
     srcs_version = "PY2AND3",
     tags = [
-        # TODO(b/111881877): Enable in oss after resolving op registry issues.
-        "no_oss",
-        "no_windows",
+        "no_mac",  # TODO(b/159077703): Enable Python API Flex support on MacOS.
+        "no_windows",  # TODO(b/159077703): Enable Python API Flex support on Windows.
     ],
     deps = [
         ":lite",
diff --git a/tensorflow/lite/python/lite_flex_test.py b/tensorflow/lite/python/lite_flex_test.py
index 26bee206d27..ffc157c2128 100644
--- a/tensorflow/lite/python/lite_flex_test.py
+++ b/tensorflow/lite/python/lite_flex_test.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from absl.testing import parameterized
+import numpy as np
 
 from tensorflow.lite.python import lite
 from tensorflow.lite.python.interpreter import Interpreter
@@ -41,8 +42,7 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       ('DisableMlirConverter', False))  # disable mlir
   def testFlexMode(self, enable_mlir):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(
-          shape=[1, 16, 16, 3], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -54,19 +54,22 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
   def testDeprecatedFlags(self):
     with ops.Graph().as_default():
-      in_tensor = array_ops.placeholder(
-          shape=[1, 16, 16, 3], dtype=dtypes.float32)
+      in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
       out_tensor = in_tensor + in_tensor
       sess = session.Session()
 
@@ -83,14 +86,18 @@ class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     tflite_model = converter.convert()
     self.assertTrue(tflite_model)
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
 
 class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
@@ -114,14 +121,18 @@ class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
     converter.experimental_new_converter = enable_mlir
     tflite_model = converter.convert()
 
-    # Ensures the model contains TensorFlow ops.
-    # TODO(nupurgarg): Check values once there is a Python delegate interface.
+    # Check the model works with TensorFlow ops.
     interpreter = Interpreter(model_content=tflite_model)
-    with self.assertRaises(RuntimeError) as error:
-      interpreter.allocate_tensors()
-    self.assertIn(
-        'Regular TensorFlow ops are not supported by this interpreter.',
-        str(error.exception))
+    interpreter.allocate_tensors()
+    input_details = interpreter.get_input_details()
+    test_input = np.array([4.0], dtype=np.float32)
+    interpreter.set_tensor(input_details[0]['index'], test_input)
+    interpreter.invoke()
+
+    output_details = interpreter.get_output_details()
+    expected_output = np.array([24.0], dtype=np.float32)
+    output_data = interpreter.get_tensor(output_details[0]['index'])
+    self.assertTrue((expected_output == output_data).all())
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index d141b719aef..f53859b2915 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -6058,7 +6058,12 @@ pywrap_tensorflow_macro(
         "@ngraph_tf//:ngraph_tf",
     ]) + if_xla_available([
         "//tensorflow/compiler/aot:tfcompile_lib",
-    ]),
+    ]) + select({
+        "//tensorflow:windows": [],  # TODO(b/159077703): Enable Flex on Windows
+        "//conditions:default": [
+            "//tensorflow/lite/delegates/flex:delegate",
+        ],
+    }),
 )
 
 # ** Targets for Windows build (start) **

From fabcd8f89cd5975331994049705e15cb75f32e0c Mon Sep 17 00:00:00 2001
From: Marat Dukhan 
Date: Thu, 18 Jun 2020 00:06:13 -0700
Subject: [PATCH 0478/1390] Update XNNPACK dependency

Bring in fix for x86 builds on Android NDK r20

PiperOrigin-RevId: 317046250
Change-Id: I493cf294ea4e51c91a68b9bc8b062f6cecf4da7f
---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index fab6faf62d1..78f7e0ce03e 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -164,11 +164,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
 
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "dfcc7b2894c5c6bc570d65ff22b371a28e0fcc672e75705cc3f1ccc1264c3f8b",
-        strip_prefix = "XNNPACK-cac25227b5d8046170f875ad80545696be908ee7",
+        sha256 = "4af883fea0a6ada106867f29670a6c0b7af74bee85d74a2e04356a670814a3d4",
+        strip_prefix = "XNNPACK-69a6a7667d96a84c596b0f4e00632b2037c17723",
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/cac25227b5d8046170f875ad80545696be908ee7.zip",
-            "https://github.com/google/XNNPACK/archive/cac25227b5d8046170f875ad80545696be908ee7.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/69a6a7667d96a84c596b0f4e00632b2037c17723.zip",
+            "https://github.com/google/XNNPACK/archive/69a6a7667d96a84c596b0f4e00632b2037c17723.zip",
         ],
     )
 

From b38d5d1889b398ca45aa7f2e4f9f0184c77c6c55 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Thu, 18 Jun 2020 00:52:32 -0700
Subject: [PATCH 0479/1390] Remove dynamic dimension of strided slice grad if
 input to strided slice is static.

If we slice a dynamic shaped tensor from a static tensor, the output
of the gradient should still be static. Unfortunately this cannot be
deduced alone by xla, so extra information is needed from the tf2xla
bridge.

PiperOrigin-RevId: 317051543
Change-Id: I7a8113c47a4aed145dfba7f7d12992ca35a9cf19
---
 .../tf2xla/kernels/strided_slice_op.cc        | 22 ++++++++++++++++
 tensorflow/compiler/xla/client/xla_builder.cc | 23 ++++++++++++++++
 tensorflow/compiler/xla/client/xla_builder.h  |  6 +++++
 .../compiler/xla/client/xla_builder_test.cc   | 26 +++++++++++++++++++
 4 files changed, 77 insertions(+)

diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
index 2684c982600..51764018df1 100644
--- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
@@ -350,6 +350,28 @@ class StridedSliceGradOp : public XlaOpKernel {
       grad = xla::Rev(grad, dimensions_to_reverse);
     }
     grad = xla::Pad(grad, zero, padding_config);
+
+    xla::XlaOp dynamic_shape = ctx->Input(0);
+    xla::Shape grad_shape = ctx->builder()->GetShape(grad).ValueOrDie();
+    ctx->set_dynamic_dimension_is_minus_one(true);
+    std::vector dynamic_size;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(0, &dynamic_size));
+    // Input of strided_slice_op has to have the same shape as output.
+    DCHECK_EQ(grad_shape.rank(), input_shape.dims());
+    for (int64 dim = 0; dim < input_shape.dims(); ++dim) {
+      DCHECK_EQ(grad_shape.dimensions(dim), input_shape.dim_size(dim));
+      if (dynamic_size[dim] == -1) {
+        // Input is a dynamic dimension, set the same dynamic dimension size in
+        // the output.
+        auto dim_size = xla::Slice(dynamic_shape, {dim}, {dim + 1}, {1});
+        grad = xla::SetDimensionSize(grad, dim_size, dim);
+      } else if (grad_shape.is_dynamic_dimension(dim)) {
+        // Input is static but output is dynamic, respect input and remove any
+        // dynamic dim in the output.
+        grad = xla::RemoveDynamicDimension(grad, dim);
+      }
+    }
+
     ctx->SetOutput(0, grad);
   }
 
diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index bfba48862f6..c7b6a7f9491 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -2727,6 +2727,25 @@ XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64 dimension) {
   });
 }
 
+XlaOp XlaBuilder::RemoveDynamicDimension(XlaOp operand, int64 dimension) {
+  return ReportErrorOrReturn([&]() -> StatusOr {
+    HloInstructionProto instr;
+    TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+
+    Shape shape = *operand_shape;
+    shape.set_dynamic_dimension(dimension, false);
+    // Setting an op's dynamic dimension to its static size removes the dynamic
+    // dimension.
+    XlaOp static_size =
+        ConstantR0(this, operand_shape->dimensions(dimension));
+
+    *instr.mutable_shape() = shape.ToProto();
+    instr.add_dimensions(dimension);
+    return AddInstruction(std::move(instr), HloOpcode::kSetDimensionSize,
+                          {operand, static_size});
+  });
+}
+
 XlaOp XlaBuilder::SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension) {
   return ReportErrorOrReturn([&]() -> StatusOr {
     HloInstructionProto instr;
@@ -3827,4 +3846,8 @@ XlaOp SetDimensionSize(const XlaOp operand, const XlaOp val, int64 dimension) {
   return operand.builder()->SetDimensionSize(operand, val, dimension);
 }
 
+XlaOp RemoveDynamicDimension(const XlaOp operand, int64 dimension) {
+  return operand.builder()->RemoveDynamicDimension(operand, dimension);
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h
index ffa6a7c3439..b8af180b83e 100644
--- a/tensorflow/compiler/xla/client/xla_builder.h
+++ b/tensorflow/compiler/xla/client/xla_builder.h
@@ -704,6 +704,8 @@ class XlaBuilder {
 
   XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
+  XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
+
   StatusOr AddInstruction(HloInstructionProto&& instr, HloOpcode opcode,
                                  absl::Span operands = {});
 
@@ -1151,6 +1153,7 @@ class XlaBuilder {
 
   friend XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
   friend XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
+  friend XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
 
  protected:
   // Returns OK status if the given op was built using this builder. Otherwise,
@@ -2149,6 +2152,9 @@ XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
 
 XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
+// Returns the same op but with dynamic dimension removed.
+XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
+
 // Implementation details below this point.
 //
 
diff --git a/tensorflow/compiler/xla/client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_builder_test.cc
index 4fa47077fca..7011c946203 100644
--- a/tensorflow/compiler/xla/client/xla_builder_test.cc
+++ b/tensorflow/compiler/xla/client/xla_builder_test.cc
@@ -556,6 +556,32 @@ TEST_F(XlaBuilderTest, DynamicParameter) {
   EXPECT_TRUE(param_shape.is_dynamic_dimension(0));
 }
 
+TEST_F(XlaBuilderTest, SetDimensionSize) {
+  XlaBuilder b(TestName());
+  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
+  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
+  auto set_dim_size = SetDimensionSize(p0, p1, 0);
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          BuildHloModule(&b, /*root=*/set_dim_size));
+  const Shape& root_shape =
+      module->entry_computation()->root_instruction()->shape();
+  EXPECT_TRUE(root_shape.is_dynamic_dimension(0));
+}
+
+TEST_F(XlaBuilderTest, RemoveDimensionSize) {
+  XlaBuilder b(TestName());
+  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
+  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
+  auto set_dim_size = SetDimensionSize(p0, p1, 0);
+  auto remove_dim_size = RemoveDynamicDimension(set_dim_size, 0);
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          BuildHloModule(&b, /*root=*/remove_dim_size));
+  const Shape& root_shape =
+      module->entry_computation()->root_instruction()->shape();
+  // Dynamic dimension has been removed.
+  EXPECT_FALSE(root_shape.is_dynamic_dimension(0));
+}
+
 TEST_F(XlaBuilderTest, DynamicUnary) {
   XlaBuilder b(TestName());
   Shape tuple_param_shape = ShapeUtil::MakeTupleShape(

From 6a8268d74e64e3bf1e445a9b1adb62e04bcdea34 Mon Sep 17 00:00:00 2001
From: Felix Johnny 
Date: Wed, 3 Jun 2020 08:27:57 +0200
Subject: [PATCH 0480/1390] CMSIS-NN wrapper update for interface changes

Background: CMSIS-NN int8  APIs are changed
where the pass by value function arguments
are replaced by pass by reference struct
arguments.

The following changes are done in TFL micro

1. Update int8 cmsis-nn's depthwise conv,
   fully connected and average pooling
   wrapper files to use the new  interface.

2. For the above mentioned operators,
   updates from the reference implementation
   (e.g, CalculateOpData() in Prepare() instead
   of Eval()) are ported as well.
---
 .../micro/kernels/cmsis-nn/depthwise_conv.cc  | 394 ++++++++++--------
 .../micro/kernels/cmsis-nn/fully_connected.cc | 221 ++++++----
 .../lite/micro/kernels/cmsis-nn/pooling.cc    | 195 ++++-----
 .../lite/micro/tools/make/ext_libs/cmsis.inc  |   1 +
 .../tools/make/third_party_downloads.inc      |   4 +-
 5 files changed, 442 insertions(+), 373 deletions(-)

diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc b/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc
index 4d398855abc..eb5fad0a08a 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/depthwise_conv.cc
@@ -1,4 +1,4 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -36,7 +36,6 @@ constexpr int kInputTensor = 0;
 constexpr int kFilterTensor = 1;
 constexpr int kBiasTensor = 2;
 constexpr int kOutputTensor = 0;
-constexpr int kMaxChannels = 256;
 
 // Depthwise conv is quantized along dimension 3:
 // https://www.tensorflow.org/lite/performance/quantization_spec
@@ -50,14 +49,14 @@ struct OpData {
   int output_shift;
 
   // Per channel output multiplier and shift.
-  // TODO(b/141139247): Allocate these dynamically when possible.
-  int32_t per_channel_output_multiplier[kMaxChannels];
-  int32_t per_channel_output_shift[kMaxChannels];
-
+  int32_t* per_channel_output_multiplier;
+  int32_t* per_channel_output_shift;
   // The range of the fused activation layer. For example for kNone and
   // uint8_t these would be 0 and 255.
   int32_t output_activation_min;
   int32_t output_activation_max;
+  // Index to buffer for optimizations if applicable.
+  int buffer_idx;
 };
 
 TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
@@ -70,6 +69,8 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
   TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
 
   int unused_output_height, unused_output_width;
+  // Set buffer index to a reset value
+  data->buffer_idx = -1;
   data->padding = ComputePaddingHeightWidth(
       params->stride_height, params->stride_width, 1, 1, height, width,
       filter_height, filter_width, params->padding, &unused_output_height,
@@ -85,12 +86,12 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
     TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
     int num_channels = filter->dims->data[kDepthwiseConvQuantizedDimension];
 
-    TF_LITE_ENSURE_STATUS(tflite::PopulateConvolutionQuantizationParams(
+    return tflite::PopulateConvolutionQuantizationParams(
         context, input, filter, bias, output, params->activation,
         &data->output_multiplier, &data->output_shift,
         &data->output_activation_min, &data->output_activation_max,
         data->per_channel_output_multiplier,
-        reinterpret_cast(data->per_channel_output_shift), num_channels));
+        reinterpret_cast(data->per_channel_output_shift), num_channels);
   }
   return kTfLiteOk;
 }
@@ -98,47 +99,117 @@ TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
 }  // namespace
 
 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
-  void* raw;
-  context->AllocatePersistentBuffer(context, sizeof(int), &raw);
-  return raw;
+  TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
+  void* data = nullptr;
+  if (context->AllocatePersistentBuffer(context, sizeof(OpData), &data) ==
+      kTfLiteError) {
+    return nullptr;
+  }
+  return data;
 }
 
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
+  TFLITE_DCHECK(node->user_data != nullptr);
+  TFLITE_DCHECK(node->builtin_data != nullptr);
+
+  OpData* data = static_cast(node->user_data);
   auto* params =
       reinterpret_cast(node->builtin_data);
 
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   const TfLiteTensor* filter = GetInput(context, node, kFilterTensor);
 
-  const int filter_width = SizeOfDimension(filter, 2);
-  const int filter_height = SizeOfDimension(filter, 1);
+  const TfLiteType data_type = input->type;
+  int width = SizeOfDimension(input, 2);
+  int height = SizeOfDimension(input, 1);
+  int filter_width = SizeOfDimension(filter, 2);
+  int filter_height = SizeOfDimension(filter, 1);
 
-  RuntimeShape input_shape = GetTensorShape(input);
-  const int input_depth = input_shape.Dims(3);
+  if (input->type == kTfLiteInt8) {
+    // Allocate memory for per-channel quantization parameters
+    const int num_channels =
+        filter->dims->data[kDepthwiseConvQuantizedDimension];
+    // Dynamically allocate per-channel quantization parameters.
+    TF_LITE_ENSURE_STATUS(context->AllocatePersistentBuffer(
+        context, num_channels * sizeof(int32_t),
+        reinterpret_cast(&data->per_channel_output_multiplier)));
+    TF_LITE_ENSURE_STATUS(context->AllocatePersistentBuffer(
+        context, num_channels * sizeof(int32_t),
+        reinterpret_cast(&data->per_channel_output_shift)));
+    TF_LITE_ENSURE_EQ(context, filter->quantization.type,
+                      kTfLiteAffineQuantization);
 
-  int* buffer_idx = reinterpret_cast(node->user_data);
+    // All per-channel quantized tensors need valid zero point and scale arrays.
+    const auto* affine_quantization =
+        reinterpret_cast(
+            filter->quantization.params);
+    TF_LITE_ENSURE(context, affine_quantization);
+    TF_LITE_ENSURE(context, affine_quantization->scale);
+    TF_LITE_ENSURE(context, affine_quantization->zero_point);
+    TF_LITE_ENSURE(
+        context, affine_quantization->scale->size == 1 ||
+                     affine_quantization->scale->size ==
+                         filter->dims->data[kDepthwiseConvQuantizedDimension]);
+    TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size,
+                      affine_quantization->zero_point->size);
+  }
 
-  *buffer_idx = -1;
-  node->user_data = buffer_idx;
+  TF_LITE_ENSURE_STATUS(CalculateOpData(context, node, params, width, height,
+                                        filter_width, filter_height, data_type,
+                                        data));
 
-  if (params->depth_multiplier == 1) {
-    const int32_t buf_size = arm_depthwise_conv_s8_opt_get_buffer_size(
-        input_depth, filter_width, filter_height);
+  if (input->type == kTfLiteInt8) {
+    const TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+    RuntimeShape input_shape = GetTensorShape(input);
+    RuntimeShape output_shape = GetTensorShape(output);
+    RuntimeShape filter_shape = GetTensorShape(filter);
+    TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+    TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+    TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+    const int batch_size = MatchingDim(input_shape, 0, output_shape, 0);
+    const int output_depth = MatchingDim(output_shape, 3, filter_shape, 3);
+    TFLITE_DCHECK_EQ(batch_size, 1); /* Only batch = 1 is supported */
+
+    cmsis_nn_dims input_dims;
+    input_dims.n = batch_size;
+    input_dims.h = height;
+    input_dims.w = width;
+    input_dims.c = input_shape.Dims(3);
+
+    cmsis_nn_dims filter_dims;
+    filter_dims.n = 1;
+    filter_dims.h = filter_height;
+    filter_dims.w = filter_width;
+    filter_dims.c = output_depth;
+
+    cmsis_nn_dims output_dims;
+    output_dims.n = batch_size;
+    output_dims.h = output_shape.Dims(1);
+    output_dims.w = output_shape.Dims(2);
+    output_dims.c = output_depth;
+
+    cmsis_nn_dw_conv_params dw_conv_params;
+    dw_conv_params.padding.h = data->padding.height;
+    dw_conv_params.padding.w = data->padding.width;
+
+    const int32_t buf_size = arm_depthwise_conv_wrapper_s8_get_buffer_size(
+        &dw_conv_params, &input_dims, &filter_dims, &output_dims);
 
     if (buf_size > 0) {
-      TF_LITE_ENSURE_STATUS(
-          context->RequestScratchBufferInArena(context, buf_size, buffer_idx));
+      TF_LITE_ENSURE_STATUS(context->RequestScratchBufferInArena(
+          context, buf_size, &data->buffer_idx));
+    } else {
+      data->buffer_idx = -1;
     }
   }
-#endif
   return kTfLiteOk;
 }
 
-TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
-                       TfLiteDepthwiseConvParams* params, OpData* data,
-                       const TfLiteTensor* input, const TfLiteTensor* filter,
-                       const TfLiteTensor* bias, TfLiteTensor* output) {
+void EvalFloat(TfLiteContext* context, TfLiteNode* node,
+               TfLiteDepthwiseConvParams* params, const OpData* data,
+               const TfLiteTensor* input, const TfLiteTensor* filter,
+               const TfLiteTensor* bias, TfLiteTensor* output) {
   float output_activation_min, output_activation_max;
   CalculateActivationRange(params->activation, &output_activation_min,
                            &output_activation_max);
@@ -150,8 +221,8 @@ TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
   op_params.padding_values.height = data->padding.height;
   op_params.stride_width = params->stride_width;
   op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = 1;
-  op_params.dilation_height_factor = 1;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
   op_params.depth_multiplier = params->depth_multiplier;
   op_params.float_activation_min = output_activation_min;
   op_params.float_activation_max = output_activation_max;
@@ -161,106 +232,120 @@ TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
       GetTensorShape(filter), GetTensorData(filter),
       GetTensorShape(bias), GetTensorData(bias), GetTensorShape(output),
       GetTensorData(output));
-  return kTfLiteOk;
 }
 
-TfLiteStatus EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node,
-                                     TfLiteDepthwiseConvParams* params,
-                                     OpData* data, const TfLiteTensor* input,
-                                     const TfLiteTensor* filter,
-                                     const TfLiteTensor* bias,
-                                     TfLiteTensor* output) {
-  DepthwiseParams op_params;
-  op_params.padding_type = PaddingType::kSame;
-  op_params.padding_values.width = data->padding.width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.stride_width = params->stride_width;
-  op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = params->dilation_width_factor;
-  op_params.dilation_height_factor = params->dilation_height_factor;
-  op_params.depth_multiplier = params->depth_multiplier;
-  op_params.input_offset = -input->params.zero_point;
-  op_params.weights_offset = 0;
-  op_params.output_offset = output->params.zero_point;
-  // TODO(b/130439627): Use calculated value for clamping.
-  op_params.quantized_activation_min = std::numeric_limits::min();
-  op_params.quantized_activation_max = std::numeric_limits::max();
+void EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node,
+                             TfLiteDepthwiseConvParams* params,
+                             const OpData* data, const TfLiteTensor* input,
+                             const TfLiteTensor* filter,
+                             const TfLiteTensor* bias, TfLiteTensor* output) {
+  cmsis_nn_dw_conv_params dw_conv_params;
+  dw_conv_params.dilation.h = params->dilation_height_factor;
+  dw_conv_params.dilation.w = params->dilation_width_factor;
+  // Call to reference implementation can be removed when dilation is supported
+  // in the optimized implementations.
+  if (1 == dw_conv_params.dilation.h && 1 == dw_conv_params.dilation.w) {
+    dw_conv_params.input_offset = -input->params.zero_point;
+    dw_conv_params.output_offset = output->params.zero_point;
+    dw_conv_params.stride.h = params->stride_height;
+    dw_conv_params.stride.w = params->stride_width;
+    dw_conv_params.padding.h = data->padding.height;
+    dw_conv_params.padding.w = data->padding.width;
+    // TODO(b/130439627): Use calculated value for clamping.
+    dw_conv_params.activation.min = std::numeric_limits::min();
+    dw_conv_params.activation.max = std::numeric_limits::max();
+    dw_conv_params.ch_mult = params->depth_multiplier;
 
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
-  RuntimeShape filter_shape = GetTensorShape(filter);
-  const int filter_height = filter_shape.Dims(1);
-  const int filter_width = filter_shape.Dims(2);
-  RuntimeShape input_shape = GetTensorShape(input);
-  const int input_height = input_shape.Dims(1);
-  const int input_width = input_shape.Dims(2);
-  const int input_depth = input_shape.Dims(3);
-  RuntimeShape output_shape = GetTensorShape(output);
-  const int output_height = output_shape.Dims(1);
-  const int output_width = output_shape.Dims(2);
-  RuntimeShape bias_shape = GetTensorShape(bias);
+    cmsis_nn_per_channel_quant_params quant_params;
+    quant_params.multiplier = data->per_channel_output_multiplier;
+    quant_params.shift = data->per_channel_output_shift;
 
-  if (op_params.depth_multiplier == 1) {
-    int16_t* buf = nullptr;
-    auto* buffer_idx = reinterpret_cast(node->user_data);
-    if (*buffer_idx > -1) {
-      void* raw = context->GetScratchBuffer(context, *buffer_idx);
-      buf = reinterpret_cast(raw);
+    RuntimeShape filter_shape = GetTensorShape(filter);
+    RuntimeShape input_shape = GetTensorShape(input);
+    RuntimeShape output_shape = GetTensorShape(output);
+    RuntimeShape bias_shape = GetTensorShape(bias);
+
+    TFLITE_DCHECK_LE(dw_conv_params.activation.min,
+                     dw_conv_params.activation.max);
+
+    const int batch_size = MatchingDim(input_shape, 0, output_shape, 0);
+    const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
+
+    if (GetTensorData(bias)) {
+      TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
     }
 
-    TF_LITE_ENSURE_EQ(
-        context,
-        arm_depthwise_conv_s8_opt(
-            GetTensorData(input), input_width, input_height,
-            input_depth, GetTensorData(filter), input_depth,
-            filter_width, filter_height, op_params.padding_values.width,
-            op_params.padding_values.height, op_params.stride_width,
-            op_params.stride_height, GetTensorData(bias),
-            GetTensorData(output), data->per_channel_output_shift,
-            data->per_channel_output_multiplier, output_width, output_height,
-            op_params.output_offset, op_params.input_offset,
-            op_params.quantized_activation_min,
-            op_params.quantized_activation_max, op_params.dilation_width_factor,
-            op_params.dilation_height_factor, buf),
-        ARM_MATH_SUCCESS);
+    cmsis_nn_dims input_dims;
+    input_dims.n = batch_size;
+    input_dims.h = input_shape.Dims(1);
+    input_dims.w = input_shape.Dims(2);
+    input_dims.c = input_shape.Dims(3);
+
+    cmsis_nn_dims filter_dims;
+    filter_dims.n = filter_shape.Dims(0);
+    filter_dims.h = filter_shape.Dims(1);
+    filter_dims.w = filter_shape.Dims(2);
+    filter_dims.c = output_depth;
+
+    cmsis_nn_dims bias_dims;
+    bias_dims.n = 1;
+    bias_dims.h = 1;
+    bias_dims.w = 1;
+    bias_dims.c = output_depth;
+
+    cmsis_nn_dims output_dims;
+    output_dims.n = batch_size;
+    output_dims.h = output_shape.Dims(1);
+    output_dims.w = output_shape.Dims(2);
+    output_dims.c = output_depth;
+
+    cmsis_nn_context ctx;
+    ctx.buf = nullptr;
+    /* 'size' is unused */
+    ctx.size = 0;
+
+    if (data->buffer_idx > -1) {
+      ctx.buf = context->GetScratchBuffer(context, data->buffer_idx);
+    }
+
+    TFLITE_DCHECK_EQ(arm_depthwise_conv_wrapper_s8(
+                         &ctx, &dw_conv_params, &quant_params, &input_dims,
+                         GetTensorData(input), &filter_dims,
+                         GetTensorData(filter), &bias_dims,
+                         GetTensorData(bias), &output_dims,
+                         GetTensorData(output)),
+                     ARM_MATH_SUCCESS);
   } else {
-    TF_LITE_ENSURE_EQ(
-        context,
-        arm_depthwise_conv_s8(
-            GetTensorData(input), input_width, input_height,
-            input_depth, GetTensorData(filter),
-            op_params.depth_multiplier * input_depth,
-            op_params.depth_multiplier, filter_width, filter_height,
-            op_params.padding_values.width, op_params.padding_values.height,
-            op_params.stride_width, op_params.stride_height,
-            GetTensorData(bias), GetTensorData(output),
-            data->per_channel_output_shift, data->per_channel_output_multiplier,
-            output_width, output_height, op_params.output_offset,
-            op_params.input_offset, op_params.quantized_activation_min,
-            op_params.quantized_activation_max, op_params.dilation_width_factor,
-            op_params.dilation_height_factor, nullptr),
-        ARM_MATH_SUCCESS);
+    DepthwiseParams op_params;
+    op_params.padding_type = PaddingType::kSame;
+    op_params.padding_values.width = data->padding.width;
+    op_params.padding_values.height = data->padding.height;
+    op_params.stride_width = params->stride_width;
+    op_params.stride_height = params->stride_height;
+    op_params.dilation_width_factor = params->dilation_width_factor;
+    op_params.dilation_height_factor = params->dilation_height_factor;
+    op_params.depth_multiplier = params->depth_multiplier;
+    op_params.input_offset = -input->params.zero_point;
+    op_params.weights_offset = 0;
+    op_params.output_offset = output->params.zero_point;
+    // TODO(b/130439627): Use calculated value for clamping.
+    op_params.quantized_activation_min = std::numeric_limits::min();
+    op_params.quantized_activation_max = std::numeric_limits::max();
+
+    reference_integer_ops::DepthwiseConvPerChannel(
+        op_params, data->per_channel_output_multiplier,
+        data->per_channel_output_shift, GetTensorShape(input),
+        GetTensorData(input), GetTensorShape(filter),
+        GetTensorData(filter), GetTensorShape(bias),
+        GetTensorData(bias), GetTensorShape(output),
+        GetTensorData(output));
   }
-#else
-#pragma message( \
-    "CMSIS-NN optimization for depthwise_conv not available for this target. Using reference kernel.")
-
-  reference_integer_ops::DepthwiseConvPerChannel(
-      op_params, data->per_channel_output_multiplier,
-      data->per_channel_output_shift, GetTensorShape(input),
-      GetTensorData(input), GetTensorShape(filter),
-      GetTensorData(filter), GetTensorShape(bias),
-      GetTensorData(bias), GetTensorShape(output),
-      GetTensorData(output));
-
-#endif
-  return kTfLiteOk;
 }
 
-TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
-                           TfLiteDepthwiseConvParams* params, OpData* data,
-                           const TfLiteTensor* input,
-                           const TfLiteTensor* filter, const TfLiteTensor* bias,
-                           TfLiteTensor* output) {
+void EvalQuantized(TfLiteContext* context, TfLiteNode* node,
+                   TfLiteDepthwiseConvParams* params, const OpData* data,
+                   const TfLiteTensor* input, const TfLiteTensor* filter,
+                   const TfLiteTensor* bias, TfLiteTensor* output) {
   const int32_t input_offset = -input->params.zero_point;
   const int32_t filter_offset = -filter->params.zero_point;
   const int32_t output_offset = output->params.zero_point;
@@ -272,8 +357,8 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
   op_params.padding_values.height = data->padding.height;
   op_params.stride_width = params->stride_width;
   op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = 1;
-  op_params.dilation_height_factor = 1;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
   op_params.depth_multiplier = params->depth_multiplier;
   op_params.quantized_activation_min = data->output_activation_min;
   op_params.quantized_activation_max = data->output_activation_max;
@@ -284,13 +369,11 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
   // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
   op_params.output_shift = -data->output_shift;
 
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
-  // optimizations utilize loop unrolling which requires the following power
-  // of two kernel dimensions
-  RuntimeShape filter_shape = GetTensorShape(filter);
-  const int filter_height = filter_shape.Dims(1);
-  const int filter_width = filter_shape.Dims(2);
-  if (0 == op_params.depth_multiplier % 2 && 0 == filter_width % 2) {
+  if (1 == op_params.dilation_width_factor &&
+      1 == op_params.dilation_height_factor) {
+    RuntimeShape filter_shape = GetTensorShape(filter);
+    const int filter_height = filter_shape.Dims(1);
+    const int filter_width = filter_shape.Dims(2);
     RuntimeShape input_shape = GetTensorShape(input);
     const int input_height = input_shape.Dims(1);
     const int input_width = input_shape.Dims(2);
@@ -310,22 +393,22 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
         output_height, op_params.quantized_activation_min,
         op_params.quantized_activation_max, op_params.output_shift,
         op_params.output_multiplier);
-  } else
-#endif
-
-  {
+  } else {
     tflite::reference_ops::DepthwiseConv(
         op_params, GetTensorShape(input), GetTensorData(input),
         GetTensorShape(filter), GetTensorData(filter),
         GetTensorShape(bias), GetTensorData(bias),
         GetTensorShape(output), GetTensorData(output));
   }
-  return kTfLiteOk;
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  TFLITE_DCHECK(node->user_data != nullptr);
+  TFLITE_DCHECK(node->builtin_data != nullptr);
+
   auto* params =
       reinterpret_cast(node->builtin_data);
+  const OpData& data = *(static_cast(node->user_data));
 
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
@@ -333,51 +416,18 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* bias =
       (NumInputs(node) == 3) ? GetInput(context, node, kBiasTensor) : nullptr;
 
-  const TfLiteType data_type = input->type;
-  int width = SizeOfDimension(input, 2);
-  int height = SizeOfDimension(input, 1);
-  int filter_width = SizeOfDimension(filter, 2);
-  int filter_height = SizeOfDimension(filter, 1);
-
-  OpData data;
-
-  // All per-channel quantized tensors need valid zero point and scale arrays.
-  if (input->type == kTfLiteInt8) {
-    TF_LITE_ENSURE_EQ(context, filter->quantization.type,
-                      kTfLiteAffineQuantization);
-
-    const auto* affine_quantization =
-        reinterpret_cast(
-            filter->quantization.params);
-    TF_LITE_ENSURE(context, affine_quantization);
-    TF_LITE_ENSURE(context, affine_quantization->scale);
-    TF_LITE_ENSURE(context, affine_quantization->zero_point);
-    TF_LITE_ENSURE(
-        context, affine_quantization->scale->size == 1 ||
-                     affine_quantization->scale->size ==
-                         filter->dims->data[kDepthwiseConvQuantizedDimension]);
-    TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size,
-                      affine_quantization->zero_point->size);
-  }
-
-  TF_LITE_ENSURE_STATUS(CalculateOpData(context, node, params, width, height,
-                                        filter_width, filter_height, data_type,
-                                        &data));
-
   // TODO(aselle): Consider whether float conv and quantized conv should be
   // separate ops to avoid dispatch overhead here.
   switch (input->type) {  // Already know in/out types are same.
     case kTfLiteFloat32:
-      return EvalFloat(context, node, params, &data, input, filter, bias,
-                       output);
+      EvalFloat(context, node, params, &data, input, filter, bias, output);
       break;
     case kTfLiteInt8:
-      return EvalQuantizedPerChannel(context, node, params, &data, input,
-                                     filter, bias, output);
+      EvalQuantizedPerChannel(context, node, params, &data, input, filter, bias,
+                              output);
       break;
     case kTfLiteUInt8:
-      return EvalQuantized(context, node, params, &data, input, filter, bias,
-                           output);
+      EvalQuantized(context, node, params, &data, input, filter, bias, output);
       break;
     default:
       TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.",
diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
index 6ae3a14bc96..3fac250f504 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/fully_connected.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -41,6 +41,8 @@ struct OpData {
   int32_t output_activation_max;
   // The index of the temporary tensor where the quantized inputs are cached.
   int input_quantized_index;
+  // Index to buffer for optimizations if applicable.
+  int buffer_idx;
 };
 
 constexpr int kInputTensor = 0;
@@ -49,12 +51,14 @@ constexpr int kBiasTensor = 2;
 constexpr int kOutputTensor = 0;
 
 TfLiteStatus CalculateOpData(TfLiteContext* context,
-                             TfLiteFullyConnectedParams* params,
+                             TfLiteFusedActivation activation,
                              TfLiteType data_type, const TfLiteTensor* input,
                              const TfLiteTensor* filter,
                              const TfLiteTensor* bias, TfLiteTensor* output,
                              OpData* data) {
   TfLiteStatus status = kTfLiteOk;
+  // Set buffer index to a reset value
+  data->buffer_idx = -1;
   if (data_type != kTfLiteFloat32) {
     double real_multiplier = 0.0;
     TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler(
@@ -63,7 +67,7 @@ TfLiteStatus CalculateOpData(TfLiteContext* context,
     QuantizeMultiplier(real_multiplier, &data->output_multiplier, &exponent);
     data->output_shift = -exponent;
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
-        context, params->activation, output, &data->output_activation_min,
+        context, activation, output, &data->output_activation_min,
         &data->output_activation_max));
   }
   return status;
@@ -72,96 +76,148 @@ TfLiteStatus CalculateOpData(TfLiteContext* context,
 }  // namespace
 
 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
-  void* raw;
-  context->AllocatePersistentBuffer(context, sizeof(int), &raw);
-  return raw;
+  TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
+  void* data = nullptr;
+  if (context->AllocatePersistentBuffer(context, sizeof(OpData), &data) ==
+      kTfLiteError) {
+    return nullptr;
+  }
+  return data;
 }
 
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TFLITE_DCHECK(node->user_data != nullptr);
+  TFLITE_DCHECK(node->builtin_data != nullptr);
+
+  OpData* data = static_cast(node->user_data);
+  const auto params =
+      static_cast(node->builtin_data);
+
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor);
+  const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
   TF_LITE_ENSURE_MSG(context, input->type == filter->type,
                      "Hybrid models are not supported on TFLite Micro.");
+  TF_LITE_ENSURE_STATUS(CalculateOpData(context, params->activation,
+                                        input->type, input, filter, bias,
+                                        output, data));
 
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
-  RuntimeShape filter_shape = GetTensorShape(filter);
-  const int filter_dim_count = filter_shape.DimensionsCount();
-  const int accum_depth = filter_shape.Dims(filter_dim_count - 1);
-  const int32_t buf_size = arm_fully_connected_s8_get_buffer_size(accum_depth);
+  if (input->type == kTfLiteInt8 && nullptr != GetTensorData(bias)) {
+    RuntimeShape filter_shape = GetTensorShape(filter);
+    RuntimeShape output_shape = GetTensorShape(output);
 
-  int* buffer_idx = reinterpret_cast(node->user_data);
+    TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 2);
+    const int filter_dim_count = filter_shape.DimensionsCount();
+    cmsis_nn_dims filter_dims;
+    filter_dims.n = filter_shape.Dims(filter_dim_count - 1);
+    filter_dims.h = 1;
+    filter_dims.w = 1;
+    filter_dims.c = output_shape.Dims(1);
 
-  node->user_data = buffer_idx;
-  if (buf_size > 0) {
-    TF_LITE_ENSURE_STATUS(
-        context->RequestScratchBufferInArena(context, buf_size, buffer_idx));
-  } else {
-    *buffer_idx = -1;
+    const int32_t buf_size =
+        arm_fully_connected_s8_get_buffer_size(&filter_dims);
+
+    if (buf_size > 0) {
+      TF_LITE_ENSURE_STATUS(context->RequestScratchBufferInArena(
+          context, buf_size, &data->buffer_idx));
+    } else {
+      data->buffer_idx = -1;
+    }
   }
-#endif
-
   return kTfLiteOk;
 }
 
 TfLiteStatus EvalQuantizedInt8(TfLiteContext* context, TfLiteNode* node,
-                               TfLiteFullyConnectedParams* params, OpData* data,
-                               const TfLiteTensor* input,
+                               const OpData& data, const TfLiteTensor* input,
                                const TfLiteTensor* filter,
                                const TfLiteTensor* bias, TfLiteTensor* output) {
-  RuntimeShape output_shape = GetTensorShape(output);
-  const int batches = output_shape.Dims(0);
-  const int output_depth = output_shape.Dims(1);
-  RuntimeShape filter_shape = GetTensorShape(filter);
-  const int filter_dim_count = filter_shape.DimensionsCount();
-  const int accum_depth = filter_shape.Dims(filter_dim_count - 1);
+  // The 'if' condition can be removed when null handling of bias is added to
+  // arm_fully_connected_s8
+  if (nullptr != GetTensorData(bias)) {
+    RuntimeShape output_shape = GetTensorShape(output);
+    TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 2);
+    const int batches = output_shape.Dims(0);
+    const int output_depth = output_shape.Dims(1);
+    const RuntimeShape filter_shape = GetTensorShape(filter);
+    const int filter_dim_count = filter_shape.DimensionsCount();
+    const int accum_depth = filter_shape.Dims(filter_dim_count - 1);
+    const RuntimeShape input_shape = GetTensorShape(input);
 
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
-  int16_t* buf = nullptr;
+    cmsis_nn_fc_params fc_params;
+    fc_params.input_offset = -input->params.zero_point;
+    fc_params.filter_offset = -filter->params.zero_point;
+    fc_params.output_offset = output->params.zero_point;
+    fc_params.activation.min = data.output_activation_min;
+    fc_params.activation.max = data.output_activation_max;
 
-  auto* buffer_idx = reinterpret_cast(node->user_data);
-  if (*buffer_idx > -1) {
-    void* raw = context->GetScratchBuffer(context, *buffer_idx);
-    buf = reinterpret_cast(raw);
+    cmsis_nn_per_tensor_quant_params quant_params;
+    quant_params.multiplier = data.output_multiplier;
+    // TODO(b/138810107): Figure out whether output shift should be inverted
+    quant_params.shift = -data.output_shift;
+
+    cmsis_nn_dims input_dims;
+    input_dims.n = batches;
+    input_dims.h = input_shape.Dims(1);
+    input_dims.w = input_shape.Dims(2);
+    input_dims.c = input_shape.Dims(3);
+
+    cmsis_nn_dims filter_dims;
+    filter_dims.n = accum_depth;
+    filter_dims.h = 1;
+    filter_dims.w = 1;
+    filter_dims.c = output_depth;
+
+    cmsis_nn_dims bias_dims;
+    bias_dims.n = 1;
+    bias_dims.h = 1;
+    bias_dims.w = 1;
+    bias_dims.c = output_depth;
+
+    cmsis_nn_dims output_dims;
+    output_dims.n = batches;
+    output_dims.h = 1;
+    output_dims.w = 1;
+    output_dims.c = output_depth;
+
+    cmsis_nn_context ctx;
+    ctx.buf = nullptr;
+    ctx.size = 0;
+
+    if (data.buffer_idx > -1) {
+      ctx.buf = context->GetScratchBuffer(context, data.buffer_idx);
+    }
+
+    TF_LITE_ENSURE_EQ(context, arm_fully_connected_s8(
+                                   &ctx, &fc_params, &quant_params, &input_dims,
+                                   GetTensorData(input), &filter_dims,
+                                   GetTensorData(filter), &bias_dims,
+                                   GetTensorData(bias), &output_dims,
+                                   GetTensorData(output)),
+                      ARM_MATH_SUCCESS);
+  } else {
+    tflite::FullyConnectedParams op_params;
+    op_params.input_offset = -input->params.zero_point;
+    op_params.weights_offset = -filter->params.zero_point;
+    op_params.output_offset = output->params.zero_point;
+    op_params.output_multiplier = data.output_multiplier;
+    // TODO(b/138810107): Figure out whether output shift should be inverted
+    op_params.output_shift = -data.output_shift;
+    op_params.quantized_activation_min = data.output_activation_min;
+    op_params.quantized_activation_max = data.output_activation_max;
+
+    reference_integer_ops::FullyConnected(
+        op_params, GetTensorShape(input), GetTensorData(input),
+        GetTensorShape(filter), GetTensorData(filter),
+        GetTensorShape(bias), GetTensorData(bias),
+        GetTensorShape(output), GetTensorData(output));
   }
-
-  TF_LITE_ENSURE_EQ(
-      context,
-      arm_fully_connected_s8(
-          GetTensorData(input), GetTensorData(filter),
-          accum_depth, output_depth, batches, -input->params.zero_point,
-          -filter->params.zero_point, data->output_multiplier,
-          -data->output_shift, output->params.zero_point,
-          GetTensorData(bias), GetTensorData(output),
-          data->output_activation_min, data->output_activation_max, buf),
-      ARM_MATH_SUCCESS);
-#else
-#pragma message( \
-    "CMSIS-NN optimization for fully_connected not available for this target. Using reference kernel.")
-
-  FullyConnectedParams op_params;
-  op_params.input_offset = -input->params.zero_point;
-  op_params.weights_offset = -filter->params.zero_point;
-  op_params.output_offset = output->params.zero_point;
-  op_params.output_multiplier = data->output_multiplier;
-  // TODO(b/138810107): Figure out whether output shift should be inverted
-  op_params.output_shift = -data->output_shift;
-  op_params.quantized_activation_min = data->output_activation_min;
-  op_params.quantized_activation_max = data->output_activation_max;
-
-  reference_integer_ops::FullyConnected(
-      op_params, GetTensorShape(input), GetTensorData(input),
-      GetTensorShape(filter), GetTensorData(filter),
-      GetTensorShape(bias), GetTensorData(bias),
-      GetTensorShape(output), GetTensorData(output));
-#endif
   return kTfLiteOk;
 }
 
 TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
-                           TfLiteFullyConnectedParams* params, OpData* data,
-                           const TfLiteTensor* input,
+                           const OpData& data, const TfLiteTensor* input,
                            const TfLiteTensor* filter, const TfLiteTensor* bias,
                            TfLiteTensor* output) {
   const int32_t input_offset = -input->params.zero_point;
@@ -172,11 +228,11 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
   op_params.input_offset = input_offset;
   op_params.weights_offset = filter_offset;
   op_params.output_offset = output_offset;
-  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_multiplier = data.output_multiplier;
   // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
-  op_params.output_shift = -data->output_shift;
-  op_params.quantized_activation_min = data->output_activation_min;
-  op_params.quantized_activation_max = data->output_activation_max;
+  op_params.output_shift = -data.output_shift;
+  op_params.quantized_activation_min = data.output_activation_min;
+  op_params.quantized_activation_max = data.output_activation_max;
 
 #define TF_LITE_FULLY_CONNECTED(output_data_type)                      \
   reference_ops::FullyConnected(                                       \
@@ -201,11 +257,11 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
 }
 
 TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
-                       TfLiteFullyConnectedParams* params, OpData* data,
+                       TfLiteFusedActivation activation,
                        const TfLiteTensor* input, const TfLiteTensor* filter,
                        const TfLiteTensor* bias, TfLiteTensor* output) {
   float output_activation_min, output_activation_max;
-  CalculateActivationRange(params->activation, &output_activation_min,
+  CalculateActivationRange(activation, &output_activation_min,
                            &output_activation_max);
   tflite::FullyConnectedParams op_params;
   op_params.float_activation_min = output_activation_min;
@@ -219,32 +275,29 @@ TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
-  auto* params =
-      reinterpret_cast(node->builtin_data);
+  TFLITE_DCHECK(node->builtin_data != nullptr);
+  const auto* params =
+      static_cast(node->builtin_data);
 
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor);
   const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TfLiteType data_type = input->type;
-  OpData local_data_object;
-  OpData* data = &local_data_object;
-  TF_LITE_ENSURE_STATUS(CalculateOpData(context, params, data_type, input,
-                                        filter, bias, output, data));
+  TFLITE_DCHECK(node->user_data != nullptr);
+  const OpData& data = *(static_cast(node->user_data));
 
   // Checks in Prepare ensure input, output and filter types are all the same.
   switch (input->type) {
     case kTfLiteFloat32:
-      return EvalFloat(context, node, params, data, input, filter, bias,
+      return EvalFloat(context, node, params->activation, input, filter, bias,
                        output);
     case kTfLiteInt8:
-      return EvalQuantizedInt8(context, node, params, data, input, filter, bias,
+      return EvalQuantizedInt8(context, node, data, input, filter, bias,
                                output);
 
     case kTfLiteUInt8:
-      return EvalQuantized(context, node, params, data, input, filter, bias,
-                           output);
+      return EvalQuantized(context, node, data, input, filter, bias, output);
 
     default:
       TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.",
diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
index 42b3c2e52ff..94f8e928868 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/pooling.cc
@@ -1,4 +1,4 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -76,13 +76,14 @@ void AverageEvalFloat(const TfLiteContext* context, const TfLiteNode* node,
       GetTensorShape(output), GetTensorData(output));
 }
 
-void AverageEvalUint8(TfLiteContext* context, const TfLiteNode* node,
-                      const TfLitePoolParams* params, const OpData* data,
-                      const TfLiteTensor* input, TfLiteTensor* output) {
+void AverageEvalQuantized(TfLiteContext* context, const TfLiteNode* node,
+                          const TfLitePoolParams* params, const OpData* data,
+                          const TfLiteTensor* input, TfLiteTensor* output) {
+  TFLITE_DCHECK(input->type == kTfLiteUInt8 || input->type == kTfLiteInt8);
+
   int32_t activation_min, activation_max;
   (void)CalculateActivationRangeQuantized(context, params->activation, output,
                                           &activation_min, &activation_max);
-
   PoolParams op_params;
   op_params.stride_height = params->stride_height;
   op_params.stride_width = params->stride_width;
@@ -92,76 +93,62 @@ void AverageEvalUint8(TfLiteContext* context, const TfLiteNode* node,
   op_params.padding_values.width = data->padding.width;
   op_params.quantized_activation_min = activation_min;
   op_params.quantized_activation_max = activation_max;
-  reference_ops::AveragePool(
-      op_params, GetTensorShape(input), GetTensorData(input),
-      GetTensorShape(output), GetTensorData(output));
-}
 
-TfLiteStatus AverageEvalInt8(TfLiteContext* context, const TfLiteNode* node,
-                             const TfLitePoolParams* params, const OpData* data,
-                             TfLiteTensor* input, TfLiteTensor* output) {
-  int32_t activation_min, activation_max;
-  (void)CalculateActivationRangeQuantized(context, params->activation, output,
-                                          &activation_min, &activation_max);
+  if (input->type == kTfLiteUInt8) {
+    reference_ops::AveragePool(
+        op_params, GetTensorShape(input), GetTensorData(input),
+        GetTensorShape(output), GetTensorData(output));
+  } else {
+    TFLITE_DCHECK_LE(activation_min, activation_max);
 
-  TFLITE_DCHECK_LE(activation_min, activation_max);
+    RuntimeShape input_shape = GetTensorShape(input);
+    TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
 
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
-  RuntimeShape input_shape = GetTensorShape(input);
-  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+    RuntimeShape output_shape = GetTensorShape(output);
+    TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
 
-  RuntimeShape output_shape = GetTensorShape(output);
-  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+    const int depth = MatchingDim(input_shape, 3, output_shape, 3);
 
-  const int depth = MatchingDim(input_shape, 3, output_shape, 3);
-  const int input_height = input_shape.Dims(1);
-  const int input_width = input_shape.Dims(2);
-  const int output_height = output_shape.Dims(1);
-  const int output_width = output_shape.Dims(2);
-  const int stride_height = params->stride_height;
-  const int stride_width = params->stride_width;
+    cmsis_nn_dims input_dims;
+    input_dims.n = 1;
+    input_dims.h = input_shape.Dims(1);
+    input_dims.w = input_shape.Dims(2);
+    input_dims.c = depth;
 
-  const int filter_height = params->filter_height;
-  const int filter_width = params->filter_width;
-  const int padding_height = data->padding.height;
-  const int padding_width = data->padding.width;
+    cmsis_nn_dims output_dims;
+    output_dims.n = 1;
+    output_dims.h = output_shape.Dims(1);
+    output_dims.w = output_shape.Dims(2);
+    output_dims.c = depth;
 
-  int16_t* scratch_buffer = nullptr;
+    cmsis_nn_pool_params pool_params;
+    pool_params.stride.h = params->stride_height;
+    pool_params.stride.w = params->stride_width;
+    pool_params.padding.h = data->padding.height;
+    pool_params.padding.w = data->padding.width;
+    pool_params.activation.min = activation_min;
+    pool_params.activation.max = activation_max;
 
-  auto* buffer_idx = reinterpret_cast(node->user_data);
+    cmsis_nn_dims filter_dims;
+    filter_dims.n = 1;
+    filter_dims.h = params->filter_height;
+    filter_dims.w = params->filter_width;
+    filter_dims.c = 1;
 
-  if (*buffer_idx > -1) {
-    void* raw = context->GetScratchBuffer(context, *buffer_idx);
-    scratch_buffer = reinterpret_cast(raw);
+    auto* buffer_idx = reinterpret_cast(node->user_data);
+    cmsis_nn_context ctx;
+    ctx.buf = nullptr;
+    ctx.size = 0;
+    if (*buffer_idx > -1) {
+      ctx.buf = context->GetScratchBuffer(context, *buffer_idx);
+    }
+
+    TFLITE_DCHECK_EQ(
+        arm_avgpool_s8(&ctx, &pool_params, &input_dims,
+                       GetTensorData(input), &filter_dims, &output_dims,
+                       GetTensorData(output)),
+        ARM_MATH_SUCCESS);
   }
-
-  TF_LITE_ENSURE_EQ(
-      context,
-      arm_avgpool_s8(input_height, input_width, output_height, output_width,
-                     stride_height, stride_width, filter_height, filter_width,
-                     padding_height, padding_width, activation_min,
-                     activation_max, depth, GetTensorData(input),
-                     scratch_buffer, GetTensorData(output)),
-      ARM_MATH_SUCCESS);
-#else
-#pragma message( \
-    "CMSIS-NN optimization for avg_pool not available for this target. Using reference kernel.")
-
-  PoolParams op_params;
-  op_params.stride_height = params->stride_height;
-  op_params.stride_width = params->stride_width;
-  op_params.filter_height = params->filter_height;
-  op_params.filter_width = params->filter_width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.padding_values.width = data->padding.width;
-  op_params.quantized_activation_min = activation_min;
-  op_params.quantized_activation_max = activation_max;
-  reference_integer_ops::AveragePool(
-      op_params, GetTensorShape(input), GetTensorData(input),
-      GetTensorShape(output), GetTensorData(output));
-
-#endif
-  return kTfLiteOk;
 }
 
 void MaxEvalFloat(TfLiteContext* context, TfLiteNode* node,
@@ -215,7 +202,6 @@ TfLiteStatus MaxEvalInt8(TfLiteContext* context, const TfLiteNode* node,
 
   TFLITE_DCHECK_LE(activation_min, activation_max);
 
-#if defined(__ARM_FEATURE_DSP)
   RuntimeShape input_shape = GetTensorShape(input);
   TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
 
@@ -245,32 +231,14 @@ TfLiteStatus MaxEvalInt8(TfLiteContext* context, const TfLiteNode* node,
   }
 
   TF_LITE_ENSURE_EQ(
-      context,
-      arm_max_pool_s8_opt(input_height, input_width, output_height,
-                          output_width, stride_height, stride_width,
-                          filter_height, filter_width, padding_height,
-                          padding_width, activation_min, activation_max, depth,
-                          GetTensorData(input), scratch_buffer,
-                          GetTensorData(output)),
+      context, arm_max_pool_s8_opt(
+                   input_height, input_width, output_height, output_width,
+                   stride_height, stride_width, filter_height, filter_width,
+                   padding_height, padding_width, activation_min,
+                   activation_max, depth, GetTensorData(input),
+                   scratch_buffer, GetTensorData(output)),
       ARM_MATH_SUCCESS);
-#else
-#pragma message( \
-    "CMSIS-NN optimization for max_pool not available for this target. Using reference kernel.")
 
-  PoolParams op_params;
-  op_params.stride_height = params->stride_height;
-  op_params.stride_width = params->stride_width;
-  op_params.filter_height = params->filter_height;
-  op_params.filter_width = params->filter_width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.padding_values.width = data->padding.width;
-  op_params.quantized_activation_min = activation_min;
-  op_params.quantized_activation_max = activation_max;
-  reference_integer_ops::MaxPool(
-      op_params, GetTensorShape(input), GetTensorData(input),
-      GetTensorShape(output), GetTensorData(output));
-
-#endif
   return kTfLiteOk;
 }
 
@@ -283,32 +251,33 @@ void* Init(TfLiteContext* context, const char* buffer, size_t length) {
 }
 
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
-#if defined(__ARM_FEATURE_DSP) || defined(__ARM_FEATURE_MVE)
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  const TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  RuntimeShape input_shape = GetTensorShape(input);
-  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  if (input->type == kTfLiteInt8) {
+    const TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  RuntimeShape output_shape = GetTensorShape(output);
-  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+    RuntimeShape input_shape = GetTensorShape(input);
+    TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
 
-  const int depth = MatchingDim(input_shape, 3, output_shape, 3);
-  const int output_width = output_shape.Dims(2);
+    RuntimeShape output_shape = GetTensorShape(output);
+    TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
 
-  const int32_t buffer_size =
-      arm_avgpool_s8_get_buffer_size(output_width, depth);
+    const int depth = MatchingDim(input_shape, 3, output_shape, 3);
+    const int output_width = output_shape.Dims(2);
 
-  int* buffer_idx = reinterpret_cast(node->user_data);
+    const int32_t buffer_size =
+        arm_avgpool_s8_get_buffer_size(output_width, depth);
 
-  node->user_data = buffer_idx;
-  if (buffer_size > 0) {
-    TF_LITE_ENSURE_STATUS(
-        context->RequestScratchBufferInArena(context, buffer_size, buffer_idx));
-  } else {
-    *buffer_idx = -1;
+    int* buffer_idx = reinterpret_cast(node->user_data);
+
+    node->user_data = buffer_idx;
+    if (buffer_size > 0) {
+      TF_LITE_ENSURE_STATUS(context->RequestScratchBufferInArena(
+          context, buffer_size, buffer_idx));
+    } else {
+      *buffer_idx = -1;
+    }
   }
-#endif
   return kTfLiteOk;
 }
 
@@ -316,9 +285,7 @@ TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
   auto* params = reinterpret_cast(node->builtin_data);
   OpData data;
 
-  // Todo: make 'input' const once CMSIS-reuse is fixed
-  TfLiteTensor* input = &context->tensors[flatbuffers::EndianScalar(
-      node->inputs->data[kInputTensor])];
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   TF_LITE_ENSURE_STATUS(CalculateOpData(context, params, input, output, &data));
@@ -329,10 +296,8 @@ TfLiteStatus AverageEval(TfLiteContext* context, TfLiteNode* node) {
       AverageEvalFloat(context, node, params, &data, input, output);
       break;
     case kTfLiteUInt8:
-      AverageEvalUint8(context, node, params, &data, input, output);
-      break;
     case kTfLiteInt8:
-      return AverageEvalInt8(context, node, params, &data, input, output);
+      AverageEvalQuantized(context, node, params, &data, input, output);
       break;
     default:
       TF_LITE_KERNEL_LOG(context, "Input type %s is not currently supported",
@@ -387,7 +352,7 @@ TfLiteRegistration* Register_AVERAGE_POOL_2D() {
 TfLiteRegistration* Register_MAX_POOL_2D() {
   static TfLiteRegistration r = {/*init=*/pooling::Init,
                                  /*free=*/nullptr,
-                                 /*prepare=*/pooling::Prepare,
+                                 /*prepare=*/nullptr,
                                  /*invoke=*/pooling::MaxEval,
                                  /*profiling_string=*/nullptr,
                                  /*builtin_code=*/0,
diff --git a/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc b/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc
index d445a6f7b37..32874f13dbf 100644
--- a/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc
+++ b/tensorflow/lite/micro/tools/make/ext_libs/cmsis.inc
@@ -55,6 +55,7 @@ ifneq ($(filter cmsis-nn,$(ALL_TAGS)),)
       $(CMSIS_PATH)/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_1_x_n_s8.c \
       $(CMSIS_PATH)/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_HWC_q7_basic_nonsquare.c \
       $(CMSIS_PATH)/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_wrapper_s8.c \
+      $(CMSIS_PATH)/CMSIS/NN/Source/ConvolutionFunctions/arm_depthwise_conv_wrapper_s8.c \
       $(CMSIS_PATH)/CMSIS/NN/Source/PoolingFunctions/arm_pool_q7_HWC.c \
       $(CMSIS_PATH)/CMSIS/NN/Source/PoolingFunctions/arm_max_pool_s8_opt.c \
       $(CMSIS_PATH)/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c \
diff --git a/tensorflow/lite/micro/tools/make/third_party_downloads.inc b/tensorflow/lite/micro/tools/make/third_party_downloads.inc
index 28d3f3ab529..2f2f9396dc0 100644
--- a/tensorflow/lite/micro/tools/make/third_party_downloads.inc
+++ b/tensorflow/lite/micro/tools/make/third_party_downloads.inc
@@ -28,8 +28,8 @@ LEON_BCC2_MD5 := "cdf78082be4882da2a92c9baa82fe765"
 TSIM_URL := "https://www.gaisler.com/anonftp/tsim/tsim-eval-2.0.63.tar.gz"
 TSIM_MD5 := "afa0095d3ed989a949e1467f94e41d2f"
 
-CMSIS_URL := "https://github.com/ARM-software/CMSIS_5/archive/1150e71e07c79b538efd842aba5b210a31827ae5.zip"
-CMSIS_MD5 := "e05f4222ef58825193910b41a0871dcb"
+CMSIS_URL := "https://github.com/ARM-software/CMSIS_5/archive/0f1587564506b385d57a58baed8c2c6a1e2b959d.zip"
+CMSIS_MD5 := "b7bf586417df9ed586d50cb9b885509f"
 
 AM_SDK_URL := "http://s3.asia.ambiqmicro.com/downloads/AmbiqSuite-Rel2.2.0.zip"
 AM_SDK_MD5 := "7605fa2d4d97e6bb7a1190c92b66b597"

From e4af590df8ff7685de1679f2a6eb01a677d84775 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Thu, 18 Jun 2020 01:35:40 -0700
Subject: [PATCH 0481/1390] Properly configure block and grid dimensions when
 launching generated kernels.

Also prepare the possibility to specify unrolling. This is not enabled yet
because there are some LLVM changes required.

PiperOrigin-RevId: 317056534
Change-Id: I3de5dda52d80b528c4bd0026a5e160fda4296c32
---
 .../mlir_generated_cwise_op_gpu_tanh.cu.cc    | 71 ++++++++++++++-----
 1 file changed, 53 insertions(+), 18 deletions(-)

diff --git a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
index 40dd7c7e49e..70de777239f 100644
--- a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
+++ b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include 
+#include 
 
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
@@ -45,9 +46,41 @@ Status CreateKernel(absl::string_view kernel_name, uint64_t num_args,
   return stream_exec->GetKernel(loader_spec, kernel_base.get());
 }
 
-class MlirGenerateTanhOp : public OpKernel {
+struct LaunchConfig {
+  se::BlockDim blockDim;
+  se::ThreadDim threadDim;
+};
+
+LaunchConfig GetLaunchConfiguration(std::vector tile_sizes,
+                                    std::vector unrolling_factors,
+                                    std::vector shape) {
+  LaunchConfig result;
+  // Ensure the vectors are length 3 and pad with ones.
+  tile_sizes.resize(3, 1);
+  unrolling_factors.resize(3, 1);
+  shape.resize(3, 1);
+  // The number of threads is given by the tiling size.
+  result.threadDim = se::ThreadDim(tile_sizes[0], tile_sizes[1], tile_sizes[2]);
+  // We know that the kernel was generated by mapping the three outer-most
+  // dimensions to x,y,z dimensions. So we only need to compute those.
+  std::vector block_dims(3);
+  for (int i = 0; i < 3; ++i) {
+    // Compute the number of grids. We use ceildiv here as we have to allocate
+    // an extra thread/block if the division is not even. The kernel contains
+    // code to handle the boundaries.
+    int number_of_threads =
+        (shape[i] + unrolling_factors[i] - 1) / unrolling_factors[i];
+    int number_of_grids =
+        (number_of_threads + tile_sizes[i] - 1) / tile_sizes[i];
+    block_dims[i] = number_of_grids;
+  }
+  result.blockDim = se::BlockDim(block_dims[0], block_dims[1], block_dims[2]);
+  return result;
+}
+
+class MlirGeneratedTanhOp : public OpKernel {
  public:
-  explicit MlirGenerateTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  explicit MlirGeneratedTanhOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
     auto* stream = ctx->op_device_context()->stream();
@@ -88,11 +121,13 @@ class MlirGenerateTanhOp : public OpKernel {
     args.add_argument(inp.NumElements());
     args.add_argument(1);
 
-    // TODO(b/158649746): Choose block size and thread dim according to the
-    // number of input elements. For now, this supports at most 1024 elements.
+    // This has to be aligned with the configuration that was used when building
+    // the kernels. See the corresponding build rules in `cubin_headers/BUILD`.
+    LaunchConfig config = GetLaunchConfiguration(
+        {256}, {}, {static_cast(inp.NumElements())});
     OP_REQUIRES_OK(
-        ctx, stream->parent()->Launch(stream, se::ThreadDim(inp.NumElements()),
-                                      se::BlockDim(1), *kernel, args));
+        ctx, stream->parent()->Launch(stream, config.threadDim, config.blockDim,
+                                      *kernel, args));
   }
 
  protected:
@@ -103,26 +138,26 @@ class MlirGenerateTanhOp : public OpKernel {
   std::mutex mu_;
 };
 
-class MlirGenerateTanhF16Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF16Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF16Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF16Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF16Kernel;
   }
 };
 
-class MlirGenerateTanhF32Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF32Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF32Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF32Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF32Kernel;
   }
 };
 
-class MlirGenerateTanhF64Op : public MlirGenerateTanhOp {
+class MlirGeneratedTanhF64Op : public MlirGeneratedTanhOp {
  public:
-  explicit MlirGenerateTanhF64Op(OpKernelConstruction* ctx)
-      : MlirGenerateTanhOp(ctx) {
+  explicit MlirGeneratedTanhF64Op(OpKernelConstruction* ctx)
+      : MlirGeneratedTanhOp(ctx) {
     cubin_data_ = kTanhF64Kernel;
   }
 };
@@ -130,11 +165,11 @@ class MlirGenerateTanhF64Op : public MlirGenerateTanhOp {
 
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF16Op);
+    MlirGeneratedTanhF16Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF32Op);
+    MlirGeneratedTanhF32Op);
 REGISTER_KERNEL_BUILDER(
     Name("Tanh").Device(DEVICE_GPU).TypeConstraint("T"),
-    MlirGenerateTanhF64Op);
+    MlirGeneratedTanhF64Op);
 }  // namespace tensorflow

From ad1434444bda1e8321dcc09965a5dce8da847eed Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 01:51:06 -0700
Subject: [PATCH 0482/1390] Remove dynamic dimension of strided slice grad if
 input to strided slice is static.

If we slice a dynamic shaped tensor from a static tensor, the output
of the gradient should still be static. Unfortunately this cannot be
deduced alone by xla, so extra information is needed from the tf2xla
bridge.

PiperOrigin-RevId: 317058146
Change-Id: I33e4895e169c238ad3d73a57ada11c4984d11dfb
---
 .../tf2xla/kernels/strided_slice_op.cc        | 22 ----------------
 tensorflow/compiler/xla/client/xla_builder.cc | 23 ----------------
 tensorflow/compiler/xla/client/xla_builder.h  |  6 -----
 .../compiler/xla/client/xla_builder_test.cc   | 26 -------------------
 4 files changed, 77 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
index 51764018df1..2684c982600 100644
--- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
@@ -350,28 +350,6 @@ class StridedSliceGradOp : public XlaOpKernel {
       grad = xla::Rev(grad, dimensions_to_reverse);
     }
     grad = xla::Pad(grad, zero, padding_config);
-
-    xla::XlaOp dynamic_shape = ctx->Input(0);
-    xla::Shape grad_shape = ctx->builder()->GetShape(grad).ValueOrDie();
-    ctx->set_dynamic_dimension_is_minus_one(true);
-    std::vector dynamic_size;
-    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(0, &dynamic_size));
-    // Input of strided_slice_op has to have the same shape as output.
-    DCHECK_EQ(grad_shape.rank(), input_shape.dims());
-    for (int64 dim = 0; dim < input_shape.dims(); ++dim) {
-      DCHECK_EQ(grad_shape.dimensions(dim), input_shape.dim_size(dim));
-      if (dynamic_size[dim] == -1) {
-        // Input is a dynamic dimension, set the same dynamic dimension size in
-        // the output.
-        auto dim_size = xla::Slice(dynamic_shape, {dim}, {dim + 1}, {1});
-        grad = xla::SetDimensionSize(grad, dim_size, dim);
-      } else if (grad_shape.is_dynamic_dimension(dim)) {
-        // Input is static but output is dynamic, respect input and remove any
-        // dynamic dim in the output.
-        grad = xla::RemoveDynamicDimension(grad, dim);
-      }
-    }
-
     ctx->SetOutput(0, grad);
   }
 
diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index c7b6a7f9491..bfba48862f6 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -2727,25 +2727,6 @@ XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64 dimension) {
   });
 }
 
-XlaOp XlaBuilder::RemoveDynamicDimension(XlaOp operand, int64 dimension) {
-  return ReportErrorOrReturn([&]() -> StatusOr {
-    HloInstructionProto instr;
-    TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
-
-    Shape shape = *operand_shape;
-    shape.set_dynamic_dimension(dimension, false);
-    // Setting an op's dynamic dimension to its static size removes the dynamic
-    // dimension.
-    XlaOp static_size =
-        ConstantR0(this, operand_shape->dimensions(dimension));
-
-    *instr.mutable_shape() = shape.ToProto();
-    instr.add_dimensions(dimension);
-    return AddInstruction(std::move(instr), HloOpcode::kSetDimensionSize,
-                          {operand, static_size});
-  });
-}
-
 XlaOp XlaBuilder::SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension) {
   return ReportErrorOrReturn([&]() -> StatusOr {
     HloInstructionProto instr;
@@ -3846,8 +3827,4 @@ XlaOp SetDimensionSize(const XlaOp operand, const XlaOp val, int64 dimension) {
   return operand.builder()->SetDimensionSize(operand, val, dimension);
 }
 
-XlaOp RemoveDynamicDimension(const XlaOp operand, int64 dimension) {
-  return operand.builder()->RemoveDynamicDimension(operand, dimension);
-}
-
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h
index b8af180b83e..ffa6a7c3439 100644
--- a/tensorflow/compiler/xla/client/xla_builder.h
+++ b/tensorflow/compiler/xla/client/xla_builder.h
@@ -704,8 +704,6 @@ class XlaBuilder {
 
   XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
-  XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
-
   StatusOr AddInstruction(HloInstructionProto&& instr, HloOpcode opcode,
                                  absl::Span operands = {});
 
@@ -1153,7 +1151,6 @@ class XlaBuilder {
 
   friend XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
   friend XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
-  friend XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
 
  protected:
   // Returns OK status if the given op was built using this builder. Otherwise,
@@ -2152,9 +2149,6 @@ XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
 
 XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
-// Returns the same op but with dynamic dimension removed.
-XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
-
 // Implementation details below this point.
 //
 
diff --git a/tensorflow/compiler/xla/client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_builder_test.cc
index 7011c946203..4fa47077fca 100644
--- a/tensorflow/compiler/xla/client/xla_builder_test.cc
+++ b/tensorflow/compiler/xla/client/xla_builder_test.cc
@@ -556,32 +556,6 @@ TEST_F(XlaBuilderTest, DynamicParameter) {
   EXPECT_TRUE(param_shape.is_dynamic_dimension(0));
 }
 
-TEST_F(XlaBuilderTest, SetDimensionSize) {
-  XlaBuilder b(TestName());
-  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
-  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
-  auto set_dim_size = SetDimensionSize(p0, p1, 0);
-  TF_ASSERT_OK_AND_ASSIGN(auto module,
-                          BuildHloModule(&b, /*root=*/set_dim_size));
-  const Shape& root_shape =
-      module->entry_computation()->root_instruction()->shape();
-  EXPECT_TRUE(root_shape.is_dynamic_dimension(0));
-}
-
-TEST_F(XlaBuilderTest, RemoveDimensionSize) {
-  XlaBuilder b(TestName());
-  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
-  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
-  auto set_dim_size = SetDimensionSize(p0, p1, 0);
-  auto remove_dim_size = RemoveDynamicDimension(set_dim_size, 0);
-  TF_ASSERT_OK_AND_ASSIGN(auto module,
-                          BuildHloModule(&b, /*root=*/remove_dim_size));
-  const Shape& root_shape =
-      module->entry_computation()->root_instruction()->shape();
-  // Dynamic dimension has been removed.
-  EXPECT_FALSE(root_shape.is_dynamic_dimension(0));
-}
-
 TEST_F(XlaBuilderTest, DynamicUnary) {
   XlaBuilder b(TestName());
   Shape tuple_param_shape = ShapeUtil::MakeTupleShape(

From 50f0fcc0005d3afec88b281a07563a603dbe0c7c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 02:01:45 -0700
Subject: [PATCH 0483/1390] compat: Update forward compatibility horizon to
 2020-06-18

PiperOrigin-RevId: 317059247
Change-Id: I61e0e6659103a5b413a0013fb0782ba0d470d285
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 8a42b3dfdd3..32545ac8463 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 17)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 18)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 39c6a3b2cd0193eae45a051344921a3c18a2e691 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 02:01:54 -0700
Subject: [PATCH 0484/1390] Update GraphDef version to 436.

PiperOrigin-RevId: 317059270
Change-Id: I2eb1760fec06cae5cf58c41effa3263a0bdde142
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 3e4e3888d87..546d86e58fa 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 435  // Updated: 2020/6/17
+#define TF_GRAPH_DEF_VERSION 436  // Updated: 2020/6/18
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From dd49e65c5b68c4b8113dfe5aadb988fdeb2abd57 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer 
Date: Thu, 18 Jun 2020 02:38:26 -0700
Subject: [PATCH 0485/1390] [XLA:CPU] Fusion: Only check for reuse on expensive
 instructions

The reuse condition walks over all instructions in the fusion, the fusion pass
walks over all instructions, making this essentially quadratic. Moving the
is_expensive check up doesn't completely avoid this behavior, but makes it much
more unlikely.

PiperOrigin-RevId: 317063582
Change-Id: I22459aa922e6d65c6c639ed81208d1d441a132bc
---
 .../compiler/xla/service/cpu/cpu_instruction_fusion.cc       | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc
index 97e0a518499..9460cc55e10 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc
@@ -94,9 +94,8 @@ bool CpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
 
   // Cost condition: not fuse (simple, expensive producers) and (consumers who
   // reuse operand elements).
-  if (producer->opcode() != HloOpcode::kFusion &&
-      consumer->ReusesOperandElements(operand_index) &&
-      is_expensive(*producer)) {
+  if (producer->opcode() != HloOpcode::kFusion && is_expensive(*producer) &&
+      consumer->ReusesOperandElements(operand_index)) {
     VLOG(2) << "Fusion is not profitable.";
     return false;
   }

From c21328bd992fe359b585452df92fc82976d27557 Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Thu, 18 Jun 2020 12:06:40 +0200
Subject: [PATCH 0486/1390] Don't skip distributed test case in eager mode

---
 .../mixed_precision/experimental/autocast_variable_test.py    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 940bd07c813..2fa7c103258 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -304,8 +304,8 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         self.assertAllClose(3., self.evaluate(x.assign_sub(3.)))
 
         # Assign multiple times
-        # This currently doesn't work in graph mode
-        if context.executing_eagerly() or ops.inside_function():
+        # This currently doesn't work in graph mode if a strategy is used
+        if not ds_context.has_strategy() or context.executing_eagerly():
           assign = x.assign(1.)
           self.assertAllClose(1., self.evaluate(assign))
           self.assertAllClose(0., self.evaluate(assign.assign(0.)))

From de5620b74cb1688dcc81f3e50ad8fc24b18beb9b Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Thu, 18 Jun 2020 04:33:29 -0700
Subject: [PATCH 0487/1390] Add option in kernel_lowering to use the tanh
 approximation.

PiperOrigin-RevId: 317077338
Change-Id: I5cc7bfb84d6defba05439186377f90e422247d68
---
 tensorflow/compiler/xla/service/mlir_gpu/BUILD              | 1 +
 tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc | 5 +++++
 tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h  | 1 +
 3 files changed, 7 insertions(+)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/BUILD b/tensorflow/compiler/xla/service/mlir_gpu/BUILD
index ce45d937424..efe69450846 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/BUILD
+++ b/tensorflow/compiler/xla/service/mlir_gpu/BUILD
@@ -167,6 +167,7 @@ cc_library(
         "//tensorflow/compiler/mlir/xla:lhlo_legalize_to_affine",
         "//tensorflow/compiler/mlir/xla:lhlo_legalize_to_gpu",
         "//tensorflow/compiler/mlir/xla:xla_dialect_registration",
+        "//tensorflow/compiler/mlir/xla:xla_legalize_tanh_to_approximation",
         "//tensorflow/compiler/mlir/xla:xla_legalize_to_linalg",
         "//tensorflow/compiler/xla:status",
         "//tensorflow/compiler/xla:statusor",
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index b0cbddcdb92..9d5b52df010 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -505,6 +505,11 @@ Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   // Some basic cleanup.
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCanonicalizerPass());
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCSEPass());
+  // Approximate of requested.
+  if (options.use_approximations) {
+    pm.addNestedPass<::mlir::FuncOp>(
+        ::mlir::xla::createLegalizeTanhToApproximationPass());
+  }
   // Move scalar operations into the launch to ensure smaller signatures.
   pm.addPass(absl::make_unique());
   // Take launches to launches with kernels.
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
index 77cf75b9e47..bd633bb06cb 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.h
@@ -28,6 +28,7 @@ struct LowerLHLOToGPUOptions {
   llvm::ArrayRef unroll_factors = {};
   bool collapse_parallel_loops = true;
   bool rewrite_signature = true;
+  bool use_approximations = false;
 };
 
 Status LowerLHLOToGPU(mlir::ModuleOp module,

From 16e6f9e792741b98d27b4f1463057313d04acdd8 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Thu, 18 Jun 2020 05:00:40 -0700
Subject: [PATCH 0488/1390] Enable unrolling for the generated tanh kernel.

PiperOrigin-RevId: 317080498
Change-Id: Idaac4b2efee78cd2cd44c68ede4ac56f1a4cde40
---
 tensorflow/core/kernels/cubin_headers/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/cubin_headers/BUILD b/tensorflow/core/kernels/cubin_headers/BUILD
index b8ba164fbc3..ec8b44050db 100644
--- a/tensorflow/core/kernels/cubin_headers/BUILD
+++ b/tensorflow/core/kernels/cubin_headers/BUILD
@@ -37,4 +37,5 @@ gen_kernel_library(
         "f32",
         "f64",
     ],
+    unroll_factors = "4",
 )

From 0a541ad1cc89f1eead6a47dc1676bd86dc810937 Mon Sep 17 00:00:00 2001
From: Christian Sigg 
Date: Thu, 18 Jun 2020 05:07:41 -0700
Subject: [PATCH 0489/1390] Remove intermediate relocatable code stored in
 __nv_relfatbin sections, if objcopy is at least version 2.26 (which added
 support for --update-sections).

The intermediate code is a result of separate compilation and linking, and removing it reduces TF's GPU wheel size.

PiperOrigin-RevId: 317081343
Change-Id: I603477b4499344aeec653765be78de11f392eac6
---
 third_party/nccl/build_defs.bzl.tpl | 103 ++++++++++++++++++++++++----
 1 file changed, 88 insertions(+), 15 deletions(-)

diff --git a/third_party/nccl/build_defs.bzl.tpl b/third_party/nccl/build_defs.bzl.tpl
index 9268af7c890..b520f71d0f1 100644
--- a/third_party/nccl/build_defs.bzl.tpl
+++ b/third_party/nccl/build_defs.bzl.tpl
@@ -169,35 +169,94 @@ _device_link = rule(
 )
 """Links device code and generates source code for kernel registration."""
 
+def _prune_relocatable_code_impl(ctx):
+    """Clears __nv_relfatbin section containing relocatable device code."""
+    empty_file = ctx.actions.declare_file(ctx.attr.name + "__nv_relfatbin")
+    ctx.actions.write(empty_file, "")
+
+    # Parse 'objcopy --version' and update section if it's at least v2.26.
+    # Otherwise, simply copy the file without changing it.
+    # TODO(csigg): version parsing is brittle, can we do better?
+    command = r"""
+        objcopy=$1                                         \
+        section=$2                                         \
+        input=$3                                           \
+        output=$4                                          \
+        args=""                                            \
+        pattern='([0-9])\.([0-9]+)';                       \
+        if [[ $($objcopy --version) =~ $pattern ]] && {    \
+            [ ${BASH_REMATCH[1]} -gt 2 ] ||                \
+            [ ${BASH_REMATCH[2]} -ge 26 ]; }; then         \
+          args="--update-section __nv_relfatbin=$section"; \
+        fi;                                                \
+        $objcopy $args $input $output
+    """
+    cc_toolchain = find_cpp_toolchain(ctx)
+    outputs = []
+    for src in ctx.files.srcs:
+        out = ctx.actions.declare_file("pruned_" + src.basename, sibling = src)
+        ctx.actions.run_shell(
+            inputs = [empty_file] + ctx.files.srcs,  # + ctx.files._crosstool,
+            outputs = [out],
+            arguments = [
+                cc_toolchain.objcopy_executable,
+                empty_file.path,
+                src.path,
+                out.path,
+            ],
+            command = command,
+        )
+        outputs.append(out)
+    return DefaultInfo(files = depset(outputs))
+
+_prune_relocatable_code = rule(
+    implementation = _prune_relocatable_code_impl,
+    attrs = {
+        "srcs": attr.label_list(mandatory = True, allow_files = True),
+        "_cc_toolchain": attr.label(
+            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+        ),
+        # "_crosstool": attr.label_list(
+        #     cfg = "host",
+        #     default = ["@bazel_tools//tools/cpp:crosstool"]
+        # ),
+    },
+)
+
 def _merge_archive_impl(ctx):
     # Generate an mri script to the merge archives in srcs and pass it to 'ar'.
     # See https://stackoverflow.com/a/23621751.
     files = _pic_only(ctx.files.srcs)
     mri_script = "create " + ctx.outputs.out.path
     for f in files:
-        mri_script += "\\naddlib " + f.path
-    mri_script += "\\nsave\\nend"
+        mri_script += r"\naddlib " + f.path
+    mri_script += r"\nsave\nend"
 
     cc_toolchain = find_cpp_toolchain(ctx)
     ctx.actions.run_shell(
         inputs = ctx.files.srcs,  # + ctx.files._crosstool,
         outputs = [ctx.outputs.out],
-        command = "printf \"%s\" | %s -M" % (mri_script, cc_toolchain.ar_executable),
+        command = "echo -e \"%s\" | %s -M" % (mri_script, cc_toolchain.ar_executable),
     )
 
 _merge_archive = rule(
     implementation = _merge_archive_impl,
     attrs = {
         "srcs": attr.label_list(mandatory = True, allow_files = True),
-        "_cc_toolchain": attr.label(default = "@bazel_tools//tools/cpp:current_cc_toolchain"),
-        # "_crosstool": attr.label_list(cfg = "host", default = ["@bazel_tools//tools/cpp:crosstool"]),
+        "_cc_toolchain": attr.label(
+            default = "@bazel_tools//tools/cpp:current_cc_toolchain",
+        ),
+        # "_crosstool": attr.label_list(
+        #     cfg = "host",
+        #     default = ["@bazel_tools//tools/cpp:crosstool"]
+        # ),
     },
     outputs = {"out": "lib%{name}.a"},
 )
 """Merges srcs into a single archive."""
 
 def cuda_rdc_library(name, hdrs = None, copts = None, linkstatic = True, **kwargs):
-    """Produces a cuda_library using separate compilation and linking.
+    r"""Produces a cuda_library using separate compilation and linking.
 
     CUDA separate compilation and linking allows device function calls across
     translation units. This is different from the normal whole program
@@ -239,17 +298,24 @@ def cuda_rdc_library(name, hdrs = None, copts = None, linkstatic = True, **kwarg
 
     The steps marked with '*' are implemented in the _device_link rule.
 
+    The intermediate relocatable device code in xy.a is no longer needed at
+    this point and the corresponding section is replaced with an empty one using
+    objcopy. We do not remove the section completely because it is referenced by
+    relocations, and removing those as well breaks fatbin registration.
+
     The object files in both xy.a and dlink.a reference symbols defined in the
     other archive. The separate archives are a side effect of using two
     cc_library targets to implement a single compilation trajectory. We could
     fix this once bazel supports C++ sandwich. For now, we just merge the two
     archives to avoid unresolved symbols:
 
-    xy.a      dlink.a
-        \    /           merge archive
-      xy_dlink.a
-           |             cc_library (or alternatively, cc_import)
-     final target
+                    xy.a
+                     |         objcopy --update-section __nv_relfatbin=''
+    dlink.a     xy_pruned.a
+         \           /         merge archive
+          xy_merged.a
+              |                cc_library (or alternatively, cc_import)
+         final target
 
     Another complication is that cc_library produces (depending on the
     configuration) both PIC and non-PIC archives, but the distinction
@@ -313,19 +379,26 @@ def cuda_rdc_library(name, hdrs = None, copts = None, linkstatic = True, **kwarg
         linkstatic = linkstatic,
     )
 
+    # Remove intermediate relocatable device code.
+    pruned = name + "_pruned"
+    _prune_relocatable_code(
+        name = pruned,
+        srcs = [lib],
+    )
+
     # Repackage the two libs into a single archive. This is required because
     # both libs reference symbols defined in the other one. For details, see
     # https://eli.thegreenplace.net/2013/07/09/library-order-in-static-linking
-    archive = name + "_a"
+    merged = name + "_merged"
     _merge_archive(
-        name = archive,
-        srcs = [lib, dlink],
+        name = merged,
+        srcs = [pruned, dlink],
     )
 
     # Create cc target from archive.
     native.cc_library(
         name = name,
-        srcs = [archive],
+        srcs = [merged],
         hdrs = hdrs,
         linkstatic = linkstatic,
     )

From 7378fabf90f489387e7dbfb111e20ddf03b1910a Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Thu, 18 Jun 2020 05:19:27 -0700
Subject: [PATCH 0490/1390] Emit tuple at the end of the Sort emitter.

If sort has more than one operand, the result is a tuple. So far, we didn't
emit the tuple at the end of the emitter.

PiperOrigin-RevId: 317082573
Change-Id: I7bec31302ba2e40556b17654daa081428871a00e
---
 .../xla/service/gpu/ir_emitter_unnested.cc    |  7 ++++
 tensorflow/compiler/xla/tests/tuple_test.cc   | 32 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index 937a0ea5bbc..74aad5f5bd5 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -1418,6 +1418,13 @@ Status IrEmitterUnnested::HandleSort(HloInstruction* sort) {
 
   AddThunkToThunkSequence(
       absl::make_unique(std::move(thunks), sort));
+  if (sort->operand_count() > 1) {
+    // Emit the tuple as part of the last stage of sorting.
+    // We are currently in the block sorted.in_bounds.after.
+    b_.SetInsertPoint(b_.GetInsertBlock()->getTerminator());
+    llvm_ir::EmitTuple(GetIrArray(*sort, *sort),
+                       ConstructIrArrayForOutputs(*sort), &b_);
+  }
   return Status::OK();
 }
 
diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc
index 9ef589e5511..b6ad44497e6 100644
--- a/tensorflow/compiler/xla/tests/tuple_test.cc
+++ b/tensorflow/compiler/xla/tests/tuple_test.cc
@@ -577,5 +577,37 @@ XLA_TEST_F(TupleHloTest,
   EXPECT_TRUE(LiteralTestUtil::Equal(expected, literal));
 }
 
+XLA_TEST_F(TupleHloTest, TupleSelectOfSort) {
+  const char* testcase = R"(
+    HloModule sort
+
+    compare {
+      p.1.lhs = s32[] parameter(2)
+      p.1.rhs = s32[] parameter(3)
+      p.0.lhs = f32[] parameter(0)
+      p.0.rhs = f32[] parameter(1)
+      ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT
+    }
+
+    ENTRY Sort {
+      keys = f32[2]{0} iota(), iota_dimension=0
+      values = s32[2]{0} iota(), iota_dimension=0
+      preds = pred[] constant(true)
+      alt = (f32[2], s32[2]) parameter(0)
+
+      sorted = (f32[2]{0}, s32[2]{0}) sort(keys, values), dimensions={0},
+               to_apply=compare
+      ROOT selected = (f32[2], s32[2]) tuple-select(preds, sorted, alt)
+    }
+  )";
+  auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie();
+  auto param = LiteralUtil::MakeTupleOwned(LiteralUtil::CreateR1({2, 3}),
+                                           LiteralUtil::CreateR1({3, 4}));
+  auto expected = LiteralUtil::MakeTupleOwned(
+      LiteralUtil::CreateR1({0, 1}), LiteralUtil::CreateR1({0, 1}));
+  auto result = ExecuteAndTransfer(std::move(module), {¶m});
+  EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
+}
+
 }  // namespace
 }  // namespace xla

From 78f6d51b8d9a2f3542918f004a29a6fbba232a40 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Thu, 18 Jun 2020 05:26:02 -0700
Subject: [PATCH 0491/1390] Also set the unroll factor in the C++ integration
 code for tanh.

PiperOrigin-RevId: 317083225
Change-Id: I7c2c26d664c15cbc967188da4b3012161edbcf49
---
 tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
index 70de777239f..a122c5112e6 100644
--- a/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
+++ b/tensorflow/core/kernels/mlir_generated_cwise_op_gpu_tanh.cu.cc
@@ -124,7 +124,7 @@ class MlirGeneratedTanhOp : public OpKernel {
     // This has to be aligned with the configuration that was used when building
     // the kernels. See the corresponding build rules in `cubin_headers/BUILD`.
     LaunchConfig config = GetLaunchConfiguration(
-        {256}, {}, {static_cast(inp.NumElements())});
+        {256}, {4}, {static_cast(inp.NumElements())});
     OP_REQUIRES_OK(
         ctx, stream->parent()->Launch(stream, config.threadDim, config.blockDim,
                                       *kernel, args));

From 2ad57928530e3e10ba89e91ff4df3c6fa9feafa0 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Thu, 18 Jun 2020 05:29:17 -0700
Subject: [PATCH 0492/1390] Bump open source llvm revision to
 e3836fe1a5562875396705369353078ab07cf07a

PiperOrigin-RevId: 317083604
Change-Id: Ifd810846a890ad137a31b4db5af5e049450765d3
---
 tensorflow/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 78f7e0ce03e..1c165ba5aba 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "4799fb63b5513f655ca8e85416ec8fe35df49bae"
-    LLVM_SHA256 = "f401a61bd7f5b05bd8a3ffdfb1f32e9379cae2c8e988f3ae6772b588ad97c84a"
+    LLVM_COMMIT = "e3836fe1a5562875396705369353078ab07cf07a"
+    LLVM_SHA256 = "6a78815e2c71c560a11c8c1740d31b88a607d82b7ccc61dc142bef0f1f3fbde8"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),

From 4f5c65d4494b4e4831d016176d506227c011f01b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 06:50:27 -0700
Subject: [PATCH 0493/1390] Make linear layout more explicit.

PiperOrigin-RevId: 317093123
Change-Id: I7af437c2c8afb31683bb659b1939eac2ce851da5
---
 .../compiler/xla/service/layout_assignment.cc | 14 +++++++------
 .../xla/service/layout_assignment_test.cc     | 21 +++++++++++++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
index 307fd82069e..a35ba140e86 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -951,12 +951,7 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) {
                 if (!Shape::Equal()
                          .IgnoreDynamicDimension()
                          .MinorToMajorOnlyInLayout()(instruction_subshape,
-                                                     buffer->shape()) &&
-                    // TODO(mingyao): Use explicit linear layout tiling to
-                    // detect and allow special bitcast.
-                    instruction->opcode() != HloOpcode::kBitcast &&
-                    instruction->opcode() != HloOpcode::kGetTupleElement &&
-                    instruction->opcode() != HloOpcode::kTuple) {
+                                                     buffer->shape())) {
                   return InternalError(
                       "Layout of instruction %s at index {%s} does not match "
                       "source LogicalBuffer %s: %s vs %s",
@@ -1803,6 +1798,13 @@ Status LayoutAssignment::ClearComputationLayouts(HloComputation* computation) {
   // potential bugs in the layout assignment pass that may accidentally use the
   // existing layout.
   for (HloInstruction* instruction : computation->instructions()) {
+    if (instruction->opcode() == HloOpcode::kBitcast) {
+      // bitcasts are inherently layout sensitive and so a bitcast instruction
+      // present in the IR before layout assignment is a bug.
+      return InternalError(
+          "Unexpected bitcast operation seen during layout assignment: %s.",
+          instruction->ToString());
+    }
     // Some instructions carry mandatory layouts in their shape.
     if (instruction->opcode() != HloOpcode::kInfeed &&
         !IsLayoutConstrainedCustomCall(instruction) &&
diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc
index 6e575247e6b..304a80c7a52 100644
--- a/tensorflow/compiler/xla/service/layout_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc
@@ -814,6 +814,27 @@ TEST_F(LayoutAssignmentTest, ConditionalAsymmetricLayout) {
   EXPECT_THAT(false_result->opcode(), HloOpcode::kCopy);
 }
 
+TEST_F(LayoutAssignmentTest, InternalErrorOnBitcast) {
+  auto builder = HloComputation::Builder(TestName());
+  auto constant0 = builder.AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR2WithLayout(
+          {{1.0, 2.0}, {3.0, 4.0}}, LayoutUtil::MakeLayout({0, 1}))));
+  builder.AddInstruction(
+      HloInstruction::CreateBitcast(constant0->shape(), constant0));
+  auto m = CreateNewVerifiedModule();
+  m->AddEntryComputation(builder.Build());
+
+  ComputationLayout computation_layout(
+      m->entry_computation()->ComputeProgramShape());
+  LayoutAssignment layout_assignment(&computation_layout);
+  Status error_status = layout_assignment.Run(m.get()).status();
+  EXPECT_FALSE(error_status.ok());
+  EXPECT_THAT(
+      error_status.error_message(),
+      ::testing::HasSubstr(
+          "Unexpected bitcast operation seen during layout assignment"));
+}
+
 TEST_F(LayoutAssignmentTest, ChannelLayoutMismatch) {
   // Pin non matching layouts to parameter and root.
   const char* module_str = R"(

From c65fccf0a6671c90599d0d3426dd18597688ea3a Mon Sep 17 00:00:00 2001
From: shwetaoj 
Date: Wed, 17 Jun 2020 11:55:27 -0700
Subject: [PATCH 0494/1390] Format changes as per Google's feedback

---
 .../ci_build/linux/mkl/Dockerfile.devel-mkl   |  4 ++--
 .../ci_build/linux/mkl/build-dev-container.sh | 21 +++++++++----------
 .../linux/mkl/install_openmpi_horovod.sh      |  7 ++++---
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
index a78d13c7755..1f80cba35f0 100755
--- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
+++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
@@ -64,9 +64,9 @@ RUN bazel --bazelrc=/root/.bazelrc build -c opt \
 
 # Install OpenMPI/Horovod
 COPY install_openmpi_horovod.sh .
-RUN if [ "${ENABLE_HOROVOD}" = "yes" ]; then \
+RUN if [ "${ENABLE_HOROVOD}" == "yes" ]; then \
         chmod +x install_openmpi_horovod.sh && \
-        ./install_openmpi_horovod.sh OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} HOROVOD_VERSION=${HOROVOD_VERSION} && \
+        OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} HOROVOD_VERSION=${HOROVOD_VERSION} ./install_openmpi_horovod.sh && \
         rm -rf install_openmpi_horovod.sh; \
     fi
 
diff --git a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh
index 6e789a54e87..a0880b0e51c 100755
--- a/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh
+++ b/tensorflow/tools/ci_build/linux/mkl/build-dev-container.sh
@@ -201,21 +201,20 @@ function test_container()
   debug "ID of the running docker container: ${CONTAINER_ID}"
 
   debug "Performing basic sanity checks on the running container..."
-  TEST_CMD_1=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'from tensorflow.python import _pywrap_util_port; print(_pywrap_util_port.IsMklEnabled())'")
-  # Make TEST_CMD backward compatible with older code
-  TEST_CMD_2=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'from tensorflow.python import pywrap_tensorflow; print(pywrap_tensorflow.IsMklEnabled())'")
-
-  if [ "${TEST_CMD_1}" = "True" -o "${TEST_CMD_2}" = "True" ] ; then
-      echo "PASS: MKL enabled test in ${TEMP_IMAGE_NAME}"
-  else
-      die "FAIL: MKL enabled test in ${TEMP_IMAGE_NAME}"
-  fi
+  {
+    ${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'from tensorflow.python import _pywrap_util_port; print(_pywrap_util_port.IsMklEnabled())'"
+    echo "PASS: MKL enabled test in ${TEMP_IMAGE_NAME}"
+  } || {
+    ${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'from tensorflow.python import pywrap_tensorflow; print(pywrap_tensorflow.IsMklEnabled())'"
+    echo "PASS: Old MKL enabled in ${TEMP_IMAGE_NAME}"
+  } || {
+    die "FAIL: MKL enabled test in ${TEMP_IMAGE_NAME}"
+  }
 
   # Test to check if horovod is installed successfully
   if [[ ${ENABLE_HOROVOD} == "yes" ]]; then
       debug "Test horovod in the container..."
-      HOROVOD_TEST_CMD=$(${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'")
-      ${HOROVOD_TEST_CMD}
+      ${DOCKER_BINARY} exec ${CONTAINER_ID} bash -c "${PYTHON} -c 'import horovod.tensorflow as hvd;'"
       if [[ $? == "0" ]]; then
           echo "PASS: HOROVOD installation test in ${TEMP_IMAGE_NAME}"
       else
diff --git a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh
index aec40543a17..9bc92ca4fef 100755
--- a/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh
+++ b/tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh
@@ -54,8 +54,9 @@ echo 'OpenMPI version:'
 mpirun --version
 
 # Install OpenSSH for MPI to communicate between containers
-apt-get clean && apt-get update && apt-get install -y --no-install-recommends --fix-missing \
-    openssh-client openssh-server libnuma-dev && \
+apt-get clean && apt-get update && \
+    apt-get install -y --no-install-recommends --fix-missing \
+        openssh-client openssh-server libnuma-dev && \
     rm -rf /var/lib/apt/lists/*
 if [[ $?  == "0" ]]; then
     echo "PASS: OpenSSH installation"
@@ -70,7 +71,7 @@ else
 fi
 mkdir -p /var/run/sshd
 # Allow OpenSSH to talk to containers without asking for confirmation
-cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new
+grep -v StrictHostKeyChecking /etc/ssh/ssh_config > /etc/ssh/ssh_config.new
 echo " StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new
 mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config
 

From 55929ce6ae2bc9552526a038ef4d01d8bef4f4fd Mon Sep 17 00:00:00 2001
From: shwetaoj 
Date: Wed, 17 Jun 2020 12:52:57 -0700
Subject: [PATCH 0495/1390] Reverting == in Dockerfile

---
 tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
index 1f80cba35f0..80091e55a17 100755
--- a/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
+++ b/tensorflow/tools/ci_build/linux/mkl/Dockerfile.devel-mkl
@@ -64,7 +64,7 @@ RUN bazel --bazelrc=/root/.bazelrc build -c opt \
 
 # Install OpenMPI/Horovod
 COPY install_openmpi_horovod.sh .
-RUN if [ "${ENABLE_HOROVOD}" == "yes" ]; then \
+RUN if [ "${ENABLE_HOROVOD}" = "yes" ]; then \
         chmod +x install_openmpi_horovod.sh && \
         OPENMPI_VERSION=${OPENMPI_VERSION} OPENMPI_DOWNLOAD_URL=${OPENMPI_DOWNLOAD_URL} HOROVOD_VERSION=${HOROVOD_VERSION} ./install_openmpi_horovod.sh && \
         rm -rf install_openmpi_horovod.sh; \

From c054f40f66fa625f51085a20c48554c61d05c5fd Mon Sep 17 00:00:00 2001
From: Deven Desai 
Date: Thu, 18 Jun 2020 14:48:42 +0000
Subject: [PATCH 0496/1390] [ROCm] Fix ROCm CSB build failure - 200618

The folllowing commit introduces build error on the ROCm platform

https://github.com/tensorflow/tensorflow/commit/b6ff68822a59578f942e4fb8076757da8db278ae

build error

```
In file included from tensorflow/core/kernels/split_lib_gpu.cu.cc:27:
In file included from ./tensorflow/core/util/gpu_kernel_helper.h:25:
./tensorflow/core/util/gpu_device_functions.h:824:25: error: redefinition of 'GpuAtomicMax'
__device__ inline int64 GpuAtomicMax(int64* ptr, int64 value) {
                        ^
./tensorflow/core/util/gpu_device_functions.h:792:29: note: previous definition is here
__device__ inline long long GpuAtomicMax(long long* ptr, long long value) {
                            ^
./tensorflow/core/util/gpu_device_functions.h:894:25: error: redefinition of 'GpuAtomicMin'
__device__ inline int64 GpuAtomicMin(int64* ptr, int64 value) {
                        ^
./tensorflow/core/util/gpu_device_functions.h:862:29: note: previous definition is here
__device__ inline long long GpuAtomicMin(long long* ptr, long long value) {
                            ^
2 errors generated.
...
...
```

The cause is a combination of two things
* The condition `#if __CUDA_ARCH__ < 320` will hold true for ROCm too!
* The issue being addressed by (the build breaking commit, for CUDA) was already fixed by this commit (https://github.com/tensorflow/tensorflow/commit/307485737f46a76c97aefb51b0fc3cd264c2bb94) within a `#if TENSORFLOW_USE_ROCM` block

The fix being submitted in this PR, is to undo some of the changes introduced by the earlier ROCm commit, and combine that change with the change in the breaking commit.
---
 tensorflow/core/util/gpu_device_functions.h | 14 ++------------
 1 file changed, 2 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h
index d4e09a7fc98..a8158a1ab08 100644
--- a/tensorflow/core/util/gpu_device_functions.h
+++ b/tensorflow/core/util/gpu_device_functions.h
@@ -789,11 +789,6 @@ __device__ inline double GpuAtomicMax(double* ptr, double value) {
       ptr, [value](double a) { return fmax(a, value); });
 }
 
-__device__ inline long long GpuAtomicMax(long long* ptr, long long value) {
-  return detail::GpuAtomicCasHelper(
-      ptr, [value](long long a) { return max(a, value); });
-}
-
 #else
 
 __device__ inline float GpuAtomicMax(float* ptr, float value) {
@@ -814,7 +809,7 @@ __device__ inline Eigen::half GpuAtomicMax(Eigen::half* ptr,
       ptr, [value](Eigen::half a) { return max(a, value); });
 }
 
-#if __CUDA_ARCH__ < 320
+#if TENSORFLOW_USE_ROCM || (__CUDA_ARCH__ < 320)
 __device__ inline tensorflow::uint64 GpuAtomicMax(tensorflow::uint64* ptr,
                                                   tensorflow::uint64 value) {
   return detail::GpuAtomicCasHelper(
@@ -859,11 +854,6 @@ __device__ inline double GpuAtomicMin(double* ptr, double value) {
       ptr, [value](double a) { return fmin(a, value); });
 }
 
-__device__ inline long long GpuAtomicMin(long long* ptr, long long value) {
-  return detail::GpuAtomicCasHelper(
-      ptr, [value](long long a) { return min(a, value); });
-}
-
 #else
 
 __device__ inline float GpuAtomicMin(float* ptr, float value) {
@@ -884,7 +874,7 @@ __device__ inline Eigen::half GpuAtomicMin(Eigen::half* ptr,
       ptr, [value](Eigen::half a) { return min(a, value); });
 }
 
-#if __CUDA_ARCH__ < 320
+#if TENSORFLOW_USE_ROCM || (__CUDA_ARCH__ < 320)
 __device__ inline tensorflow::uint64 GpuAtomicMin(tensorflow::uint64* ptr,
                                                   tensorflow::uint64 value) {
   return detail::GpuAtomicCasHelper(

From 49e3e63dd7d6d7cf700013000efc8c8da80662ab Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 08:20:28 -0700
Subject: [PATCH 0497/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/92d8ad02e92f

PiperOrigin-RevId: 317105768
Change-Id: I76b13f63a85dd4b8bbcaff3fd1ca624b82411079
---
 tensorflow/workspace.bzl |  4 ++--
 third_party/mlir/BUILD   | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 1c165ba5aba..f5b0b7537dc 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "e3836fe1a5562875396705369353078ab07cf07a"
-    LLVM_SHA256 = "6a78815e2c71c560a11c8c1740d31b88a607d82b7ccc61dc142bef0f1f3fbde8"
+    LLVM_COMMIT = "92d8ad02e92fed3884169ba5d98056fe4fa5660d"
+    LLVM_SHA256 = "a4995ace7ddaef0c49293dc65771f58ef1fea96ebe1f39aa0a2d6d75d07f6cc7"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index db75b27e78b..a0d10667d23 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -738,12 +738,32 @@ cc_library(
         ":Pass",
         ":SCFDialect",
         ":Shape",
+        ":ShapeToStandardPatternsIncGen",
         ":StandardOps",
         ":Support",
         ":Transforms",
     ],
 )
 
+gentbl(
+    name = "ShapeToStandardPatternsIncGen",
+    strip_include_prefix = "include/mlir/Conversion/ShapeToStandard",
+    tbl_outs = [
+        (
+            "-gen-rewriters",
+            "include/mlir/Conversion/ShapeToStandard/ShapeToStandardPatterns.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td",
+    td_srcs = [
+        ":StdOpsTdFiles",
+        "include/mlir/Dialect/Shape/IR/ShapeBase.td",
+        "include/mlir/Dialect/Shape/IR/ShapeOps.td",
+        "include/mlir/Interfaces/InferTypeOpInterface.td",
+    ],
+)
+
 cc_library(
     name = "ShapeToSCF",
     srcs = glob([

From 3cea671a74251d62549280fbb6444ffc2cdc4f03 Mon Sep 17 00:00:00 2001
From: Zhenyu Tan 
Date: Thu, 18 Jun 2020 08:28:08 -0700
Subject: [PATCH 0498/1390] set `_batch_input_shape` before build for kpl.

PiperOrigin-RevId: 317107061
Change-Id: Idd7f5dcbce9d15b2e3d708dd7ea3b2c3e5c1be7e
---
 .../python/keras/engine/base_preprocessing_layer.py  |  2 ++
 .../keras/layers/preprocessing/normalization_test.py | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/tensorflow/python/keras/engine/base_preprocessing_layer.py b/tensorflow/python/keras/engine/base_preprocessing_layer.py
index 08df07e33e3..b2ab0880422 100644
--- a/tensorflow/python/keras/engine/base_preprocessing_layer.py
+++ b/tensorflow/python/keras/engine/base_preprocessing_layer.py
@@ -190,6 +190,8 @@ class CombinerPreprocessingLayer(PreprocessingLayer):
           shape = data_element.shape
         except AttributeError:
           shape = None
+        # TODO (b/159261555): move this to base layer build.
+        self._batch_input_shape = shape
         self.build(shape)
 
       # Once we have built the Layer, we can process the input data. We do so
diff --git a/tensorflow/python/keras/layers/preprocessing/normalization_test.py b/tensorflow/python/keras/layers/preprocessing/normalization_test.py
index f5f68d9c51a..f97b8db50ec 100644
--- a/tensorflow/python/keras/layers/preprocessing/normalization_test.py
+++ b/tensorflow/python/keras/layers/preprocessing/normalization_test.py
@@ -318,6 +318,18 @@ class NormalizationTest(keras_parameterized.TestCase,
     layer.adapt(data)
     self.assertAllClose(expect, layer(data))
 
+  def test_model_summary_after_layer_adapt(self):
+    data = np.array([[[0., 1., 2.], [0., 2., 6.]],
+                     [[2., 3., 4.], [3., 6., 10.]]])
+    cls = get_layer_class()
+    layer = cls(axis=-1)
+    layer.adapt(data)
+    model = keras.Sequential(
+        [layer,
+         keras.layers.Dense(64, activation="relu"),
+         keras.layers.Dense(1)])
+    model.summary()
+
 
 if __name__ == "__main__":
   test.main()

From 06bc84b12edce4c4ce616d0dbad5e0d5178218b7 Mon Sep 17 00:00:00 2001
From: Jian Li 
Date: Thu, 18 Jun 2020 08:48:06 -0700
Subject: [PATCH 0499/1390] Create fully integer ResizeBilinear kernel.

PiperOrigin-RevId: 317110515
Change-Id: I233b89083dda9f9e0150c97a88391a154fb32d50
---
 .../internal/reference/reference_ops.h        | 103 ++++++++++++++++++
 1 file changed, 103 insertions(+)

diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index e991a21e3bd..5208b21eb4d 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -1645,6 +1645,109 @@ inline void ResizeBilinear(const tflite::ResizeBilinearParams& op_params,
   }
 }
 
+inline void ComputeInterpolationValues(const int32 value, const int32 scale_10,
+                                       const bool half_pixel_centers,
+                                       int32 input_size, int32* scaled_value,
+                                       int32* lower_bound, int32* upper_bound) {
+  if (half_pixel_centers) {
+    *scaled_value = value * scale_10 + scale_10 / 2 - (1 << 9);
+  } else {
+    *scaled_value = value * scale_10;
+  }
+  *lower_bound = std::max(*scaled_value / (1 << 10), 0);
+  *upper_bound = std::min(*scaled_value / (1 << 10) + 1, input_size - 1);
+}
+
+// Same as above but takes int8 as input and output.
+inline void ResizeBilinear(const tflite::ResizeBilinearParams& op_params,
+                           const RuntimeShape& unextended_input_shape,
+                           const int8_t* input_data,
+                           const RuntimeShape& unextended_output_size_shape,
+                           const int32* output_size_data,
+                           const RuntimeShape& unextended_output_shape,
+                           int8_t* output_data) {
+  // If half_pixel_centers is True, align_corners must be False.
+  TFLITE_DCHECK(!op_params.half_pixel_centers || !op_params.align_corners);
+  TFLITE_DCHECK_LE(unextended_input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_LE(unextended_output_size_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4);
+  const RuntimeShape input_shape =
+      RuntimeShape::ExtendedShape(4, unextended_input_shape);
+  const RuntimeShape output_size_shape =
+      RuntimeShape::ExtendedShape(4, unextended_output_size_shape);
+  const RuntimeShape output_shape =
+      RuntimeShape::ExtendedShape(4, unextended_output_shape);
+
+  const int32 batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int32 input_height = input_shape.Dims(1);
+  const int32 input_width = input_shape.Dims(2);
+  const int32 depth = MatchingDim(input_shape, 3, output_shape, 3);
+
+  TFLITE_DCHECK_EQ(output_size_shape.Dims(0), 1);
+  TFLITE_DCHECK_EQ(output_size_shape.Dims(1), 1);
+  TFLITE_DCHECK_EQ(output_size_shape.Dims(2), 1);
+  TFLITE_DCHECK_EQ(output_size_shape.Dims(3), 2);
+  const int32 output_height =
+      output_size_data[Offset(output_size_shape, 0, 0, 0, 0)];
+  const int32 output_width =
+      output_size_data[Offset(output_size_shape, 0, 0, 0, 1)];
+
+  int32 height_scale_10 =
+      ((1 << 10) * input_height + output_height / 2) / output_height;
+  int32 width_scale_10 =
+      ((1 << 10) * input_width + output_width / 2) / output_width;
+  if (op_params.align_corners && output_height > 1) {
+    height_scale_10 =
+        ((1 << 10) * (input_height - 1) + (output_height - 1) / 2) /
+        (output_height - 1);
+  }
+  if (op_params.align_corners && output_width > 1) {
+    width_scale_10 = ((1 << 10) * (input_width - 1) + (output_width - 1) / 2) /
+                     (output_width - 1);
+  }
+
+  for (int b = 0; b < batches; ++b) {
+    for (int y = 0; y < output_height; ++y) {
+      int32 input_y, y0, y1;
+      ComputeInterpolationValues(y, height_scale_10,
+                                 op_params.half_pixel_centers, input_height,
+                                 &input_y, &y0, &y1);
+      for (int x = 0; x < output_width; ++x) {
+        int32 input_x, x0, x1;
+        ComputeInterpolationValues(x, width_scale_10,
+                                   op_params.half_pixel_centers, input_width,
+                                   &input_x, &x0, &x1);
+        for (int c = 0; c < depth; ++c) {
+          const int64_t output_20_ll =
+              static_cast(
+                  input_data[Offset(input_shape, b, y0, x0, c)]) *
+              ((1 << 10) - (input_y - (1 << 10) * y0)) *
+              ((1 << 10) - (input_x - (1 << 10) * x0));
+          const int64_t output_20_lu =
+              static_cast(
+                  input_data[Offset(input_shape, b, y1, x0, c)]) *
+              (input_y - (1 << 10) * y0) *
+              ((1 << 10) - (input_x - (1 << 10) * x0));
+          const int64_t output_20_rl =
+              static_cast(
+                  input_data[Offset(input_shape, b, y0, x1, c)]) *
+              ((1 << 10) - (input_y - (1 << 10) * y0)) *
+              (input_x - (1 << 10) * x0);
+          const int64_t output_20_ru =
+              static_cast(
+                  input_data[Offset(input_shape, b, y1, x1, c)]) *
+              (input_y - (1 << 10) * y0) * (input_x - (1 << 10) * x0);
+          const int64_t output_20 =
+              output_20_ll + output_20_lu + output_20_rl + output_20_ru;
+          const int8_t interpolation =
+              static_cast((output_20 + (1 << 19)) / (1 << 20));
+          output_data[Offset(output_shape, b, y, x, c)] = interpolation;
+        }
+      }
+    }
+  }
+}
+
 template 
 inline void SpaceToBatchND(
     const SpaceToBatchParams& params,

From eda7d05793ec75227069eb0c3f49e0377f33c963 Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Thu, 18 Jun 2020 23:02:21 +0700
Subject: [PATCH 0500/1390] Add NewWritableFile

---
 .../experimental/filesystem/plugins/gcs/BUILD |  1 +
 .../filesystem/plugins/gcs/gcs_filesystem.cc  | 50 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
index c9fee433589..05fd371088c 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
@@ -24,6 +24,7 @@ cc_library(
         "//tensorflow:windows": get_win_copts(),
     }),
     deps = [
+        "//tensorflow/c:env",
         "//tensorflow/c:tf_status",
         "//tensorflow/c/experimental/filesystem:filesystem_interface",
         "@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
index 8c54bc85439..4ddc8548486 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
@@ -15,8 +15,11 @@ limitations under the License.
 #include 
 #include 
 
+#include 
+
 #include "absl/strings/string_view.h"
 #include "google/cloud/storage/client.h"
+#include "tensorflow/c/env.h"
 #include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
 #include "tensorflow/c/tf_status.h"
 
@@ -75,6 +78,25 @@ static void ParseGCSPath(absl::string_view fname, bool object_empty_ok,
   strcpy(*object, object_view.data());
 }
 
+class TempFile : public std::fstream {
+ public:
+  // We should specify openmode each time we call TempFile.
+  TempFile(const char* temp_file_name, std::ios::openmode mode)
+      : std::fstream(temp_file_name, mode), name(temp_file_name) {}
+  TempFile(TempFile&& rhs) : std::fstream(std::move(rhs)), name(rhs.name) {
+    rhs.name = nullptr;
+  }
+  ~TempFile() {
+    std::fstream::close();
+    std::remove(name);
+    plugin_memory_free(const_cast(name));
+  }
+  const char* getName() { return name; }
+
+ private:
+  const char* name;
+};
+
 // SECTION 1. Implementation for `TF_RandomAccessFile`
 // ----------------------------------------------------------------------------
 namespace tf_random_access_file {
@@ -86,6 +108,20 @@ namespace tf_random_access_file {
 // SECTION 2. Implementation for `TF_WritableFile`
 // ----------------------------------------------------------------------------
 namespace tf_writable_file {
+typedef struct GCSFile {
+  const char* bucket;
+  const char* object;
+  gcs::Client* gcs_client;  // not owned
+  TempFile outfile;
+  bool sync_need;
+} GCSFile;
+
+static void Cleanup(TF_WritableFile* file) {
+  auto gcs_file = static_cast(file->plugin_file);
+  plugin_memory_free(const_cast(gcs_file->bucket));
+  plugin_memory_free(const_cast(gcs_file->object));
+  delete gcs_file;
+}
 
 // TODO(vnvo2409): Implement later
 
@@ -119,6 +155,20 @@ static void Init(TF_Filesystem* filesystem, TF_Status* status) {
 
 // TODO(vnvo2409): Implement later
 
+static void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
+                            TF_WritableFile* file, TF_Status* status) {
+  char* bucket;
+  char* object;
+  ParseGCSPath(path, false, &bucket, &object, status);
+  if (TF_GetCode(status) != TF_OK) return;
+
+  auto gcs_client = static_cast(filesystem->plugin_filesystem);
+  TempFile outfile(TF_GetTempFileName(""), std::ios::binary | std::ios::out);
+  file->plugin_file = new tf_writable_file::GCSFile(
+      {bucket, object, gcs_client, std::move(outfile), true});
+  TF_SetStatus(status, TF_OK, "");
+}
+
 }  // namespace tf_gcs_filesystem
 
 static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,

From 2e3310031296f9232f7a58a5cfca2ee03b9a7c91 Mon Sep 17 00:00:00 2001
From: Jian Li 
Date: Thu, 18 Jun 2020 09:26:13 -0700
Subject: [PATCH 0501/1390] Add quantization test for transpose.

PiperOrigin-RevId: 317117545
Change-Id: I81c4f9583f29205bcbdaae175eac59439cd19047
---
 tensorflow/lite/tools/optimize/BUILD          |   1 +
 .../tools/optimize/quantize_model_test.cc     |  44 ++++++++++++++++++
 tensorflow/lite/tools/optimize/test_util.cc   |   2 +
 tensorflow/lite/tools/optimize/test_util.h    |   3 ++
 .../tools/optimize/testdata/transpose.bin     | Bin 0 -> 544 bytes
 5 files changed, 50 insertions(+)
 create mode 100644 tensorflow/lite/tools/optimize/testdata/transpose.bin

diff --git a/tensorflow/lite/tools/optimize/BUILD b/tensorflow/lite/tools/optimize/BUILD
index 3011c01cdeb..c10d4465e5c 100644
--- a/tensorflow/lite/tools/optimize/BUILD
+++ b/tensorflow/lite/tools/optimize/BUILD
@@ -296,6 +296,7 @@ tf_cc_test(
         "//tensorflow/lite/tools/optimize:testdata/split.bin",
         "//tensorflow/lite/tools/optimize:testdata/svdf_calibrated.bin",
         "//tensorflow/lite/tools/optimize:testdata/svdf_quantized.bin",
+        "//tensorflow/lite/tools/optimize:testdata/transpose.bin",
         "//tensorflow/lite/tools/optimize:testdata/unpack.bin",
     ],
     tags = [
diff --git a/tensorflow/lite/tools/optimize/quantize_model_test.cc b/tensorflow/lite/tools/optimize/quantize_model_test.cc
index f8f1a9d4113..36b35af0065 100644
--- a/tensorflow/lite/tools/optimize/quantize_model_test.cc
+++ b/tensorflow/lite/tools/optimize/quantize_model_test.cc
@@ -1454,6 +1454,50 @@ TEST_F(QuantizeUnpackTest, VerifyUnpack) {
                   unpack_output_1->quantization->zero_point[0]);
 }
 
+class QuantizeTransposeTest : public QuantizeModelTest {
+ protected:
+  QuantizeTransposeTest() {
+    input_model_ = ReadModel(internal::kModelWithTranspose);
+    readonly_model_ = input_model_->GetModel();
+    readonly_model_->UnPackTo(&model_);
+  }
+};
+
+TEST_F(QuantizeTransposeTest, VerifyTranspose) {
+  auto status = QuantizeModel(&builder_, &model_, &error_reporter_);
+
+  ASSERT_EQ(kTfLiteOk, status);
+
+  const auto subgraph = model_.subgraphs[0].get();
+  auto op = subgraph->operators[1].get();
+
+  auto float_graph = readonly_model_->subgraphs()->Get(0);
+
+  ASSERT_EQ(model_.operator_codes[op->opcode_index].get()->builtin_code,
+            BuiltinOperator_TRANSPOSE);
+
+  // The model should only have one input and one outputs.
+  EXPECT_EQ(subgraph->inputs.size(), 1);
+  EXPECT_EQ(subgraph->outputs.size(), 1);
+
+  // Get transpose input and output tensors
+  auto transpose_input = subgraph->tensors[op->inputs[0]].get();
+  auto transpose_output = subgraph->tensors[op->outputs[0]].get();
+
+  // Verify transpose input is quantized.
+  ASSERT_EQ(float_graph->tensors()->Get(op->inputs[0])->type(),
+            TensorType_FLOAT32);
+  EXPECT_EQ(transpose_input->type, TensorType_INT8);
+
+  // Ensure quantization parameters before and after transpose
+  // are preserved after quantization for all outputs of
+  // transpose.
+  EXPECT_FLOAT_EQ(transpose_input->quantization->scale[0],
+                  transpose_output->quantization->scale[0]);
+  EXPECT_EQ(transpose_input->quantization->zero_point[0],
+            transpose_output->quantization->zero_point[0]);
+}
+
 }  // namespace
 }  // namespace optimize
 }  // namespace tflite
diff --git a/tensorflow/lite/tools/optimize/test_util.cc b/tensorflow/lite/tools/optimize/test_util.cc
index 7d5e9d65f06..61e82ed3e34 100644
--- a/tensorflow/lite/tools/optimize/test_util.cc
+++ b/tensorflow/lite/tools/optimize/test_util.cc
@@ -61,6 +61,8 @@ const char* kModelWithMaximumOp = "maximum.bin";
 const char* kLstmCalibrated2 = "lstm_calibrated2.bin";
 const char* kLstmQuantized2 = "lstm_quantized2.bin";
 
+const char* kModelWithTranspose = "transpose.bin";
+
 const char* kSvdfCalibrated = "svdf_calibrated.bin";
 const char* kSvdfQuantized = "svdf_quantized.bin";
 
diff --git a/tensorflow/lite/tools/optimize/test_util.h b/tensorflow/lite/tools/optimize/test_util.h
index abcdbc21d36..4d2eadf283f 100644
--- a/tensorflow/lite/tools/optimize/test_util.h
+++ b/tensorflow/lite/tools/optimize/test_util.h
@@ -98,6 +98,9 @@ extern const char* kModelWithMaximumOp;
 extern const char* kLstmCalibrated2;
 extern const char* kLstmQuantized2;
 
+// Test model with a transpose op.
+extern const char* kModelWithTranspose;
+
 // Test model with SVDF op.
 extern const char* kSvdfCalibrated;
 extern const char* kSvdfQuantized;
diff --git a/tensorflow/lite/tools/optimize/testdata/transpose.bin b/tensorflow/lite/tools/optimize/testdata/transpose.bin
new file mode 100644
index 0000000000000000000000000000000000000000..a76886e5b473b8de04253089e2acc931f5b6ec86
GIT binary patch
literal 544
zcmb1OU|TQp3Q&zyifC3=9kw
z3=9k)y&PaMHn1E60~-Se0}BHiSPY_`fuX^k0SWy1|NlP=0|+uO@Po|tclK9s&d)1L
zEh$Xt+K1_lQQ
z2+hC%3XA;Gl7i9_{YsE|AiVAW|No${0;z$Su>!-418_4yVFxh-68<24APkBhP>6tR
z0mTp~ejFGW7(i~;0f!ev1{4w?cc8ejp~eMbhQdKs1_p-Ayn@masJ}qsFboPWNF1_*
K
Date: Thu, 18 Jun 2020 23:46:41 +0700
Subject: [PATCH 0502/1390] Move TempFile  to gcs_helper

---
 .../experimental/filesystem/plugins/gcs/BUILD | 12 ++++++-
 .../filesystem/plugins/gcs/gcs_filesystem.cc  | 20 +----------
 .../filesystem/plugins/gcs/gcs_helper.cc      | 19 +++++++++++
 .../filesystem/plugins/gcs/gcs_helper.h       | 33 +++++++++++++++++++
 4 files changed, 64 insertions(+), 20 deletions(-)
 create mode 100644 tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
 create mode 100644 tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h

diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
index 05fd371088c..d104181b264 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
@@ -24,10 +24,20 @@ cc_library(
         "//tensorflow:windows": get_win_copts(),
     }),
     deps = [
-        "//tensorflow/c:env",
+        ":gcs_helper",
         "//tensorflow/c:tf_status",
         "//tensorflow/c/experimental/filesystem:filesystem_interface",
         "@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
         "@com_google_absl//absl/strings",
     ],
 )
+
+cc_library(
+    name = "gcs_helper",
+    srcs = ["gcs_helper.cc"],
+    hdrs = ["gcs_helper.h"],
+    linkstatic = 1,
+    deps = [
+        "//tensorflow/c:env",
+    ],
+)
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
index 4ddc8548486..2793194e0a8 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "google/cloud/storage/client.h"
 #include "tensorflow/c/env.h"
 #include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
 #include "tensorflow/c/tf_status.h"
 
 // Implementation of a filesystem for GCS environments.
@@ -78,25 +79,6 @@ static void ParseGCSPath(absl::string_view fname, bool object_empty_ok,
   strcpy(*object, object_view.data());
 }
 
-class TempFile : public std::fstream {
- public:
-  // We should specify openmode each time we call TempFile.
-  TempFile(const char* temp_file_name, std::ios::openmode mode)
-      : std::fstream(temp_file_name, mode), name(temp_file_name) {}
-  TempFile(TempFile&& rhs) : std::fstream(std::move(rhs)), name(rhs.name) {
-    rhs.name = nullptr;
-  }
-  ~TempFile() {
-    std::fstream::close();
-    std::remove(name);
-    plugin_memory_free(const_cast(name));
-  }
-  const char* getName() { return name; }
-
- private:
-  const char* name;
-};
-
 // SECTION 1. Implementation for `TF_RandomAccessFile`
 // ----------------------------------------------------------------------------
 namespace tf_random_access_file {
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
new file mode 100644
index 00000000000..139579c53ae
--- /dev/null
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
@@ -0,0 +1,19 @@
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
+
+#include 
+
+#include 
+#include 
+
+TempFile::TempFile(const char* temp_file_name, std::ios::openmode mode)
+    : std::fstream(temp_file_name, mode), name(temp_file_name) {}
+
+TempFile::TempFile(TempFile&& rhs)
+    : std::fstream(std::move(rhs)), name(std::move(rhs.name)) {}
+
+TempFile::~TempFile() {
+  std::fstream::close();
+  std::remove(name.c_str());
+}
+
+const std::string TempFile::getName() const { return name; }
\ No newline at end of file
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
new file mode 100644
index 00000000000..437cbe560d6
--- /dev/null
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+
+#include 
+#include 
+
+class TempFile : public std::fstream {
+ public:
+  // We should specify openmode each time we call TempFile.
+  TempFile(const char* temp_file_name, std::ios::openmode mode);
+  TempFile(TempFile&& rhs);
+  ~TempFile();
+  const std::string getName() const;
+
+ private:
+  const std::string name;
+};
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_

From 18f54c42c62191b60edeff7e1308ed7b3305b0eb Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Thu, 18 Jun 2020 09:45:59 -0700
Subject: [PATCH 0503/1390] Disable special_math_ops_test on asan build.

PiperOrigin-RevId: 317121277
Change-Id: I902f6c046c2766638b5dbbb90a0ec2942a3c536f
---
 tensorflow/python/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index f53859b2915..a4e72bf2460 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -5277,6 +5277,7 @@ cuda_py_test(
     shard_count = 10,
     tags = [
         "no_windows_gpu",
+        "noasan",  # b/159332048
         "nomsan",  # b/148630708
     ],
     deps = [

From 89b80c5fb98ba57175626a9df490dbdd2bc8776b Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Thu, 18 Jun 2020 09:59:14 -0700
Subject: [PATCH 0504/1390] Add banded triangular solve op.

PiperOrigin-RevId: 317124054
Change-Id: I54f090d7583b21fa18788a2deb02262d9c8231be
---
 .../api_def_BandedTriangularSolve.pbtxt       |   4 +
 .../api_def_BandedTriangularSolve.pbtxt       |   4 +
 tensorflow/core/kernels/BUILD                 |  27 ++
 .../kernels/banded_triangular_solve_op.cc     | 293 ++++++++++++++++++
 .../banded_triangular_solve_op_test.cc        | 180 +++++++++++
 tensorflow/core/ops/linalg_ops.cc             |  54 ++++
 tensorflow/python/kernel_tests/BUILD          |  11 +
 .../banded_triangular_solve_op_test.py        | 232 ++++++++++++++
 .../python/kernel_tests/linalg_grad_test.py   |  52 ++++
 tensorflow/python/ops/linalg/linalg_impl.py   |  96 ++++++
 tensorflow/python/ops/linalg_grad.py          |  33 ++
 .../api/golden/v1/tensorflow.raw_ops.pbtxt    |   4 +
 .../api/golden/v2/tensorflow.linalg.pbtxt     |   4 +
 .../api/golden/v2/tensorflow.raw_ops.pbtxt    |   4 +
 14 files changed, 998 insertions(+)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/api_def/python_api/api_def_BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/kernels/banded_triangular_solve_op.cc
 create mode 100644 tensorflow/core/kernels/banded_triangular_solve_op_test.cc
 create mode 100644 tensorflow/python/kernel_tests/banded_triangular_solve_op_test.py

diff --git a/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt b/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
new file mode 100644
index 00000000000..ba5e1bdcaf2
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BandedTriangularSolve"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_BandedTriangularSolve.pbtxt b/tensorflow/core/api_def/python_api/api_def_BandedTriangularSolve.pbtxt
new file mode 100644
index 00000000000..ba5e1bdcaf2
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_BandedTriangularSolve.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BandedTriangularSolve"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index b2b54adbcf9..b4730dad96c 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3577,6 +3577,7 @@ tf_cc_tests(
 cc_library(
     name = "linalg",
     deps = [
+        ":banded_triangular_solve_op",
         ":cholesky_grad",
         ":cholesky_op",
         ":determinant_op",
@@ -3750,6 +3751,12 @@ tf_kernel_library(
     deps = LINALG_DEPS,
 )
 
+tf_kernel_library(
+    name = "banded_triangular_solve_op",
+    prefix = "banded_triangular_solve_op",
+    deps = LINALG_DEPS + [":fill_functor"],
+)
+
 tf_kernel_library(
     name = "matrix_triangular_solve_op",
     hdrs = ["matrix_triangular_solve_op_impl.h"],
@@ -4425,6 +4432,26 @@ tf_cuda_cc_test(
     ],
 )
 
+tf_cuda_cc_test(
+    name = "banded_triangular_solve_op_test",
+    size = "small",
+    srcs = ["banded_triangular_solve_op_test.cc"],
+    deps = [
+        ":banded_triangular_solve_op",
+        ":matrix_set_diag_op",
+        ":matrix_triangular_solve_op",
+        ":ops_testutil",
+        ":ops_util",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
 tf_cuda_cc_test(
     name = "matrix_triangular_solve_op_test",
     size = "small",
diff --git a/tensorflow/core/kernels/banded_triangular_solve_op.cc b/tensorflow/core/kernels/banded_triangular_solve_op.cc
new file mode 100644
index 00000000000..d01a015502a
--- /dev/null
+++ b/tensorflow/core/kernels/banded_triangular_solve_op.cc
@@ -0,0 +1,293 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+// See docs in ../ops/linalg_ops.cc.
+
+#include "third_party/eigen3/Eigen/Core"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/kernels/fill_functor.h"
+#include "tensorflow/core/kernels/linalg_ops_common.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/matmul_bcast.h"
+
+namespace tensorflow {
+
+typedef Eigen::ThreadPoolDevice CPUDevice;
+
+template 
+Scalar eigen_conj(const Scalar& scalar) {
+  return Eigen::numext::conj(scalar);
+}
+
+// Sequential batch matrix triangular solve kernel that calls Eigen's
+// matrix triangular solve.
+template 
+struct SequentialBandedTriangularSolveKernel {
+  using Matrix =
+      Eigen::Matrix;
+  using ConstMatrixMap = Eigen::Map;
+  using MatrixMap = Eigen::Map;
+  using RealScalar = typename Eigen::NumTraits::Real;
+
+  static ConstMatrixMap ConstTensorSliceToEigenMatrix(const Tensor& t,
+                                                      int slice) {
+    return ConstMatrixMap(
+        t.flat().data() + slice * t.dim_size(1) * t.dim_size(2),
+        t.dim_size(1), t.dim_size(2));
+  }
+
+  static MatrixMap TensorSliceToEigenMatrix(Tensor* t, int slice) {
+    return MatrixMap(
+        t->flat().data() + slice * t->dim_size(1) * t->dim_size(2),
+        t->dim_size(1), t->dim_size(2));
+  }
+
+  static void Run(const Tensor& in_x, const Tensor& in_y, bool lower,
+                  bool adjoint, const MatMulBCast& bcast, Tensor* out,
+                  int start, int limit) {
+    const bool should_bcast = bcast.IsBroadcastingRequired();
+    const auto& x_batch_indices = bcast.x_batch_indices();
+    const auto& y_batch_indices = bcast.y_batch_indices();
+    int num_bands = in_x.dim_size(1);
+    int matrix_size = in_x.dim_size(2);
+
+    for (int64 i = start; i < limit; ++i) {
+      const int64 x_batch_index = should_bcast ? x_batch_indices[i] : i;
+      const int64 y_batch_index = should_bcast ? y_batch_indices[i] : i;
+      auto matrix = ConstTensorSliceToEigenMatrix(in_x, x_batch_index);
+      auto rhs = ConstTensorSliceToEigenMatrix(in_y, y_batch_index);
+      auto output = TensorSliceToEigenMatrix(out, i);
+      // Below, we use the standard algorithm for computing a triangular solve,
+      // except we band limit it.
+      // Given A x = b, where A is lower triangular,
+      // x_i = (b_i - sum a_ij * x_j) / a_ii, where the sum is from
+      // j = 0 to i - 1.
+      //
+      // Now, in a banded triangular matrix, when i exceeds the band size,
+      // then the sum goes from j = i - band_size to i - 1, since the other
+      // elements are zero.
+      //
+      // Finally, given the band storage format, we'll need to change the
+      // indexing.
+      if (lower) {
+        if (!adjoint) {
+          output.row(0) = rhs.row(0) / matrix(0, 0);
+          for (int i = 1; i < matrix_size; ++i) {
+            if (i < num_bands) {
+              output.row(i).noalias() =
+                  (rhs.row(i) - matrix.block(1, i, i, 1).reverse().transpose() *
+                                    output.topRows(i)) /
+                  matrix(0, i);
+            } else {
+              output.row(i).noalias() =
+                  (rhs.row(i) -
+                   matrix.block(1, i, num_bands - 1, 1).reverse().transpose() *
+                       output.middleRows(i - (num_bands - 1), num_bands - 1)) /
+                  matrix(0, i);
+            }
+          }
+        } else {
+          // In the adjoint case, here and below, we now have an upper (lower)
+          // triangular matrix, and thus need to work through with the other
+          // case. We can't simply conjugate `matrix` and use the upper (lower)
+          // algorithm because the band storage format for upper and lower
+          // triangular matrices are different (in the lower case, we pad
+          // entries on the left, and in the upper case we pad entries on the
+          // right.
+          output.row(matrix_size - 1) =
+              rhs.row(matrix_size - 1) / eigen_conj(matrix(0, matrix_size - 1));
+          for (int i = matrix_size - 1; i >= 0; --i) {
+            output.row(i).noalias() = rhs.row(i);
+            for (int j = i + 1; j < std::min(matrix_size, i + num_bands); ++j) {
+              output.row(i).noalias() -=
+                  eigen_conj(matrix(j - i, j)) * output.row(j);
+            }
+            output.row(i) /= eigen_conj(matrix(0, i));
+          }
+        }
+      } else {
+        if (!adjoint) {
+          output.row(matrix_size - 1) =
+              rhs.row(matrix_size - 1) / matrix(num_bands - 1, matrix_size - 1);
+          for (int i = 1; i < matrix_size; ++i) {
+            int k = matrix_size - 1 - i;
+            if (i < num_bands) {
+              output.row(k).noalias() =
+                  (rhs.row(k) - matrix.block(num_bands - 1 - i, k, i, 1)
+                                        .reverse()
+                                        .transpose() *
+                                    output.bottomRows(i)) /
+                  matrix(num_bands - 1, k);
+            } else {
+              output.row(k).noalias() =
+                  (rhs.row(k) -
+                   matrix.block(0, k, num_bands - 1, 1).reverse().transpose() *
+                       output.middleRows(k + 1, num_bands - 1)) /
+                  matrix(num_bands - 1, k);
+            }
+          }
+        } else {
+          output.row(0) = rhs.row(0) / eigen_conj(matrix(num_bands - 1, 0));
+          for (int i = 1; i < matrix_size; ++i) {
+            output.row(i).noalias() = rhs.row(i);
+            for (int j = std::max(0, i - (num_bands - 1)); j < i; ++j) {
+              output.row(i).noalias() -=
+                  eigen_conj(matrix(num_bands - 1 - (i - j), j)) *
+                  output.row(j);
+            }
+            output.row(i) /= eigen_conj(matrix(num_bands - 1, i));
+          }
+        }
+      }
+    }
+  }
+};
+
+template 
+struct LaunchBatchBandedTriangularSolve;
+
+template 
+struct LaunchBatchBandedTriangularSolve {
+  static void Launch(OpKernelContext* context, const Tensor& in_x,
+                     const Tensor& in_y, bool adjoint, bool lower,
+                     const MatMulBCast& bcast, Tensor* out) {
+    // Number of banded matrix triangular solves i.e. size of the batch.
+    const int64 batch_size = bcast.output_batch_size();
+    const int64 cost_per_unit =
+        in_x.dim_size(1) * in_x.dim_size(2) * in_y.dim_size(2);
+    auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());
+
+    using Matrix =
+        Eigen::Matrix;
+    using ConstMatrixMap = Eigen::Map;
+    using RealScalar = typename Eigen::NumTraits::Real;
+    // Check diagonal before doing any solves. This is the first row in the
+    // lower case and else is the last row.
+    auto matrix = ConstMatrixMap(in_x.flat().data(), in_x.dim_size(1),
+                                 in_x.dim_size(2));
+    RealScalar min_abs_pivot;
+    if (lower) {
+      min_abs_pivot = matrix.row(0).cwiseAbs().minCoeff();
+    } else {
+      min_abs_pivot = matrix.row(in_x.dim_size(1) - 1).cwiseAbs().minCoeff();
+    }
+    OP_REQUIRES(context, min_abs_pivot > RealScalar(0),
+                errors::InvalidArgument("Input matrix is not invertible."));
+
+    Shard(worker_threads.num_threads, worker_threads.workers, batch_size,
+          cost_per_unit,
+          [&in_x, &in_y, adjoint, lower, &bcast, out](int start, int limit) {
+            SequentialBandedTriangularSolveKernel::Run(
+                in_x, in_y, lower, adjoint, bcast, out, start, limit);
+          });
+  }
+};
+
+template 
+class BandedTriangularSolveOpCpu : public OpKernel {
+ public:
+  explicit BandedTriangularSolveOpCpu(OpKernelConstruction* context)
+      : OpKernel(context) {
+    OP_REQUIRES_OK(context, context->GetAttr("lower", &lower_));
+    OP_REQUIRES_OK(context, context->GetAttr("adjoint", &adjoint_));
+  }
+
+  ~BandedTriangularSolveOpCpu() override {}
+
+  void Compute(OpKernelContext* ctx) override {
+    const Tensor& in0 = ctx->input(0);
+    const Tensor& in1 = ctx->input(1);
+
+    ValidateInputTensors(ctx, in0, in1);
+
+    MatMulBCast bcast(in0.shape().dim_sizes(), in1.shape().dim_sizes());
+    OP_REQUIRES(
+        ctx, bcast.IsValid(),
+        errors::InvalidArgument(
+            "In[0] and In[1] must have compatible batch dimensions: ",
+            in0.shape().DebugString(), " vs. ", in1.shape().DebugString()));
+
+    TensorShape out_shape = bcast.output_batch_shape();
+    auto batch_size = bcast.output_batch_size();
+    auto d0 = in0.dim_size(in0.dims() - 2);  // Band size.
+    auto d1 = in0.dim_size(in0.dims() - 1);
+    Tensor in0_reshaped;
+    OP_REQUIRES(
+        ctx,
+        in0_reshaped.CopyFrom(in0, TensorShape({bcast.x_batch_size(), d0, d1})),
+        errors::Internal("Failed to reshape In[0] from ",
+                         in0.shape().DebugString()));
+    auto d2 = in1.dim_size(in1.dims() - 2);
+    auto d3 = in1.dim_size(in1.dims() - 1);
+    Tensor in1_reshaped;
+    OP_REQUIRES(
+        ctx,
+        in1_reshaped.CopyFrom(in1, TensorShape({bcast.y_batch_size(), d2, d3})),
+        errors::Internal("Failed to reshape In[1] from ",
+                         in1.shape().DebugString()));
+    OP_REQUIRES(ctx, d1 == d2,
+                errors::InvalidArgument(
+                    "In[0] mismatch In[1] shape: ", d1, " vs. ", d2, ": ",
+                    in0.shape().DebugString(), " ", in1.shape().DebugString(),
+                    " ", lower_, " ", adjoint_));
+    out_shape.AddDim(d1);
+    out_shape.AddDim(d3);
+    Tensor* out = nullptr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out));
+    if (out->NumElements() == 0) {
+      return;
+    }
+    Tensor out_reshaped;
+    OP_REQUIRES(ctx,
+                out_reshaped.CopyFrom(*out, TensorShape({batch_size, d1, d3})),
+                errors::Internal("Failed to reshape output from ",
+                                 out->shape().DebugString()));
+    LaunchBatchBandedTriangularSolve::Launch(
+        ctx, in0_reshaped, in1_reshaped, adjoint_, lower_, bcast,
+        &out_reshaped);
+  }
+
+ private:
+  void ValidateInputTensors(OpKernelContext* ctx, const Tensor& in0,
+                            const Tensor& in1) {
+    OP_REQUIRES(
+        ctx, in0.dims() >= 2,
+        errors::InvalidArgument("In[0] ndims must be >= 2: ", in0.dims()));
+
+    OP_REQUIRES(
+        ctx, in1.dims() >= 2,
+        errors::InvalidArgument("In[1] ndims must be >= 2: ", in1.dims()));
+  }
+  bool lower_;
+  bool adjoint_;
+};
+
+#define REGISTER_BANDED_TRIANGULAR_SOLVE_CPU(TYPE)        \
+  REGISTER_KERNEL_BUILDER(Name("BandedTriangularSolve")   \
+                              .Device(DEVICE_CPU)         \
+                              .TypeConstraint("T"), \
+                          BandedTriangularSolveOpCpu);
+
+REGISTER_BANDED_TRIANGULAR_SOLVE_CPU(float);
+REGISTER_BANDED_TRIANGULAR_SOLVE_CPU(double);
+REGISTER_BANDED_TRIANGULAR_SOLVE_CPU(complex64);
+REGISTER_BANDED_TRIANGULAR_SOLVE_CPU(complex128);
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/kernels/banded_triangular_solve_op_test.cc b/tensorflow/core/kernels/banded_triangular_solve_op_test.cc
new file mode 100644
index 00000000000..37e904a3e0e
--- /dev/null
+++ b/tensorflow/core/kernels/banded_triangular_solve_op_test.cc
@@ -0,0 +1,180 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/graph/testlib.h"
+#include "tensorflow/core/kernels/matrix_set_diag_op.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+namespace {
+
+Node* SetDiag(int num_bands, Graph* g, Node* bands, Node* triangular) {
+  Node* ret;
+  Tensor bandwidth(DT_INT32, TensorShape({2}));
+  bandwidth.flat()(0) = -(num_bands - 1);
+  bandwidth.flat()(1) = 0;
+  TF_CHECK_OK(NodeBuilder(g->NewName("n"), "MatrixSetDiagV3")
+                  .Input(triangular)
+                  .Input(bands)
+                  .Input(test::graph::Constant(g, bandwidth))
+                  .Attr("align", "RIGHT_LEFT")
+                  .Finalize(g, &ret));
+  return ret;
+}
+
+Node* BandedTriangularSolve(Graph* g, Node* in0, Node* in1) {
+  Node* ret;
+  TF_CHECK_OK(NodeBuilder(g->NewName("n"), "BandedTriangularSolve")
+                  .Input(in0)
+                  .Input(in1)
+                  .Attr("lower", true)
+                  .Attr("adjoint", false)
+                  .Finalize(g, &ret));
+  return ret;
+}
+
+Node* MatrixTriangularSolve(Graph* g, Node* in0, Node* in1) {
+  Node* ret;
+  TF_CHECK_OK(NodeBuilder(g->NewName("n"), "MatrixTriangularSolve")
+                  .Input(in0)
+                  .Input(in1)
+                  .Attr("lower", true)
+                  .Attr("adjoint", false)
+                  .Finalize(g, &ret));
+  return ret;
+}
+
+template 
+static Graph* BandedTriangularSolve(int64 num_bands, int64 n, int64 m,
+                                    bool use_banded_solver, DataType type) {
+  Graph* g = new Graph(OpRegistry::Global());
+  Tensor in0(type, TensorShape({num_bands, n}));
+  // Set diagonal to nonzero to guarantee invertibility.
+  in0.flat().setRandom();
+  in0.flat() =
+      in0.flat().abs() + in0.flat().constant(static_cast(0.5));
+  Tensor in1(type, TensorShape({n, m}));
+  in1.flat().setRandom();
+  if (use_banded_solver) {
+    BandedTriangularSolve(g, test::graph::Constant(g, in0),
+                          test::graph::Constant(g, in1));
+  } else {
+    // Create a zero tensor.
+    Tensor in2(type, TensorShape({n, n}));
+    in2.flat().setZero();
+    Node* triangular_matrix =
+        SetDiag(num_bands, g, test::graph::Constant(g, in0),
+                test::graph::Constant(g, in2));
+    MatrixTriangularSolve(g, triangular_matrix, test::graph::Constant(g, in1));
+  }
+  return g;
+}
+
+// Macro arguments names: --------------------------------------------------- //
+//   K: Number of bands
+//   N: Inner dimension of LHS, Inner dimension of RHS.
+//   M: Outer dimensions of RHS
+//   BS: boolean indicating whether to use the banded solver
+//    T: C++ type of scalars (e.g. float, std::complex)
+//   TT: TensorFlow type of scalars (e.g. DT_FLOAT, DT_COMPLEX128
+#define BM_BandedTriangularSolveDev(K, N, M, BS, T, TT, D)                     \
+  static void BM_BandedTriangularSolve##_##K##_##N##_##M##_##BS##_##TT(        \
+      int iters) {                                                             \
+    testing::UseRealTime();                                                    \
+    testing::ItemsProcessed(static_cast(iters) * K * N + N * M);        \
+    test::Benchmark(#D, BandedTriangularSolve(K, N, M, BS, TT)).Run(iters); \
+  }                                                                            \
+  BENCHMARK(BM_BandedTriangularSolve##_##K##_##N##_##M##_##BS##_##TT);
+
+#define BM_BandedTriangularSolve(K, N, M, BS, D)                \
+  BM_BandedTriangularSolveDev(K, N, M, BS, float, DT_FLOAT, D); \
+  BM_BandedTriangularSolveDev(K, N, M, BS, double, DT_DOUBLE, D);
+
+// Small number of bands, few rhs
+BM_BandedTriangularSolve(2, 32, 1, true, cpu);
+BM_BandedTriangularSolve(2, 32, 1, false, cpu);
+BM_BandedTriangularSolve(4, 32, 1, true, cpu);
+BM_BandedTriangularSolve(4, 32, 1, false, cpu);
+BM_BandedTriangularSolve(8, 32, 1, true, cpu);
+BM_BandedTriangularSolve(8, 32, 1, false, cpu);
+BM_BandedTriangularSolve(16, 32, 1, true, cpu);
+BM_BandedTriangularSolve(16, 32, 1, false, cpu);
+BM_BandedTriangularSolve(2, 128, 1, true, cpu);
+BM_BandedTriangularSolve(2, 128, 1, false, cpu);
+BM_BandedTriangularSolve(4, 128, 1, true, cpu);
+BM_BandedTriangularSolve(4, 128, 1, false, cpu);
+BM_BandedTriangularSolve(8, 128, 1, true, cpu);
+BM_BandedTriangularSolve(8, 128, 1, false, cpu);
+BM_BandedTriangularSolve(16, 128, 1, true, cpu);
+BM_BandedTriangularSolve(16, 128, 1, false, cpu);
+BM_BandedTriangularSolve(2, 512, 1, true, cpu);
+BM_BandedTriangularSolve(2, 512, 1, false, cpu);
+BM_BandedTriangularSolve(4, 512, 1, true, cpu);
+BM_BandedTriangularSolve(4, 512, 1, false, cpu);
+BM_BandedTriangularSolve(8, 512, 1, true, cpu);
+BM_BandedTriangularSolve(8, 512, 1, false, cpu);
+BM_BandedTriangularSolve(16, 512, 1, true, cpu);
+BM_BandedTriangularSolve(16, 512, 1, false, cpu);
+
+// Larger # rhs
+BM_BandedTriangularSolve(2, 32, 32, true, cpu);
+BM_BandedTriangularSolve(2, 32, 32, false, cpu);
+BM_BandedTriangularSolve(4, 32, 32, true, cpu);
+BM_BandedTriangularSolve(4, 32, 32, false, cpu);
+BM_BandedTriangularSolve(8, 32, 32, true, cpu);
+BM_BandedTriangularSolve(8, 32, 32, false, cpu);
+BM_BandedTriangularSolve(16, 32, 32, true, cpu);
+BM_BandedTriangularSolve(16, 32, 32, false, cpu);
+BM_BandedTriangularSolve(2, 128, 128, true, cpu);
+BM_BandedTriangularSolve(2, 128, 128, false, cpu);
+BM_BandedTriangularSolve(4, 128, 128, true, cpu);
+BM_BandedTriangularSolve(4, 128, 128, false, cpu);
+BM_BandedTriangularSolve(8, 128, 128, true, cpu);
+BM_BandedTriangularSolve(8, 128, 128, false, cpu);
+BM_BandedTriangularSolve(16, 128, 128, true, cpu);
+BM_BandedTriangularSolve(16, 128, 128, false, cpu);
+BM_BandedTriangularSolve(2, 512, 512, true, cpu);
+BM_BandedTriangularSolve(2, 512, 512, false, cpu);
+BM_BandedTriangularSolve(4, 512, 512, true, cpu);
+BM_BandedTriangularSolve(4, 512, 512, false, cpu);
+BM_BandedTriangularSolve(8, 512, 512, true, cpu);
+BM_BandedTriangularSolve(8, 512, 512, false, cpu);
+BM_BandedTriangularSolve(16, 512, 512, true, cpu);
+BM_BandedTriangularSolve(16, 512, 512, false, cpu);
+
+BM_BandedTriangularSolve(2, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(2, 2048, 2048, false, cpu);
+BM_BandedTriangularSolve(4, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(4, 2048, 2048, false, cpu);
+BM_BandedTriangularSolve(8, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(8, 2048, 2048, false, cpu);
+BM_BandedTriangularSolve(16, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(16, 2048, 2048, false, cpu);
+BM_BandedTriangularSolve(32, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(32, 2048, 2048, false, cpu);
+BM_BandedTriangularSolve(64, 2048, 2048, true, cpu);
+BM_BandedTriangularSolve(64, 2048, 2048, false, cpu);
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/ops/linalg_ops.cc b/tensorflow/core/ops/linalg_ops.cc
index 75340b28eb0..a05231834b7 100644
--- a/tensorflow/core/ops/linalg_ops.cc
+++ b/tensorflow/core/ops/linalg_ops.cc
@@ -47,6 +47,49 @@ Status BatchUnchangedSquareShapeFn(InferenceContext* c) {
   return Status::OK();
 }
 
+// The first input is [...,K,M] and second input is [...,M,N].
+Status BandedTriangularSolveShapeFn(InferenceContext* c) {
+  ShapeHandle lhs;
+  ShapeHandle rhs;
+
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &lhs));
+  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 2, &rhs));
+
+  // Check K > 0.
+  DimensionHandle num_bands = c->Dim(lhs, -2);
+  DimensionHandle m = c->Dim(lhs, -1);
+  if (c->ValueKnown(num_bands) && c->Value(num_bands) <= 0) {
+    return errors::InvalidArgument("Number of bands must be positive, but is ",
+                                   c->Value(num_bands));
+  }
+  if (c->ValueKnown(num_bands) && c->ValueKnown(m) &&
+      c->Value(num_bands) > c->Value(m)) {
+    return errors::InvalidArgument("Number of bands ", c->Value(num_bands),
+                                   " cannot exceed the size of the matrix ",
+                                   c->Value(m));
+  }
+
+  ShapeHandle lhs_batch_shape;
+  ShapeHandle rhs_batch_shape;
+  ShapeHandle output_batch_shape;
+  // Make the common batch subshape.
+  TF_RETURN_IF_ERROR(c->Subshape(lhs, 0, -2, &lhs_batch_shape));
+  TF_RETURN_IF_ERROR(c->Subshape(rhs, 0, -2, &rhs_batch_shape));
+  TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
+      c, lhs_batch_shape, rhs_batch_shape, true, &output_batch_shape));
+
+  // lhs and rhs have the same value for M to be compatible.
+  TF_RETURN_IF_ERROR(c->Merge(m, c->Dim(rhs, -2), &m));
+
+  // Build final shape (batch_shape + m + n) in .
+  ShapeHandle out;
+  TF_RETURN_IF_ERROR(
+      c->Concatenate(output_batch_shape, c->Matrix(m, c->Dim(rhs, -1)), &out));
+
+  c->set_output(0, out);
+  return Status::OK();
+}
+
 // The first input is [...,M,N] and second input is either [...,M,K] or [...,M].
 // Output is [...,N,K] or [...,N]. If , then input is [...,M,M].
 Status MatrixSolveShapeFn(InferenceContext* c, bool square) {
@@ -446,6 +489,17 @@ REGISTER_OP("MatrixSolve")
       return MatrixSolveShapeFn(c, true /* square (*/);
     });
 
+REGISTER_OP("BandedTriangularSolve")
+    .Input("matrix: T")
+    .Input("rhs: T")
+    .Output("output: T")
+    .Attr("lower: bool = True")
+    .Attr("adjoint: bool = False")
+    .Attr("T: {double, float, half, complex64, complex128}")
+    .SetShapeFn([](InferenceContext* c) {
+      return BandedTriangularSolveShapeFn(c);
+    });
+
 REGISTER_OP("MatrixTriangularSolve")
     .Input("matrix: T")
     .Input("rhs: T")
diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD
index 846c582737f..f2c614974f5 100644
--- a/tensorflow/python/kernel_tests/BUILD
+++ b/tensorflow/python/kernel_tests/BUILD
@@ -762,6 +762,17 @@ cuda_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "banded_triangular_solve_op_test",
+    size = "small",
+    srcs = ["banded_triangular_solve_op_test.py"],
+    deps = [
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:linalg_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
 cuda_py_test(
     name = "matrix_triangular_solve_op_test",
     size = "medium",
diff --git a/tensorflow/python/kernel_tests/banded_triangular_solve_op_test.py b/tensorflow/python/kernel_tests/banded_triangular_solve_op_test.py
new file mode 100644
index 00000000000..bd0fdae03c5
--- /dev/null
+++ b/tensorflow/python/kernel_tests/banded_triangular_solve_op_test.py
@@ -0,0 +1,232 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for tensorflow.ops.math_ops.banded_triangular_solve."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.platform import test
+
+
+class BandedTriangularSolveOpTest(test.TestCase):
+
+  def _verifySolveAllWays(self, x, y, dtypes, batch_dims=None):
+    for lower in (False,):
+      for adjoint in (False, True):
+        for use_placeholder in True, False:
+          self._verifySolve(
+              x,
+              y,
+              lower=lower,
+              adjoint=adjoint,
+              batch_dims=batch_dims,
+              use_placeholder=use_placeholder,
+              dtypes=dtypes)
+
+  def _verifySolveAllWaysReal(self, x, y, batch_dims=None):
+    self._verifySolveAllWays(x, y, (np.float32, np.float64), batch_dims)
+
+  def _verifySolveAllWaysComplex(self, x, y, batch_dims=None):
+    self._verifySolveAllWays(x, y, (np.complex64, np.complex128), batch_dims)
+
+  def _verifySolve(self,
+                   x,
+                   y,
+                   lower=True,
+                   adjoint=False,
+                   batch_dims=None,
+                   use_placeholder=False,
+                   dtypes=(np.float32, np.float64)):
+    for np_type in dtypes:
+      a = x.astype(np_type)
+      b = y.astype(np_type)
+
+      # Now we need to convert a to a dense triangular matrix.
+      def make_diags(diags, lower=True):
+        n = len(diags[0])
+        a = np.zeros(n * n, dtype=diags.dtype)
+        if lower:
+          for i, diag in enumerate(diags):
+            a[n * i:n * n:n + 1] = diag[i:]
+        else:
+          diags_flip = np.flip(diags, 0)
+          for i, diag in enumerate(diags_flip):
+            a[i:(n - i) * n:n + 1] = diag[:(n - i)]
+        return a.reshape(n, n)
+
+      # For numpy.solve we have to explicitly zero out the strictly
+      # upper or lower triangle.
+      if a.size > 0:
+        a_np = make_diags(a, lower=lower)
+      else:
+        a_np = a
+      if adjoint:
+        a_np = np.conj(np.transpose(a_np))
+
+      if batch_dims is not None:
+        a = np.tile(a, batch_dims + [1, 1])
+        a_np = np.tile(a_np, batch_dims + [1, 1])
+        b = np.tile(b, batch_dims + [1, 1])
+
+      with self.cached_session(use_gpu=True):
+        a_tf = a
+        b_tf = b
+        if use_placeholder:
+          a_tf = array_ops.placeholder_with_default(a_tf, shape=None)
+          b_tf = array_ops.placeholder_with_default(b_tf, shape=None)
+        tf_ans = linalg_ops.banded_triangular_solve(
+            a_tf, b_tf, lower=lower, adjoint=adjoint)
+        tf_val = self.evaluate(tf_ans)
+        np_ans = np.linalg.solve(a_np, b)
+        self.assertEqual(np_ans.shape, tf_val.shape)
+        self.assertAllClose(np_ans, tf_val)
+
+  @test_util.run_deprecated_v1
+  def testSolve(self):
+    # 1x1 matrix, single rhs.
+    matrix = np.array([[0.1]])
+    rhs0 = np.array([[1.]])
+    self._verifySolveAllWaysReal(matrix, rhs0)
+    # 2x2 matrix with 2 bands, single right-hand side.
+    # Corresponds to the lower triangular
+    # [[1., 0.], [3., 4.]]
+    # and upper triangular
+    # [[2., 1.], [0., 3.]]
+    matrix = np.array([[1., 4.], [2., 3.]])
+    rhs0 = np.array([[1.], [1.]])
+    self._verifySolveAllWaysReal(matrix, rhs0)
+    # 2x2 matrix with 2 bands, 3 right-hand sides.
+    rhs1 = np.array([[1., 0., 1.], [0., 1., 1.]])
+    self._verifySolveAllWaysReal(matrix, rhs1)
+    # 4 x 4 matrix with 2 bands, 3 right hand sides.
+    # Corresponds to the lower triangular
+    # [[1.,  0., 0., 0.],
+    #  [-1., 2., 0., 0.],
+    #  [0., -2., 3., 0.],
+    #  [0., 0., -3., 4.]]
+    # and upper triangular
+    # [[1.,  1., 0., 0.],
+    #  [0., -1., 2., 0.],
+    #  [0., 0., -2., 3.],
+    #  [0., 0., 0., -3.]]
+    matrix = np.array([[1., 2., 3., 4.], [1., -1., -2., -3.]])
+    rhs0 = np.array([[1., 0., 1.], [0., 1., 1.], [-1., 2., 1.], [0., -1., -1.]])
+    self._verifySolveAllWaysReal(matrix, rhs0)
+
+  def testSolveBandSizeSmaller(self):
+    rhs0 = np.random.randn(6, 4)
+
+    # 6 x 6 matrix with 2 bands. Ensure all non-zero entries.
+    matrix = 2. * np.random.uniform(size=[3, 6]) + 1.
+    self._verifySolveAllWaysReal(matrix, rhs0)
+
+    # 6 x 6 matrix with 3 bands. Ensure all non-zero entries.
+    matrix = 2. * np.random.uniform(size=[3, 6]) + 1.
+    self._verifySolveAllWaysReal(matrix, rhs0)
+
+  @test_util.run_deprecated_v1
+  def testSolveComplex(self):
+    if test.is_built_with_rocm():
+      self.skipTest("ROCm does not support BLAS operations for complex types")
+    # 1x1 matrix, single rhs.
+    matrix = np.array([[0.1 + 1j * 0.1]])
+    rhs0 = np.array([[1. + 1j]])
+    self._verifySolveAllWaysComplex(matrix, rhs0)
+    # 2x2 matrix with 2 bands, single right-hand side.
+    # Corresponds to
+    # [[1. + 1j, 0.], [4 + 1j, 2 + 1j]]
+    matrix = np.array([[1., 2.], [3., 4.]]).astype(np.complex64)
+    matrix += 1j * matrix
+    rhs0 = np.array([[1.], [1.]]).astype(np.complex64)
+    rhs0 += 1j * rhs0
+    self._verifySolveAllWaysComplex(matrix, rhs0)
+    # 2x2 matrix with 2 bands, 3 right-hand sides.
+    rhs1 = np.array([[1., 0., 1.], [0., 1., 1.]]).astype(np.complex64)
+    rhs1 += 1j * rhs1
+    self._verifySolveAllWaysComplex(matrix, rhs1)
+
+  @test_util.run_deprecated_v1
+  def testSolveBatch(self):
+    matrix = np.array([[1., 2.], [3., 4.]])
+    rhs = np.array([[1., 0., 1.], [0., 1., 1.]])
+    # Batch of 2x3x2x2 matrices, 2x3x2x3 right-hand sides.
+    self._verifySolveAllWaysReal(matrix, rhs, batch_dims=[2, 3])
+    # Batch of 3x2x2x2 matrices, 3x2x2x3 right-hand sides.
+    self._verifySolveAllWaysReal(matrix, rhs, batch_dims=[3, 2])
+
+    matrix = np.array([[1., 2., 3., 4.], [-1., -2., -3., -4.],
+                       [-1., 1., 2., 3.]])
+    rhs = np.array([[-1., 2.], [1., 1.], [0., 1.], [2., 3.]])
+    # Batch of 2x3x4x4 matrices with 3 bands, 2x3x4x2 right-hand sides.
+    self._verifySolveAllWaysReal(matrix, rhs, batch_dims=[2, 3])
+    # Batch of 3x2x4x4 matrices with 3 bands, 3x2x4x2 right-hand sides.
+    self._verifySolveAllWaysReal(matrix, rhs, batch_dims=[3, 2])
+
+  @test_util.run_deprecated_v1
+  def testSolveBatchComplex(self):
+    if test.is_built_with_rocm():
+      self.skipTest("ROCm does not support BLAS operations for complex types")
+    matrix = np.array([[1., 2.], [3., 4.]]).astype(np.complex64)
+    matrix += 1j * matrix
+    rhs = np.array([[1., 0., 1.], [0., 1., 1.]]).astype(np.complex64)
+    rhs += 1j * rhs
+    # Batch of 2x3x2x2 matrices, 2x3x2x3 right-hand sides.
+    self._verifySolveAllWaysComplex(matrix, rhs, batch_dims=[2, 3])
+    # Batch of 3x2x2x2 matrices, 3x2x2x3 right-hand sides.
+    self._verifySolveAllWaysComplex(matrix, rhs, batch_dims=[3, 2])
+
+  @test_util.run_deprecated_v1
+  def testWrongDimensions(self):
+    # The matrix should have the same number of rows as the
+    # right-hand sides.
+    matrix = np.array([[1., 1.], [1., 1.]])
+    rhs = np.array([[1., 0.]])
+    with self.cached_session(use_gpu=True):
+      with self.assertRaises(ValueError):
+        self._verifySolve(matrix, rhs)
+      with self.assertRaises(ValueError):
+        self._verifySolve(matrix, rhs, batch_dims=[2, 3])
+
+    # Number of bands exceeds the dimension of the matrix.
+    matrix = np.ones((6, 4))
+    rhs = np.ones((4, 2))
+    with self.cached_session(use_gpu=True):
+      with self.assertRaises(ValueError):
+        self._verifySolve(matrix, rhs)
+      with self.assertRaises(ValueError):
+        self._verifySolve(matrix, rhs, batch_dims=[2, 3])
+
+  @test_util.run_deprecated_v1
+  @test_util.disable_xla("XLA cannot throw assertion errors during a kernel.")
+  def testNotInvertible(self):
+    # The input should be invertible.
+    # The matrix is singular because it has a zero on the diagonal.
+    # FIXME(rmlarsen): The GPU kernel does not check for singularity.
+    singular_matrix = np.array([[1., 0., -1.], [-1., 0., 1.], [0., -1., 1.]])
+    with self.cached_session():
+      with self.assertRaisesOpError("Input matrix is not invertible."):
+        self._verifySolve(singular_matrix, singular_matrix)
+      with self.assertRaisesOpError("Input matrix is not invertible."):
+        self._verifySolve(singular_matrix, singular_matrix, batch_dims=[2, 3])
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/python/kernel_tests/linalg_grad_test.py b/tensorflow/python/kernel_tests/linalg_grad_test.py
index 36e58bee829..3aceddf4d5f 100644
--- a/tensorflow/python/kernel_tests/linalg_grad_test.py
+++ b/tensorflow/python/kernel_tests/linalg_grad_test.py
@@ -132,6 +132,44 @@ def _GetMatrixBinaryFunctorGradientTest(functor_,
   return Test
 
 
+def _GetBandedTriangularSolveGradientTest(
+    functor_,
+    dtype_,
+    shape_,
+    float32_tol_fudge=1.0,  # pylint: disable=redefined-outer-name
+    **kwargs_):
+
+  @test_util.run_in_graph_and_eager_modes(use_gpu=True)
+  def Test(self):
+    n = shape_[-1]
+
+    np.random.seed(1)
+    # Make sure invertible.
+    a_np = np.random.uniform(low=1.0, high=2.0, size=shape_).astype(dtype_)
+    a = constant_op.constant(a_np)
+
+    b_np = np.random.uniform(low=-1.0, high=1.0, size=[n, n]).astype(dtype_)
+    b = constant_op.constant(b_np)
+
+    epsilon = np.finfo(dtype_).eps
+    delta = epsilon**(1.0 / 3.0)
+    # tolerance obtained by looking at actual differences using
+    # np.linalg.norm(theoretical-numerical, np.inf) on -mavx build
+    tol = 1e-6 if dtype_ == np.float64 else float32_tol_fudge * 0.05
+
+    # check gradient w.r.t. left argument.
+    theoretical, numerical = gradient_checker_v2.compute_gradient(
+        lambda x: functor_(x, b, **kwargs_), [a], delta=delta)
+    self.assertAllClose(theoretical, numerical, atol=tol, rtol=tol)
+
+    # check gradient w.r.t. right argument.
+    theoretical, numerical = gradient_checker_v2.compute_gradient(
+        lambda y: functor_(a, y, **kwargs_), [b], delta=delta)
+    self.assertAllClose(theoretical, numerical, atol=tol, rtol=tol)
+
+  return Test
+
+
 if __name__ == '__main__':
   # Tests for gradients of binary matrix operations.
   for dtype in np.float32, np.float64:
@@ -166,6 +204,20 @@ if __name__ == '__main__':
                          adjoint=adjoint,
                          lower=lower))
 
+            band_shape = extra + (size // 2 + 1, size)
+            name = '%s_%s_adj_%s_low_%s' % (dtype.__name__, '_'.join(
+                map(str, band_shape)), str(adjoint), lower)
+            _AddTest(
+                MatrixBinaryFunctorGradientTest,
+                'BandedTriangularSolveGradient', name,
+                _GetBandedTriangularSolveGradientTest(
+                    linalg_ops.banded_triangular_solve,
+                    dtype,
+                    band_shape,
+                    float32_tol_fudge=4.0,
+                    adjoint=adjoint,
+                    lower=lower))
+
   # Tests for gradients of unary matrix operations.
   for dtype in np.float32, np.float64:
     for size in 2, 5, 10:
diff --git a/tensorflow/python/ops/linalg/linalg_impl.py b/tensorflow/python/ops/linalg/linalg_impl.py
index 82acd09caec..9ddf7b5e8b8 100644
--- a/tensorflow/python/ops/linalg/linalg_impl.py
+++ b/tensorflow/python/ops/linalg/linalg_impl.py
@@ -340,6 +340,102 @@ def matrix_exponential(input, name=None):  # pylint: disable=redefined-builtin
     return array_ops.reshape(result, batch_shape.concatenate(result.shape[-2:]))
 
 
+@tf_export('linalg.banded_triangular_solve', v1=[])
+def banded_triangular_solve(
+    bands,
+    rhs,
+    lower=True,
+    adjoint=False,  # pylint: disable=redefined-outer-name
+    name=None):
+  r"""Solve triangular systems of equations with a banded solver.
+
+  `bands` is a tensor of shape `[..., K, M]`, where `K` represents the number
+  of bands stored. This corresponds to a batch of `M` by `M` matrices, whose
+  `K` subdiagonals (when `lower` is `True`) are stored.
+
+  This operator broadcasts the batch dimensions of `bands` and the batch
+  dimensions of `rhs`.
+
+
+  Examples:
+
+  Storing 2 bands of a 3x3 matrix.
+  Note that first element in the second row is ignored due to
+  the 'LEFT_RIGHT' padding.
+
+  >>> x = [[2., 3., 4.], [1., 2., 3.]]
+  >>> x2 = [[2., 3., 4.], [10000., 2., 3.]]
+  >>> y = tf.zeros([3, 3])
+  >>> z = tf.linalg.set_diag(y, x, align='LEFT_RIGHT', k=(-1, 0))
+  >>> z
+  
+  >>> soln = tf.linalg.banded_triangular_solve(x, tf.ones([3, 1]))
+  >>> soln
+  
+  >>> are_equal = soln == tf.linalg.banded_triangular_solve(x2, tf.ones([3, 1]))
+  >>> tf.reduce_all(are_equal).numpy()
+  True
+  >>> are_equal = soln == tf.linalg.triangular_solve(z, tf.ones([3, 1]))
+  >>> tf.reduce_all(are_equal).numpy()
+  True
+
+  Storing 2 superdiagonals of a 4x4 matrix. Because of the 'LEFT_RIGHT' padding
+  the last element of the first row is ignored.
+
+  >>> x = [[2., 3., 4., 5.], [-1., -2., -3., -4.]]
+  >>> y = tf.zeros([4, 4])
+  >>> z = tf.linalg.set_diag(y, x, align='LEFT_RIGHT', k=(0, 1))
+  >>> z
+  
+  >>> soln = tf.linalg.banded_triangular_solve(x, tf.ones([4, 1]), lower=False)
+  >>> soln
+  
+  >>> are_equal = (soln == tf.linalg.triangular_solve(
+  ...   z, tf.ones([4, 1]), lower=False))
+  >>> tf.reduce_all(are_equal).numpy()
+  True
+
+
+  Args:
+    bands: A `Tensor` describing the bands of the left hand side, with shape
+      `[..., K, M]`. The `K` rows correspond to the diagonal to the `K - 1`-th
+      diagonal (the diagonal is the top row) when `lower` is `True` and
+      otherwise the `K - 1`-th superdiagonal to the diagonal (the diagonal is
+      the bottom row) when `lower` is `False`. The bands are stored with
+      'LEFT_RIGHT' alignment, where the superdiagonals are padded on the right
+      and subdiagonals are padded on the left. This is the alignment cuSPARSE
+      uses.  See  `tf.linalg.set_diag` for more details.
+    rhs: A `Tensor` of shape [..., M] or [..., M, N] and with the same dtype as
+      `diagonals`. Note that if the shape of `rhs` and/or `diags` isn't known
+      statically, `rhs` will be treated as a matrix rather than a vector.
+    lower: An optional `bool`. Defaults to `True`. Boolean indicating whether
+      `bands` represents a lower or upper triangular matrix.
+    adjoint: An optional `bool`. Defaults to `False`. Boolean indicating whether
+      to solve with the matrix's block-wise adjoint.
+    name:  A name to give this `Op` (optional).
+
+  Returns:
+    A `Tensor` of shape [..., M] or [..., M, N] containing the solutions.
+  """
+  with ops.name_scope(name, 'banded_triangular_solve', [bands, rhs]):
+    return gen_linalg_ops.banded_triangular_solve(
+        bands, rhs, lower=lower, adjoint=adjoint)
+
+
 @tf_export('linalg.tridiagonal_solve')
 @dispatch.add_dispatch_support
 def tridiagonal_solve(diagonals,
diff --git a/tensorflow/python/ops/linalg_grad.py b/tensorflow/python/ops/linalg_grad.py
index 8d3664144a1..437e28e7e6b 100644
--- a/tensorflow/python/ops/linalg_grad.py
+++ b/tensorflow/python/ops/linalg_grad.py
@@ -607,6 +607,39 @@ def _MatrixSolveLsGrad(op, grad):
                                  lambda: _Underdetermined(op, grad))
 
 
+@ops.RegisterGradient("BandedTriangularSolve")
+def _BandedTriangularSolveGrad(op, grad):
+  """Gradient for BandedTriangularSolve."""
+  a = op.inputs[0]
+  b = op.inputs[1]
+  num_bands = array_ops.shape(a)[-2]
+  adjoint_a = op.get_attr("adjoint")
+  lower_a = op.get_attr("lower")
+  c = op.outputs[0]
+  grad_b = linalg_ops.banded_triangular_solve(
+      a, grad, lower=lower_a, adjoint=not adjoint_a)
+  if adjoint_a:
+    grad_a = -math_ops.matmul(c, grad_b, adjoint_b=True)
+  else:
+    grad_a = -math_ops.matmul(grad_b, c, adjoint_b=True)
+  if lower_a:
+    grad_a = array_ops.matrix_diag_part(
+        grad_a, k=(-(num_bands - 1), 0), align="LEFT_RIGHT")
+  else:
+    grad_a = array_ops.matrix_diag_part(
+        grad_a, k=(0, num_bands - 1), align="LEFT_RIGHT")
+  # If the static batch shapes are equal, we don't need to unbroadcast.
+  if (a.shape.is_fully_defined() and b.shape.is_fully_defined() and
+      a.shape[:-2] == b.shape[:-2]):
+    return grad_a, grad_b
+  a_shape = array_ops.shape(a)
+  b_shape = array_ops.shape(b)
+  ra, rb = array_ops.broadcast_gradient_args(a_shape[:-2], b_shape[:-2])
+  grad_a = array_ops.reshape(math_ops.reduce_sum(grad_a, axis=ra), a_shape)
+  grad_b = array_ops.reshape(math_ops.reduce_sum(grad_b, axis=rb), b_shape)
+  return grad_a, grad_b
+
+
 @ops.RegisterGradient("MatrixTriangularSolve")
 def _MatrixTriangularSolveGrad(op, grad):
   """Gradient for MatrixTriangularSolve."""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
index 3d298e928e9..62969b5a0dd 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
@@ -284,6 +284,10 @@ tf_module {
     name: "AvgPoolGrad"
     argspec: "args=[\'orig_input_shape\', \'grad\', \'ksize\', \'strides\', \'padding\', \'data_format\', \'name\'], varargs=None, keywords=None, defaults=[\'NHWC\', \'None\'], "
   }
+  member_method {
+    name: "BandedTriangularSolve"
+    argspec: "args=[\'matrix\', \'rhs\', \'lower\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'False\', \'None\'], "
+  }
   member_method {
     name: "Barrier"
     argspec: "args=[\'component_types\', \'shapes\', \'capacity\', \'container\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'[]\', \'-1\', \'\', \'\', \'None\'], "
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
index 734837b99cb..4f62af20dc0 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.linalg.pbtxt
@@ -96,6 +96,10 @@ tf_module {
     name: "band_part"
     argspec: "args=[\'input\', \'num_lower\', \'num_upper\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
+  member_method {
+    name: "banded_triangular_solve"
+    argspec: "args=[\'bands\', \'rhs\', \'lower\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'False\', \'None\'], "
+  }
   member_method {
     name: "cholesky"
     argspec: "args=[\'input\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
index 3d298e928e9..62969b5a0dd 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
@@ -284,6 +284,10 @@ tf_module {
     name: "AvgPoolGrad"
     argspec: "args=[\'orig_input_shape\', \'grad\', \'ksize\', \'strides\', \'padding\', \'data_format\', \'name\'], varargs=None, keywords=None, defaults=[\'NHWC\', \'None\'], "
   }
+  member_method {
+    name: "BandedTriangularSolve"
+    argspec: "args=[\'matrix\', \'rhs\', \'lower\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'True\', \'False\', \'None\'], "
+  }
   member_method {
     name: "Barrier"
     argspec: "args=[\'component_types\', \'shapes\', \'capacity\', \'container\', \'shared_name\', \'name\'], varargs=None, keywords=None, defaults=[\'[]\', \'-1\', \'\', \'\', \'None\'], "

From 455750f3623b15a3b5d46c11d4c5102e9388dbda Mon Sep 17 00:00:00 2001
From: Srinivas Vasudevan 
Date: Thu, 18 Jun 2020 10:23:23 -0700
Subject: [PATCH 0505/1390] Add StatelessParameterizedTruncatedNormal sampler.

This sampler supports broadcasting of its input parameters as well as puts the # samples at the left of the output shape, rather than the right.

PiperOrigin-RevId: 317129622
Change-Id: I4b62ad2e89a9637ae8b30b73af4b662ad0caa943
---
 ...tatelessParameterizedTruncatedNormal.pbtxt |  54 +++
 tensorflow/core/kernels/BUILD                 |   1 +
 .../parameterized_truncated_normal_op.cc      | 435 +++++++++++++++++-
 .../parameterized_truncated_normal_op.h       |  16 +
 tensorflow/core/ops/stateless_random_ops.cc   |  35 ++
 .../eager/pywrap_gradient_exclusions.cc       |   3 +-
 tensorflow/python/kernel_tests/BUILD          |  18 -
 tensorflow/python/kernel_tests/random/BUILD   |  18 +
 .../parameterized_truncated_normal_op_test.py | 198 ++++++--
 tensorflow/python/ops/random_grad.py          | 120 +++++
 tensorflow/python/ops/stateless_random_ops.py |  70 +++
 .../api/golden/v1/tensorflow.random.pbtxt     |   4 +
 .../api/golden/v1/tensorflow.raw_ops.pbtxt    |   4 +
 .../api/golden/v2/tensorflow.random.pbtxt     |   4 +
 .../api/golden/v2/tensorflow.raw_ops.pbtxt    |   4 +
 15 files changed, 913 insertions(+), 71 deletions(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
 rename tensorflow/python/kernel_tests/{ => random}/parameterized_truncated_normal_op_test.py (63%)

diff --git a/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt b/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
new file mode 100644
index 00000000000..15bd4670cef
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
@@ -0,0 +1,54 @@
+op {
+  graph_op_name: "StatelessParameterizedTruncatedNormal"
+  visibility: HIDDEN
+  in_arg {
+    name: "shape"
+    description: < {
     const T kStdDevsInsideBoundsToUseRandnSampler = T(1.3);
     auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());
 
-    auto DoWork = [samples_per_batch, num_elements, &ctx, &means, &stddevs,
-                   &minvals, &maxvals, &gen, &output,
-                   kStdDevsInsideBoundsToUseRandnSampler](int start_batch,
-                                                          int limit_batch) {
+    auto do_work = [samples_per_batch, num_elements, &ctx, &means, &stddevs,
+                    &minvals, &maxvals, &gen, &output,
+                    kStdDevsInsideBoundsToUseRandnSampler](int start_batch,
+                                                           int limit_batch) {
       // Capturing "gen" by-value would only make a copy for the _shared_
       // lambda.  Since we want to let each worker have its own copy, we pass
       // "gen" by reference and explicitly do a copy assignment here.
@@ -80,9 +81,9 @@ struct TruncatedNormalFunctor {
       // The sample from each iteration uses 2 random numbers.
       gen_copy.Skip(start_batch * 2 * kMaxIterations * (samples_per_batch + 3) /
                     4);
-      typedef random::UniformDistribution Uniform;
+      using Uniform = random::UniformDistribution;
       Uniform dist;
-      typedef random::NormalDistribution Normal;
+      using Normal = random::NormalDistribution;
       Normal normal_dist;
 
       // Vectorized intermediate calculations for uniform rejection sampling.
@@ -112,7 +113,7 @@ struct TruncatedNormalFunctor {
                          Eigen::numext::isfinite(maxval)),
                     errors::InvalidArgument("Invalid parameters"));
 
-        int numIterations = 0;
+        int num_iterations = 0;
 
         // If possible, make one-sided bound be the lower bound, or make both
         // bounds positive. Otherwise, the bounds are on either side of the
@@ -160,10 +161,10 @@ struct TruncatedNormalFunctor {
                 if (sample >= limit_sample) {
                   break;
                 }
-                numIterations = 0;
+                num_iterations = 0;
               } else {
-                numIterations++;
-                if (numIterations > kMaxIterations) {
+                num_iterations++;
+                if (num_iterations > kMaxIterations) {
                   // This should never occur because this sampler should
                   // (by the selection criteria above) be used if at least 3
                   // standard deviations of one side of the distribution
@@ -201,7 +202,7 @@ struct TruncatedNormalFunctor {
             const auto u = dist(&gen_copy);
             for (int i = 0; i < size; i++) {
               auto accept = u[i] <= Eigen::numext::exp(g[i]);
-              if (accept || numIterations + 1 >= kMaxIterations) {
+              if (accept || num_iterations + 1 >= kMaxIterations) {
                 // Accept the sample z.
                 // If we run out of iterations, just use the current uniform
                 // sample, but emit a warning.
@@ -223,9 +224,9 @@ struct TruncatedNormalFunctor {
                 if (sample >= limit_sample) {
                   break;
                 }
-                numIterations = 0;
+                num_iterations = 0;
               } else {
-                numIterations++;
+                num_iterations++;
               }
             }
           }
@@ -248,7 +249,7 @@ struct TruncatedNormalFunctor {
               const T u = rand[i];
               i++;
               auto accept = (u <= g && z < normMax);
-              if (accept || numIterations + 1 >= kMaxIterations) {
+              if (accept || num_iterations + 1 >= kMaxIterations) {
                 if (!accept) {
                   LOG(ERROR) << "TruncatedNormal exponential distribution "
                              << "rejection sampler exceeds max iterations. "
@@ -263,9 +264,9 @@ struct TruncatedNormalFunctor {
                 if (sample >= limit_sample) {
                   break;
                 }
-                numIterations = 0;
+                num_iterations = 0;
               } else {
-                numIterations++;
+                num_iterations++;
               }
             }
           }
@@ -305,7 +306,297 @@ struct TruncatedNormalFunctor {
     const int64 batchCost =
         batchInitCost + uniformRejectionSamplingCost * 2 * samples_per_batch;
     Shard(worker_threads.num_threads, worker_threads.workers, num_batches,
-          batchCost, DoWork);
+          batchCost, do_work);
+  }
+};
+
+template 
+struct TruncatedNormalFunctorV2 {
+  void operator()(OpKernelContext* ctx, const CPUDevice& d, int64 num_batches,
+                  int64 samples_per_batch, int64 num_elements,
+                  const BCastList<4>& bcast,
+                  typename TTypes::ConstFlat means,
+                  typename TTypes::ConstFlat stddevs,
+                  typename TTypes::ConstFlat minvals,
+                  typename TTypes::ConstFlat maxvals,
+                  const random::PhiloxRandom& gen,
+                  typename TTypes::Flat output) {
+    // The randn rejection sampling is used when the mean and at least this many
+    // standard deviations are inside the bounds.
+    // The uniform proposal samplers become less efficient as the bounds are
+    // further from the mean, the reverse is true for the randn sampler.
+    // This number was chosen by empirical benchmarking. If modified, the
+    // benchmarks in parameterized_truncated_normal_op_test should also be
+    // changed.
+    const T kStdDevsInsideBoundsToUseRandnSampler = T(1.3);
+    auto worker_threads = *(ctx->device()->tensorflow_cpu_worker_threads());
+
+    auto do_work = [num_batches, samples_per_batch, &ctx, &bcast, &means,
+                    &stddevs, &minvals, &maxvals, &gen, &output,
+                    kStdDevsInsideBoundsToUseRandnSampler](int start_output,
+                                                           int limit_output) {
+      // Capturing "gen" by-value would only make a copy for the _shared_
+      // lambda.  Since we want to let each worker have its own copy, we pass
+      // "gen" by reference and explicitly do a copy assignment here.
+      random::PhiloxRandom gen_copy = gen;
+      using Uniform = random::UniformDistribution;
+      Uniform dist;
+      using Normal = random::NormalDistribution;
+      Normal normal_dist;
+      // Skip takes units of 128 bits. The Uniform::kResultElementCount - 1
+      // is so rounding doesn't lead to
+      // us using the same state in different workloads.
+      // The sample from each iteration uses 2 random numbers.
+      gen_copy.Skip((start_output * 2 * kMaxIterations +
+                     Uniform::kResultElementCount - 1) /
+                    Uniform::kResultElementCount);
+
+      // Vectorized intermediate calculations for uniform rejection sampling.
+      // We always generate at most 4 samples.
+      Eigen::array z;
+      Eigen::array g;
+
+      const bool should_bcast = bcast.IsBroadcastingRequired();
+      const auto& means_batch_indices = bcast.batch_indices(0);
+      const auto& stddevs_batch_indices = bcast.batch_indices(1);
+      const auto& minvals_batch_indices = bcast.batch_indices(2);
+      const auto& maxvals_batch_indices = bcast.batch_indices(3);
+      auto output_flat = output.data();
+
+      // We partition work across batches and then across samples
+      // per batch member, to avoid extra work.
+      for (int64 output_idx = start_output; output_idx < limit_output;
+           // output_idx is incremented with the inner loops below.
+      ) {
+        int64 batch_idx = output_idx / samples_per_batch;
+        // The output layout is [samples_per_batch, num_batches]. Thus
+        // the output address is sample_idx * num_batches + batch_idx.
+        // Below, code will index at output_batch_offset[sample_idx *
+        // num_batches] matching this.
+        T* const output_batch_offset = output_flat + batch_idx;
+        // Generate batch counts from BCast, as it has the right indices to loop
+        // over.
+        T mean, stddev, minval, maxval;
+        if (should_bcast) {
+          mean = means(means_batch_indices[batch_idx]);
+          stddev = stddevs(stddevs_batch_indices[batch_idx]);
+          minval = minvals(minvals_batch_indices[batch_idx]);
+          maxval = maxvals(maxvals_batch_indices[batch_idx]);
+        } else {
+          mean = means(batch_idx);
+          stddev = stddevs(batch_idx);
+          minval = minvals(batch_idx);
+          maxval = maxvals(batch_idx);
+        }
+
+        // On GPU, this check will just fill samples with NAN if it fails.
+        OP_REQUIRES(ctx,
+                    stddev > T(0) && minval < maxval &&
+                        (Eigen::numext::isfinite(minval) ||
+                         Eigen::numext::isfinite(maxval)),
+                    errors::InvalidArgument("Invalid parameters"));
+
+        int num_iterations = 0;
+
+        // If possible, make one-sided bound be the lower bound, or make both
+        // bounds positive. Otherwise, the bounds are on either side of the
+        // mean.
+        if ((Eigen::numext::isinf(minval) && minval < T(0)) || maxval < mean) {
+          // Reverse all calculations. normMin and normMax will be flipped.
+          std::swap(minval, maxval);
+          stddev = -stddev;
+        }
+
+        // Calculate normalized samples, then convert them.
+        const T normMin = (minval - mean) / stddev;
+        const T normMax = (maxval - mean) / stddev;
+
+        // Determine the method to use.
+        const T sqrtFactor = Eigen::numext::sqrt((normMin * normMin) + T(4));
+        const T cutoff =
+            T(2) *
+            Eigen::numext::exp(T(0.5) +
+                               (normMin * (normMin - sqrtFactor)) / T(4)) /
+            (normMin + sqrtFactor);
+        const T diff = normMax - normMin;
+
+        if (((normMin < -kStdDevsInsideBoundsToUseRandnSampler) &&
+             (normMax >= T(0.))) ||
+            ((normMax > kStdDevsInsideBoundsToUseRandnSampler) &&
+             (normMin <= T(0.)))) {
+          // If the bounds are a least 3 standard deviations from the mean
+          // on at least one side then we rejection sample by sampling
+          // from the normal distribution and rejecting samples outside
+          // the bounds.
+          // Under this condition the acceptance rate per iteration should
+          // always be ~ 50%. This sampler is more efficient (and more
+          // numerically stable when one or both bounds is far from the mean).
+          for (int64 sample_idx = output_idx % samples_per_batch;
+               sample_idx < samples_per_batch && output_idx < limit_output;) {
+            const auto randn_sample = normal_dist(&gen_copy);
+            const int size = randn_sample.size();
+            for (int i = 0; i < size; ++i) {
+              if ((randn_sample[i] >= normMin) &&
+                  (randn_sample[i] <= normMax)) {
+                output_batch_offset[sample_idx * num_batches] =
+                    randn_sample[i] * stddev + mean;
+                ++sample_idx;
+                ++output_idx;
+                if (sample_idx >= samples_per_batch ||
+                    output_idx >= limit_output) {
+                  break;
+                }
+                num_iterations = 0;
+              } else {
+                ++num_iterations;
+                if (num_iterations > kMaxIterations) {
+                  // This should never occur because this sampler should
+                  // (by the selection criteria above) be used if at least 3
+                  // standard deviations of one side of the distribution
+                  // is within the limits (so acceptance probability per
+                  // iterations >~ 1/2 per iteration).
+                  LOG(ERROR) << "TruncatedNormal randn rejection sampler "
+                             << "exceeded maximum iterations for "
+                             << "normMin=" << normMin << " normMax=" << normMax
+                             << " kMaxIterations=" << kMaxIterations;
+                  ctx->SetStatus(errors::Internal(
+                      "TruncatedNormal randn rejection sampler failed to accept"
+                      " a sample."));
+                  return;
+                }
+              }
+            }
+          }
+        } else if (diff < cutoff) {
+          // Sample from a uniform distribution on [normMin, normMax].
+
+          const T plusFactor = (normMin < T(0)) ? T(0) : normMin * normMin;
+
+          for (int64 sample_idx = output_idx % samples_per_batch;
+               sample_idx < samples_per_batch && output_idx < limit_output;) {
+            const auto rand = dist(&gen_copy);
+            const int size = rand.size();
+            // NOTE(ringwalt): These loops seem to only generate packed AVX
+            // instructions for float32.
+            for (int i = 0; i < size; i++) {
+              z[i] = rand[i] * diff + normMin;
+              g[i] = (plusFactor - z[i] * z[i]) / T(2.0);
+            }
+
+            const auto u = dist(&gen_copy);
+            for (int i = 0; i < size; i++) {
+              auto accept = u[i] <= Eigen::numext::exp(g[i]);
+              if (accept || num_iterations + 1 >= kMaxIterations) {
+                // Accept the sample z.
+                // If we run out of iterations, just use the current uniform
+                // sample, but emit a warning.
+                // TODO(jjhunt) For small entropies (relative to the bounds),
+                // this sampler is poor and may take many iterations since
+                // the proposal distribution is the uniform distribution
+                // U(lower_bound, upper_bound).
+                if (!accept) {
+                  LOG(ERROR) << "TruncatedNormal uniform rejection sampler "
+                             << "exceeded max iterations. Sample may contain "
+                             << "outliers.";
+                  ctx->SetStatus(errors::Internal(
+                      "TruncatedNormal uniform rejection sampler failed to "
+                      " accept a sample."));
+                  return;
+                }
+                output_batch_offset[sample_idx * num_batches] =
+                    z[i] * stddev + mean;
+                ++sample_idx;
+                ++output_idx;
+                if (sample_idx >= samples_per_batch ||
+                    output_idx >= limit_output) {
+                  break;
+                }
+                num_iterations = 0;
+              } else {
+                num_iterations++;
+              }
+            }
+          }
+        } else {
+          // Sample from an exponential distribution with alpha maximizing
+          // acceptance probability, offset by normMin from the origin.
+          // Accept only if less than normMax.
+          const T alpha =
+              (normMin + Eigen::numext::sqrt((normMin * normMin) + T(4))) /
+              T(2);
+          for (int64 sample_idx = output_idx % samples_per_batch;
+               sample_idx < samples_per_batch && output_idx < limit_output;) {
+            auto rand = dist(&gen_copy);
+            const int size = rand.size();
+            int i = 0;
+            while (i < size) {
+              const T z = -Eigen::numext::log(rand[i]) / alpha + normMin;
+              i++;
+              const T x = normMin < alpha ? alpha - z : normMin - alpha;
+              const T g = Eigen::numext::exp(-x * x / T(2.0));
+              const T u = rand[i];
+              i++;
+              auto accept = (u <= g && z < normMax);
+              if (accept || num_iterations + 1 >= kMaxIterations) {
+                if (!accept) {
+                  LOG(ERROR) << "TruncatedNormal exponential distribution "
+                             << "rejection sampler exceeds max iterations. "
+                             << "Sample may contain outliers.";
+                  ctx->SetStatus(errors::Internal(
+                      "TruncatedNormal exponential distribution rejection"
+                      " sampler failed to accept a sample."));
+                  return;
+                }
+                output_batch_offset[sample_idx * num_batches] =
+                    z * stddev + mean;
+                ++sample_idx;
+                ++output_idx;
+                if (sample_idx >= samples_per_batch ||
+                    output_idx >= limit_output) {
+                  break;
+                }
+                num_iterations = 0;
+              } else {
+                num_iterations++;
+              }
+            }
+          }
+        }
+      }
+    };
+    // The cost of the initial calculations for the batch.
+    const int64 batchInitCost =
+        // normMin, normMax
+        (Eigen::TensorOpCost::AddCost() +
+         Eigen::TensorOpCost::MulCost()) *
+            2
+        // sqrtFactor
+        + Eigen::TensorOpCost::AddCost() +
+        Eigen::TensorOpCost::MulCost() +
+        Eigen::internal::functor_traits<
+            Eigen::internal::scalar_sqrt_op>::Cost
+        // cutoff
+        + Eigen::TensorOpCost::MulCost() * 4 +
+        Eigen::internal::functor_traits>::Cost
+        // diff
+        + Eigen::TensorOpCost::AddCost();
+    const int64 uniformSampleCost =
+        random::PhiloxRandom::kElementCost +
+        random::UniformDistribution::kElementCost;
+    // The cost of a single uniform sampling round.
+    const int64 uniformRejectionSamplingCost =
+        uniformSampleCost + Eigen::TensorOpCost::MulCost() +
+        Eigen::TensorOpCost::AddCost() +
+        Eigen::TensorOpCost::MulCost() * 2 +
+        Eigen::TensorOpCost::AddCost() + uniformSampleCost +
+        Eigen::internal::functor_traits<
+            Eigen::internal::scalar_exp_op>::Cost +
+        Eigen::TensorOpCost::MulCost() + Eigen::TensorOpCost::AddCost();
+    // Estimate the cost for an entire batch.
+    // Assume we use uniform sampling, and accept the 2nd sample on average.
+    const int64 batchCost = batchInitCost + uniformRejectionSamplingCost * 2;
+    Shard(worker_threads.num_threads, worker_threads.workers, num_elements,
+          batchCost, do_work);
   }
 };
 
@@ -436,13 +727,113 @@ class ParameterizedTruncatedNormalOp : public OpKernel {
   TF_DISALLOW_COPY_AND_ASSIGN(ParameterizedTruncatedNormalOp);
 };
 
+// Samples from a truncated normal distribution, using the given parameters.
+template 
+class StatelessParameterizedTruncatedNormal : public OpKernel {
+  // Reshape batches so each batch is this size if possible.
+  static const int32 kDesiredBatchSize = 100;
+
+ public:
+  explicit StatelessParameterizedTruncatedNormal(OpKernelConstruction* context)
+      : OpKernel(context) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    const Tensor& shape_tensor = ctx->input(0);
+    const Tensor& seed_tensor = ctx->input(1);
+    const Tensor& means_tensor = ctx->input(2);
+    const Tensor& stddevs_tensor = ctx->input(3);
+    const Tensor& minvals_tensor = ctx->input(4);
+    const Tensor& maxvals_tensor = ctx->input(5);
+
+    OP_REQUIRES(ctx, seed_tensor.dims() == 1 && seed_tensor.dim_size(0) == 2,
+                errors::InvalidArgument("seed must have shape [2], not ",
+                                        seed_tensor.shape().DebugString()));
+
+    tensorflow::BCastList<4> bcast(
+        {means_tensor.shape().dim_sizes(), stddevs_tensor.shape().dim_sizes(),
+         minvals_tensor.shape().dim_sizes(),
+         maxvals_tensor.shape().dim_sizes()},
+        /*fewer_dims_optimization=*/false,
+        /*return_flattened_batch_indices=*/true);
+
+    OP_REQUIRES(ctx, bcast.IsValid(),
+                errors::InvalidArgument(
+                    "means, stddevs, minvals, maxvals must have compatible "
+                    "batch dimensions: ",
+                    means_tensor.shape().DebugString(), " vs. ",
+                    stddevs_tensor.shape().DebugString(), " vs. ",
+                    minvals_tensor.shape().DebugString(), " vs. ",
+                    maxvals_tensor.shape().DebugString()));
+
+    // Let's check that the shape tensor dominates the broadcasted tensor.
+    TensorShape bcast_shape = BCast::ToShape(bcast.output_shape());
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsVector(shape_tensor.shape()),
+        errors::InvalidArgument("Input shape should be a vector, got shape: ",
+                                shape_tensor.shape().DebugString()));
+    TensorShape output_shape;
+    if (shape_tensor.dtype() == DataType::DT_INT32) {
+      OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(shape_tensor.vec(),
+                                                      &output_shape));
+    } else {
+      OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(shape_tensor.vec(),
+                                                      &output_shape));
+    }
+    OP_REQUIRES(ctx, TensorShapeUtils::EndsWith(output_shape, bcast_shape),
+                errors::InvalidArgument(
+                    "Shape passed in must end with broadcasted shape."));
+
+    int64 samples_per_batch = 1;
+    const int64 num_sample_dims =
+        (shape_tensor.dim_size(0) - bcast.output_shape().size());
+    for (int64 i = 0; i < num_sample_dims; ++i) {
+      samples_per_batch *= output_shape.dim_size(i);
+    }
+    int64 num_batches = 1;
+    for (int64 i = num_sample_dims; i < shape_tensor.dim_size(0); ++i) {
+      num_batches *= output_shape.dim_size(i);
+    }
+    const int64 num_elements = num_batches * samples_per_batch;
+
+    Tensor* samples_tensor;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &samples_tensor));
+
+    auto truncFunctor = functor::TruncatedNormalFunctorV2();
+    // Each worker has the same fudge factor, so use it here.
+    random::PhiloxRandom::Key key;
+    random::PhiloxRandom::ResultType counter;
+    OP_REQUIRES_OK(ctx, GenerateKey(seed_tensor, &key, &counter));
+
+    auto philox = random::PhiloxRandom(counter, key);
+
+    truncFunctor(ctx, ctx->eigen_device(), num_batches,
+                 samples_per_batch, num_elements, bcast, means_tensor.flat(),
+                 stddevs_tensor.flat(), minvals_tensor.flat(),
+                 maxvals_tensor.flat(), philox, samples_tensor->flat());
+  }
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(StatelessParameterizedTruncatedNormal);
+};
+
 }  // namespace
 
-#define REGISTER(TYPE)                                         \
-  REGISTER_KERNEL_BUILDER(Name("ParameterizedTruncatedNormal") \
-                              .Device(DEVICE_CPU)              \
-                              .TypeConstraint("dtype"),  \
-                          ParameterizedTruncatedNormalOp)
+#define REGISTER(TYPE)                                                     \
+  REGISTER_KERNEL_BUILDER(Name("ParameterizedTruncatedNormal")             \
+                              .Device(DEVICE_CPU)                          \
+                              .TypeConstraint("dtype"),              \
+                          ParameterizedTruncatedNormalOp) \
+  REGISTER_KERNEL_BUILDER(                                                 \
+      Name("StatelessParameterizedTruncatedNormal")                        \
+          .HostMemory("shape")                                             \
+          .HostMemory("seed")                                              \
+          .HostMemory("means")                                             \
+          .HostMemory("stddevs")                                           \
+          .HostMemory("minvals")                                           \
+          .HostMemory("maxvals")                                           \
+          .Device(DEVICE_CPU)                                              \
+          .TypeConstraint("dtype"),                                  \
+      StatelessParameterizedTruncatedNormal)
 
 TF_CALL_half(REGISTER);
 TF_CALL_float(REGISTER);
diff --git a/tensorflow/core/kernels/parameterized_truncated_normal_op.h b/tensorflow/core/kernels/parameterized_truncated_normal_op.h
index c919a22c7b0..ee7fb7bf605 100644
--- a/tensorflow/core/kernels/parameterized_truncated_normal_op.h
+++ b/tensorflow/core/kernels/parameterized_truncated_normal_op.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/lib/random/random_distributions.h"
+#include "tensorflow/core/util/bcast.h"
 
 namespace tensorflow {
 
@@ -44,6 +45,21 @@ struct TruncatedNormalFunctor {
                   typename TTypes::Flat output);
 };
 
+// This version supports broadcasting of the arguments, as well as puts
+// the sample dimension on the left.
+template 
+struct TruncatedNormalFunctorV2 {
+  void operator()(OpKernelContext* ctx, const Device& d, int64 num_batches,
+                  int64 samples_per_batch, int64 num_elements,
+                  const BCastList<4>& bcast,
+                  typename TTypes::ConstFlat means,
+                  typename TTypes::ConstFlat stddevs,
+                  typename TTypes::ConstFlat minvals,
+                  typename TTypes::ConstFlat maxvals,
+                  const random::PhiloxRandom& gen,
+                  typename TTypes::Flat output);
+};
+
 }  // namespace functor
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/ops/stateless_random_ops.cc b/tensorflow/core/ops/stateless_random_ops.cc
index d540b9a04d9..e1820ea4feb 100644
--- a/tensorflow/core/ops/stateless_random_ops.cc
+++ b/tensorflow/core/ops/stateless_random_ops.cc
@@ -124,6 +124,41 @@ REGISTER_OP("StatelessRandomBinomial")
     .Attr("dtype: {half, float, double, int32, int64} = DT_INT64")
     .SetShapeFn(StatelessShape);
 
+REGISTER_OP("StatelessParameterizedTruncatedNormal")
+    .Input("shape: S")
+    .Input("seed: Tseed")
+    .Input("means: dtype")
+    .Input("stddevs: dtype")
+    .Input("minvals: dtype")
+    .Input("maxvals: dtype")
+    .Output("output: dtype")
+    .Attr("S: {int32, int64}")
+    .Attr("Tseed: {int32, int64} = DT_INT64")
+    .Attr("dtype: {float16, float32, float64}")
+    .SetShapeFn([](InferenceContext* c) {
+      // Check seed shape
+      ShapeHandle seed;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &seed));
+      DimensionHandle unused_dim;
+      TF_RETURN_IF_ERROR(c->WithValue(c->Dim(seed, 0), 2, &unused_dim));
+
+      ShapeHandle bcast_means_stddevs;
+      ShapeHandle bcast_except_maxvals;
+      ShapeHandle bcast_all;
+      TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
+          c, c->input(2), c->input(3), true, &bcast_means_stddevs));
+      TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
+          c, c->input(4), bcast_means_stddevs, true, &bcast_except_maxvals));
+      TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
+          c, c->input(5), bcast_except_maxvals, true, &bcast_all));
+
+      // Set output shape
+      ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out));
+      c->set_output(0, out);
+      return Status::OK();
+    });
+
 REGISTER_OP("StatelessRandomPoisson")
     .Input("shape: T")
     .Input("seed: Tseed")
diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
index a7c7ab7abc7..7da45e36118 100644
--- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc
+++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
@@ -50,7 +50,7 @@ auto OpGradientInfoInit(const T &a) {
 
 absl::optional> OpGradientUnusedInputIndices(
     const tensorflow::string &op_name) {
-  static std::array a = {{
+  static std::array a = {{
       {"Acosh"},
       {"AllToAll", 1, {0}},
       {"ApproximateEqual"},
@@ -326,6 +326,7 @@ absl::optional> OpGradientUnusedInputIndices(
       {"StackPop"},
       {"StackPush"},
       {"StatelessMultinomial"},
+      {"StatelessParameterizedTruncatedNormal", 1, {1}},
       {"StatelessRandomBinomial"},
       {"StatelessRandomGammaV2", 1, {1}},
       {"StatelessRandomNormal"},
diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD
index f2c614974f5..f93bf5cd1ae 100644
--- a/tensorflow/python/kernel_tests/BUILD
+++ b/tensorflow/python/kernel_tests/BUILD
@@ -785,24 +785,6 @@ cuda_py_test(
     ],
 )
 
-cuda_py_test(
-    name = "parameterized_truncated_normal_op_test",
-    size = "medium",
-    srcs = ["parameterized_truncated_normal_op_test.py"],
-    deps = [
-        "//tensorflow/core:protos_all_py",
-        "//tensorflow/python:client",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:control_flow_ops",
-        "//tensorflow/python:framework",
-        "//tensorflow/python:framework_for_generated_wrappers",
-        "//tensorflow/python:platform",
-        "//tensorflow/python:random_ops",
-        "//third_party/py/numpy",
-        "@absl_py//absl/testing:parameterized",
-    ],
-)
-
 tf_py_test(
     name = "parsing_ops_test",
     size = "medium",
diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index b5d291d2973..6e404b4cd5f 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -20,6 +20,24 @@ py_library(
     ],
 )
 
+cuda_py_test(
+    name = "parameterized_truncated_normal_op_test",
+    size = "medium",
+    srcs = ["parameterized_truncated_normal_op_test.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:client",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:control_flow_ops",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:random_ops",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 tf_py_test(
     name = "random_shuffle_queue_test",
     size = "small",
diff --git a/tensorflow/python/kernel_tests/parameterized_truncated_normal_op_test.py b/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
similarity index 63%
rename from tensorflow/python/kernel_tests/parameterized_truncated_normal_op_test.py
rename to tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
index ac8ad7a2bd4..309c3e404db 100644
--- a/tensorflow/python/kernel_tests/parameterized_truncated_normal_op_test.py
+++ b/tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
@@ -27,11 +27,15 @@ from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import stateless_random_ops as stateless
+from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging
 
@@ -91,13 +95,8 @@ class TruncatedNormalMoments(object):
 
 def calculate_moments(samples, max_moment):
   moments = [0.0] * (max_moment + 1)
-  for sample in samples:
-    value = 1.0
-    for k in range(len(moments)):
-      moments[k] += value
-      value *= sample
-  for i in range(len(moments)):
-    moments[i] /= len(samples)
+  for k in range(len(moments)):
+    moments[k] = np.mean(samples**k, axis=0)
   return moments
 
 
@@ -118,16 +117,31 @@ class ParameterizedTruncatedNormalTest(test.TestCase):
   # Stop at moment 10 to avoid numerical errors in the theoretical moments.
   max_moment = 10
 
-  def validateMoments(self, shape, mean, stddev, minval, maxval, seed=1618):
+  def validateMoments(self,
+                      shape,
+                      mean,
+                      stddev,
+                      minval,
+                      maxval,
+                      use_stateless=False,
+                      seed=1618):
     try:
       # TruncatedNormalMoments requires scipy.stats.
       # Give up early if we are unable to import it.
-      import scipy.stats  # pylint: disable=g-import-not-at-top,unused-variable
       random_seed.set_random_seed(seed)
       with self.cached_session(use_gpu=True):
-        samples = random_ops.parameterized_truncated_normal(shape, mean, stddev,
-                                                            minval,
-                                                            maxval).eval()
+        if use_stateless:
+          # Generate a seed that stateless ops can use.
+          new_seed = random_ops.random_uniform([2],
+                                               seed=seed,
+                                               minval=0,
+                                               maxval=(2**31 - 1),
+                                               dtype=np.int32)
+          samples = stateless.stateless_parameterized_truncated_normal(
+              shape, new_seed, mean, stddev, minval, maxval).eval()
+        else:
+          samples = random_ops.parameterized_truncated_normal(
+              shape, mean, stddev, minval, maxval).eval()
         assert (~np.isnan(samples)).all()
       moments = calculate_moments(samples, self.max_moment)
       expected_moments = TruncatedNormalMoments(mean, stddev, minval, maxval)
@@ -144,14 +158,24 @@ class ParameterizedTruncatedNormalTest(test.TestCase):
                                 stddev,
                                 minval,
                                 maxval,
+                                use_stateless=False,
                                 seed=1618):
     try:
       import scipy.stats  # pylint: disable=g-import-not-at-top
       random_seed.set_random_seed(seed)
       with self.cached_session(use_gpu=True):
-        samples = random_ops.parameterized_truncated_normal(shape, mean, stddev,
-                                                            minval,
-                                                            maxval).eval()
+        if use_stateless:
+          new_seed = random_ops.random_uniform([2],
+                                               seed=seed,
+                                               minval=0,
+                                               maxval=(2**31 - 1),
+                                               dtype=np.int32)
+          samples = stateless.stateless_parameterized_truncated_normal(
+              shape, new_seed, mean, stddev, minval, maxval).eval()
+        else:
+          samples = random_ops.parameterized_truncated_normal(
+              shape, mean, stddev, minval, maxval).eval()
+
       assert (~np.isnan(samples)).all()
       minval = max(mean - stddev * 10, minval)
       maxval = min(mean + stddev * 10, maxval)
@@ -169,61 +193,160 @@ class ParameterizedTruncatedNormalTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testDefaults(self):
-    self.validateMoments([10**5], 0.0, 1.0, -2.0, 2.0)
+    self.validateMoments([int(1e5)], 0.0, 1.0, -2.0, 2.0)
+    self.validateMoments([int(1e5)], 0.0, 1.0, -2.0, 2.0, use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testShifted(self):
-    self.validateMoments([10**5], -1.0, 1.0, -2.0, 2.0)
+    self.validateMoments([int(1e5)], -1.0, 1.0, -2.0, 2.0)
+    self.validateMoments([int(1e5)], -1.0, 1.0, -2.0, 2.0, use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testRightTail(self):
-    self.validateMoments([10**5], 0.0, 1.0, 4.0, np.infty)
+    self.validateMoments([int(1e5)], 0.0, 1.0, 4.0, np.infty)
+    self.validateMoments([int(1e5)],
+                         0.0,
+                         1.0,
+                         4.0,
+                         np.infty,
+                         use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testLeftTail(self):
-    self.validateMoments([10**5], 0.0, 1.0, -np.infty, -4.0)
+    self.validateMoments([int(1e5)], 0.0, 1.0, -np.infty, -4.0)
+    self.validateMoments([int(1e5)],
+                         0.0,
+                         1.0,
+                         -np.infty,
+                         -4.0,
+                         use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testLeftTailTwoSidedBounds(self):
-    self.validateMoments([10**5], 0.0, 1.0, -6.0, -3.0)
+    self.validateMoments([int(1e5)], 0.0, 1.0, -6.0, -3.0)
+    self.validateMoments([int(1e5)], 0.0, 1.0, -6.0, -3.0, use_stateless=True)
 
   @test_util.run_deprecated_v1
   @test_util.disable_xla("Low probability region")
   def testTwoSidedLeftTailShifted(self):
-    self.validateKolmogorovSmirnov([10**5], 6.0, 1.0, -1.0, 1.0)
+    self.validateKolmogorovSmirnov([int(1e5)], 6.0, 1.0, -1.0, 1.0)
+    self.validateKolmogorovSmirnov([int(1e5)],
+                                   6.0,
+                                   1.0,
+                                   -1.0,
+                                   1.0,
+                                   use_stateless=True)
 
   @test_util.run_deprecated_v1
   @test_util.disable_xla("Low probability region")
   def testRightTailShifted(self):
-    self.validateMoments([10**5], -5.0, 1.0, 2.0, np.infty)
+    self.validateMoments([int(1e5)], -5.0, 1.0, 2.0, np.infty)
+    self.validateMoments([int(1e5)],
+                         -5.0,
+                         1.0,
+                         2.0,
+                         np.infty,
+                         use_stateless=True)
 
   # Take the normal distribution around the mean, but truncating the left tail
   # far from the mean.
   @test_util.run_deprecated_v1
   def testTruncateOnLeft_entireTailOnRight(self):
-    self.validateKolmogorovSmirnov([10**5], 10.0, 1.0, 4.0, np.infty)
+    self.validateKolmogorovSmirnov([int(1e5)], 10.0, 1.0, 4.0, np.infty)
+    self.validateKolmogorovSmirnov([int(1e5)],
+                                   10.0,
+                                   1.0,
+                                   4.0,
+                                   np.infty,
+                                   use_stateless=True)
 
   # Take the normal distribution around the mean, but truncating the right tail.
   @test_util.run_deprecated_v1
   def testTruncateOnRight_entireTailOnLeft(self):
-    self.validateKolmogorovSmirnov([10**5], -8, 1.0, -np.infty, -4.0)
+    self.validateKolmogorovSmirnov([int(1e5)], -8, 1.0, -np.infty, -4.0)
+    self.validateKolmogorovSmirnov([int(1e5)],
+                                   -8.,
+                                   1.0,
+                                   -np.infty,
+                                   -4.0,
+                                   use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testSmallStddev(self):
-    self.validateKolmogorovSmirnov([10**5], 0.0, 0.1, 0.05, 0.10)
+    self.validateKolmogorovSmirnov([int(1e5)], 0.0, 0.1, 0.05, 0.10)
+    self.validateKolmogorovSmirnov([int(1e5)],
+                                   0.0,
+                                   0.1,
+                                   0.05,
+                                   0.10,
+                                   use_stateless=True)
 
   @test_util.run_deprecated_v1
   def testSamplingWithSmallStdDevFarFromBound(self):
     sample_op = random_ops.parameterized_truncated_normal(
         shape=(int(1e5),), means=0.8, stddevs=0.05, minvals=-1., maxvals=1.)
+    new_seed = random_ops.random_uniform([2],
+                                         seed=1234,
+                                         minval=0,
+                                         maxval=(2**31 - 1),
+                                         dtype=np.int32)
+    sample_op_stateless = stateless.stateless_parameterized_truncated_normal(
+        shape=(int(1e5),),
+        seed=new_seed,
+        means=0.8,
+        stddevs=0.05,
+        minvals=-1.,
+        maxvals=1.)
 
     with self.session(use_gpu=True) as sess:
-      samples = sess.run(sample_op)
+      samples, samples_stateless = sess.run([sample_op, sample_op_stateless])
       # 0. is more than 16 standard deviations from the mean, and
       # should have a likelihood < 1e-57.
       assert (~np.isnan(samples)).all()
-      no_neg_samples = np.sum(samples < 0.)
-      self.assertEqual(no_neg_samples, 0.)
+      assert (~np.isnan(samples_stateless)).all()
+      self.assertAllGreater(samples, 0.)
+      self.assertAllGreater(samples_stateless, 0.)
+
+  def testStatelessParameterizedTruncatedNormalHasGrads(self):
+    mean = variables.Variable(0.01)
+    stddev = variables.Variable(1.)
+    minval = variables.Variable(-1.)
+    maxval = variables.Variable(1.)
+
+    with self.cached_session(use_gpu=True) as sess:
+      with backprop.GradientTape(persistent=True) as tape:
+        samples = stateless.stateless_parameterized_truncated_normal(
+            [1], [1, 2], mean, stddev, minval, maxval)
+
+      sess.run(variables.variables_initializer([mean, stddev, minval, maxval]))
+      [mean_grad, std_grad], mean_actual_grad, std_actual_grad = sess.run([
+          tape.gradient(samples, [mean, stddev]),
+          array_ops.ones_like(mean),
+          (samples - mean) / stddev])
+      self.assertAllClose(mean_grad, mean_actual_grad)
+      self.assertAllClose(std_grad, std_actual_grad[0])
+
+      try:
+        import scipy.stats  # pylint:disable=g-import-not-at-top
+        truncnorm = scipy.stats.truncnorm(a=-1., b=1., loc=0., scale=1.)
+        samples_np, [minval_grad, maxval_grad] = sess.run([
+            samples, tape.gradient(samples, [minval, maxval])])
+
+        sample_cdf = truncnorm.cdf(samples_np)
+        # These come from the implicit reparameterization trick.
+        scipy_maxval_grad = np.exp(
+            0.5 * (samples_np ** 2 - ((1. - 0.01) / 1.) ** 2) +
+            np.log(sample_cdf))
+
+        scipy_minval_grad = np.exp(
+            0.5 * (samples_np ** 2 - ((-1. - 0.01) / 1.) ** 2) +
+            np.log1p(-sample_cdf))
+
+        self.assertAllClose(minval_grad, scipy_minval_grad[0], rtol=1e-2)
+        self.assertAllClose(maxval_grad, scipy_maxval_grad[0], rtol=1e-2)
+
+      except ImportError as e:
+        tf_logging.warn("Cannot test truncated normal op: %s" % str(e))
 
   @test_util.run_deprecated_v1
   def testSamplingAtRandnSwitchover(self):
@@ -239,18 +362,33 @@ class ParameterizedTruncatedNormalTest(test.TestCase):
 
     epsilon = 0.001
     self.validateMoments(
-        shape=[10**6],
+        shape=[int(1e6)],
         mean=0.,
         stddev=1.0,
         minval=-epsilon,
         maxval=stddev_inside_bounds_before_using_randn - epsilon)
     self.validateMoments(
-        shape=[10**6],
+        shape=[int(1e6)],
         mean=0.,
         stddev=1.0,
         minval=-epsilon,
         maxval=stddev_inside_bounds_before_using_randn + epsilon)
 
+    self.validateMoments(
+        shape=[int(1e6)],
+        mean=0.,
+        stddev=1.0,
+        minval=-epsilon,
+        maxval=stddev_inside_bounds_before_using_randn - epsilon,
+        use_stateless=True)
+    self.validateMoments(
+        shape=[int(1e6)],
+        mean=0.,
+        stddev=1.0,
+        minval=-epsilon,
+        maxval=stddev_inside_bounds_before_using_randn + epsilon,
+        use_stateless=True)
+
 
 # Benchmarking code
 def parameterized_vs_naive(shape, num_iters, use_gpu=False):
diff --git a/tensorflow/python/ops/random_grad.py b/tensorflow/python/ops/random_grad.py
index 771980932cb..3caa08d96f9 100644
--- a/tensorflow/python/ops/random_grad.py
+++ b/tensorflow/python/ops/random_grad.py
@@ -18,9 +18,14 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gen_random_ops
 from tensorflow.python.ops import math_ops
 
@@ -114,3 +119,118 @@ def _StatelessRandomGammaV2Grad(op, grad):  # pylint: disable=invalid-name
     return (None, None,
             math_ops.reduce_sum(
                 grad * partial_a, axis=math_ops.range(num_sample_dimensions)))
+
+
+def _Ndtr(x):
+  """Normal distribution function."""
+  half_sqrt_2 = constant_op.constant(
+      0.5 * np.sqrt(2.), dtype=x.dtype, name="half_sqrt_2")
+  w = x * half_sqrt_2
+  z = math_ops.abs(w)
+  y = array_ops.where(
+      z < half_sqrt_2,
+      1. + math_ops.erf(w),
+      array_ops.where(
+          w > 0., 2. - math_ops.erfc(z), math_ops.erfc(z)))
+  return 0.5 * y
+
+
+@ops.RegisterGradient("StatelessParameterizedTruncatedNormal")
+def _StatelessParameterizedTruncatedNormalGrad(op, grad):  # pylint: disable=invalid-name
+  """Returns the gradient of a TruncatedNormal sample w.r.t. parameters.
+
+  The gradient is computed using implicit differentiation
+  (Figurnov et al., 2018).
+
+  Args:
+    op: A `StatelessParameterizedTruncatedNormal` operation. We assume that the
+      inputs to the operation are `shape`, `seed`, `mean`, `stddev`, `minval`,
+      and `maxval` tensors, and the output is the `sample` tensor.
+    grad: The incoming gradient `dloss / dsample` of the same shape as
+      `op.outputs[0]`.
+
+  Returns:
+    A list of `Tensor` with derivates with respect to each parameter.
+
+  References:
+    Implicit Reparameterization Gradients:
+      [Figurnov et al., 2018]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients)
+      ([pdf]
+      (http://papers.nips.cc/paper/7326-implicit-reparameterization-gradients.pdf))
+  """
+  shape = op.inputs[0]
+  mean = op.inputs[2]
+  stddev = op.inputs[3]
+  minval = op.inputs[4]
+  maxval = op.inputs[5]
+  sample = op.outputs[0]
+
+  with ops.control_dependencies([grad]):
+    minval_std = (minval - mean) / stddev
+    maxval_std = (maxval - mean) / stddev
+    sample_std = (sample - mean) / stddev
+
+    cdf_sample = (_Ndtr(sample_std) - _Ndtr(minval_std)) / (
+        _Ndtr(maxval_std) - _Ndtr(minval_std))
+
+    # Clip to avoid zero argument for log_cdf expression
+    tiny = np.finfo(mean.dtype.as_numpy_dtype).tiny
+    eps = np.finfo(mean.dtype.as_numpy_dtype).eps
+    cdf_sample = clip_ops.clip_by_value(cdf_sample, tiny, 1 - eps)
+
+    dmaxval = math_ops.exp(0.5 * (sample_std ** 2 - maxval_std ** 2) +
+                           math_ops.log(cdf_sample))
+    dminval = math_ops.exp(0.5 * (sample_std ** 2 - minval_std ** 2) +
+                           math_ops.log1p(-cdf_sample))
+    dmean = array_ops.ones_like(sample_std)
+    dstddev = sample_std
+
+    # Reduce over extra dimensions caused by `shape`. We need to get the
+    # difference in rank from shape vs. the broadcasted rank.
+
+    mean_shape = array_ops.shape(mean)
+    stddev_shape = array_ops.shape(stddev)
+    minval_shape = array_ops.shape(minval)
+    maxval_shape = array_ops.shape(maxval)
+
+    broadcast_shape = array_ops.broadcast_dynamic_shape(
+        mean_shape, stddev_shape)
+    broadcast_shape = array_ops.broadcast_dynamic_shape(
+        minval_shape, broadcast_shape)
+    broadcast_shape = array_ops.broadcast_dynamic_shape(
+        maxval_shape, broadcast_shape)
+    extra_dims = math_ops.range(
+        array_ops.size(shape) - array_ops.size(broadcast_shape))
+
+    grad_mean = math_ops.reduce_sum(grad * dmean, axis=extra_dims)
+    grad_stddev = math_ops.reduce_sum(grad * dstddev, axis=extra_dims)
+    grad_minval = math_ops.reduce_sum(grad * dminval, axis=extra_dims)
+    grad_maxval = math_ops.reduce_sum(grad * dmaxval, axis=extra_dims)
+
+    _, rmean = gen_array_ops.broadcast_gradient_args(
+        broadcast_shape, mean_shape)
+    _, rstddev = gen_array_ops.broadcast_gradient_args(
+        broadcast_shape, stddev_shape)
+    _, rminval = gen_array_ops.broadcast_gradient_args(
+        broadcast_shape, minval_shape)
+    _, rmaxval = gen_array_ops.broadcast_gradient_args(
+        broadcast_shape, maxval_shape)
+
+    grad_mean = array_ops.reshape(
+        math_ops.reduce_sum(grad_mean, axis=rmean, keepdims=True), mean_shape)
+
+    grad_stddev = array_ops.reshape(
+        math_ops.reduce_sum(grad_stddev, axis=rstddev, keepdims=True),
+        stddev_shape)
+
+    grad_minval = array_ops.reshape(
+        math_ops.reduce_sum(grad_minval, axis=rminval, keepdims=True),
+        minval_shape)
+
+    grad_maxval = array_ops.reshape(
+        math_ops.reduce_sum(grad_maxval, axis=rmaxval, keepdims=True),
+        maxval_shape)
+
+    # The first two inputs are shape.
+    return (None, None, grad_mean, grad_stddev, grad_minval, grad_maxval)
diff --git a/tensorflow/python/ops/stateless_random_ops.py b/tensorflow/python/ops/stateless_random_ops.py
index 25fefcc514c..3e825cc4775 100644
--- a/tensorflow/python/ops/stateless_random_ops.py
+++ b/tensorflow/python/ops/stateless_random_ops.py
@@ -618,3 +618,73 @@ def stateless_multinomial_categorical_impl(logits, num_samples, dtype, seed):
   logits = ops.convert_to_tensor(logits, name="logits")
   return gen_stateless_random_ops.stateless_multinomial(
       logits, num_samples, seed, output_dtype=dtype)
+
+
+@dispatch.add_dispatch_support
+@tf_export("random.stateless_parameterized_truncated_normal")
+def stateless_parameterized_truncated_normal(shape,
+                                             seed,
+                                             means=0.0,
+                                             stddevs=1.0,
+                                             minvals=-2.0,
+                                             maxvals=2.0,
+                                             name=None):
+  """Outputs random values from a truncated normal distribution.
+
+  The generated values follow a normal distribution with specified mean and
+  standard deviation, except that values whose magnitude is more than 2 standard
+  deviations from the mean are dropped and re-picked.
+
+
+  Examples:
+
+  Sample from a Truncated normal, with deferring shape parameters that
+  broadcast.
+
+  >>> means = 0.
+  >>> stddevs = tf.math.exp(tf.random.uniform(shape=[2, 3]))
+  >>> minvals = [-1., -2., -1000.]
+  >>> maxvals = [[10000.], [1.]]
+  >>> y = tf.random.stateless_parameterized_truncated_normal(
+  ...   shape=[10, 2, 3], seed=[7, 17],
+  ...   means=means, stddevs=stddevs, minvals=minvals, maxvals=maxvals)
+  >>> y.shape
+  TensorShape([10, 2, 3])
+
+  Args:
+    shape: A 1-D integer `Tensor` or Python array. The shape of the output
+      tensor.
+    seed: A shape [2] Tensor, the seed to the random number generator. Must have
+      dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
+    means: A `Tensor` or Python value of type `dtype`. The mean of the truncated
+      normal distribution. This must broadcast with `stddevs`, `minvals` and
+      `maxvals`, and the broadcasted shape must be dominated by `shape`.
+    stddevs: A `Tensor` or Python value of type `dtype`. The standard deviation
+      of the truncated normal distribution. This must broadcast with `means`,
+      `minvals` and `maxvals`, and the broadcasted shape must be dominated by
+      `shape`.
+    minvals: A `Tensor` or Python value of type `dtype`. The minimum value of
+      the truncated normal distribution. This must broadcast with `means`,
+      `stddevs` and `maxvals`, and the broadcasted shape must be dominated by
+      `shape`.
+    maxvals: A `Tensor` or Python value of type `dtype`. The maximum value of
+      the truncated normal distribution. This must broadcast with `means`,
+      `stddevs` and `minvals`, and the broadcasted shape must be dominated by
+      `shape`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A tensor of the specified shape filled with random truncated normal values.
+  """
+  with ops.name_scope(name, "stateless_parameterized_truncated_normal",
+                      [shape, means, stddevs, minvals, maxvals]) as name:
+    shape_tensor = tensor_util.shape_tensor(shape)
+    means_tensor = ops.convert_to_tensor(means, name="means")
+    stddevs_tensor = ops.convert_to_tensor(stddevs, name="stddevs")
+    minvals_tensor = ops.convert_to_tensor(minvals, name="minvals")
+    maxvals_tensor = ops.convert_to_tensor(maxvals, name="maxvals")
+    rnd = gen_stateless_random_ops.stateless_parameterized_truncated_normal(
+        shape_tensor, seed, means_tensor, stddevs_tensor, minvals_tensor,
+        maxvals_tensor)
+    tensor_util.maybe_set_static_shape(rnd, shape)
+    return rnd
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
index 9c6fa7154a3..f5963f1324c 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt
@@ -92,6 +92,10 @@ tf_module {
     name: "stateless_normal"
     argspec: "args=[\'shape\', \'seed\', \'mean\', \'stddev\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"\", \'None\'], "
   }
+  member_method {
+    name: "stateless_parameterized_truncated_normal"
+    argspec: "args=[\'shape\', \'seed\', \'means\', \'stddevs\', \'minvals\', \'maxvals\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \'-2.0\', \'2.0\', \'None\'], "
+  }
   member_method {
     name: "stateless_poisson"
     argspec: "args=[\'shape\', \'seed\', \'lam\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
index 62969b5a0dd..8e5303cbea4 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
@@ -4492,6 +4492,10 @@ tf_module {
     name: "StatelessMultinomial"
     argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "
   }
+  member_method {
+    name: "StatelessParameterizedTruncatedNormal"
+    argspec: "args=[\'shape\', \'seed\', \'means\', \'stddevs\', \'minvals\', \'maxvals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
   member_method {
     name: "StatelessRandomBinomial"
     argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
index e3a11ee4610..d1b8c90bfae 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt
@@ -80,6 +80,10 @@ tf_module {
     name: "stateless_normal"
     argspec: "args=[\'shape\', \'seed\', \'mean\', \'stddev\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \"\", \'None\'], "
   }
+  member_method {
+    name: "stateless_parameterized_truncated_normal"
+    argspec: "args=[\'shape\', \'seed\', \'means\', \'stddevs\', \'minvals\', \'maxvals\', \'name\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \'-2.0\', \'2.0\', \'None\'], "
+  }
   member_method {
     name: "stateless_poisson"
     argspec: "args=[\'shape\', \'seed\', \'lam\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
index 62969b5a0dd..8e5303cbea4 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
@@ -4492,6 +4492,10 @@ tf_module {
     name: "StatelessMultinomial"
     argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "
   }
+  member_method {
+    name: "StatelessParameterizedTruncatedNormal"
+    argspec: "args=[\'shape\', \'seed\', \'means\', \'stddevs\', \'minvals\', \'maxvals\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
   member_method {
     name: "StatelessRandomBinomial"
     argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "

From 81041bcd8267056741c8b4766bd8018a7270f454 Mon Sep 17 00:00:00 2001
From: Lu Wang 
Date: Thu, 18 Jun 2020 10:29:58 -0700
Subject: [PATCH 0506/1390] Update the documentation for Model metadata

PiperOrigin-RevId: 317131186
Change-Id: Ida1cd7fc729d7d601c32d646bd650fe1aa539591
---
 tensorflow/lite/g3doc/convert/metadata.md | 72 +++++++++++++++++++----
 1 file changed, 60 insertions(+), 12 deletions(-)

diff --git a/tensorflow/lite/g3doc/convert/metadata.md b/tensorflow/lite/g3doc/convert/metadata.md
index 29b2c5ce2b3..cd86333b305 100644
--- a/tensorflow/lite/g3doc/convert/metadata.md
+++ b/tensorflow/lite/g3doc/convert/metadata.md
@@ -9,7 +9,8 @@ input / output information. The metadata consists of both
 *   human readable parts which convey the best practice when using the model,
     and
 *   machine readable parts that can be leveraged by code generators, such as
-    [the TensorFlow Lite Android code generator](../guide/codegen.md).
+    [the TensorFlow Lite Android code generator](../guide/codegen.md) and
+    [the Android Studio ML Binding feature](https://developer.android.com/studio/preview/features#tensor-flow-lite-models).
 
 ## Setup the metadata tools
 
@@ -32,11 +33,21 @@ There are three parts to the
 [model metadata](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs):
 
 1.  **Model information** - Overall description of the model as well as items
-    such as licence terms.
+    such as licence terms. See
+    [ModelMetadata](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L515).
 2.  **Input information** - Description of the inputs and pre-processing
-    required such as normalization.
+    required such as normalization. See
+    [SubGraphMetadata.input_tensor_metadata](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L500).
 3.  **Output information** - Description of the output and post-processing
-    required such as mapping to labels.
+    required such as mapping to labels. See
+    [SubGraphMetadata.output_tensor_metadata](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L509).
+
+Since TensorFlow Lite only supports single subgraph at this point, the
+[TensorFlow Lite code generator](../guide/codegen.md) and
+[the Android Studio ML Binding feature](https://developer.android.com/studio/preview/features#tensor-flow-lite-models)
+will use `ModelMetadata.name` and `ModelMetadata.description`, instead of
+`SubGraphMetadata.name` and `SubGraphMetadata.description`, when displaying
+metadata and generating code.
 
 ### Supported Input / Output types
 
@@ -51,6 +62,29 @@ Lite metadata:
 *   Bounding box - Rectangular shape bounding boxes. The schema supports
     [a variety of numbering schemes](https://github.com/tensorflow/tensorflow/blob/268853ee81edab09e07f455cc918f7ef9a421485/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L165).
 
+### Pack the associated files
+
+TensorFlow Lite models may come with different associated files. For example,
+natural language models usually have vocab files that map word pieces to word
+IDs; classification models may have label files that indicate object categories.
+Without the associated files (if there are), a model will not function well.
+
+The associated files can now be bundled with the model through the metadata
+Python library. The new TensorFlow Lite model becomes a zip file that contains
+both the model and the associated files. It can be unpacked with common zip
+tools. This new model format keeps using the same file extension, `.tflite`. It
+is compatible with existing TFLite framework and Interpreter. See
+[Pack mtadata and associated files into the model](#pack-metadata-and-associated-files-into-the-model)
+for more details.
+
+The associate file information can be recored in the metadata. Depending on the
+file type and where the file is attached to (i.e. `ModelMetadata`,
+`SubGraphMetadata`, and `TensorMetadata`),
+[the TensorFlow Lite Android code generator](../guide/codegen.md) may apply
+corresponding pre/post processing automatically to the object. See
+[the \ section of each associate file type](https://github.com/tensorflow/tensorflow/blob/268853ee81edab09e07f455cc918f7ef9a421485/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L37-L77)
+in the schema for more details.
+
 ### Examples
 
 Note: The export directory specified has to exist before you run the script; it
@@ -63,7 +97,9 @@ types of models here:
 
 Download the script
 [here](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/metadata/metadata_writer_for_image_classifier.py)
-and run the script like this:
+, which populates metadata to
+[mobilenet_v1_0.75_160_quantized.tflite](https://tfhub.dev/tensorflow/lite-model/mobilenet_v1_0.75_160_quantized/1/default/1).
+Run the script like this:
 
 ```sh
 python ./metadata_writer_for_image_classifier.py \
@@ -72,8 +108,11 @@ python ./metadata_writer_for_image_classifier.py \
     --export_directory=model_with_metadata
 ```
 
-The rest of this guide will highlight some of the key sections in the image
-classification example to illustrate the key elements.
+To populate metadata for other image classification models, add the model specs
+like
+[this](https://github.com/tensorflow/examples/blob/master/lite/examples/image_classification/metadata/metadata_writer_for_image_classifier.py#L63-L74)
+into the script. The rest of this guide will highlight some of the key sections
+in the image classification example to illustrate the key elements.
 
 ### Deep dive into the image classification example
 
@@ -173,7 +212,7 @@ label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
 output_meta.associatedFiles = [label_file]
 ```
 
-#### Put it all together
+#### Create the metadata Flatbuffers
 
 The following code combines the model information with the input and output
 information:
@@ -192,8 +231,10 @@ b.Finish(
 metadata_buf = b.Output()
 ```
 
-Once the data structure is ready, the metadata is written into the TFLite file
-via the `populate` method:
+#### Pack metadata and associated files into the model
+
+Once the metadata Flatbuffers is created, the metadata and the label file are
+written into the TFLite file via the `populate` method:
 
 ```python
 populator = _metadata.MetadataPopulator.with_model_file(model_file)
@@ -202,9 +243,16 @@ populator.load_associated_files(["your_path_to_label_file"])
 populator.populate()
 ```
 
-#### Verify the metadata
+You can pack as many associated files as you want into the model through
+`load_associated_files`. However, it is required to pack at least those files
+documented in the metadata. In this example, packing the lable file is
+mandatory.
 
-You can read the metadata in a TFLite file using the `MetadataDisplayer`:
+### Visualize the metadata
+
+You can use [Netron](https://github.com/lutzroeder/netron) to visualize your
+metadata, or you can read the metadata from a TensorFlow Lite model into a json
+format using the `MetadataDisplayer`:
 
 ```python
 displayer = _metadata.MetadataDisplayer.with_model_file(export_model_path)

From 47582983cb1064b5bb81233db4f0adeeaa10b74d Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Thu, 18 Jun 2020 10:30:56 -0700
Subject: [PATCH 0507/1390] Fix save model issue for stateless ConvLSTM2D
 layer.

The root cause is that ConvLSTM2D.state is a tuple rather than a list. When converting the state for save_model, the tuple is not converted to trackable objects since the states are (None, None). On the other hand, save_model requires all objects to be trackable when saving.

We didn't hit this issue for keras.LSTM since its state is a list, rather than tuple. The list is auto convert to ListWrapper since list itself is mutable.

This should fix https://github.com/tensorflow/tensorflow/issues/40328 and partly https://github.com/tensorflow/tensorflow/issues/38220

PiperOrigin-RevId: 317131403
Change-Id: I202d4dbdb29accc7a047d5f5a2fef08d24d05c7c
---
 .../saving/saved_model/layer_serialization.py | 10 ++++++++--
 .../saving/saved_model/saved_model_test.py    | 20 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/keras/saving/saved_model/layer_serialization.py b/tensorflow/python/keras/saving/saved_model/layer_serialization.py
index 559b6158d87..4216457bf28 100644
--- a/tensorflow/python/keras/saving/saved_model/layer_serialization.py
+++ b/tensorflow/python/keras/saving/saved_model/layer_serialization.py
@@ -159,6 +159,12 @@ class RNNSavedModelSaver(LayerSavedModelSaver):
     objects, functions = (
         super(RNNSavedModelSaver, self)._get_serialized_attributes_internal(
             serialization_cache))
-
-    objects['states'] = data_structures.wrap_or_unwrap(self.obj.states)
+    states = data_structures.wrap_or_unwrap(self.obj.states)
+    # Force the tuple into TupleWrapper which is a trackable object. The
+    # save/load code requires all the objects to be trackable.
+    # Tuple is not converted to TupleWrapper by data_structures.wrap_or_unwrap()
+    # if it doesn't contains any trackable objects.
+    if isinstance(states, tuple):
+      states = data_structures._TupleWrapper(states)  # pylint: disable=protected-access
+    objects['states'] = states
     return objects, functions
diff --git a/tensorflow/python/keras/saving/saved_model/saved_model_test.py b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
index 8d4d27e2357..3f55d5f40b5 100644
--- a/tensorflow/python/keras/saving/saved_model/saved_model_test.py
+++ b/tensorflow/python/keras/saving/saved_model/saved_model_test.py
@@ -773,6 +773,26 @@ class TestModelSavingAndLoadingV2(keras_parameterized.TestCase):
     self.assertAllClose(layer.states, loaded_layer.states)
     self.assertAllClose(model(input_arr), loaded(input_arr))
 
+  def testSaveStatelessConvLSTM2D(self):
+    data_format = 'channels_first'
+    batch, timesteps, channels, rows, cols = 12, 10, 8, 4, 4
+    input_arr = np.ones(
+        (batch, timesteps, channels, rows, cols)).astype('float32')
+    layer = keras.layers.ConvLSTM2D(
+        filters=16, kernel_size=(1, 1), data_format=data_format)
+    x = keras.Input(batch_shape=(batch, timesteps, channels, rows, cols))
+    y = layer(x)
+    model = keras.Model(x, y)
+
+    predict_1 = model(input_arr)
+    saved_model_dir = self._save_model_dir()
+    tf_save.save(model, saved_model_dir)
+    del model
+
+    loaded = keras_load.load(saved_model_dir)
+    predict_2 = loaded(input_arr)
+    self.assertAllClose(predict_1, predict_2)
+
   def testSaveWithRaggedInputs(self):
 
     class EmbeddingMerger(keras.layers.Layer):

From 69467b92630bb37c077eda3e411903ccddd906ff Mon Sep 17 00:00:00 2001
From: Sachin Joglekar 
Date: Thu, 18 Jun 2020 10:32:53 -0700
Subject: [PATCH 0508/1390] Make control_dep test mlir-only, since tf2 behavior
 introduces control-flow ops that are not supported by TOCO

PiperOrigin-RevId: 317131879
Change-Id: I8cf51927a3d4985cec7d2b03930d2f553a6a3286
---
 tensorflow/lite/build_def.bzl                 |  1 -
 .../lite/testing/generate_examples_lib.py     |  1 -
 .../lite/testing/op_tests/control_dep.py      | 61 -------------------
 3 files changed, 63 deletions(-)
 delete mode 100644 tensorflow/lite/testing/op_tests/control_dep.py

diff --git a/tensorflow/lite/build_def.bzl b/tensorflow/lite/build_def.bzl
index e6c92691b15..ad43b56743a 100644
--- a/tensorflow/lite/build_def.bzl
+++ b/tensorflow/lite/build_def.bzl
@@ -251,7 +251,6 @@ def generated_test_models():
         "ceil",
         "concat",
         "constant",
-        # "control_dep", # b/150647401
         "conv",
         "conv_relu",
         "conv_relu1",
diff --git a/tensorflow/lite/testing/generate_examples_lib.py b/tensorflow/lite/testing/generate_examples_lib.py
index fc92991bd57..fce2beabf45 100644
--- a/tensorflow/lite/testing/generate_examples_lib.py
+++ b/tensorflow/lite/testing/generate_examples_lib.py
@@ -52,7 +52,6 @@ from tensorflow.lite.testing.op_tests.cast import make_cast_tests
 from tensorflow.lite.testing.op_tests.ceil import make_ceil_tests
 from tensorflow.lite.testing.op_tests.concat import make_concat_tests
 from tensorflow.lite.testing.op_tests.constant import make_constant_tests
-from tensorflow.lite.testing.op_tests.control_dep import make_control_dep_tests
 from tensorflow.lite.testing.op_tests.conv import make_conv_tests
 from tensorflow.lite.testing.op_tests.conv2d_transpose import make_conv2d_transpose_tests
 from tensorflow.lite.testing.op_tests.conv_activation import make_conv_relu_tests, make_conv_relu1_tests, make_conv_relu6_tests
diff --git a/tensorflow/lite/testing/op_tests/control_dep.py b/tensorflow/lite/testing/op_tests/control_dep.py
deleted file mode 100644
index bd9e369303b..00000000000
--- a/tensorflow/lite/testing/op_tests/control_dep.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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 configs for control_dep."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow.compat.v1 as tf
-from tensorflow.lite.testing.zip_test_utils import create_tensor_data
-from tensorflow.lite.testing.zip_test_utils import ExtraTocoOptions
-from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests
-from tensorflow.lite.testing.zip_test_utils import register_make_test_function
-
-TEST_INPUT_DEPTH = 3
-
-
-@register_make_test_function()
-def make_control_dep_tests(options):
-  """Make a set of tests that use control dependencies."""
-
-  test_parameters = [{
-      "input_shape": [[], [1, 1, 1, 1], [1, 15, 14, 1], [3, 15, 14, 3]],
-  }]
-
-  def build_graph(parameters):
-    input_tensor = tf.compat.v1.placeholder(
-        dtype=tf.float32, name="input", shape=parameters["input_shape"])
-    filter_value = tf.zeros((3, 3, TEST_INPUT_DEPTH, 8), tf.float32)
-    assert_op = tf.compat.v1.assert_greater_equal(input_tensor,
-                                                  input_tensor - 1)
-    with tf.control_dependencies([assert_op]):
-      out = tf.nn.conv2d(
-          input_tensor, filter_value, strides=(1, 1, 1, 1), padding="SAME")
-      return [input_tensor], [out]
-
-  def build_inputs(parameters, sess, inputs, outputs):
-    input_values = create_tensor_data(tf.float32, parameters["input_shape"])
-    return [input_values], sess.run(
-        outputs, feed_dict=dict(zip(inputs, [input_values])))
-
-  extra_toco_options = ExtraTocoOptions()
-  extra_toco_options.drop_control_dependency = True
-  make_zip_of_tests(
-      options,
-      test_parameters,
-      build_graph,
-      build_inputs,
-      extra_toco_options,
-      expected_tf_failures=3)

From d1157c976bd605b36506151493439ebd2f80bded Mon Sep 17 00:00:00 2001
From: Brian Zhao 
Date: Thu, 18 Jun 2020 10:34:15 -0700
Subject: [PATCH 0509/1390] Adding missing kernel dependency to fix test for
 windows.

PiperOrigin-RevId: 317132189
Change-Id: Ia59020db1c96301af4b9723170eb7808b66ef5b5
---
 tensorflow/c/experimental/saved_model/core/ops/BUILD | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/c/experimental/saved_model/core/ops/BUILD b/tensorflow/c/experimental/saved_model/core/ops/BUILD
index 8c4c41c6d75..332b92bec45 100644
--- a/tensorflow/c/experimental/saved_model/core/ops/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/ops/BUILD
@@ -78,14 +78,12 @@ tf_cc_test(
     srcs = [
         "variable_ops_test.cc",
     ],
-    tags = [
-        "no_windows",  # TODO(b/159210739): Remove this tag after fixing the bug.
-    ],
     deps = [
         ":owned_eager_context",
         ":owned_tensor",
         ":owned_tensor_handle",
         ":variable_ops",
+        "//tensorflow/core:all_kernels",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",

From cf00e559d7e9649714be26173c44f9d6697e702e Mon Sep 17 00:00:00 2001
From: Frank Chen 
Date: Thu, 18 Jun 2020 10:39:58 -0700
Subject: [PATCH 0510/1390] Add TPU configuration ops to default TensorFlow
 build

PiperOrigin-RevId: 317133514
Change-Id: I33bc6d7fdbba5915bd0d1291d4e086139c07eb14
---
 .bazelrc                                      |   4 +
 tensorflow/BUILD                              |   7 +
 tensorflow/core/BUILD                         |   3 +
 tensorflow/core/tpu/BUILD                     |  14 ++
 tensorflow/core/tpu/kernels/BUILD             |  30 +++-
 .../core/tpu/kernels/tpu_compile_c_api.h      |  27 +--
 .../core/tpu/kernels/tpu_compile_op_common.cc |   2 +-
 .../core/tpu/kernels/tpu_configuration_ops.cc |  29 +--
 .../core/tpu/kernels/tpu_mesh_state_c_api.h   |  15 +-
 .../tpu/kernels/tpu_mesh_state_interface.h    |  11 +-
 tensorflow/core/tpu/kernels/tpu_util_c_api.h  |   6 +
 tensorflow/core/tpu/tpu_library_init_fns.inc  | 166 ++++++++++++++++++
 tensorflow/core/tpu/tpu_library_loader.cc     |  75 ++++----
 tensorflow/core/tpu/tpu_library_loader.h      |  16 ++
 .../core/tpu/tpu_library_loader_windows.cc    |  10 ++
 tensorflow/stream_executor/tpu/BUILD          |  14 +-
 .../stream_executor/tpu/tpu_executor.cc       | 129 ++++++++------
 .../stream_executor/tpu/tpu_executor_c_api.h  |  92 +++++++++-
 .../stream_executor/tpu/tpu_node_context.cc   |  14 +-
 .../tpu/tpu_node_context_c_api.h              |  12 ++
 .../stream_executor/tpu/tpu_platform.cc       |  64 ++++---
 tensorflow/stream_executor/tpu/tpu_platform.h |   6 +-
 tensorflow/stream_executor/tpu/tpu_stream.h   |  25 ++-
 tensorflow/stream_executor/tpu/tpu_timer.h    |  13 +-
 .../tpu/tpu_transfer_manager.cc               |  20 ++-
 tensorflow/tensorflow.bzl                     |   7 +
 26 files changed, 627 insertions(+), 184 deletions(-)
 create mode 100644 tensorflow/core/tpu/tpu_library_init_fns.inc

diff --git a/.bazelrc b/.bazelrc
index e21a1a32917..e67c3eecc3b 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -39,6 +39,7 @@
 #
 # Feature and Third party library support options:
 #     xla:          Build TF with XLA
+#     tpu:          Build TF with TPU support
 #     using_cuda:   CUDA is available to build system.
 #     cuda:         Build with full cuda support.
 #     rocm:         Build with AMD GPU support (rocm).
@@ -180,6 +181,9 @@ build:dbg --cxxopt -DTF_LITE_DISABLE_X86_NEON
 # AWS SDK must be compiled in release mode. see: https://github.com/tensorflow/tensorflow/issues/37498
 build:dbg --copt -DDEBUG_BUILD
 
+# Config to build TPU backend
+build:tpu --define=with_tpu_support=true
+
 build:tensorrt --action_env TF_NEED_TENSORRT=1
 
 build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
diff --git a/tensorflow/BUILD b/tensorflow/BUILD
index bd0619b0c05..d00608ccc98 100644
--- a/tensorflow/BUILD
+++ b/tensorflow/BUILD
@@ -467,6 +467,13 @@ config_setting(
     visibility = ["//visibility:public"],
 )
 
+# This flag enables experimental TPU support
+config_setting(
+    name = "with_tpu_support",
+    values = {"define": "with_tpu_support=true"},
+    visibility = ["//visibility:public"],
+)
+
 # Specifies via a config setting if this is a mobile build or not, makes
 # it easier to combine settings later.
 selects.config_setting_group(
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 50f1f2527a5..7f1c1bd549b 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -72,6 +72,7 @@ load(
     "if_ios",
     "if_mobile",
     "if_not_windows",
+    "if_tpu",
     "tf_android_core_proto_headers",
     "tf_cc_test",
     "tf_cc_test_mkl",
@@ -1093,6 +1094,8 @@ cc_library(
     ]) + if_tensorrt([
         "//tensorflow/compiler/tf2tensorrt:trt_engine_resource_op_kernels",
         "//tensorflow/compiler/tf2tensorrt:trt_op_kernels",
+    ]) + if_tpu([
+        "//tensorflow/core/tpu/kernels",
     ]),
 )
 
diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD
index d4bcdfd52c5..9e89cd69235 100644
--- a/tensorflow/core/tpu/BUILD
+++ b/tensorflow/core/tpu/BUILD
@@ -103,6 +103,7 @@ cc_library(
         ":libtftpu_header",
         "//tensorflow/c:tf_status",
     ],
+    alwayslink = True,
 )
 
 cc_library(
@@ -116,7 +117,20 @@ cc_library(
     deps = [
         ":libtftpu_header",
         ":tpu_config_c_api",
+        ":tpu_library_init_fns",
         "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:status",
+        "//tensorflow/core/tpu/kernels:tpu_compile_c_api_hdrs",
+        "//tensorflow/core/tpu/kernels:tpu_mesh_state_c_api_hdrs",
+        "//tensorflow/core/tpu/kernels:tpu_util_c_api_hdrs",
+        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
+        "//tensorflow/stream_executor/tpu:tpu_node_context_c_api_hdrs",
+        "//tensorflow/stream_executor/tpu:tpu_platform_hdrs",
     ],
 )
+
+cc_library(
+    name = "tpu_library_init_fns",
+    hdrs = ["tpu_library_init_fns.inc"],
+    visibility = ["//visibility:public"],
+)
diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index 9ba9ad61aa0..94d3c8edf2b 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -3,6 +3,10 @@ load(
     "//tensorflow/core/platform:build_config.bzl",
     "tf_proto_library_cc",
 )
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_kernel_library",
+)
 
 package(
     default_visibility = [
@@ -12,6 +16,12 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+tf_kernel_library(
+    name = "kernels",
+    visibility = ["//visibility:public"],
+    deps = [":tpu_configuration_ops"],
+)
+
 cc_library(
     name = "tpu_compile_op_common",
     srcs = ["tpu_compile_op_common.cc"],
@@ -50,7 +60,7 @@ cc_library(
     hdrs = ["tpu_compile_op_options.h"],
 )
 
-cc_library(
+tf_kernel_library(
     name = "tpu_configuration_ops",
     srcs = ["tpu_configuration_ops.cc"],
     hdrs = ["tpu_configuration_ops.h"],
@@ -75,12 +85,13 @@ cc_library(
     name = "tpu_compile_c_api_hdrs",
     hdrs = ["tpu_compile_c_api.h"],
     deps = [
-        ":tpu_mesh_state_c_api",
+        ":tpu_mesh_state_c_api_hdrs",
         ":tpu_ops_common_c_api_hdrs",
         ":tpu_program_c_api_hdrs",
-        "//tensorflow/c:tf_datatype",
+        "//tensorflow/core/tpu:libtftpu_header",
         "//tensorflow/stream_executor/tpu:proto_helper",
     ],
+    alwayslink = True,
 )
 
 tf_proto_library_cc(
@@ -197,8 +208,10 @@ cc_library(
 )
 
 cc_library(
-    name = "tpu_mesh_state_c_api",
+    name = "tpu_mesh_state_c_api_hdrs",
     hdrs = ["tpu_mesh_state_c_api.h"],
+    deps = ["//tensorflow/core/tpu:libtftpu_header"],
+    alwayslink = True,
 )
 
 cc_library(
@@ -207,12 +220,11 @@ cc_library(
     hdrs = ["tpu_mesh_state_interface.h"],
     deps = [
         ":tpu_compile_c_api_hdrs",
-        ":tpu_mesh_state_c_api",
+        ":tpu_mesh_state_c_api_hdrs",
         "//tensorflow/compiler/xla/service",
         "//tensorflow/core:framework",
-        "//tensorflow/core/platform:errors",
         "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
-        "//tensorflow/core/tpu:tpu_config_c_api",
+        "//tensorflow/core/tpu:tpu_library_loader",
     ],
 )
 
@@ -371,13 +383,16 @@ cc_library(
     name = "tpu_util_c_api_hdrs",
     hdrs = ["tpu_util_c_api.h"],
     deps = [
+        "//tensorflow/core/tpu:libtftpu_header",
         "//tensorflow/stream_executor/tpu:proto_helper",
     ],
+    alwayslink = True,
 )
 
 cc_library(
     name = "tpu_ops_common_c_api_hdrs",
     hdrs = ["tpu_ops_common_c_api.h"],
+    alwayslink = True,
 )
 
 cc_library(
@@ -387,6 +402,7 @@ cc_library(
         ":tpu_ops_common_c_api_hdrs",
         "//tensorflow/stream_executor/tpu:proto_helper",
     ],
+    alwayslink = True,
 )
 
 cc_library(
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
index 70e3a7d2340..d1546ed9610 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_ops_common_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_program_c_api.h"
+#include "tensorflow/core/tpu/libtftpu.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"
 
 enum TpuCoreTypeEnum {
@@ -44,35 +45,41 @@ struct CompilationCacheKeyProperty {
 extern "C" {
 
 // Returns the number of available TPU core count.
-int TpuTopology_AvailableCoreCount(const XLA_TpuMeshState* mesh_state,
-                                   TpuCoreTypeEnum tpu_core_type);
+TFTPU_CAPI_EXPORT int TpuTopology_AvailableCoreCount(
+    const XLA_TpuMeshState* mesh_state, TpuCoreTypeEnum tpu_core_type);
 
 // Creates a unique compilation cache `key` used for `put` and `get` operations.
 // Returned buffer is heap-allocated and must be owned.
-const char* TpuCompile_CreateCompilationCacheKey(
+TFTPU_CAPI_EXPORT const char* TpuCompile_CreateCompilationCacheKey(
     CompilationCacheKeyProperty property);
 
 // Creates a guaranteed const fingerprint. Guarantee const is normally used in
 // TPU inference to avoid re-copying unchanged variables onto the TPU device.
 // It promises the value is identical for every execution in the same session
 // even if the actual value changes in later executions.
-uint64_t TpuCompile_CreateGuaranteedConstFingerprint(uint64_t fingerprint,
-                                                     const char* data,
-                                                     size_t size);
+TFTPU_CAPI_EXPORT uint64_t TpuCompile_CreateGuaranteedConstFingerprint(
+    uint64_t fingerprint, const char* data, size_t size);
 
 // Executes the computations using XLA TPU compiler and returns TPU programs
 // ready for execution.
-void TpuCompile_CompileAheadOfTime(
-    TpuSerializedProto aot_compilation_request,
-    XLA_TpuProgram** tpu_programs[],
+TFTPU_CAPI_EXPORT void TpuCompile_CompileAheadOfTime(
+    TpuSerializedProto aot_compilation_request, XLA_TpuProgram** tpu_programs[],
     size_t* count, SE_Status* status);
 
 // Builds `DeviceAssignment` from `TpuCompileMetadata` serialized proto.
-void TpuCompile_BuildXLADeviceAssignment(
+TFTPU_CAPI_EXPORT void TpuCompile_BuildXLADeviceAssignment(
     TpuSerializedProto serialized_tpu_compile_metadata,
     const XLA_TpuMeshState* mesh_state,
     TpuSerializedProto* serialized_device_assignment, SE_Status* status);
 
+struct TfTpu_CompileApiFn {
+  TFTPU_ADD_FN_IN_STRUCT(TpuTopology_AvailableCoreCount);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CreateCompilationCacheKey);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CreateGuaranteedConstFingerprint);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CompileAheadOfTime);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_BuildXLADeviceAssignment);
+};
+
 }  // extern "C"
 
 #endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_COMPILE_C_API_H_
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
index 7ab1c9b8027..92d1fa1337e 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
@@ -353,7 +353,7 @@ Status TpuCompileOpKernelCommon::CompileTFFunctionToHlo(
     return;
   }
 
-  LogAndExit(42);
+  std::quick_exit(42);
 }
 
 /* static */ Status TpuCompileOpKernelCommon::GetDynamicShapes(
diff --git a/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc b/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
index 7fa345d735c..12a3256a44f 100644
--- a/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
+++ b/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/tpu/tpu_config_c_api.h"
 #include "tensorflow/core/tpu/tpu_configuration.h"
 #include "tensorflow/core/tpu/tpu_defs.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"
 
 namespace tensorflow {
@@ -97,13 +98,14 @@ void ConfigureDistributedTpuOp::Compute(OpKernelContext* ctx) {
   OP_REQUIRES_OK(ctx, DeleteIfExists(
                           rmgr, tpu::kTpuMeshCommonStateResourceName));
 
-  ConfigureDistributedTpuOp_DoWork(
+  tpu::ConfigApiFn()->ConfigureDistributedTpuOp_DoWorkFn(
       num_devices_per_host.size(), num_devices_per_host.data(),
       &host_config_output_size, &host_config_output, status);
 
-  OP_REQUIRES_OK(ctx, rmgr->Create(rmgr->default_container(),
-                                   tpu::kTpuMeshCommonStateResourceName,
-                                   tpu::TpuMeshStateInterface::Create()));
+  auto* tpu_mesh = tpu::TpuMeshStateInterface::Create();
+  OP_REQUIRES_OK(ctx,
+                 rmgr->Create(rmgr->default_container(),
+                              tpu::kTpuMeshCommonStateResourceName, tpu_mesh));
 
   Tensor* ctx_output;
   OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &ctx_output));
@@ -112,7 +114,8 @@ void ConfigureDistributedTpuOp::Compute(OpKernelContext* ctx) {
 
   OP_REQUIRES_OK(ctx, StatusFromTF_Status(status));
   TF_DeleteStatus(status);
-  TpuConfigurationApi_FreeCharArray(host_config_output);
+
+  tpu::ConfigApiFn()->TpuConfigurationApi_FreeCharArrayFn(host_config_output);
 
   VLOG(1) << "ConfigureDistributedTpuOp done";
 }
@@ -171,7 +174,7 @@ void WaitForDistributedTpuOp::Compute(OpKernelContext* ctx) {
   OP_REQUIRES_OK(ctx, GetTpuMeshStateInterface(rmgr, &mesh_state));
   core::ScopedUnref mesh_state_unref(mesh_state);
 
-  WaitForDistributedTpuOp_DoWork(
+  tpu::ConfigApiFn()->WaitForDistributedTpuOp_DoWorkFn(
       num_hosts, num_devices_per_host,
       const_cast(mapping_arg.data()), mesh_state,
       &tpu_topology_output_size, &tpu_topology_output, status);
@@ -183,7 +186,7 @@ void WaitForDistributedTpuOp::Compute(OpKernelContext* ctx) {
 
   OP_REQUIRES_OK(ctx, StatusFromTF_Status(status));
   TF_DeleteStatus(status);
-  TpuConfigurationApi_FreeCharArray(tpu_topology_output);
+  tpu::ConfigApiFn()->TpuConfigurationApi_FreeCharArrayFn(tpu_topology_output);
 
   VLOG(1) << "WaitForDistributedTpuOp done";
 }
@@ -196,7 +199,7 @@ void ShutdownDistributedTpuOp::Compute(OpKernelContext* ctx) {
   OP_REQUIRES_OK(ctx, DeleteIfExists(
                           GetTPUConfigResourceMgr(),
                           tpu::kTpuMeshCommonStateResourceName));
-  ShutdownDistributedTpuOp_DoWork(status);
+  tpu::ConfigApiFn()->ShutdownDistributedTpuOp_DoWorkFn(status);
   OP_REQUIRES_OK(ctx, StatusFromTF_Status(status));
   TF_DeleteStatus(status);
 
@@ -213,7 +216,7 @@ void InitializeHostForDistributedTpuOp::Compute(OpKernelContext* ctx) {
   int32_t* device_id_output;
   TF_Status* status = TF_NewStatus();
 
-  InitializeHostForDistributedTpuOp_DoWork(
+  tpu::ConfigApiFn()->InitializeHostForDistributedTpuOp_DoWorkFn(
       tpu_host_config.size(), tpu_host_config.data(),
       enable_whole_mesh_compilations_, &device_id_output_size,
       &device_id_output, status);
@@ -230,7 +233,7 @@ void InitializeHostForDistributedTpuOp::Compute(OpKernelContext* ctx) {
 
   OP_REQUIRES_OK(ctx, StatusFromTF_Status(status));
   TF_DeleteStatus(status);
-  TpuConfigurationApi_FreeInt32Array(device_id_output);
+  tpu::ConfigApiFn()->TpuConfigurationApi_FreeInt32ArrayFn(device_id_output);
 
   VLOG(1) << "InitializeHostForDistributedTpuOp done";
 }
@@ -242,7 +245,8 @@ void SetGlobalTPUArrayOp::Compute(OpKernelContext* ctx) {
   auto tpu_topology = ctx->input(0).scalar()();
   TF_Status* status = TF_NewStatus();
 
-  SetGlobalTPUArrayOp_DoWork(tpu_topology.size(), tpu_topology.data(), status);
+  tpu::ConfigApiFn()->SetGlobalTPUArrayOp_DoWorkFn(tpu_topology.size(),
+                                                   tpu_topology.data(), status);
 
   OP_REQUIRES_OK(ctx, StatusFromTF_Status(status));
   TF_DeleteStatus(status);
@@ -257,7 +261,8 @@ void DisconnectDistributedTpuChipsOp::Compute(OpKernelContext* ctx) {
   TF_Status* status = TF_NewStatus();
   int32_t number_of_chips_output = 0;
 
-  DisconnectDistributedTpuChipsOp_DoWork(&number_of_chips_output, status);
+  tpu::ConfigApiFn()->DisconnectDistributedTpuChipsOp_DoWorkFn(
+      &number_of_chips_output, status);
 
   Tensor* ctx_output;
   OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &ctx_output));
diff --git a/tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h b/tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h
index 3ed65fe5cc4..a6434d7d2fd 100644
--- a/tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h
@@ -15,20 +15,29 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_MESH_STATE_C_API_H_
 #define TENSORFLOW_CORE_TPU_KERNELS_TPU_MESH_STATE_C_API_H_
 
+#include "tensorflow/core/tpu/libtftpu.h"
+
 typedef struct XLA_TpuMeshState XLA_TpuMeshState;
 
 extern "C" {
 
 // Creates a new TPU mesh state object.
-XLA_TpuMeshState* TpuMeshState_Create();
+TFTPU_CAPI_EXPORT XLA_TpuMeshState* TpuMeshState_Create();
 
 // Deletes the given TPU `mesh_state` object. Once deleted the object is
 // unusable.
-void TpuMeshState_Free(XLA_TpuMeshState* mesh_state);
+TFTPU_CAPI_EXPORT void TpuMeshState_Free(XLA_TpuMeshState* mesh_state);
 
 // Returns a pointer to an opaque mesh data structure used internally.
-void* TpuMeshState_MeshCommonState(XLA_TpuMeshState* mesh_state);
+TFTPU_CAPI_EXPORT void* TpuMeshState_MeshCommonState(
+    XLA_TpuMeshState* mesh_state);
 
 }  // extern "C"
 
+struct TfTpu_MeshStateApiFn {
+  TFTPU_ADD_FN_IN_STRUCT(TpuMeshState_Create);
+  TFTPU_ADD_FN_IN_STRUCT(TpuMeshState_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuMeshState_MeshCommonState);
+};
+
 #endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_MESH_STATE_C_API_H_
diff --git a/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h b/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
index 34202a78718..3eff3be4915 100644
--- a/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
+++ b/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 
 namespace tensorflow {
 
@@ -38,19 +39,19 @@ class TpuMeshStateInterface : public tensorflow::ResourceBase {
 
   ~TpuMeshStateInterface() override {
     if (mesh_state_ != nullptr) {
-      TpuMeshState_Free(mesh_state_);
+      MeshStateApiFn()->TpuMeshState_FreeFn(mesh_state_);
     }
   }
 
   static TpuMeshStateInterface* Create() {
-    return new TpuMeshStateInterface(TpuMeshState_Create());
+    return new TpuMeshStateInterface(MeshStateApiFn()->TpuMeshState_CreateFn());
   }
 
   const XLA_TpuMeshState* data() const { return mesh_state_; }
 
   tensorflow::TpuMeshCommonState* mesh_common_state() const {
     return static_cast(
-        TpuMeshState_MeshCommonState(mesh_state_));
+        MeshStateApiFn()->TpuMeshState_MeshCommonStateFn(mesh_state_));
   }
 
   // Returns whether we should include the device assignment as a static field
@@ -62,8 +63,8 @@ class TpuMeshStateInterface : public tensorflow::ResourceBase {
     // Static device assignment enables XLA to perform certain optimization when
     // all cores are used in the replicated computation.
     return metadata.num_cores_per_replica() * metadata.num_replicas() ==
-           TpuTopology_AvailableCoreCount(mesh_state_,
-                                          tpu_core_type);
+           CompileApiFn()->TpuTopology_AvailableCoreCountFn(mesh_state_,
+                                                            tpu_core_type);
   }
 
   string DebugString() const override { return "TpuMeshStateInterface"; }
diff --git a/tensorflow/core/tpu/kernels/tpu_util_c_api.h b/tensorflow/core/tpu/kernels/tpu_util_c_api.h
index 4d992449cfc..32b946d56c9 100644
--- a/tensorflow/core/tpu/kernels/tpu_util_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_util_c_api.h
@@ -15,6 +15,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
 #define TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
 
+#include "tensorflow/core/tpu/libtftpu.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"
 
 typedef struct SE_Status SE_Status;
@@ -32,4 +33,9 @@ void TpuCompile_ToTpuShapeRepresentation(
 
 }  // extern "C"
 
+struct TfTpu_UtilApiFn {
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_IsTpuCompilationEnabled);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_ToTpuShapeRepresentation);
+};
+
 #endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_UTIL_C_API_H_
diff --git a/tensorflow/core/tpu/tpu_library_init_fns.inc b/tensorflow/core/tpu/tpu_library_init_fns.inc
new file mode 100644
index 00000000000..e21d7f195ad
--- /dev/null
+++ b/tensorflow/core/tpu/tpu_library_init_fns.inc
@@ -0,0 +1,166 @@
+namespace {
+
+tensorflow::Status SetTpuConfigStructFns(void* library_handle) {
+  auto* config_fn = tensorflow::tpu::ConfigApiFn();
+
+  TFTPU_SET_FN(config_fn, ConfigureDistributedTpuOp_DoWork);
+  TFTPU_SET_FN(config_fn, WaitForDistributedTpuOp_DoWork);
+  TFTPU_SET_FN(config_fn, ShutdownDistributedTpuOp_DoWork);
+  TFTPU_SET_FN(config_fn, InitializeHostForDistributedTpuOp_DoWork);
+  TFTPU_SET_FN(config_fn, SetGlobalTPUArrayOp_DoWork);
+  TFTPU_SET_FN(config_fn, DisconnectDistributedTpuChipsOp_DoWork);
+  TFTPU_SET_FN(config_fn, TpuConfigurationApi_FreeCharArray);
+  TFTPU_SET_FN(config_fn, TpuConfigurationApi_FreeInt32Array);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetTpuMeshStateStructFns(void* library_handle) {
+  auto* mesh_state_fn = tensorflow::tpu::MeshStateApiFn();
+
+  TFTPU_SET_FN(mesh_state_fn, TpuMeshState_Create);
+  TFTPU_SET_FN(mesh_state_fn, TpuMeshState_Free);
+  TFTPU_SET_FN(mesh_state_fn, TpuMeshState_MeshCommonState);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetCompileStructFn(void* library_handle) {
+  auto* compile_fn = tensorflow::tpu::CompileApiFn();
+
+  TFTPU_SET_FN(compile_fn, TpuTopology_AvailableCoreCount);
+  TFTPU_SET_FN(compile_fn, TpuCompile_CreateCompilationCacheKey);
+  TFTPU_SET_FN(compile_fn, TpuCompile_CreateGuaranteedConstFingerprint);
+  TFTPU_SET_FN(compile_fn, TpuCompile_CompileAheadOfTime);
+  TFTPU_SET_FN(compile_fn, TpuCompile_BuildXLADeviceAssignment);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetExecutorStructFn(void* library_handle) {
+  auto* executor_fn = tensorflow::tpu::ExecutorApiFn();
+
+  TFTPU_SET_FN(executor_fn, TpuPlatform_New);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_Free);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_Initialize);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_Initialized);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_GetExecutor);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_Id);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_VisibleDeviceCount);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_TpuMemoryLimit);
+  TFTPU_SET_FN(executor_fn, TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy);
+
+  TFTPU_SET_FN(executor_fn, TpuExecutor_Init);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_Free);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_PlatformDeviceCount);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_Allocate);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_Deallocate);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_GetAllocatorStats);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_DeviceMemoryUsage);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateStream);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateStream);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_CreateStreamDependency);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_GetStatus);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateEvent);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateEvent);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_PollForEventStatus);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_RecordEvent);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForEvent);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_AllocateTimer);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_DeallocateTimer);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_StartTimer);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_StopTimer);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyToHost);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronousMemcpyFromHost);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyToHost);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_MemcpyFromHost);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_EnqueueInfeed);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_DequeueOutfeed);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForInfeedReady);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_WaitForOutfeedReady);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_BlockHostUntilDone);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_BlockUntilDoneOrFailed);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_SyncAndForgetFailedStreams);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_SynchronizeAllActivity);
+
+  TFTPU_SET_FN(executor_fn, TpuStream_New);
+  TFTPU_SET_FN(executor_fn, TpuStream_Free);
+  TFTPU_SET_FN(executor_fn, TpuStream_Stream);
+  TFTPU_SET_FN(executor_fn, TpuStream_Status);
+  TFTPU_SET_FN(executor_fn, TpuStream_IsSameSharedMemoryLocation);
+  TFTPU_SET_FN(executor_fn, TpuStream_TpuEnqueueOnDeviceSendRecvLocal);
+
+  TFTPU_SET_FN(executor_fn, TpuEvent_New);
+  TFTPU_SET_FN(executor_fn, TpuEvent_Free);
+
+  TFTPU_SET_FN(executor_fn, TpuTimer_New);
+  TFTPU_SET_FN(executor_fn, TpuTimer_Free);
+  TFTPU_SET_FN(executor_fn, TpuTimer_Nanoseconds);
+  TFTPU_SET_FN(executor_fn, TpuTimer_Microseconds);
+
+  TFTPU_SET_FN(executor_fn, TpuStatus_New);
+  TFTPU_SET_FN(executor_fn, TpuStatus_Create);
+  TFTPU_SET_FN(executor_fn, TpuStatus_Free);
+  TFTPU_SET_FN(executor_fn, TpuStatus_Message);
+  TFTPU_SET_FN(executor_fn, TpuStatus_Code);
+  TFTPU_SET_FN(executor_fn, TpuStatus_Ok);
+
+  TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Default);
+  TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_SetOrdinal);
+  TFTPU_SET_FN(executor_fn, TpuStreamExecutorConfig_Free);
+
+  TFTPU_SET_FN(executor_fn, TpuDeviceDescription_New);
+  TFTPU_SET_FN(executor_fn, TpuDeviceDescription_Free);
+
+  TFTPU_SET_FN(executor_fn, TpuExecutor_CreateDeviceDescription);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_NewDeviceOptions);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_FreeDeviceOptions);
+  TFTPU_SET_FN(executor_fn, TpuExecutor_HostCallback);
+
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_New);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_Free);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_PlatformId);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_HostShapeToDeviceShape);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralToDeviceAsync);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_TransferLiteralFromDevice);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_GetByteSizeRequirement);
+  TFTPU_SET_FN(executor_fn, TpuTransferManager_WriteSingleTupleIndexTable);
+
+  TFTPU_SET_FN(executor_fn, TpuComputationPlacer_New);
+  TFTPU_SET_FN(executor_fn, TpuComputationPlacer_Free);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetTpuNodeContextStructFns(void* library_handle) {
+  auto* node_context_fn = tensorflow::tpu::NodeContextApiFn();
+
+  TFTPU_SET_FN(node_context_fn, TpuNodeContext_Create);
+  TFTPU_SET_FN(node_context_fn, TpuNodeContext_Free);
+  TFTPU_SET_FN(node_context_fn, TpuNodeContext_StopChipHeartbeats);
+  TFTPU_SET_FN(node_context_fn, TpuNodeContext_CloseTpuHost);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetTpuUtilStructFns(void* library_handle) {
+  auto* util_fn = tensorflow::tpu::UtilApiFn();
+
+  TFTPU_SET_FN(util_fn, TpuCompile_IsTpuCompilationEnabled);
+  TFTPU_SET_FN(util_fn, TpuCompile_ToTpuShapeRepresentation);
+
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status InitializeTpuStructFns(void* library_handle) {
+  TF_RETURN_IF_ERROR(SetTpuConfigStructFns(library_handle));
+  TF_RETURN_IF_ERROR(SetTpuMeshStateStructFns(library_handle));
+  TF_RETURN_IF_ERROR(SetCompileStructFn(library_handle));
+  TF_RETURN_IF_ERROR(SetExecutorStructFn(library_handle));
+  TF_RETURN_IF_ERROR(SetTpuNodeContextStructFns(library_handle));
+  TF_RETURN_IF_ERROR(SetTpuUtilStructFns(library_handle));
+
+  return tensorflow::Status::OK();
+}
+
+}  // namespace
\ No newline at end of file
diff --git a/tensorflow/core/tpu/tpu_library_loader.cc b/tensorflow/core/tpu/tpu_library_loader.cc
index c89de142a9f..834b86e68a7 100644
--- a/tensorflow/core/tpu/tpu_library_loader.cc
+++ b/tensorflow/core/tpu/tpu_library_loader.cc
@@ -13,16 +13,23 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+// TODO(frankchn): Rename to `tpu_api_dlsym_initializer` or similar.
+
 #include "tensorflow/core/tpu/tpu_library_loader.h"
 
 #include 
 
-#define TFTPU_SET_FN(Struct, FnName) \
-  Struct->FnName##Fn =               \
-      reinterpret_cast(dlsym(library_handle, #FnName));
-
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_platform.h"
+
+#define TFTPU_SET_FN(Struct, FnName)                                       \
+  Struct->FnName##Fn =                                                     \
+      reinterpret_cast(dlsym(library_handle, #FnName)); \
+  if (!(Struct->FnName##Fn)) {                                             \
+    LOG(ERROR) << #FnName " not available in this library.";               \
+  }
 
 // Reminder: Update tpu_library_loader_windows.cc if you are adding new publicly
 // visible methods.
@@ -30,28 +37,7 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
-Status SetTpuInitializeStructFns(void* library_handle) {
-  auto* base_fn = InitializeApiFn();
-
-  TFTPU_SET_FN(base_fn, TfTpu_Initialize);
-
-  return Status::OK();
-}
-
-Status SetTpuConfigStructFns(void* library_handle) {
-  auto* config_fn = ConfigApiFn();
-
-  TFTPU_SET_FN(config_fn, ConfigureDistributedTpuOp_DoWork);
-  TFTPU_SET_FN(config_fn, WaitForDistributedTpuOp_DoWork);
-  TFTPU_SET_FN(config_fn, ShutdownDistributedTpuOp_DoWork);
-  TFTPU_SET_FN(config_fn, InitializeHostForDistributedTpuOp_DoWork);
-  TFTPU_SET_FN(config_fn, SetGlobalTPUArrayOp_DoWork);
-  TFTPU_SET_FN(config_fn, DisconnectDistributedTpuChipsOp_DoWork);
-  TFTPU_SET_FN(config_fn, TpuConfigurationApi_FreeCharArray);
-  TFTPU_SET_FN(config_fn, TpuConfigurationApi_FreeInt32Array);
-
-  return Status::OK();
-}
+#include "tensorflow/core/tpu/tpu_library_init_fns.inc"
 
 TfTpu_BaseFn* InitializeApiFn() {
   static TfTpu_BaseFn base_fn;
@@ -63,19 +49,48 @@ TfTpu_ConfigApiFn* ConfigApiFn() {
   return &config_api_fn;
 }
 
+TfTpu_MeshStateApiFn* MeshStateApiFn() {
+  static TfTpu_MeshStateApiFn mesh_state_api_fn;
+  return &mesh_state_api_fn;
+}
+
+TfTpu_CompileApiFn* CompileApiFn() {
+  static TfTpu_CompileApiFn compile_api_fn;
+  return &compile_api_fn;
+}
+
+TfTpu_ExecutorApiFn* ExecutorApiFn() {
+  static TfTpu_ExecutorApiFn executor_api_fn;
+  return &executor_api_fn;
+}
+
+TfTpu_NodeContextApiFn* NodeContextApiFn() {
+  static TfTpu_NodeContextApiFn node_context_api_fn;
+  return &node_context_api_fn;
+}
+
+TfTpu_UtilApiFn* UtilApiFn() {
+  static TfTpu_UtilApiFn util_api_fn;
+  return &util_api_fn;
+}
+
 Status InitializeTpuLibrary(void* library_handle) {
   bool shared_object_loaded = true;
   if (library_handle == nullptr) {
-    library_handle = dlopen(nullptr, RTLD_LAZY);
+    library_handle = dlopen(nullptr, RTLD_NOW);
     shared_object_loaded = false;
   }
 
-  TF_RETURN_IF_ERROR(SetTpuInitializeStructFns(library_handle));
-  TF_RETURN_IF_ERROR(SetTpuConfigStructFns(library_handle));
+  TF_RETURN_IF_ERROR(InitializeTpuStructFns(library_handle));
 
   if (shared_object_loaded) {
+    // TODO(frankchn): Make initialization actually work
     // Initialize TPU platform when the platform code is loaded from a library.
-    InitializeApiFn()->TfTpu_InitializeFn();
+    // InitializeApiFn()->TfTpu_InitializeFn();
+
+    // We should only register the TPU platform when the library is loaded.
+    // TODO(frankchn): Resolve the circular dependency and register the platform
+    // RegisterTpuPlatform();
   }
 
   return Status::OK();
diff --git a/tensorflow/core/tpu/tpu_library_loader.h b/tensorflow/core/tpu/tpu_library_loader.h
index a51948cf719..ba6c324707d 100644
--- a/tensorflow/core/tpu/tpu_library_loader.h
+++ b/tensorflow/core/tpu/tpu_library_loader.h
@@ -17,8 +17,13 @@ limitations under the License.
 #define TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
 
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
+#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
+#include "tensorflow/core/tpu/kernels/tpu_util_c_api.h"
 #include "tensorflow/core/tpu/libtftpu.h"
 #include "tensorflow/core/tpu/tpu_config_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
 
 // LINT.IfChange
 namespace tensorflow {
@@ -26,10 +31,21 @@ namespace tpu {
 
 Status InitializeTpuLibrary(void* library_handle);
 
+// TODO(frankchn): Separate out API functions from the loader.
 TfTpu_BaseFn* InitializeApiFn();
 
 TfTpu_ConfigApiFn* ConfigApiFn();
 
+TfTpu_MeshStateApiFn* MeshStateApiFn();
+
+TfTpu_CompileApiFn* CompileApiFn();
+
+TfTpu_ExecutorApiFn* ExecutorApiFn();
+
+TfTpu_NodeContextApiFn* NodeContextApiFn();
+
+TfTpu_UtilApiFn* UtilApiFn();
+
 }  // namespace tpu
 }  // namespace tensorflow
 // LINT.ThenChange(//tensorflow/core/tpu/tpu_library_loader_windows.cc)
diff --git a/tensorflow/core/tpu/tpu_library_loader_windows.cc b/tensorflow/core/tpu/tpu_library_loader_windows.cc
index e7c25df415e..7cf1b5cdb1d 100644
--- a/tensorflow/core/tpu/tpu_library_loader_windows.cc
+++ b/tensorflow/core/tpu/tpu_library_loader_windows.cc
@@ -27,6 +27,16 @@ TfTpu_BaseFn* InitializeApiFn() { return nullptr; }
 
 TfTpu_ConfigApiFn* ConfigApiFn() { return nullptr; }
 
+TfTpu_MeshStateApiFn* MeshStateApiFn() { return nullptr; }
+
+TfTpu_CompileApiFn* CompileApiFn() { return nullptr; }
+
+TfTpu_ExecutorApiFn* ExecutorApiFn() { return nullptr; }
+
+TfTpu_NodeContextApiFn* NodeContextApiFn() { return nullptr; }
+
+TfTpu_UtilApiFn* UtilApiFn() { return nullptr; }
+
 Status InitializeTpuLibrary(void* library_handle) {
   return errors::Unimplemented(
       "Loading TPU library is not supported on Windows.");
diff --git a/tensorflow/stream_executor/tpu/BUILD b/tensorflow/stream_executor/tpu/BUILD
index 964f36b82c7..bf88e9809d0 100644
--- a/tensorflow/stream_executor/tpu/BUILD
+++ b/tensorflow/stream_executor/tpu/BUILD
@@ -11,20 +11,25 @@ package(
 cc_library(
     name = "tpu_executor_c_api_hdrs",
     hdrs = ["tpu_executor_c_api.h"],
+    visibility = ["//visibility:public"],
     deps = [
         "//tensorflow/c:tf_attrtype",
-        "//tensorflow/c:tf_datatype",
         "//tensorflow/c:tf_status",
+        "//tensorflow/core/tpu:libtftpu_header",
         "//tensorflow/core/tpu/kernels:tpu_ops_common_c_api_hdrs",
     ],
+    alwayslink = True,
 )
 
 cc_library(
     name = "tpu_node_context_c_api_hdrs",
     hdrs = ["tpu_node_context_c_api.h"],
+    visibility = ["//visibility:public"],
     deps = [
         ":tpu_executor_c_api_hdrs",
+        "//tensorflow/core/tpu:libtftpu_header",
     ],
+    alwayslink = True,
 )
 
 cc_library(
@@ -65,6 +70,7 @@ cc_library(
         ":status_helper",
         ":tpu_executor_c_api_hdrs",
         ":tpu_stream_interface",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:stream",
     ],
 )
@@ -75,6 +81,7 @@ cc_library(
     deps = [
         ":tpu_executor_c_api_hdrs",
         "//tensorflow/core/platform:types",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:stream",
     ],
 )
@@ -94,6 +101,7 @@ cc_library(
         ":tpu_timer",
         "//tensorflow/c:tf_status",
         "//tensorflow/core:lib",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:stream",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -143,6 +151,7 @@ cc_library(
         "//tensorflow/compiler/xla/service:stream_pool",
         "//tensorflow/compiler/xla/service:transfer_manager",
         "//tensorflow/core:framework",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:device_memory_allocator",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/memory",
@@ -160,6 +169,7 @@ cc_library(
         ":tpu_platform_interface",
         "//tensorflow/c:tf_status",
         "//tensorflow/core/platform:types",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:stream",
         "@com_google_absl//absl/container:flat_hash_map",
     ],
@@ -191,6 +201,7 @@ cc_library(
         "//tensorflow/compiler/xla:xla_data_proto_cc",
         "//tensorflow/compiler/xla/service:shaped_buffer",
         "//tensorflow/compiler/xla/service:transfer_manager",
+        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor:stream",
     ],
 )
@@ -217,6 +228,7 @@ cc_library(
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor:multi_platform_manager",
         "//tensorflow/stream_executor:stream_executor_headers",
+        "//tensorflow/stream_executor:stream_executor_pimpl",
     ],
 )
 
diff --git a/tensorflow/stream_executor/tpu/tpu_executor.cc b/tensorflow/stream_executor/tpu/tpu_executor.cc
index 03cab5801e6..cb1410880eb 100644
--- a/tensorflow/stream_executor/tpu/tpu_executor.cc
+++ b/tensorflow/stream_executor/tpu/tpu_executor.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/c/tf_status.h"
 #include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/device_memory.h"
 #include "tensorflow/stream_executor/lib/status.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
@@ -33,63 +34,68 @@ namespace {
 using ::stream_executor::port::Status;
 }  // namespace
 
-TpuExecutor::~TpuExecutor() { TpuExecutor_Free(executor_); }
+TpuExecutor::~TpuExecutor() {
+  tpu::ExecutorApiFn()->TpuExecutor_FreeFn(executor_);
+}
 
 Status TpuExecutor::Init(int device_ordinal,
                          ::stream_executor::DeviceOptions device_options) {
   StatusHelper status;
   SE_DeviceOptions* options =
-      TpuExecutor_NewDeviceOptions(device_options.flags());
-  TpuExecutor_Init(executor_, device_ordinal, options, status.c_status);
-  TpuExecutor_FreeDeviceOptions(options);
+      tpu::ExecutorApiFn()->TpuExecutor_NewDeviceOptionsFn(
+          device_options.flags());
+  tpu::ExecutorApiFn()->TpuExecutor_InitFn(executor_, device_ordinal, options,
+                                           status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_FreeDeviceOptionsFn(options);
   return status.status();
 }
 
 int TpuExecutor::PlatformDeviceCount() {
-  return TpuExecutor_PlatformDeviceCount(executor_);
+  return tpu::ExecutorApiFn()->TpuExecutor_PlatformDeviceCountFn(executor_);
 }
 
 void TpuExecutor::SyncAndForgetFailedStreams() {
-  TpuExecutor_SyncAndForgetFailedStreams(executor_);
+  tpu::ExecutorApiFn()->TpuExecutor_SyncAndForgetFailedStreamsFn(executor_);
 }
 
 bool TpuExecutor::SynchronizeAllActivity() {
-  return TpuExecutor_SynchronizeAllActivity(executor_);
+  return tpu::ExecutorApiFn()->TpuExecutor_SynchronizeAllActivityFn(executor_);
 }
 
 Status TpuExecutor::BlockHostUntilDone(Stream* stream) {
   StatusHelper status;
-  TpuExecutor_BlockHostUntilDone(
+  tpu::ExecutorApiFn()->TpuExecutor_BlockHostUntilDoneFn(
       executor_, stream_map().at(stream->implementation()), status.c_status);
   return status.status();
 }
 
 Status TpuExecutor::BlockUntilDoneOrFailed() {
   StatusHelper status;
-  TpuExecutor_BlockUntilDoneOrFailed(executor_, status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_BlockUntilDoneOrFailedFn(executor_,
+                                                             status.c_status);
   return status.status();
 }
 
 Status TpuExecutor::GetStatus(Stream* stream) {
   StatusHelper status;
-  TpuExecutor_GetStatus(executor_, stream_map().at(stream->implementation()),
-                        status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_GetStatusFn(
+      executor_, stream_map().at(stream->implementation()), status.c_status);
   return status.status();
 }
 
 bool TpuExecutor::AllocateStream(Stream* stream) {
-  return TpuExecutor_AllocateStream(executor_,
-                                    stream_map().at(stream->implementation()));
+  return tpu::ExecutorApiFn()->TpuExecutor_AllocateStreamFn(
+      executor_, stream_map().at(stream->implementation()));
 }
 
 void TpuExecutor::DeallocateStream(Stream* stream) {
-  TpuExecutor_DeallocateStream(executor_,
-                               stream_map().at(stream->implementation()));
+  tpu::ExecutorApiFn()->TpuExecutor_DeallocateStreamFn(
+      executor_, stream_map().at(stream->implementation()));
   stream_map().erase(stream->implementation());
 }
 
 bool TpuExecutor::CreateStreamDependency(Stream* dependent, Stream* other) {
-  return TpuExecutor_CreateStreamDependency(
+  return tpu::ExecutorApiFn()->TpuExecutor_CreateStreamDependencyFn(
       executor_, stream_map().at(dependent->implementation()),
       stream_map().at(other->implementation()));
 }
@@ -104,15 +110,15 @@ bool TpuExecutor::AllocateTimer(Timer* timer) { return true; }
 void TpuExecutor::DeallocateTimer(Timer* timer) {}
 
 bool TpuExecutor::StartTimer(Stream* stream, ::stream_executor::Timer* timer) {
-  return TpuExecutor_StartTimer(executor_,
-                                stream_map().at(stream->implementation()),
-                                timer_map_.at(timer->implementation()));
+  return tpu::ExecutorApiFn()->TpuExecutor_StartTimerFn(
+      executor_, stream_map().at(stream->implementation()),
+      timer_map_.at(timer->implementation()));
 }
 
 bool TpuExecutor::StopTimer(Stream* stream, ::stream_executor::Timer* timer) {
-  return TpuExecutor_StopTimer(executor_,
-                               stream_map().at(stream->implementation()),
-                               timer_map_.at(timer->implementation()));
+  return tpu::ExecutorApiFn()->TpuExecutor_StopTimerFn(
+      executor_, stream_map().at(stream->implementation()),
+      timer_map_.at(timer->implementation()));
 }
 
 stream_executor::Event::Status TpuExecutor::PollForEventStatus(
@@ -148,7 +154,7 @@ Status TpuExecutor::WaitForEvent(Stream* stream,
 // Called by Timer::Timer
 std::unique_ptr<::stream_executor::internal::TimerInterface>
 TpuExecutor::GetTimerImplementation() {
-  SE_Timer* tpu_timer = TpuTimer_New(executor_);
+  SE_Timer* tpu_timer = tpu::ExecutorApiFn()->TpuTimer_NewFn(executor_);
   auto ptr = absl::make_unique(tpu_timer);
   timer_map_[ptr.get()] = tpu_timer;
   return ptr;
@@ -157,7 +163,7 @@ TpuExecutor::GetTimerImplementation() {
 // Called by Stream::Stream
 std::unique_ptr<::stream_executor::internal::StreamInterface>
 TpuExecutor::GetStreamImplementation() {
-  SE_Stream* tpu_stream = TpuStream_New(executor_);
+  SE_Stream* tpu_stream = tpu::ExecutorApiFn()->TpuStream_NewFn(executor_);
   auto ptr = absl::make_unique(tpu_stream);
   stream_map()[ptr.get()] = tpu_stream;
   return ptr;
@@ -166,34 +172,35 @@ TpuExecutor::GetStreamImplementation() {
 // Called by Event::Event
 std::unique_ptr<::stream_executor::internal::EventInterface>
 TpuExecutor::CreateEventImplementation() {
-  SE_Event* tpu_event = TpuEvent_New(executor_);
+  SE_Event* tpu_event = tpu::ExecutorApiFn()->TpuEvent_NewFn(executor_);
   auto ptr = absl::make_unique(tpu_event);
   event_map()[ptr.get()] = tpu_event;
   return ptr;
 }
 
 DeviceMemoryBase TpuExecutor::Allocate(uint64 size, int64 memory_space) {
-  SE_DeviceMemoryBase se_base =
-      TpuExecutor_Allocate(executor_, size, memory_space);
+  SE_DeviceMemoryBase se_base = tpu::ExecutorApiFn()->TpuExecutor_AllocateFn(
+      executor_, size, memory_space);
   return TpuConversions::SE_DeviceMemoryBaseToDeviceMemoryBase(se_base);
 }
 
 void TpuExecutor::Deallocate(const DeviceMemoryBase& memory) {
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(memory);
-  TpuExecutor_Deallocate(executor_, &se_base);
+  tpu::ExecutorApiFn()->TpuExecutor_DeallocateFn(executor_, &se_base);
 }
 
 void TpuExecutor::Deallocate(DeviceMemoryBase* memory) {
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(*memory);
-  TpuExecutor_Deallocate(executor_, &se_base);
+  tpu::ExecutorApiFn()->TpuExecutor_DeallocateFn(executor_, &se_base);
 }
 
 bool TpuExecutor::DeviceMemoryUsage(int64* free, int64* total) const {
   int64_t _free;
   int64_t _total;
-  if (TpuExecutor_DeviceMemoryUsage(executor_, &_free, &_total)) {
+  if (tpu::ExecutorApiFn()->TpuExecutor_DeviceMemoryUsageFn(executor_, &_free,
+                                                            &_total)) {
     *free = _free;
     *total = _total;
     return true;
@@ -204,7 +211,8 @@ bool TpuExecutor::DeviceMemoryUsage(int64* free, int64* total) const {
 absl::optional
 TpuExecutor::GetAllocatorStats() {
   SE_AllocatorStats c_stats;
-  if (TpuExecutor_GetAllocatorStats(executor_, &c_stats)) {
+  if (tpu::ExecutorApiFn()->TpuExecutor_GetAllocatorStatsFn(executor_,
+                                                            &c_stats)) {
     ::stream_executor::AllocatorStats stats;
     stats.num_allocs = c_stats.num_allocs;
     stats.bytes_in_use = c_stats.bytes_in_use;
@@ -226,31 +234,33 @@ TpuExecutor::GetAllocatorStats() {
 
 Status TpuExecutor::WaitForInfeedReady(int32 infeed_queue_index) {
   StatusHelper status;
-  TpuExecutor_WaitForInfeedReady(executor_, infeed_queue_index,
-                                 status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_WaitForInfeedReadyFn(
+      executor_, infeed_queue_index, status.c_status);
   return status.status();
 }
 
 Status TpuExecutor::WaitForOutfeedReady(int32 outfeed_queue_index) {
   StatusHelper status;
-  TpuExecutor_WaitForOutfeedReady(executor_, outfeed_queue_index,
-                                  status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_WaitForOutfeedReadyFn(
+      executor_, outfeed_queue_index, status.c_status);
   return status.status();
 }
 
 void TpuExecutor::DequeueOutfeed(int32 outfeed_queue_index,
                                  absl::Span bytes, StatusCallback done) {
   StatusHelper status;
-  TpuExecutor_DequeueOutfeed(executor_, outfeed_queue_index, bytes.data(),
-                             bytes.size(), status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_DequeueOutfeedFn(
+      executor_, outfeed_queue_index, bytes.data(), bytes.size(),
+      status.c_status);
   done(status.status());
 }
 
 Status TpuExecutor::EnqueueInfeed(int32 infeed_queue_index,
                                   absl::Span bytes) {
   StatusHelper status;
-  TpuExecutor_EnqueueInfeed(executor_, infeed_queue_index, bytes.data(),
-                            bytes.size(), status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_EnqueueInfeedFn(
+      executor_, infeed_queue_index, bytes.data(), bytes.size(),
+      status.c_status);
   return status.status();
 }
 
@@ -259,9 +269,9 @@ bool TpuExecutor::Memcpy(Stream* stream, void* host_dst,
                          uint64 size) {
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(device_src);
-  return TpuExecutor_MemcpyToHost(executor_,
-                                  stream_map().at(stream->implementation()),
-                                  host_dst, &se_base, size);
+  return tpu::ExecutorApiFn()->TpuExecutor_MemcpyToHostFn(
+      executor_, stream_map().at(stream->implementation()), host_dst, &se_base,
+      size);
 }
 
 bool TpuExecutor::Memcpy(Stream* stream,
@@ -269,9 +279,9 @@ bool TpuExecutor::Memcpy(Stream* stream,
                          const void* host_src, uint64 size) {
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(*device_dst);
-  return TpuExecutor_MemcpyFromHost(executor_,
-                                    stream_map().at(stream->implementation()),
-                                    &se_base, host_src, size);
+  return tpu::ExecutorApiFn()->TpuExecutor_MemcpyFromHostFn(
+      executor_, stream_map().at(stream->implementation()), &se_base, host_src,
+      size);
 }
 
 Status TpuExecutor::SynchronousMemcpy(
@@ -280,8 +290,8 @@ Status TpuExecutor::SynchronousMemcpy(
   StatusHelper status;
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(*device_dst);
-  TpuExecutor_SynchronousMemcpyFromHost(executor_, &se_base, host_src, size,
-                                        status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_SynchronousMemcpyFromHostFn(
+      executor_, &se_base, host_src, size, status.c_status);
   return status.status();
 }
 
@@ -291,8 +301,8 @@ Status TpuExecutor::SynchronousMemcpy(
   StatusHelper status;
   SE_DeviceMemoryBase se_base =
       TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(device_src);
-  TpuExecutor_SynchronousMemcpyToHost(executor_, host_dst, &se_base, size,
-                                      status.c_status);
+  tpu::ExecutorApiFn()->TpuExecutor_SynchronousMemcpyToHostFn(
+      executor_, host_dst, &se_base, size, status.c_status);
   return status.status();
 }
 
@@ -316,8 +326,8 @@ struct HostCallbackContext {
 SE_Status* HostCallbackTrampoline(void* ctx) {
   HostCallbackContext* host_ctx = reinterpret_cast(ctx);
   Status status = host_ctx->callback();
-  SE_Status* c_status =
-      TpuStatus_Create(status.code(), status.error_message().c_str());
+  SE_Status* c_status = tpu::ExecutorApiFn()->TpuStatus_CreateFn(
+      status.code(), status.error_message().c_str());
   delete host_ctx;
   return c_status;
 }
@@ -325,18 +335,21 @@ SE_Status* HostCallbackTrampoline(void* ctx) {
 bool TpuExecutor::HostCallback(Stream* stream,
                                std::function callback) {
   HostCallbackContext* ctx = new HostCallbackContext{callback};
-  return TpuExecutor_HostCallback(executor_,
-                                  stream_map().at(stream->implementation()),
-                                  &HostCallbackTrampoline, ctx);
+  return tpu::ExecutorApiFn()->TpuExecutor_HostCallbackFn(
+      executor_, stream_map().at(stream->implementation()),
+      &HostCallbackTrampoline, ctx);
 }
 
 TpuExecutor::StatusOr>
 TpuExecutor::CreateDeviceDescription() const {
   StatusHelper status;
-  SE_DeviceDescription* description = TpuDeviceDescription_New();
-  auto cleanup = tensorflow::gtl::MakeCleanup(
-      [description]() { TpuDeviceDescription_Free(description); });
-  TpuExecutor_CreateDeviceDescription(executor_, description, status.c_status);
+  SE_DeviceDescription* description =
+      tpu::ExecutorApiFn()->TpuDeviceDescription_NewFn();
+  auto cleanup = tensorflow::gtl::MakeCleanup([description]() {
+    tpu::ExecutorApiFn()->TpuDeviceDescription_FreeFn(description);
+  });
+  tpu::ExecutorApiFn()->TpuExecutor_CreateDeviceDescriptionFn(
+      executor_, description, status.c_status);
   if (status.status().ok()) {
     stream_executor::internal::DeviceDescriptionBuilder builder;
     CHECK_NE(description->device_vendor, nullptr);
diff --git a/tensorflow/stream_executor/tpu/tpu_executor_c_api.h b/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
index e77e09bb911..eee69a35b23 100644
--- a/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
+++ b/tensorflow/stream_executor/tpu/tpu_executor_c_api.h
@@ -20,9 +20,9 @@ limitations under the License.
 #include 
 
 #include "tensorflow/c/tf_attrtype.h"
-#include "tensorflow/c/tf_datatype.h"
 #include "tensorflow/c/tf_status.h"
 #include "tensorflow/core/tpu/kernels/tpu_ops_common_c_api.h"
+#include "tensorflow/core/tpu/libtftpu.h"
 
 typedef struct SE_Platform SE_Platform;
 typedef struct SE_StreamExecutor SE_StreamExecutor;
@@ -292,6 +292,96 @@ void TpuTransferManager_WriteSingleTupleIndexTable(
 
 XLA_ComputationPlacer* TpuComputationPlacer_New();
 void TpuComputationPlacer_Free(XLA_ComputationPlacer* placer);
+
+struct TfTpu_ExecutorApiFn {
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_Initialize);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_Initialized);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_GetExecutor);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_Id);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_VisibleDeviceCount);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_TpuMemoryLimit);
+  TFTPU_ADD_FN_IN_STRUCT(TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_Init);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_PlatformDeviceCount);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_Allocate);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_Deallocate);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetAllocatorStats);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeviceMemoryUsage);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_AllocateStream);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateStream);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_CreateStreamDependency);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_GetStatus);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_AllocateEvent);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateEvent);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_PollForEventStatus);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_RecordEvent);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_WaitForEvent);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_AllocateTimer);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DeallocateTimer);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_StartTimer);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_StopTimer);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_SynchronousMemcpyToHost);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_SynchronousMemcpyFromHost);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_MemcpyToHost);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_MemcpyFromHost);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_EnqueueInfeed);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_DequeueOutfeed);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_WaitForInfeedReady);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_WaitForOutfeedReady);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_BlockHostUntilDone);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_BlockUntilDoneOrFailed);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_SyncAndForgetFailedStreams);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_SynchronizeAllActivity);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_Stream);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_Status);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_IsSameSharedMemoryLocation);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStream_TpuEnqueueOnDeviceSendRecvLocal);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuEvent_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuEvent_Free);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuTimer_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTimer_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTimer_Nanoseconds);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTimer_Microseconds);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_Create);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_Message);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_Code);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStatus_Ok);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuStreamExecutorConfig_Default);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStreamExecutorConfig_SetOrdinal);
+  TFTPU_ADD_FN_IN_STRUCT(TpuStreamExecutorConfig_Free);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuDeviceDescription_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuDeviceDescription_Free);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_CreateDeviceDescription);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_NewDeviceOptions);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_FreeDeviceOptions);
+  TFTPU_ADD_FN_IN_STRUCT(TpuExecutor_HostCallback);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_PlatformId);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_HostShapeToDeviceShape);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_TransferLiteralToDeviceAsync);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_TransferLiteralFromDevice);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_GetByteSizeRequirement);
+  TFTPU_ADD_FN_IN_STRUCT(TpuTransferManager_WriteSingleTupleIndexTable);
+
+  TFTPU_ADD_FN_IN_STRUCT(TpuComputationPlacer_New);
+  TFTPU_ADD_FN_IN_STRUCT(TpuComputationPlacer_Free);
+};
 }
 
 // extern "C"
diff --git a/tensorflow/stream_executor/tpu/tpu_node_context.cc b/tensorflow/stream_executor/tpu/tpu_node_context.cc
index 35a9eb53bcd..356ede40fb3 100644
--- a/tensorflow/stream_executor/tpu/tpu_node_context.cc
+++ b/tensorflow/stream_executor/tpu/tpu_node_context.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/backend.h"
 #include "tensorflow/compiler/xla/service/platform_util.h"
 #include "tensorflow/compiler/xla/service/transfer_manager.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 #include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
@@ -32,15 +33,18 @@ StatusOr> TpuNodeContext::Create(
     int device_ordinal) {
   StatusHelper status;
   XLA_TpuNodeContext* node_context =
-      TpuNodeContext_Create(device_ordinal, status.c_status);
+      tpu::NodeContextApiFn()->TpuNodeContext_CreateFn(device_ordinal,
+                                                       status.c_status);
   if (!status.status().ok()) {
-    TpuNodeContext_Free(node_context);
+    tpu::NodeContextApiFn()->TpuNodeContext_FreeFn(node_context);
     return status.status();
   }
   return std::make_unique(device_ordinal, node_context);
 }
 
-TpuNodeContext::~TpuNodeContext() { TpuNodeContext_Free(node_context_); }
+TpuNodeContext::~TpuNodeContext() {
+  tpu::NodeContextApiFn()->TpuNodeContext_FreeFn(node_context_);
+}
 
 /* static */
 Status TpuNodeContext::Initialize(int device_ordinal) {
@@ -52,14 +56,14 @@ Status TpuNodeContext::Initialize(int device_ordinal) {
 /* static */
 Status TpuNodeContext::StopChipHeartbeats() {
   StatusHelper status;
-  TpuNodeContext_StopChipHeartbeats(status.c_status);
+  tpu::NodeContextApiFn()->TpuNodeContext_StopChipHeartbeatsFn(status.c_status);
   return status.status();
 }
 
 /* static */
 Status TpuNodeContext::CloseTpuHost() {
   StatusHelper status;
-  TpuNodeContext_CloseTpuHost(status.c_status);
+  tpu::NodeContextApiFn()->TpuNodeContext_CloseTpuHostFn(status.c_status);
   return status.status();
 }
 
diff --git a/tensorflow/stream_executor/tpu/tpu_node_context_c_api.h b/tensorflow/stream_executor/tpu/tpu_node_context_c_api.h
index e5092d4842b..d47fdf37a46 100644
--- a/tensorflow/stream_executor/tpu/tpu_node_context_c_api.h
+++ b/tensorflow/stream_executor/tpu/tpu_node_context_c_api.h
@@ -15,10 +15,13 @@ limitations under the License.
 #ifndef TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_NODE_CONTEXT_C_API_H_
 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_NODE_CONTEXT_C_API_H_
 
+#include "tensorflow/core/tpu/libtftpu.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 
 typedef struct XLA_TpuNodeContext XLA_TpuNodeContext;
 
+extern "C" {
+
 XLA_TpuNodeContext* TpuNodeContext_Create(int device_ordinal,
                                           SE_Status* status);
 void TpuNodeContext_Free(XLA_TpuNodeContext* node_context);
@@ -28,4 +31,13 @@ void TpuNodeContext_Initialize(int device_ordinal, SE_Status* status);
 void TpuNodeContext_StopChipHeartbeats(SE_Status* status);
 void TpuNodeContext_CloseTpuHost(SE_Status* status);
 
+}  // extern "C"
+
+struct TfTpu_NodeContextApiFn {
+  TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_Create);
+  TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_Free);
+  TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_StopChipHeartbeats);
+  TFTPU_ADD_FN_IN_STRUCT(TpuNodeContext_CloseTpuHost);
+};
+
 #endif  // TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_NODE_CONTEXT_C_API_H_
diff --git a/tensorflow/stream_executor/tpu/tpu_platform.cc b/tensorflow/stream_executor/tpu/tpu_platform.cc
index c65d8a4207a..13a845829c1 100644
--- a/tensorflow/stream_executor/tpu/tpu_platform.cc
+++ b/tensorflow/stream_executor/tpu/tpu_platform.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/tpu/tpu_platform.h"
 
 #include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/platform.h"
 #include "tensorflow/stream_executor/tpu/status_helper.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor.h"
@@ -30,7 +31,9 @@ using Status = ::stream_executor::port::Status;
 template 
 using StatusOr = ::stream_executor::port::StatusOr;
 
-TpuPlatform::TpuPlatform() { platform_ = TpuPlatform_New(); }
+TpuPlatform::TpuPlatform() {
+  platform_ = tpu::ExecutorApiFn()->TpuPlatform_NewFn();
+}
 
 TpuPlatform* TpuPlatform::GetRegisteredPlatform() {
   return tpu_registered_platform;
@@ -53,8 +56,8 @@ Status TpuPlatform::Initialize(
     i++;
   }
 
-  TpuPlatform_Initialize(platform_, options_size, options_key, options_value,
-                         status.c_status);
+  tpu::ExecutorApiFn()->TpuPlatform_InitializeFn(
+      platform_, options_size, options_key, options_value, status.c_status);
 
   free(options_key);
   free(options_value);
@@ -62,10 +65,16 @@ Status TpuPlatform::Initialize(
   return status.status();
 }
 
-TpuPlatform::~TpuPlatform() { TpuPlatform_Free(platform_); }
+bool TpuPlatform::Initialized() const {
+  return tpu::ExecutorApiFn()->TpuPlatform_InitializedFn(platform_);
+}
+
+TpuPlatform::~TpuPlatform() {
+  tpu::ExecutorApiFn()->TpuPlatform_FreeFn(platform_);
+}
 
 int TpuPlatform::VisibleDeviceCount() const {
-  return TpuPlatform_VisibleDeviceCount(platform_);
+  return tpu::ExecutorApiFn()->TpuPlatform_VisibleDeviceCountFn(platform_);
 }
 
 StatusOr<::stream_executor::StreamExecutor*> TpuPlatform::GetExecutor(
@@ -77,14 +86,16 @@ StatusOr<::stream_executor::StreamExecutor*> TpuPlatform::GetExecutor(
 StatusOr>
 TpuPlatform::GetUncachedExecutor(
     const ::stream_executor::StreamExecutorConfig& config) {
-  SE_StreamExecutorConfig* c_config = TpuStreamExecutorConfig_Default();
+  SE_StreamExecutorConfig* c_config =
+      tpu::ExecutorApiFn()->TpuStreamExecutorConfig_DefaultFn();
 
-  TpuStreamExecutorConfig_SetOrdinal(c_config, config.ordinal);
+  tpu::ExecutorApiFn()->TpuStreamExecutorConfig_SetOrdinalFn(c_config,
+                                                             config.ordinal);
 
   StatusHelper status;
-  SE_StreamExecutor* executor =
-      TpuPlatform_GetExecutor(platform_, c_config, status.c_status);
-  TpuStreamExecutorConfig_Free(c_config);
+  SE_StreamExecutor* executor = tpu::ExecutorApiFn()->TpuPlatform_GetExecutorFn(
+      platform_, c_config, status.c_status);
+  tpu::ExecutorApiFn()->TpuStreamExecutorConfig_FreeFn(c_config);
   if (!status.ok()) {
     return status.status();
   }
@@ -103,27 +114,24 @@ const std::string& TpuPlatform::Name() const {
 }
 
 int64 TpuPlatform::TpuMemoryLimit() {
-  return TpuPlatform_TpuMemoryLimit(platform_);
+  return tpu::ExecutorApiFn()->TpuPlatform_TpuMemoryLimitFn(platform_);
 }
 
 bool TpuPlatform::ShouldRegisterTpuDeviceToDeviceCopy() {
-  return TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopy(platform_);
+  return tpu::ExecutorApiFn()
+      ->TpuPlatform_ShouldRegisterTpuDeviceToDeviceCopyFn(platform_);
+}
+
+void RegisterTpuPlatform() {
+  static bool tpu_platform_registered = false;
+  if (!tpu_platform_registered) {
+    tensorflow::tpu_registered_platform = new tensorflow::TpuPlatform();
+    std::unique_ptr platform(
+        tensorflow::tpu_registered_platform);
+    SE_CHECK_OK(stream_executor::MultiPlatformManager::RegisterPlatform(
+        std::move(platform)));
+    tpu_platform_registered = true;
+  }
 }
 
 }  // namespace tensorflow
-
-void RegisterTpuPlatform() {
-  tensorflow::tpu_registered_platform = new tensorflow::TpuPlatform();
-  std::unique_ptr platform(
-      tensorflow::tpu_registered_platform);
-  SE_CHECK_OK(stream_executor::MultiPlatformManager::RegisterPlatform(
-      std::move(platform)));
-}
-
-REGISTER_MODULE_INITIALIZER(tpu_platform, RegisterTpuPlatform());
-
-// Note that module initialization sequencing is not supported in the
-// open-source project, so this will be a no-op there.
-REGISTER_MODULE_INITIALIZER_SEQUENCE(tpu_platform, multi_platform_manager);
-REGISTER_MODULE_INITIALIZER_SEQUENCE(multi_platform_manager_listener,
-                                     tpu_platform);
diff --git a/tensorflow/stream_executor/tpu/tpu_platform.h b/tensorflow/stream_executor/tpu/tpu_platform.h
index 6fdd8d15aa4..c2673ab9288 100644
--- a/tensorflow/stream_executor/tpu/tpu_platform.h
+++ b/tensorflow/stream_executor/tpu/tpu_platform.h
@@ -60,9 +60,7 @@ class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface {
 
   bool ShouldRegisterTpuDeviceToDeviceCopy() override;
 
-  bool Initialized() const override {
-    return TpuPlatform_Initialized(platform_);
-  }
+  bool Initialized() const override;
 
   Status Initialize(
       const std::map& platform_options) override;
@@ -124,6 +122,8 @@ class TpuPlatform : public ::tensorflow::tpu::TpuPlatformInterface {
   EventMap event_map_;
 };
 
+void RegisterTpuPlatform();
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_PLATFORM_H_
diff --git a/tensorflow/stream_executor/tpu/tpu_stream.h b/tensorflow/stream_executor/tpu/tpu_stream.h
index 5c71c0535f3..e1aa1164248 100644
--- a/tensorflow/stream_executor/tpu/tpu_stream.h
+++ b/tensorflow/stream_executor/tpu/tpu_stream.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
 
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/stream_executor_internal.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
 #include "tensorflow/stream_executor/tpu/status_helper.h"
@@ -27,23 +28,27 @@ class TpuStream : public tensorflow::tpu::TpuStreamInterface {
   using Status = stream_executor::port::Status;
 
   explicit TpuStream(SE_Stream* stream) : stream_(stream) {}
-  ~TpuStream() override { TpuStream_Free(stream_); }
+  ~TpuStream() override {
+    tensorflow::tpu::ExecutorApiFn()->TpuStream_FreeFn(stream_);
+  }
 
   bool IsSameSharedMemoryLocation(
       tensorflow::tpu::TpuStreamInterface* other) override {
-    return TpuStream_IsSameSharedMemoryLocation(
-        stream_, static_cast(other)->stream_);
+    return tensorflow::tpu::ExecutorApiFn()
+        ->TpuStream_IsSameSharedMemoryLocationFn(
+            stream_, static_cast(other)->stream_);
   }
 
   Status EnqueueOnTpuDeviceSendRecvLocal(
       stream_executor::DeviceMemoryBase send_buffer,
       stream_executor::DeviceMemoryBase recv_buffer) override {
     StatusHelper status;
-    TpuStream_TpuEnqueueOnDeviceSendRecvLocal(
-        stream_,
-        TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(send_buffer),
-        TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(recv_buffer),
-        status.c_status);
+    tensorflow::tpu::ExecutorApiFn()
+        ->TpuStream_TpuEnqueueOnDeviceSendRecvLocalFn(
+            stream_,
+            TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(send_buffer),
+            TpuConversions::DeviceMemoryBaseToSE_DeviceMemoryBase(recv_buffer),
+            status.c_status);
     return status.status();
   }
 
@@ -54,7 +59,9 @@ class TpuStream : public tensorflow::tpu::TpuStreamInterface {
 class TpuEvent : public ::stream_executor::internal::EventInterface {
  public:
   explicit TpuEvent(SE_Event* event) : event_(event) {}
-  ~TpuEvent() override { TpuEvent_Free(event_); }
+  ~TpuEvent() override {
+    tensorflow::tpu::ExecutorApiFn()->TpuEvent_FreeFn(event_);
+  }
 
  private:
   SE_Event* event_;
diff --git a/tensorflow/stream_executor/tpu/tpu_timer.h b/tensorflow/stream_executor/tpu/tpu_timer.h
index 246a0b7eb32..d7f8f660b37 100644
--- a/tensorflow/stream_executor/tpu/tpu_timer.h
+++ b/tensorflow/stream_executor/tpu/tpu_timer.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_TIMER_H_
 
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/stream_executor_internal.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 
@@ -25,9 +26,15 @@ namespace tensorflow {
 class TpuTimer : public ::stream_executor::internal::TimerInterface {
  public:
   explicit TpuTimer(SE_Timer* timer) : timer_(timer) {}
-  ~TpuTimer() override { TpuTimer_Free(timer_); }
-  uint64 Microseconds() const override { return TpuTimer_Microseconds(timer_); }
-  uint64 Nanoseconds() const override { return TpuTimer_Nanoseconds(timer_); }
+  ~TpuTimer() override {
+    tensorflow::tpu::ExecutorApiFn()->TpuTimer_FreeFn(timer_);
+  }
+  uint64 Microseconds() const override {
+    return tensorflow::tpu::ExecutorApiFn()->TpuTimer_MicrosecondsFn(timer_);
+  }
+  uint64 Nanoseconds() const override {
+    return tensorflow::tpu::ExecutorApiFn()->TpuTimer_NanosecondsFn(timer_);
+  }
 
  private:
   SE_Timer* timer_;
diff --git a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
index 4bedc251413..934fabbf54d 100644
--- a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
+++ b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/device_memory.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"
@@ -29,10 +30,12 @@ namespace tensorflow {
 using Status = stream_executor::port::Status;
 
 TpuTransferManager::TpuTransferManager() {
-  manager_ = TpuTransferManager_New();
+  manager_ = tpu::ExecutorApiFn()->TpuTransferManager_NewFn();
 }
 
-TpuTransferManager::~TpuTransferManager() { TpuTransferManager_Free(manager_); }
+TpuTransferManager::~TpuTransferManager() {
+  tpu::ExecutorApiFn()->TpuTransferManager_FreeFn(manager_);
+}
 
 stream_executor::Platform::Id TpuTransferManager::PlatformId() const {
   return TpuPlatform::kId;
@@ -45,8 +48,8 @@ xla::Shape TpuTransferManager::HostShapeToDeviceShape(
 
   TpuConversions::XlaShapeToCShape(host_shape, &c_host_shape);
 
-  TpuTransferManager_HostShapeToDeviceShape(manager_, &c_host_shape,
-                                            &c_device_shape);
+  tpu::ExecutorApiFn()->TpuTransferManager_HostShapeToDeviceShapeFn(
+      manager_, &c_host_shape, &c_device_shape);
   xla::Shape device_shape = TpuConversions::CShapeToXlaShape(&c_device_shape);
   TpuConversions::CShapeCleanup(&c_host_shape);
   TpuConversions::CShapeCleanup(&c_device_shape);
@@ -66,7 +69,7 @@ Status TpuTransferManager::TransferLiteralToDeviceAsync(
   TpuConversions::XLAShapedBufferToCShapedBuffer(device_buffer,
                                                  &c_device_buffer);
 
-  TpuTransferManager_TransferLiteralToDeviceAsync(
+  tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralToDeviceAsyncFn(
       manager_,
       TpuPlatform::GetRegisteredPlatform()->stream_map()->at(
           stream->implementation()),
@@ -112,7 +115,7 @@ void TpuTransferManager::TransferLiteralFromDevice(
   XLA_Literal c_literal;
   TpuConversions::XLALiteralToCLiteral(literal, &c_literal);
 
-  TpuTransferManager_TransferLiteralFromDevice(
+  tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralFromDeviceFn(
       manager_,
       TpuPlatform::GetRegisteredPlatform()->stream_map()->at(
           stream->implementation()),
@@ -127,7 +130,8 @@ int64 TpuTransferManager::GetByteSizeRequirement(
   TpuConversions::XlaShapeToCShape(shape, &c_shape);
 
   int64 size_in_bytes =
-      TpuTransferManager_GetByteSizeRequirement(manager_, &c_shape);
+      tpu::ExecutorApiFn()->TpuTransferManager_GetByteSizeRequirementFn(
+          manager_, &c_shape);
 
   TpuConversions::CShapeCleanup(&c_shape);
   return size_in_bytes;
@@ -151,7 +155,7 @@ Status TpuTransferManager::WriteSingleTupleIndexTable(
                                   region->payload()};
   StatusHelper status;
 
-  TpuTransferManager_WriteSingleTupleIndexTable(
+  tpu::ExecutorApiFn()->TpuTransferManager_WriteSingleTupleIndexTableFn(
       manager_,
       TpuPlatform::GetRegisteredPlatform()->stream_map()->at(
           stream->implementation()),
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index 5da15b0a4d6..4a4f8837867 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -2899,6 +2899,13 @@ def if_mlir(if_true, if_false = []):
         "//conditions:default": if_false,
     })
 
+def if_tpu(if_true, if_false = []):
+    """Shorthand for select()ing whether to build for TPUs."""
+    return select({
+        str(Label("//tensorflow:with_tpu_support")): if_true,
+        "//conditions:default": if_false,
+    })
+
 def tfcompile_target_cpu():
     return ""
 

From 2ff1c5a31be8a3d38e3bccd59e54cbfd5cabe5cd Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen 
Date: Thu, 18 Jun 2020 10:59:05 -0700
Subject: [PATCH 0511/1390] Import initialization graph in SignatureDef
 SavedModels as an MLIR function in TF saved model dialect.

PiperOrigin-RevId: 317137903
Change-Id: I7cbded06b3deafa30d3b3e3dad98cc8f056dd4e3
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |   4 +-
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  25 +++
 .../mlir/tensorflow/ir/tf_saved_model_ops.td  |  24 +++
 .../tests/tf_saved_model/common_v1.py         |   1 +
 .../tests/tf_saved_model/hash_table_v1.py     |  92 +++++++++++
 .../tensorflow/tests/tf_saved_model_ops.mlir  |   5 +
 .../tests/tf_saved_model_ops_invalid.mlir     |  33 ++++
 .../mlir/tensorflow/translate/import_model.cc | 149 +++++++++++++-----
 8 files changed, 290 insertions(+), 43 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 904ccb7e820..17ed0e36a28 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -661,7 +661,9 @@ cc_library(
         ":tensorflow_types",
         ":translate_utils",
         "//tensorflow/cc/saved_model:bundle_v2",
+        "//tensorflow/cc/saved_model:constants",
         "//tensorflow/cc/saved_model:loader_lite",
+        "//tensorflow/cc/saved_model:loader_util",
         "//tensorflow/compiler/jit:shape_inference_helpers",
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/tf2xla:functionalize_control_flow",
@@ -673,6 +675,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/utils:transitive_fanin",
+        "//tensorflow/core/platform:protobuf_internal",
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
@@ -682,7 +685,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:StandardOps",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index 140a778770c..6af70158e14 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -76,6 +76,23 @@ static LogicalResult Verify(GlobalTensorOp global_tensor) {
   return success();
 }
 
+static LogicalResult Verify(SessionInitializerOp session_initializer) {
+  mlir::SymbolTable symbol_table(
+      session_initializer.getParentOfType());
+
+  auto init_func_op =
+      symbol_table.lookup(session_initializer.initializer());
+  if (!init_func_op)
+    return session_initializer.emitOpError()
+           << "the initializer function does not exist";
+
+  if (!init_func_op.getType().getResults().empty())
+    return session_initializer.emitOpError()
+           << "the initializer function should have no output";
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc.inc"
 
@@ -220,6 +237,14 @@ static LogicalResult VerifySavedModelModule(
       }
     }
   }
+
+  auto session_initializers = module.getOps();
+  if (std::distance(session_initializers.begin(), session_initializers.end()) >
+      1) {
+    return (*++session_initializers.begin()).emitError()
+           << "there must be no more than one session_initializer op";
+  }
+
   SymbolTable symbol_table(module);
   auto symbol_uses = SymbolTable::getSymbolUses(&module.getBodyRegion());
   if (!symbol_uses.hasValue()) {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
index 4431a160edf..497f4d90cb9 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
@@ -128,4 +128,28 @@ def TfSavedModel_GlobalTensorOp : TfSavedModel_Op<"global_tensor"> {
   let verifier = [{ return Verify(*this); }];
 }
 
+def TfSavedModel_SessionInitializerOp: TfSavedModel_Op<"session_initializer"> {
+  let summary = "Initializes TensorFlow session state.";
+  let description = [{
+    Represents a session initializer function initializes TensorFlow session
+    state. It is used to initialize resources in the saved model before calling
+    any exported functions. There must be no more than one session initializer
+    in a saved model.
+
+    The `initializer` represents the initialization function. The function have
+    no output and this function should be only called once.
+
+    This is used, for example, to initialize hash tables stored in resources and
+    accessed by resource name (rather than as resource handles or bound inputs
+    which is how `global_tensor`s are referenced)
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$initializer
+  );
+
+
+  let verifier = [{ return Verify(*this); }];
+}
+
 #endif // SAVED_MODEL_DIALECT
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
index 7171f63bb05..51ccbeb1fbd 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
@@ -84,6 +84,7 @@ def do_test(signature_def_map, show_debug_info=False):
     builder.add_meta_graph_and_variables(
         sess, [tf.saved_model.tag_constants.SERVING],
         signature_def_map,
+        main_op=tf.tables_initializer(),
         strip_default_attrs=True)
     builder.save()
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
new file mode 100644
index 00000000000..64847434b82
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
@@ -0,0 +1,92 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# RUN: %p/hash_table_v1 | FileCheck %s
+
+# pylint: disable=missing-docstring,line-too-long
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v1 as tf
+from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
+
+# Verify that the tf.versions attribute exists. It is difficult to enforce
+# contents, since the version numbers change over time. The conversion logic
+# itself is verified in the common graphdef converter, so here just assert
+# it is being invoked.
+# CHECK: module
+# CHECK-SAME: tf.versions
+# CHECK-SAME: bad_consumers
+# CHECK-SAME: min_consumer
+# CHECK-SAME: producer
+
+# CHECK: "tf_saved_model.session_initializer"() {initializer = [[init:@.*]]} : () -> ()
+# CHECK: "tf_saved_model.global_tensor"()
+
+# CHECK:      func {{@[a-zA-Z_0-9]+}}(
+# CHECK-SAME: [[ARG0:%.*]]: tensor
+# CHECK-SAME: [[ARG1:%.*]]: tensor ()
+
   // Representation for constants: (immutable) global tensor.
   // CHECK: tf_saved_model.global_tensor
   "tf_saved_model.global_tensor"() {
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index c055c6c9f56..544600cf6b8 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -258,3 +258,36 @@ module attributes {tf_saved_model.semantics} {
   // expected-error@+1 {{'type' attribute for immutable 'tf_saved_model.global_tensor' should have a static shape}}
   "tf_saved_model.global_tensor"() { sym_name = "v", type = tensor, value = dense<1.> : tensor<1xf32> } : () -> ()
 }
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function does not exist}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function should have no output}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  // expected-error@+1 {{there must be no more than one session_initializer op}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 820d0ce31fb..3cff4217215 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -60,6 +60,8 @@ limitations under the License.
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
@@ -99,6 +101,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/protobuf_internal.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
@@ -116,6 +119,7 @@ using mlir::NamedAttrList;
 using mlir::TensorType;
 using mlir::TF::VarHandleOp;
 using mlir::tf_saved_model::GlobalTensorOp;
+using mlir::tf_saved_model::SessionInitializerOp;
 using stream_executor::port::StatusOr;
 
 namespace {
@@ -2955,6 +2959,13 @@ void SortSavedModelModule(mlir::ModuleOp module) {
     named_global_tensor.global_tensor.getOperation()->moveBefore(
         &module.getBody()->front());
   }
+
+  auto initializers = module.getOps();
+  if (!initializers.empty()) {
+    (*initializers.begin())
+        .getOperation()
+        ->moveBefore(&module.getBody()->front());
+  }
 }
 
 Status CreateSavedModelIR(
@@ -3241,17 +3252,29 @@ class SavedModelSignatureDefImporter {
                                  absl::Span exported_names,
                                  mlir::MLIRContext* context)
       : bundle_(bundle),
+        flib_def_(OpRegistry::Global(), graph_def().library()),
+        debug_info_(),
         exported_names_(exported_names),
-        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {}
+        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {
+    // debug_info might not be loaded with loader_lite.
+    if (bundle_.debug_info != nullptr) debug_info_ = *bundle_.debug_info;
+  }
 
   // Converts the SavedModel to the SavedModel dialect. Creates an MLIR function
   // for each signature.
   StatusOr ConvertSignatures();
-  Status ConvertSignature(const GraphDef& graphdef,
-                          const std::string& sig_def_key,
-                          const SignatureDef& signature_def,
-                          const GraphDebugInfo& debug_info,
-                          const FunctionLibraryDefinition& flib_def);
+  Status ConvertSignature(const std::string& sig_def_key,
+                          const SignatureDef& signature_def);
+
+  // Converts the initialization graph in the SavedModel to an MLIR function.
+  Status ConvertInitializer();
+
+  // Converts a graph with feeds and fetches to an MLIR function.
+  StatusOr ConvertGraph(
+      const std::string& name,
+      const std::vector>& inputs,
+      const std::vector>& outputs,
+      const std::vector control_outputs);
 
   // Creates GlobalTensorOp for each variable and moves each VarHandle op to
   // the enclosing function's arguments.
@@ -3273,18 +3296,62 @@ class SavedModelSignatureDefImporter {
   GraphImportConfig::InputArrays ParseInputArrays(
       const std::vector>& inputs);
 
+  const GraphDef& graph_def() const {
+    return bundle_.meta_graph_def.graph_def();
+  }
+  const FunctionLibraryDefinition& flib_def() const { return flib_def_; }
+  const GraphDebugInfo& debug_info() const { return debug_info_; }
+
   const SavedModelBundle& bundle_;
+  FunctionLibraryDefinition flib_def_;
+  GraphDebugInfo debug_info_;
   absl::Span exported_names_;
   mlir::OwningModuleRef module_;
 };
 
+Status SavedModelSignatureDefImporter::ConvertInitializer() {
+  std::vector asset_file_defs;
+  TF_RETURN_IF_ERROR(
+      internal::GetAssetFileDefs(bundle_.meta_graph_def, &asset_file_defs));
+
+  if (!asset_file_defs.empty())
+    return errors::Unimplemented(
+        absl::StrCat("Assets are not supported in signaturedef importer"));
+
+  std::string init_node_name;
+  TF_RETURN_IF_ERROR(
+      internal::GetInitOp("", bundle_.meta_graph_def, &init_node_name));
+
+  if (init_node_name.empty()) return Status::OK();
+
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(init_node_name, {}, {}, {init_node_name}));
+
+  mlir::SymbolTable symbol_table(*sub_module);
+
+  auto init_func_op = symbol_table.lookup(init_node_name);
+
+  init_func_op.removeAttr("tf.entry_function");
+
+  mlir::OpBuilder builder(module_->getBodyRegion());
+
+  builder.create(
+      module_->getLoc(), builder.getSymbolRefAttr(init_func_op.getName()));
+
+  // Move the converted functions to top level MLIR module.
+  auto* block = module_->getBody();
+  auto* sub_block = sub_module->getBody();
+  block->getOperations().splice(
+      mlir::Block::iterator(block->getTerminator()), sub_block->getOperations(),
+      sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
+
+  return Status::OK();
+}
+
 StatusOr
 SavedModelSignatureDefImporter::ConvertSignatures() {
   const auto& signatures = bundle_.GetSignatures();
-  const auto& graphdef = bundle_.meta_graph_def.graph_def();
-  PopulateTfVersions(module_.get(), graphdef.versions());
-
-  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graphdef.library());
+  PopulateTfVersions(module_.get(), graph_def().versions());
 
   // debug_info might not be loaded with loader_lite.
   GraphDebugInfo debug_info;
@@ -3307,9 +3374,10 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
       continue;
     }
 
-    TF_RETURN_IF_ERROR(ConvertSignature(graphdef, sig_def_key, signature_def,
-                                        debug_info, flib_def));
+    TF_RETURN_IF_ERROR(ConvertSignature(sig_def_key, signature_def));
   }
+
+  TF_RETURN_IF_ERROR(ConvertInitializer());
   TF_RETURN_IF_ERROR(LiftVariables());
 
   mlir::OpBuilder builder(module_->getBodyRegion());
@@ -3320,10 +3388,32 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
   return std::move(module_);
 }
 
+StatusOr SavedModelSignatureDefImporter::ConvertGraph(
+    const std::string& name,
+    const std::vector>& inputs,
+    const std::vector>& outputs,
+    const std::vector control_outputs) {
+  GraphImportConfig specs;
+  specs.prune_unused_nodes = true;
+  specs.inputs = ParseInputArrays(inputs);
+  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
+  specs.control_outputs = control_outputs;
+
+  // Convert sub-graphdef to sub-graph.
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.add_default_attributes = true;
+  Graph graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, graph_def(), &graph));
+
+  // Convert sub-graph to MLIR module.true
+  return GraphDefImporter::Convert(module_->getContext(), graph, debug_info(),
+                                   flib_def(), specs, name);
+}
+
 Status SavedModelSignatureDefImporter::ConvertSignature(
-    const GraphDef& graphdef, const std::string& sig_def_key,
-    const SignatureDef& signature_def, const GraphDebugInfo& debug_info,
-    const FunctionLibraryDefinition& flib_def) {
+    const std::string& sig_def_key, const SignatureDef& signature_def) {
   // Create local vectors for the input and output and sort them to be
   // deterministic. We don't want anyone to really depend on the order, client
   // should lookup argument/result mapping by attribute name.
@@ -3339,34 +3429,9 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
     return lhs.first.size() < rhs.first.size() || lhs.first > rhs.first;
   });
 
-  GraphImportConfig specs;
-  specs.prune_unused_nodes = true;
-  specs.inputs = ParseInputArrays(inputs);
-  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
-
-  // Remove unused nodes and create sub-graphdef.
-  GraphDef sub_graph_def;
-  TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(
-      graphdef, &sub_graph_def,
-      /*terminal_nodes=*/{specs.outputs.begin(), specs.outputs.end()}));
-
-  // Set the function library definitions in the pruned graphdef.
-  *sub_graph_def.mutable_library() = flib_def.ToProto();
-
-  // Convert sub-graphdef to sub-graph.
-  GraphConstructorOptions options;
-  options.allow_internal_ops = true;
-  options.add_default_attributes = true;
-  Graph sub_graph(OpRegistry::Global());
-
-  TF_RETURN_IF_ERROR(
-      ConvertGraphDefToGraph(options, sub_graph_def, &sub_graph));
-
   // Convert sub-graph to MLIR module.
-  TF_ASSIGN_OR_RETURN(
-      auto sub_module,
-      GraphDefImporter::Convert(module_->getContext(), sub_graph, debug_info,
-                                flib_def, specs, sig_def_key));
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(sig_def_key, inputs, outputs, {}));
   mlir::OpBuilder builder(sub_module->getBodyRegion());
 
   // Find the FuncOp which corresponds to current SignatureDef.

From b8bb250ebb52d4f037975f58cb1eadb2eff3751c Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Thu, 18 Jun 2020 10:59:05 -0700
Subject: [PATCH 0512/1390] Rollback of rollback: [TF/XLA] Only force retracing
 for non-unique XLA context ID for TPUReplicatedContext

Fixes https://github.com/tensorflow/tensorflow/issues/39872

PiperOrigin-RevId: 317137904
Change-Id: Id287e10a0ab2494b11427435d8f89a383eeaf392
---
 .../python/eager/def_function_xla_jit_test.py | 19 +++++++++++++++++++
 tensorflow/python/eager/function.py           |  7 ++++---
 tensorflow/python/ops/control_flow_ops.py     |  5 +++++
 tensorflow/python/tpu/tpu.py                  |  6 ++++++
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/eager/def_function_xla_jit_test.py b/tensorflow/python/eager/def_function_xla_jit_test.py
index b63a3b434d4..78d44a81b0b 100644
--- a/tensorflow/python/eager/def_function_xla_jit_test.py
+++ b/tensorflow/python/eager/def_function_xla_jit_test.py
@@ -29,6 +29,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import control_flow_util
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import tensor_array_ops
 from tensorflow.python.platform import test
@@ -385,6 +386,24 @@ class DefFunctionTest(test.TestCase):
     f64_input = constant_op.constant([1.1, 2.2, 3.3], dtype=dtypes.float64)
     self.assertAllClose([1.1, 3.3, 6.6], f(f64_input))
 
+  def testNoExcessiveRetracing(self):
+    inner_retracings = 0
+
+    @def_function.function(experimental_compile=True)
+    def inner(a, b):
+      nonlocal inner_retracings
+      inner_retracings += 1
+      return a * b + a
+
+    def outer(a, b):
+      return inner(a, b)
+
+    func_input = random_ops.random_normal([10, 10])
+    for _ in range(2):
+      def_function.function(outer)(func_input, func_input)
+
+    self.assertEqual(inner_retracings, 1)
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index a40eaf886b3..c02318cb814 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -2981,9 +2981,10 @@ class Function(object):
     if not executing_eagerly:
       # We want to force function retracing for each different
       # XLAControlFlowContext, so add `xla_context_id` to the cache key.
-      tpu_context = _enclosing_xla_context()
-      if tpu_context is not None:
-        xla_context_id = id(tpu_context)
+      xla_context = _enclosing_xla_context()
+      if xla_context is not None and \
+            xla_context.RequiresUniqueFunctionRetracing():
+        xla_context_id = id(xla_context)
 
       with ops.init_scope():
         # The graph, or whether we're executing eagerly, should be a part of the
diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py
index 3398308d42e..748f842a9e0 100644
--- a/tensorflow/python/ops/control_flow_ops.py
+++ b/tensorflow/python/ops/control_flow_ops.py
@@ -3682,6 +3682,11 @@ class XLAControlFlowContext(ControlFlowContext):
   def AddValue(self, x):
     return x
 
+  def RequiresUniqueFunctionRetracing(self):
+    """Returns whether the tf.function should be retraced if the context changes.
+    """
+    return False
+
 
 def from_control_flow_context_def(context_def, import_scope=None):
   """Deserializes `context_def` into the appropriate ControlFlowContext.
diff --git a/tensorflow/python/tpu/tpu.py b/tensorflow/python/tpu/tpu.py
index 28eba69b7da..ce3aaa8a058 100644
--- a/tensorflow/python/tpu/tpu.py
+++ b/tensorflow/python/tpu/tpu.py
@@ -639,6 +639,12 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
   def GetControlPivot(self):
     return self._pivot
 
+  def RequiresUniqueFunctionRetracing(self):
+    # More context: b/158152827. TPU stack uses the TPUReplicateContext to
+    # create replicated variable handles and cluster TPU computations, thus we
+    # always retrace a tf.function when the wrapped TPUReplicateContext changes.
+    return True
+
 
 class OutsideCompilationV2Context(control_flow_ops.ControlFlowContext):
   """The context for outside compilation in Tensorflow 2.0.

From 17eea4753dd71b397289f793b348fe9fe751873f Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Thu, 18 Jun 2020 11:09:28 -0700
Subject: [PATCH 0513/1390] Add OSSFuzz badge to TensorFlow.

Now that TF <-> OSSFuzz works again, add badge to show that we are fuzzing the code.

PiperOrigin-RevId: 317140301
Change-Id: I401c1ffd3da37f44910fde3ba60fb2b5c925dfcc
---
 README.md                     | 1 +
 tensorflow/security/README.md | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/README.md b/README.md
index 54c9470b04b..73a345706a4 100644
--- a/README.md
+++ b/README.md
@@ -95,6 +95,7 @@ for general questions and discussion, and please direct specific questions to
 The TensorFlow project strives to abide by generally accepted best practices in
 open-source software development:
 
+[![Fuzzing Status](https://oss-fuzz-build-logs.storage.googleapis.com/badges/tensorflow.svg)](https://bugs.chromium.org/p/oss-fuzz/issues/list?sort=-opened&can=1&q=proj:tensorflow)
 [![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/1486/badge)](https://bestpractices.coreinfrastructure.org/projects/1486)
 [![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v1.4%20adopted-ff69b4.svg)](CODE_OF_CONDUCT.md)
 
diff --git a/tensorflow/security/README.md b/tensorflow/security/README.md
index d9fa1c77a02..34f98e640d6 100644
--- a/tensorflow/security/README.md
+++ b/tensorflow/security/README.md
@@ -1,5 +1,7 @@
 # TensorFlow Security Advisories
 
+[![Fuzzing Status](https://oss-fuzz-build-logs.storage.googleapis.com/badges/tensorflow.svg)](https://bugs.chromium.org/p/oss-fuzz/issues/list?sort=-opened&can=1&q=proj:tensorflow)
+
 We regularly publish security advisories about using TensorFlow.
 
 *Note*: In conjunction with these security advisories, we strongly encourage

From a31d5da02607a2c5eb01d7c977b92001f842cc89 Mon Sep 17 00:00:00 2001
From: Bruce Fontaine 
Date: Thu, 18 Jun 2020 11:30:06 -0700
Subject: [PATCH 0514/1390] Wrap save/restore logic in tf.function when in
 eager mode. This allows parallel saving and restoring when using multiple
 devices.

PiperOrigin-RevId: 317144560
Change-Id: Iebc230589a5e2712da03c5db3f45e4fd7eeb5ff9
---
 .../grappler/optimizers/function_optimizer.cc |   8 +-
 .../parallel_device/parallel_device_test.py   |   4 +
 .../python/framework/auto_control_deps.py     |   2 +-
 tensorflow/python/training/saving/BUILD       |   1 +
 .../training/saving/functional_saver.py       | 111 ++++++++++++------
 .../training/saving/functional_saver_test.py  |  17 ++-
 6 files changed, 101 insertions(+), 42 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index a66e645e04b..0e156aaa84c 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -837,7 +837,6 @@ const bool IsExemptFromSideEffectsExecutionValidation(const string& op) {
        "ParameterizedTruncatedNormal", "TruncatedNormal", "RandomShuffle",
        "Multinomial", "RandomGamma", "RandomGammaGrad", "RandomPoisson",
        "RandomPoissonV2",
-       // LINT.ThenChange(//tensorflow/python/framework/auto_control_deps.py)
 
        // ReadVariableOp marked as stateful because it consumes DT_RESOURCE,
        // but it can't generate any observable side-effect.
@@ -851,7 +850,12 @@ const bool IsExemptFromSideEffectsExecutionValidation(const string& op) {
        // the same device_ordinal on the same host.
        "EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
        "EnqueueTPUEmbeddingSparseTensorBatch",
-       "EnqueueTPUEmbeddingRaggedTensorBatch"});
+       "EnqueueTPUEmbeddingRaggedTensorBatch",
+
+       // SaveV2 and RestoreV2 should be allowed to operate in parallel on
+       // multiple hosts.
+       "SaveV2", "RestoreV2"});
+  // LINT.ThenChange(//tensorflow/python/framework/auto_control_deps.py)
   return exemption->contains(op);
 }
 
diff --git a/tensorflow/python/distribute/parallel_device/parallel_device_test.py b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
index 8fc3dcb5816..1429c522aba 100644
--- a/tensorflow/python/distribute/parallel_device/parallel_device_test.py
+++ b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
@@ -172,6 +172,8 @@ class ParallelDeviceTests(_VirtualDeviceTestCase):
       config.set_synchronous_execution(previous)
 
   def test_checkpointing(self):
+    self.skipTest(
+        "Disable saving until SaveableObject's methods are traceable.")
     prefix = os.path.join(self.get_temp_dir(), "ckpt")
     with self.device.scope():
       different_values = self.device.pack(
@@ -263,6 +265,8 @@ class LayerTests(_VirtualDeviceTestCase):
     self.assertIn(self.device.components[1], final_kernels[1].backing_device)
 
   def test_training_loop(self):
+    self.skipTest(
+        "Disable saving until SaveableObject's methods are traceable.")
     for _ in range(5):
       layer = _Dense(5)
       checkpoint = tracking.Checkpoint(layer=layer)
diff --git a/tensorflow/python/framework/auto_control_deps.py b/tensorflow/python/framework/auto_control_deps.py
index 51dcb248b11..4b47735e0bf 100644
--- a/tensorflow/python/framework/auto_control_deps.py
+++ b/tensorflow/python/framework/auto_control_deps.py
@@ -100,7 +100,7 @@ _ORDER_INSENSITIVE_STATEFUL_OPS = [
     "CudnnRNNV2", "CudnnRNNV3", "CudnnRNNBackpropV2", "CudnnRNNBackpropV3",
     "EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
     "EnqueueTPUEmbeddingSparseTensorBatch",
-    "EnqueueTPUEmbeddingRaggedTensorBatch"
+    "EnqueueTPUEmbeddingRaggedTensorBatch", "RestoreV2", "SaveV2"
 ]
 # LINT.ThenChange(//tensorflow/core/grappler/optimizers/function_optimizer.cc)
 
diff --git a/tensorflow/python/training/saving/BUILD b/tensorflow/python/training/saving/BUILD
index 670a4c35c6f..12940840309 100644
--- a/tensorflow/python/training/saving/BUILD
+++ b/tensorflow/python/training/saving/BUILD
@@ -43,6 +43,7 @@ cuda_py_test(
         ":checkpoint_options",
         ":functional_saver",
         ":saveable_hook",
+        "//tensorflow/python/eager:remote",
         "//tensorflow/python/eager:test",
     ],
 )
diff --git a/tensorflow/python/training/saving/functional_saver.py b/tensorflow/python/training/saving/functional_saver.py
index c4334e096df..3a9b565470d 100644
--- a/tensorflow/python/training/saving/functional_saver.py
+++ b/tensorflow/python/training/saving/functional_saver.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 import uuid
 
 from tensorflow.core.protobuf import saver_pb2
+from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -161,7 +162,8 @@ class MultiDeviceSaver(object):
         self._after_restore_callbacks.append(saveable.after_restore)
 
       if is_saveable:
-        saveables_by_device.setdefault(saveable.device, []).append(saveable)
+        host_device = saveable_object_util.set_cpu0(saveable.device)
+        saveables_by_device.setdefault(host_device, []).append(saveable)
 
     self._single_device_savers = {
         device: _SingleDeviceSaver(saveables)
@@ -247,33 +249,50 @@ class MultiDeviceSaver(object):
       tmp_checkpoint_prefix = string_ops.string_join(
           [file_prefix, sharded_suffix])
 
-    num_shards = len(self._single_device_savers)
-    sharded_saves = []
-    sharded_prefixes = []
-    num_shards_tensor = constant_op.constant(num_shards, name="num_shards")
-    last_device = None
-    for shard, (device, saver) in enumerate(
-        sorted(self._single_device_savers.items())):
-      last_device = device
-      with ops.device(saveable_object_util.set_cpu0(device)):
-        shard_prefix = sharded_filename(tmp_checkpoint_prefix, shard,
-                                        num_shards_tensor)
-      sharded_prefixes.append(shard_prefix)
-      with ops.device(device):
-        # _SingleDeviceSaver will use the CPU device when necessary, but initial
-        # read operations should be placed on the SaveableObject's device.
-        sharded_saves.append(saver.save(shard_prefix, options))
+    def save_fn():
+      num_shards = len(self._single_device_savers)
+      sharded_saves = []
+      sharded_prefixes = []
+      num_shards_tensor = constant_op.constant(num_shards, name="num_shards")
+      last_device = None
+      for shard, (device, saver) in enumerate(
+          sorted(self._single_device_savers.items())):
+        last_device = device
+        with ops.device(saveable_object_util.set_cpu0(device)):
+          shard_prefix = sharded_filename(tmp_checkpoint_prefix, shard,
+                                          num_shards_tensor)
+        sharded_prefixes.append(shard_prefix)
+        with ops.device(device):
+          # _SingleDeviceSaver will use the CPU device when necessary, but
+          # initial read operations should be placed on the SaveableObject's
+          # device.
+          sharded_saves.append(saver.save(shard_prefix, options))
 
-    with ops.control_dependencies(sharded_saves):
-      # Merge on the io_device if specified, otherwise co-locates the merge op
-      # with the last device used.
-      merge_device = (options.experimental_io_device or
-                      saveable_object_util.set_cpu0(last_device))
-      with ops.device(merge_device):
-        # V2 format write path consists of a metadata merge step.  Once merged,
-        # attempts to delete the temporary directory, "_temp".
-        return gen_io_ops.merge_v2_checkpoints(
-            sharded_prefixes, file_prefix, delete_old_dirs=True)
+      with ops.control_dependencies(sharded_saves):
+        # Merge on the io_device if specified, otherwise co-locates the merge op
+        # with the last device used.
+        merge_device = (
+            options.experimental_io_device or
+            saveable_object_util.set_cpu0(last_device))
+        with ops.device(merge_device):
+          # V2 format write path consists of a metadata merge step.  Once
+          # merged, attempts to delete the temporary directory,
+          # "_temp".
+          return gen_io_ops.merge_v2_checkpoints(
+              sharded_prefixes, file_prefix, delete_old_dirs=True)
+
+    # Since this will causes a function re-trace on each save, limit this to the
+    # cases where it is needed: eager and when there are multiple tasks/single
+    # device savers. Note that the retrace is needed to ensure we pickup the
+    # latest values of options like experimental_io_device.
+    if context.executing_eagerly() and len(self._single_device_savers) > 1:
+      # Explicitly place the identity op on the first device.
+      @def_function.function(experimental_compile=False)
+      def tf_function_save():
+        save_fn()
+      tf_function_save()
+    else:
+      return save_fn()
 
   def restore(self, file_prefix, options=None):
     """Restore the saveable objects from a checkpoint with `file_prefix`.
@@ -287,12 +306,38 @@ class MultiDeviceSaver(object):
       A dictionary mapping from SaveableObject names to restore operations.
     """
     options = options or checkpoint_options.CheckpointOptions()
-    restore_ops = {}
-    # Sort by device name to avoid propagating non-deterministic dictionary
-    # ordering in some Python versions.
-    for device, saver in sorted(self._single_device_savers.items()):
-      with ops.device(device):
-        restore_ops.update(saver.restore(file_prefix, options))
+
+    def restore_fn():
+      restore_ops = {}
+      # Sort by device name to avoid propagating non-deterministic dictionary
+      # ordering in some Python versions.
+      for device, saver in sorted(self._single_device_savers.items()):
+        with ops.device(device):
+          restore_ops.update(saver.restore(file_prefix, options))
+
+      return restore_ops
+
+    # Since this will causes a function re-trace on each save, limit this to the
+    # cases where it is needed: eager and when there are multiple tasks/single
+    # device savers. Note that the retrace is needed to ensure we pickup the
+    # latest values of options like experimental_io_device.
+    if context.executing_eagerly() and len(self._single_device_savers) > 1:
+      first_device, _ = list(self._single_device_savers.items())[0]
+      @def_function.function(experimental_compile=False)
+      def tf_function_restore():
+        restore_ops = restore_fn()
+        restore_tensors = {}
+        # tf.functions must return tensors, thus we use control dependencies so
+        # that we can return a tensor which depends on the given op.
+        with ops.device(saveable_object_util.set_cpu0(first_device)):
+          for name, op in restore_ops.items():
+            with ops.control_dependencies([op]):
+              restore_tensors[name] = array_ops.identity(file_prefix)
+        return restore_tensors
+
+      restore_ops = tf_function_restore()
+    else:
+      restore_ops = restore_fn()
 
     for callback in self._after_restore_callbacks:
       callback()
diff --git a/tensorflow/python/training/saving/functional_saver_test.py b/tensorflow/python/training/saving/functional_saver_test.py
index 7db32ff72d7..8f3eef4fb9c 100644
--- a/tensorflow/python/training/saving/functional_saver_test.py
+++ b/tensorflow/python/training/saving/functional_saver_test.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 import os
 
 from tensorflow.python.eager import context
+from tensorflow.python.eager import remote
 from tensorflow.python.eager import test
 from tensorflow.python.eager import wrap_function
 from tensorflow.python.framework import config
@@ -29,6 +30,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import gfile
+from tensorflow.python.training import server_lib
 from tensorflow.python.training.saving import checkpoint_options
 from tensorflow.python.training.saving import functional_saver
 from tensorflow.python.training.saving import saveable_hook
@@ -126,13 +128,16 @@ class SaverTest(test.TestCase):
     second_saver.restore(save_path)
     self.assertEqual(2., self.evaluate(v2))
 
-  @test_util.run_in_graph_and_eager_modes
-  def test_checkpoint_is_sharded_by_device(self):
-    with ops.device("cpu:0"):
+  def test_checkpoint_is_sharded_by_task(self):
+    servers = [server_lib.Server.create_local_server() for _ in range(3)]
+    cluster_spec = server_lib.ClusterSpec({
+        "worker": [s.target[len("grpc://"):] for s in servers]})
+    remote.connect_to_cluster(cluster_spec)
+    with ops.device("/job:worker/task:0/cpu:0"):
       v0 = resource_variable_ops.ResourceVariable(0.)
-    with ops.device("cpu:1"):
+    with ops.device("/job:worker/task:1/cpu:0"):
       v1 = resource_variable_ops.ResourceVariable(1.)
-    with ops.device("cpu:2"):
+    with ops.device("/job:worker/task:2/cpu:0"):
       v2 = resource_variable_ops.ResourceVariable(2.)
 
     self.evaluate([v0.initializer, v1.initializer, v2.initializer])
@@ -167,7 +172,7 @@ class SaverTest(test.TestCase):
         list(saveable_object_util.saveable_objects_for_op(v2, "v2")))
     prefix = os.path.join(self.get_temp_dir(), "ckpt")
     self.evaluate(saver.save(constant_op.constant(prefix), self.local_options))
-    self.assertEqual(4, len(gfile.Glob(prefix + "*")))
+    self.assertEqual(2, len(gfile.Glob(prefix + "*")))
     self.evaluate(v0.assign(-1.))
     self.evaluate(v1.assign(-1.))
     self.evaluate(v2.assign(-1.))

From 4fab49e55013f1265fdd42e53522744ae01351ad Mon Sep 17 00:00:00 2001
From: Ken Franko 
Date: Thu, 18 Jun 2020 11:39:19 -0700
Subject: [PATCH 0515/1390] Add suggestion to use experimental_io_device when
 save/load can't find directory.

PiperOrigin-RevId: 317146508
Change-Id: Ia74ac5baa286c90959c7d55f23187d3db46c3c4b
---
 tensorflow/python/saved_model/load.py | 13 +++++++++++--
 tensorflow/python/saved_model/save.py | 12 +++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/saved_model/load.py b/tensorflow/python/saved_model/load.py
index 74b030a3797..fb2d01cbee2 100644
--- a/tensorflow/python/saved_model/load.py
+++ b/tensorflow/python/saved_model/load.py
@@ -28,6 +28,7 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
@@ -614,8 +615,16 @@ def load_internal(export_dir, tags=None, options=None, loader_cls=Loader):
     ckpt_options = checkpoint_options.CheckpointOptions(
         experimental_io_device=options.experimental_io_device)
     with ops.init_scope():
-      loader = loader_cls(object_graph_proto, saved_model_proto, export_dir,
-                          ckpt_options)
+      try:
+        loader = loader_cls(object_graph_proto, saved_model_proto, export_dir,
+                            ckpt_options)
+      except errors.NotFoundError as err:
+        raise FileNotFoundError(
+            str(err) + "\n If trying to load on a different device from the "
+            "computational device, consider using setting the "
+            "`experimental_io_device` option on tf.saved_model.LoadOptions "
+            "to the io_device such as '/job:localhost'."
+        )
       root = loader.get(0)
       if isinstance(loader, Loader):
         root.graph_debug_info = loader.adjust_debug_info_func_names(debug_info)
diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py
index e22b0129dda..5844c80995f 100644
--- a/tensorflow/python/saved_model/save.py
+++ b/tensorflow/python/saved_model/save.py
@@ -33,6 +33,7 @@ from tensorflow.python.eager import function as defun
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import error_interpolation
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import meta_graph
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
@@ -966,7 +967,16 @@ def save(obj, export_dir, signatures=None, options=None):
   # SavedModel. Users rely on checking saved_model_dir/saved_model.pb as an
   # indication that the SavedModel is completely written.
   if context.executing_eagerly():
-    context.async_wait()  # Ensure save operations have completed.
+    try:
+      context.async_wait()  # Ensure save operations have completed.
+    except errors.NotFoundError as err:
+      raise FileNotFoundError(
+          str(err) + "\n If trying to save on a different device from the "
+          "computational device, consider using setting the "
+          "`experimental_io_device` option on tf.saved_model.SaveOptions "
+          "to the io_device such as '/job:localhost'."
+      )
+
   path = os.path.join(
       compat.as_str(export_dir),
       compat.as_str(constants.SAVED_MODEL_FILENAME_PB))

From bd0a10139a8fa11eb49aa759eca1bb52e662ba34 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 11:50:49 -0700
Subject: [PATCH 0516/1390] Update ops-related pbtxt files.

PiperOrigin-RevId: 317148838
Change-Id: Icecbfedaea80cf2ae6a1424872df62126b7106f1
---
 .../BandedTriangularSolve.pbtxt               |  42 +++++++
 ...tatelessParameterizedTruncatedNormal.pbtxt |  65 +++++++++++
 tensorflow/core/ops/ops.pbtxt                 | 107 ++++++++++++++++++
 3 files changed, 214 insertions(+)
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt

diff --git a/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
new file mode 100644
index 00000000000..5cf85a62392
--- /dev/null
+++ b/tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
@@ -0,0 +1,42 @@
+op {
+  name: "BandedTriangularSolve"
+  input_arg {
+    name: "matrix"
+    type_attr: "T"
+  }
+  input_arg {
+    name: "rhs"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "T"
+  }
+  attr {
+    name: "lower"
+    type: "bool"
+    default_value {
+      b: true
+    }
+  }
+  attr {
+    name: "adjoint"
+    type: "bool"
+    default_value {
+      b: false
+    }
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_DOUBLE
+        type: DT_FLOAT
+        type: DT_HALF
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
new file mode 100644
index 00000000000..598125677b1
--- /dev/null
+++ b/tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
@@ -0,0 +1,65 @@
+op {
+  name: "StatelessParameterizedTruncatedNormal"
+  input_arg {
+    name: "shape"
+    type_attr: "S"
+  }
+  input_arg {
+    name: "seed"
+    type_attr: "Tseed"
+  }
+  input_arg {
+    name: "means"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "stddevs"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "minvals"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "maxvals"
+    type_attr: "dtype"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "dtype"
+  }
+  attr {
+    name: "S"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_INT32
+        type: DT_INT64
+      }
+    }
+  }
+  attr {
+    name: "Tseed"
+    type: "type"
+    default_value {
+      type: DT_INT64
+    }
+    allowed_values {
+      list {
+        type: DT_INT32
+        type: DT_INT64
+      }
+    }
+  }
+  attr {
+    name: "dtype"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index d99f8b8a479..1f1cf7444fb 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -3070,6 +3070,48 @@ op {
     }
   }
 }
+op {
+  name: "BandedTriangularSolve"
+  input_arg {
+    name: "matrix"
+    type_attr: "T"
+  }
+  input_arg {
+    name: "rhs"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "T"
+  }
+  attr {
+    name: "lower"
+    type: "bool"
+    default_value {
+      b: true
+    }
+  }
+  attr {
+    name: "adjoint"
+    type: "bool"
+    default_value {
+      b: false
+    }
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_DOUBLE
+        type: DT_FLOAT
+        type: DT_HALF
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "Barrier"
   output_arg {
@@ -48849,6 +48891,71 @@ op {
     }
   }
 }
+op {
+  name: "StatelessParameterizedTruncatedNormal"
+  input_arg {
+    name: "shape"
+    type_attr: "S"
+  }
+  input_arg {
+    name: "seed"
+    type_attr: "Tseed"
+  }
+  input_arg {
+    name: "means"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "stddevs"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "minvals"
+    type_attr: "dtype"
+  }
+  input_arg {
+    name: "maxvals"
+    type_attr: "dtype"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "dtype"
+  }
+  attr {
+    name: "S"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_INT32
+        type: DT_INT64
+      }
+    }
+  }
+  attr {
+    name: "Tseed"
+    type: "type"
+    default_value {
+      type: DT_INT64
+    }
+    allowed_values {
+      list {
+        type: DT_INT32
+        type: DT_INT64
+      }
+    }
+  }
+  attr {
+    name: "dtype"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+      }
+    }
+  }
+}
 op {
   name: "StatelessRandomBinomial"
   input_arg {

From 0957b4306325a2a1c0bded767dd0199f118c123e Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Thu, 18 Jun 2020 11:57:49 -0700
Subject: [PATCH 0517/1390] LSTM: Use int for 'count' values, fix usage of some
 appropriate intN_t for exact-size values.

PiperOrigin-RevId: 317150337
Change-Id: I5cd8bf8b8231f16cb2d130a8209c45a857b30d21
---
 tensorflow/lite/kernels/lstm.cc      | 13 +++++++------
 tensorflow/lite/kernels/lstm_eval.cc | 14 +++++++-------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index 0e0c1b9c0f0..aa6a112a022 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -78,14 +78,14 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   auto* proj_params = static_cast(
       output_tensor->quantization.params);
   if (cell_clip > 0.0) {
-    integer_lstm_param->quantized_cell_clip = static_cast(std::min(
+    integer_lstm_param->quantized_cell_clip = static_cast(std::min(
         std::max(cell_clip / cell_state_params->scale->data[0], -32768.0f),
         32767.0f));
   } else {
     integer_lstm_param->quantized_cell_clip = 0;
   }
   if (proj_clip > 0.0) {
-    integer_lstm_param->quantized_proj_clip = static_cast(std::min(
+    integer_lstm_param->quantized_proj_clip = static_cast(std::min(
         std::max(proj_clip / proj_params->scale->data[0], -128.0f), 127.0f));
   } else {
     integer_lstm_param->quantized_proj_clip = 0;
@@ -703,14 +703,15 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
       output_tensor->quantization.params);
   TF_LITE_ENSURE_EQ(context, cell_state_params->scale->data[0], 1.0 / 32768);
   if (cell_clip > 0.0 && cell_clip < 1.0) {
-    integer_lstm_param->quantized_cell_clip =
-        static_cast(cell_clip / cell_state_params->scale->data[0]);
+    integer_lstm_param->quantized_cell_clip = static_cast(std::min(
+        std::max(cell_clip / cell_state_params->scale->data[0], -32768.0f),
+        32767.0f));
   } else {
     integer_lstm_param->quantized_cell_clip = 0;
   }
   if (proj_clip > 0.0) {
-    integer_lstm_param->quantized_proj_clip =
-        proj_clip / proj_params->scale->data[0];
+    integer_lstm_param->quantized_proj_clip = static_cast(std::min(
+        std::max(proj_clip / proj_params->scale->data[0], -128.0f), 127.0f));
   } else {
     integer_lstm_param->quantized_proj_clip = 0;
   }
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 2e7f300f9a9..f45d46762bf 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -1044,8 +1044,8 @@ inline void LstmStepInteger(
     const int32_t* recurrent_to_output_effective_bias,
     const int32_t* input_to_input_effective_bias,
     const int32_t* recurrent_to_input_effective_bias,
-    const int32_t* projection_effective_bias, int32 n_batch, int32 n_cell,
-    int32 n_input, int32 n_output, int8_t* output_state_ptr,
+    const int32_t* projection_effective_bias, int n_batch, int n_cell,
+    int n_input, int n_output, int8_t* output_state_ptr,
     int32_t output_state_zp, int16_t* cell_ptr, int8_t* output_ptr,
     int16_t* scratch_0_ptr, int16_t* scratch_1_ptr, int16_t* scratch_2_ptr,
     int16_t* scratch_3_ptr, int8_t* scratch_4_ptr, int32_t* scratch_5_ptr,
@@ -1362,9 +1362,9 @@ void LstmStepInteger(
     const int32_t* cell_gate_bias_ptr, const int32_t* output_bias_ptr,
     const int32_t* proj_bias_ptr, const TfLiteLSTMParams* params,
     const int32_t* intermediate_scale_a, const int32_t* intermediate_scale_b,
-    const int32_t* intermediate_zp, int32 quantized_cell_clip,
-    int32 quantized_proj_clip, int32 n_batch, int32 n_cell, int32 n_input,
-    int32 n_output, int32 output_batch_leading_dim, int8_t* output_state_ptr,
+    const int32_t* intermediate_zp, int16_t quantized_cell_clip,
+    int8_t quantized_proj_clip, int n_batch, int n_cell, int n_input,
+    int n_output, int output_batch_leading_dim, int8_t* output_state_ptr,
     int32_t output_state_zp, int16_t* cell_ptr, int8_t* output_ptr,
     int8_t* scratch0, int8_t* scratch1, int16_t* scratch2, int16_t* scratch3,
     int16_t* scratch4, int16_t* scratch5, int16_t* scratch6,
@@ -2129,8 +2129,8 @@ TfLiteStatus EvalInteger8x8_8(
   int8_t* output_state_ptr = GetTensorData(output_state);
   int8_t* output_ptr = nullptr;
 
-  const int32 input_zp = input->params.zero_point;
-  const int32 output_state_zp = output_state->params.zero_point;
+  const int32_t input_zp = input->params.zero_point;
+  const int32_t output_state_zp = output_state->params.zero_point;
 
   // Get params for time/batch/sequence.
   const int output_batch_leading_dim =

From 0cd2551d8c0ce6f1d8e95f465fa8f3d47054eb32 Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Thu, 18 Jun 2020 12:00:17 -0700
Subject: [PATCH 0518/1390] Rewrite tf numpy's ndarray as a composite tensor.

Expand composites when taking gradients.

PiperOrigin-RevId: 317150904
Change-Id: I2fcb0eadd9797340e468bcf82afd17da0d0369e6
---
 tensorflow/python/eager/backprop.py           | 12 ++-
 tensorflow/python/eager/backprop_test.py      | 69 ++++++++++++++
 tensorflow/python/eager/function.py           |  4 +-
 tensorflow/python/ops/numpy_ops/BUILD         | 10 ++
 tensorflow/python/ops/numpy_ops/np_arrays.py  | 41 +++++++-
 .../python/ops/numpy_ops/np_arrays_test.py    | 18 ++++
 .../python/ops/numpy_ops/np_interop_test.py   | 94 +++++++++++++++++++
 7 files changed, 241 insertions(+), 7 deletions(-)
 create mode 100644 tensorflow/python/ops/numpy_ops/np_interop_test.py

diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index dc7bb7c4b11..5c2deb9c0f2 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -14,6 +14,9 @@
 # ==============================================================================
 """Code for backpropagation using the tape utilities."""
 
+# TODO(b/159343581): Properly support CompositeTensor in all functions in this
+# file.
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
@@ -1021,7 +1024,7 @@ class GradientTape(object):
             "derivatives.", 1)
 
     flat_targets = []
-    for t in nest.flatten(target):
+    for t in nest.flatten(target, expand_composites=True):
       if not backprop_util.IsTrainable(t):
         logging.vlog(
             logging.WARN, "The dtype of the target tensor must be "
@@ -1032,7 +1035,7 @@ class GradientTape(object):
           t = ops.convert_to_tensor(t)
       flat_targets.append(t)
 
-    flat_sources = nest.flatten(sources)
+    flat_sources = nest.flatten(sources, expand_composites=True)
     flat_sources_raw = flat_sources
     flat_sources = [_handle_or_self(x) for x in flat_sources]
     for t in flat_sources_raw:
@@ -1048,7 +1051,8 @@ class GradientTape(object):
 
     if output_gradients is not None:
       output_gradients = [None if x is None else ops.convert_to_tensor(x)
-                          for x in nest.flatten(output_gradients)]
+                          for x in nest.flatten(
+                              output_gradients, expand_composites=True)]
 
     flat_grad = imperative_grad.imperative_grad(
         self._tape,
@@ -1063,7 +1067,7 @@ class GradientTape(object):
       self._watched_variables = self._tape.watched_variables()
       self._tape = None
 
-    grad = nest.pack_sequence_as(sources, flat_grad)
+    grad = nest.pack_sequence_as(sources, flat_grad, expand_composites=True)
     return grad
 
   def jacobian(self,
diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py
index a0f98fc0a44..abdac526ce4 100644
--- a/tensorflow/python/eager/backprop_test.py
+++ b/tensorflow/python/eager/backprop_test.py
@@ -28,6 +28,7 @@ from tensorflow.python.eager import def_function
 from tensorflow.python.eager import function
 from tensorflow.python.eager import tape as tape_lib
 from tensorflow.python.eager import test
+from tensorflow.python.framework import composite_tensor
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
@@ -36,6 +37,7 @@ from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.framework import test_util
+from tensorflow.python.framework import type_spec
 from tensorflow.python.framework.memory_checker import MemoryChecker
 from tensorflow.python.layers.pooling import max_pooling3d
 from tensorflow.python.ops import array_ops
@@ -52,6 +54,44 @@ from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import sparse_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.training import training
+from tensorflow.python.util import nest
+
+
+# TODO(nareshmodi): This is copied from composite_tensor_test.py. Extract it out
+# to a common library to avoid duplication.
+class CTSpec(type_spec.TypeSpec):
+  """A generic CompositeTensor TypeSpec, used for constructing tests."""
+
+  def __init__(self, component_specs):
+    self.component_specs = component_specs
+
+  value_type = property(lambda self: CT)
+  _component_specs = property(lambda self: self.component_specs)
+
+  def _serialize(self):
+    return (self.component_specs,)
+
+  def _to_components(self, value):
+    return value.components
+
+  def _from_components(self, tensor_list):
+    return CT(tensor_list)
+
+
+class CT(composite_tensor.CompositeTensor):
+  """A generic CompositeTensor, used for constructing tests."""
+  _type_spec_class = CTSpec
+
+  def __init__(self, components):
+    if isinstance(components, list):
+      components = tuple(components)
+    self.components = components
+
+  @property
+  def _type_spec(self):
+    component_specs = nest.map_structure(type_spec.type_spec_from_value,
+                                         self.components)
+    return self._type_spec_class(component_specs)
 
 
 class BackpropTest(test.TestCase, parameterized.TestCase):
@@ -1581,6 +1621,35 @@ class BackpropTest(test.TestCase, parameterized.TestCase):
     memory_checker.report()
     memory_checker.assert_no_leak_if_all_possibly_except_one()
 
+  def testCompositeTensorAsSource(self):
+    t = CT([constant_op.constant(3.), constant_op.constant(2.)])
+    with backprop.GradientTape() as gt:
+      gt.watch(t)
+      y = CT([t.components[0] * 2, t.components[1] * 3])
+
+    grad = gt.gradient(y, t)
+    expected_grad = CT([constant_op.constant(2.), constant_op.constant(3.)])
+
+    flat_grads = nest.flatten(grad, expand_composites=True)
+    flat_expected_grads = nest.flatten(expected_grad, expand_composites=True)
+
+    self.assertAllClose(flat_grads, flat_expected_grads)
+
+  def testCompositeTensorAsOutputGradients(self):
+    t = CT([constant_op.constant(3.), constant_op.constant(2.)])
+    with backprop.GradientTape() as gt:
+      gt.watch(t)
+      y = CT([t.components[0] * 2, t.components[1] * 3])
+
+    output_gradients = CT([constant_op.constant(5.), constant_op.constant(10.)])
+    grad = gt.gradient(y, t, output_gradients=output_gradients)
+    expected_grad = CT([constant_op.constant(10.), constant_op.constant(30.)])
+
+    flat_grads = nest.flatten(grad, expand_composites=True)
+    flat_expected_grads = nest.flatten(expected_grad, expand_composites=True)
+
+    self.assertAllClose(flat_grads, flat_expected_grads)
+
 
 class JacobianTest(test.TestCase):
 
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index c02318cb814..ca1e60c1b7b 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -2626,7 +2626,9 @@ def _is_ndarray(value):
       # For legacy reasons we do not automatically promote Numpy strings.
       or isinstance(value, np.str_)
       # NumPy dtypes have __array__ as unbound methods.
-      or isinstance(value, type))
+      or isinstance(value, type)
+      # CompositeTensors should be flattened instead.
+      or isinstance(value, composite_tensor.CompositeTensor))
 
 
 def _convert_numpy_inputs(inputs):
diff --git a/tensorflow/python/ops/numpy_ops/BUILD b/tensorflow/python/ops/numpy_ops/BUILD
index 5879bc9f062..3f18a7b3e01 100644
--- a/tensorflow/python/ops/numpy_ops/BUILD
+++ b/tensorflow/python/ops/numpy_ops/BUILD
@@ -115,3 +115,13 @@ cuda_py_test(
         "@absl_py//absl/testing:parameterized",
     ],
 )
+
+cuda_py_test(
+    name = "np_interop_test",
+    srcs = ["np_interop_test.py"],
+    deps = [
+        ":numpy",
+        "//tensorflow/python:platform",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index e2f73100909..8bec8a469a2 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -24,10 +24,13 @@ import numbers
 import numpy as np
 import six
 
+from tensorflow.python.framework import composite_tensor
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import type_spec
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.numpy_ops import np_dtypes
@@ -175,7 +178,38 @@ def convert_to_tensor(value, dtype=None, dtype_hint=None):
   return ops.convert_to_tensor(value, dtype=dtype, dtype_hint=dtype_hint)
 
 
-class ndarray(object):  # pylint: disable=invalid-name
+class NdarraySpec(type_spec.BatchableTypeSpec):
+  """Type specification for a `tf.experiemntal.numpy.ndarray`."""
+
+  value_type = property(lambda self: ndarray)
+
+  def __init__(self, data_spec):
+    if not isinstance(data_spec, tensor_spec.TensorSpec):
+      raise ValueError('NdarraySpec.__init__ was expecting a tf.TypeSpec, '
+                       'but got a {} instead.'.format(type(data_spec)))
+    self._data_spec = data_spec
+
+  @property
+  def _component_specs(self):
+    return self._data_spec
+
+  def _to_components(self, value):
+    return value.data
+
+  def _from_components(self, data):
+    return tensor_to_ndarray(data)
+
+  def _serialize(self):
+    return (self._data_spec,)
+
+  def _batch(self, batch_size):
+    return NdarraySpec(self._data_spec.batch(batch_size))
+
+  def _unbatch(self):
+    return NdarraySpec(self._data_spec.unbatch())
+
+
+class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
   """Equivalent of numpy.ndarray backed by TensorFlow tensors.
 
   This does not support all features of NumPy ndarrays e.g. strides and
@@ -236,7 +270,10 @@ class ndarray(object):  # pylint: disable=invalid-name
     if dtype and dtype != buffer.dtype:
       buffer = array_ops.bitcast(buffer, dtype)
     self._data = buffer
-    self.base = None
+
+  @property
+  def _type_spec(self):
+    return NdarraySpec(type_spec.type_spec_from_value(self._data))
 
   @property
   def data(self):
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays_test.py b/tensorflow/python/ops/numpy_ops/np_arrays_test.py
index feced98438d..412addc0ad7 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays_test.py
@@ -29,6 +29,7 @@ from tensorflow.python.ops.numpy_ops import np_arrays
 # Required for operator overloads
 from tensorflow.python.ops.numpy_ops import np_math_ops  # pylint: disable=unused-import
 from tensorflow.python.platform import test
+from tensorflow.python.util import nest
 
 t2a = np_arrays.tensor_to_ndarray
 
@@ -182,6 +183,23 @@ class ArrayTest(test.TestCase):
     with self.assertRaisesWithPredicateMatch(TypeError, r'unhashable type'):
       hash(a)
 
+  def testFromToCompositeTensor(self):
+    tensors = [t2a(ops.convert_to_tensor(0.1)), t2a(ops.convert_to_tensor(0.2))]
+
+    flattened = nest.flatten(tensors, expand_composites=True)
+    # Each ndarray contains only one tensor, so the flattened output should be
+    # just 2 tensors in a list.
+    self.assertLen(flattened, 2)
+    self.assertIsInstance(flattened[0], ops.Tensor)
+    self.assertIsInstance(flattened[1], ops.Tensor)
+
+    repacked = nest.pack_sequence_as(tensors, flattened, expand_composites=True)
+    self.assertLen(repacked, 2)
+    self.assertIsInstance(repacked[0], np_arrays.ndarray)
+    self.assertIsInstance(repacked[1], np_arrays.ndarray)
+
+    self.assertAllClose(tensors, repacked)
+
 
 if __name__ == '__main__':
   # TODO(wangpeng): Test in graph mode as well.
diff --git a/tensorflow/python/ops/numpy_ops/np_interop_test.py b/tensorflow/python/ops/numpy_ops/np_interop_test.py
new file mode 100644
index 00000000000..052949dff9d
--- /dev/null
+++ b/tensorflow/python/ops/numpy_ops/np_interop_test.py
@@ -0,0 +1,94 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for interop between TF ops, numpy_ops, and numpy methods."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops.numpy_ops import np_array_ops
+from tensorflow.python.ops.numpy_ops import np_arrays
+from tensorflow.python.platform import test
+
+
+class InteropTest(test.TestCase):
+
+  def testGradientTapeInterop(self):
+    with backprop.GradientTape() as t:
+      x = np_array_ops.asarray(3.0)
+      y = np_array_ops.asarray(2.0)
+
+      t.watch([x, y])
+
+      xx = 2 * x
+      yy = 3 * y
+
+    dx, dy = t.gradient([xx, yy], [x, y])
+
+    # TODO(nareshmodi): Gradient tape returns tensors. Is it possible to rewrap?
+    self.assertAllClose(dx, 2.0)
+    self.assertAllClose(dy, 3.0)
+
+  def testFunctionInterop(self):
+    x = np_array_ops.asarray(3.0)
+    y = np_array_ops.asarray(2.0)
+
+    add = lambda x, y: x + y
+    add_fn = def_function.function(add)
+
+    raw_result = add(x, y)
+    fn_result = add_fn(x, y)
+
+    self.assertIsInstance(raw_result, np_arrays.ndarray)
+    self.assertIsInstance(fn_result, np_arrays.ndarray)
+    self.assertAllClose(raw_result, fn_result)
+
+  def testCondInterop(self):
+    x = np_array_ops.asarray(3.0)
+
+    def fn(x):
+      x_plus_1 = control_flow_ops.cond(x > 0, lambda: x+1, lambda: x+2)
+      x_plus_2 = control_flow_ops.cond(x < 0, lambda: x+1, lambda: x+2)
+
+      return x_plus_1, x_plus_2
+
+    raw_x_plus_1, raw_x_plus_2 = fn(x)
+    fn_x_plus_1, fn_x_plus_2 = def_function.function(fn)(x)
+
+    self.assertAllClose(raw_x_plus_1, x + 1)
+    self.assertAllClose(raw_x_plus_2, x + 2)
+
+    self.assertAllClose(fn_x_plus_1, x + 1)
+    self.assertAllClose(fn_x_plus_2, x + 2)
+
+  def testWhileInterop(self):
+    def fn():
+      x = np_array_ops.asarray(0)
+      c = lambda x: x < 10000
+      b = lambda x: [x + 1]
+      return control_flow_ops.while_loop_v2(c, b, [x], parallel_iterations=20)
+
+    self.assertEqual(10000, fn()[0])
+    self.assertEqual(10000, def_function.function(fn)()[0])
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

From 629e6077d0c4ccc342a15496cd91becedfc02bc3 Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Thu, 18 Jun 2020 12:25:56 -0700
Subject: [PATCH 0519/1390] [tf.data service] Add test that options are applied
 to the distributed dataset

PiperOrigin-RevId: 317156222
Change-Id: Ie1956a3ef474371e381618c4808fee7a8ef8ee08
---
 .../kernel_tests/data_service_ops_test.py     | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tensorflow/python/data/kernel_tests/data_service_ops_test.py b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
index 2356a866d6e..796ab328980 100644
--- a/tensorflow/python/data/kernel_tests/data_service_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
@@ -23,6 +23,7 @@ from absl.testing import parameterized
 
 from tensorflow.python.data.experimental.ops import data_service_ops
 from tensorflow.python.data.experimental.ops import distribute_options
+from tensorflow.python.data.experimental.ops import testing
 from tensorflow.python.data.experimental.service import server_lib
 from tensorflow.python.data.kernel_tests import test_base
 from tensorflow.python.data.ops import dataset_ops
@@ -30,6 +31,7 @@ from tensorflow.python.eager import def_function
 from tensorflow.python.framework import combinations
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import tensor_array_ops
 from tensorflow.python.platform import test
@@ -317,6 +319,34 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
       results.append(elem.numpy())
     self.assertCountEqual(num_repetitions * list(range(num_elements)), results)
 
+  @combinations.generate(test_base.eager_only_combinations())
+  def testApplyDeterminismOption(self):
+    elements = list(range(10))
+    master_address = self.create_cluster(1)
+
+    def dataset_fn(delay_ms):
+
+      def interleave_fn(x):
+        ds = dataset_ops.Dataset.from_tensors(x)
+        if math_ops.equal(x, 0):
+          ds = ds.apply(testing.sleep(delay_ms * 1000))
+        else:
+          ds = ds.apply(testing.sleep(0))
+        return ds
+
+      ds = dataset_ops.Dataset.from_tensor_slices(elements)
+      ds = ds.interleave(interleave_fn, cycle_length=10, num_parallel_calls=10)
+      opts = dataset_ops.Options()
+      opts.experimental_deterministic = False
+      ds = ds.with_options(opts)
+      ds = _make_distributed_dataset(ds, master_address)
+      return ds
+
+    self.checkDeterminism(
+        dataset_fn=dataset_fn,
+        expect_determinism=False,
+        expected_elements=elements)
+
   def run_stateful(self, external_state_policy):
     num_elements = 10
     ds = dataset_ops.Dataset.range(num_elements).map(

From d8bfc935fdc77ddc138e415dd8bb47f9817c4d5e Mon Sep 17 00:00:00 2001
From: Reed 
Date: Thu, 18 Jun 2020 12:42:19 -0700
Subject: [PATCH 0520/1390] Add MKL supoprt to auto_mixed_precision.

This extends the auto mixed precision grappler pass to support converting nodes to bfloat16 on MKL-supported CPUs.

Co-authored-by: Niranjan Hasabnis 
---
 tensorflow/core/grappler/optimizers/BUILD     |  10 +-
 .../optimizers/auto_mixed_precision.cc        | 229 +++---
 .../optimizers/auto_mixed_precision.h         |  20 +-
 .../optimizers/auto_mixed_precision_lists.h   | 257 +++++--
 .../optimizers/auto_mixed_precision_test.cc   | 234 ++++++-
 .../grappler/optimizers/meta_optimizer.cc     |  11 +-
 .../core/protobuf/rewriter_config.proto       |   8 +-
 .../grappler/auto_mixed_precision_test.py     | 661 ++++++++++--------
 8 files changed, 951 insertions(+), 479 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 7432e2d54ea..f9b4c6e6d81 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -1,5 +1,5 @@
 load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cc_test_mkl", "tf_kernel_library")
-load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
+load("//tensorflow:tensorflow.bzl", "tf_copts", "tf_cuda_cc_test")
 
 # Platform specific build config
 load(
@@ -7,6 +7,11 @@ load(
     "if_static",
 )
 
+load(
+    "//third_party/mkl:build_defs.bzl",
+    "mkl_deps",
+)
+
 package(
     licenses = ["notice"],  # Apache 2.0
 )
@@ -611,6 +616,7 @@ cc_library(
         "auto_mixed_precision_lists.h",
     ],
     visibility = ["//visibility:public"],
+    copts = tf_copts(),
     deps = [
         ":custom_graph_optimizer_registry",
         ":graph_optimizer",
@@ -627,7 +633,7 @@ cc_library(
         "//tensorflow/core/grappler/costs:virtual_placer",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
-    ],
+    ] + mkl_deps(),
 )
 
 tf_cuda_cc_test(
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc b/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
index fa6ca3144a5..2e0b56fd3e2 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/optimizers/auto_mixed_precision.h"
 
 #include 
+#include 
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
@@ -52,6 +53,7 @@ const std::pair kMinGPUArch = {0, 0};
 
 const char kSuffix[] = "AutoMixedPrecision";
 const char kCastToFp16[] = "CastToFp16";
+const char kCastToBf16[] = "CastToBf16";
 const char kCastToFp32[] = "CastToFp32";
 
 // Instances of this class represent unique type attribute identifiers within a
@@ -840,22 +842,6 @@ DataTypeSet AllowedDataTypes(const OpDef& op_def, const TypeAttrId& t_attr_id) {
   return AllowedDataTypes(*attr_def);
 }
 
-NodeDef BuildCastNode(const MutableGraphView::OutputPort& src, bool to_fp16,
-                      const string& device) {
-  const char* cast_string = to_fp16 ? kCastToFp16 : kCastToFp32;
-  string name = strings::StrCat(src.node->name(), "-", src.port_id, "-",
-                                cast_string, "-", kSuffix);
-  NodeDef node;
-  node.set_name(name);
-  node.set_op("Cast");
-  node.set_device(device);
-  node.add_input(strings::StrCat(src.node->name(), ":", src.port_id));
-  (*node.mutable_attr())["SrcT"].set_type(to_fp16 ? DT_FLOAT : DT_HALF);
-  (*node.mutable_attr())["DstT"].set_type(to_fp16 ? DT_HALF : DT_FLOAT);
-  (*node.mutable_attr())["Truncate"].set_b(false);
-  return node;
-}
-
 Status ValidateLists(const gtl::FlatSet& white_list,
                      const gtl::FlatSet& black_list,
                      const gtl::FlatSet& gray_list,
@@ -941,7 +927,8 @@ class AutoMixedPrecisionImpl {
  public:
   AutoMixedPrecisionImpl(Cluster* cluster,
                          const std::unordered_set& nodes_to_preserve,
-                         GraphDef* graph, string id)
+                         GraphDef* graph, string id,
+                         AutoMixedPrecisionMode mode)
       : virtual_placer_(cluster->GetDevices()),
         nodes_to_preserve_(nodes_to_preserve),
         graph_(graph),
@@ -949,23 +936,35 @@ class AutoMixedPrecisionImpl {
         id_(id),
         graph_view_(graph),
         cuda_version_(GetCudaVersion(*cluster)),
-        cudnn_version_(GetCudnnVersion(*cluster)) {}
+        cudnn_version_(GetCudnnVersion(*cluster)),
+        mode_(mode),
+        target_dtype_(mode_ == AutoMixedPrecisionMode::CUDA ? DT_HALF
+                                                            : DT_BFLOAT16) {}
 
   Status Optimize();
 
  private:
   typedef absl::flat_hash_set NodeTypeIdSet;
 
+  std::unique_ptr get_mixed_precision_lists() const {
+    switch (mode_) {
+      case AutoMixedPrecisionMode::CUDA:
+        return std::make_unique(cuda_version_,
+                                                             cudnn_version_);
+      case AutoMixedPrecisionMode::MKL:
+        return std::make_unique();
+    }
+  }
   Status PrintDebugLogs(bool preop, size_t timestamp);
   void LogSkippedNode(const NodeDef& node) const;
   bool MustPreserve(const NodeDef& node) const;
-  bool IsOnGPU(const NodeDef& node) const;
+  bool IsOnDevice(const NodeDef& node, const string& device_type) const;
   bool IsOnSuitableGPUArch(const NodeDef& node) const;
   bool ShouldProcess(const NodeDef& node) const;
-  bool NodeHasFP16KernelForTypeAttr(const NodeDef& node, TypeAttrId taid) const;
+  bool NodeHasF16KernelForTypeAttr(const NodeDef& node, TypeAttrId taid) const;
   bool NodeImplicitlyReadsNonResourceVariable(const NodeDef& node) const;
   void ConvertBatchNormOpsToV2();
-  bool SupportsFloat16(const NodeTypeId& node_type) const;
+  bool SupportsF16(const NodeTypeId& node_type) const;
   const NodeTypeId* GetTensorListFloat32NodeTypeId(const NodeDef& node) const;
   bool IsSourceOrSinkOp(const string& op) const;
   void FindFloat32TensorListOpClustersAndBlacklistUnsafe(
@@ -990,6 +989,8 @@ class AutoMixedPrecisionImpl {
       absl::flat_hash_set* white_set) const;
   void MakeCastsWhiteIfAllOutputsWhite(
       absl::flat_hash_set* white_set) const;
+  NodeDef BuildCastNode(const MutableGraphView::OutputPort& src, bool to_f16,
+                        const string& device) const;
   Status ChangeTypeAttrsAndAddCasts(const absl::flat_hash_set& white_set);
 
   VirtualPlacer virtual_placer_;
@@ -1003,21 +1004,44 @@ class AutoMixedPrecisionImpl {
   NodeTypeAttrMap node_type_map_;
   GraphTypeTopologyView graph_type_view_;
   bool force_all_fp16_;
-  gtl::FlatSet fp16_whitelist_;
-  gtl::FlatSet fp16_blacklist_;
-  gtl::FlatSet fp16_graylist_;
-  gtl::FlatSet fp16_clearlist_;
+  AutoMixedPrecisionMode mode_;
+  gtl::FlatSet f16_whitelist_;
+  gtl::FlatSet f16_blacklist_;
+  gtl::FlatSet f16_graylist_;
+  gtl::FlatSet f16_clearlist_;
   absl::flat_hash_set should_process_nodes_;
+  DataType target_dtype_;  // Either DT_HALF or DT_BFLOAT16
 };
 
-bool AutoMixedPrecisionImpl::NodeHasFP16KernelForTypeAttr(
+NodeDef AutoMixedPrecisionImpl::BuildCastNode(
+    const MutableGraphView::OutputPort& src, bool to_f16,
+    const string& device) const {
+  DataType src_type = to_f16 ? DT_FLOAT : target_dtype_;
+  DataType dst_type = to_f16 ? target_dtype_ : DT_FLOAT;
+  const char* cast_string =
+      !to_f16 ? kCastToFp32
+              : target_dtype_ == DT_HALF ? kCastToFp16 : kCastToBf16;
+  string name = strings::StrCat(src.node->name(), "-", src.port_id, "-",
+                                cast_string, "-", kSuffix);
+  NodeDef node;
+  node.set_name(name);
+  node.set_op("Cast");
+  node.set_device(device);
+  node.add_input(strings::StrCat(src.node->name(), ":", src.port_id));
+  (*node.mutable_attr())["SrcT"].set_type(src_type);
+  (*node.mutable_attr())["DstT"].set_type(dst_type);
+  (*node.mutable_attr())["Truncate"].set_b(false);
+  return node;
+}
+
+bool AutoMixedPrecisionImpl::NodeHasF16KernelForTypeAttr(
     const NodeDef& node, TypeAttrId taid) const {
   NodeDef node_copy(node);
   if (node.device().empty()) {
     string device_name = virtual_placer_.get_canonical_device_name(node);
     node_copy.set_device(device_name);
   }
-  if (!SetDataType(&node_copy, taid, DataType::DT_HALF)) {
+  if (!SetDataType(&node_copy, taid, target_dtype_)) {
     return false;
   }
   return IsKernelRegisteredForNode(node_copy).ok();
@@ -1053,21 +1077,22 @@ Status AutoMixedPrecisionImpl::PrintDebugLogs(bool preop, size_t timestamp) {
     fname = io::JoinPath(prepend_path,
                          strings::StrCat("paintbuckets", suffix, ".txt"));
     f.open(fname.c_str(), std::fstream::out);
+    std::unique_ptr mp_lists =
+        get_mixed_precision_lists();
     f << "WhiteList:\n";
-    for (const auto& x :
-         AutoMixedPrecisionLists::WhiteList(cuda_version_, cudnn_version_)) {
+    for (const auto& x : mp_lists->WhiteList()) {
       f << x << "\n";
     }
     f << "\nBlackList:\n";
-    for (const auto& x : AutoMixedPrecisionLists::BlackList()) {
+    for (const auto& x : mp_lists->BlackList()) {
       f << x << "\n";
     }
     f << "\nGrayList:\n";
-    for (const auto& x : AutoMixedPrecisionLists::GrayList()) {
+    for (const auto& x : mp_lists->GrayList()) {
       f << x << "\n";
     }
     f << "\nClearList:\n";
-    for (const auto& x : AutoMixedPrecisionLists::ClearList()) {
+    for (const auto& x : mp_lists->ClearList()) {
       f << x << "\n";
     }
     f.close();
@@ -1088,7 +1113,8 @@ bool AutoMixedPrecisionImpl::MustPreserve(const NodeDef& node) const {
   return nodes_to_preserve_.count(node.name());
 }
 
-bool AutoMixedPrecisionImpl::IsOnGPU(const NodeDef& node) const {
+bool AutoMixedPrecisionImpl::IsOnDevice(const NodeDef& node,
+                                        const string& device_type) const {
   string device_name;
   if (node.device().empty()) {
     device_name = virtual_placer_.get_canonical_device_name(node);
@@ -1099,7 +1125,7 @@ bool AutoMixedPrecisionImpl::IsOnGPU(const NodeDef& node) const {
   string not_used;
   if (DeviceNameUtils::SplitDeviceName(device_name, ¬_used, &device) &&
       absl::StrContains(absl::AsciiStrToLower(device),
-                        absl::AsciiStrToLower(DEVICE_GPU))) {
+                        absl::AsciiStrToLower(device_type))) {
     return true;
   }
   return false;
@@ -1164,15 +1190,14 @@ bool IsTensorListWriterOp(const string& op) {
   return tensor_list_writer_ops.count(op);
 }
 
-bool AutoMixedPrecisionImpl::SupportsFloat16(
-    const NodeTypeId& node_type) const {
+bool AutoMixedPrecisionImpl::SupportsF16(const NodeTypeId& node_type) const {
   const OpDef* op_def;
   Status status =
       OpRegistry::Global()->LookUpOpDef(node_type.node->op(), &op_def);
   if (!status.ok()) return false;
   return AllowedDataTypes(*op_def, node_type.type_attr)
-             .Contains(DataType::DT_HALF) &&
-         NodeHasFP16KernelForTypeAttr(*node_type.node, node_type.type_attr);
+             .Contains(target_dtype_) &&
+         NodeHasF16KernelForTypeAttr(*node_type.node, node_type.type_attr);
 }
 
 // TODO(mconley): Make this change the node's name (to aid debugging). Need to
@@ -1219,22 +1244,40 @@ Status AutoMixedPrecisionImpl::Optimize() {
       "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_LEVEL", "", &optimization_level));
   optimization_level = absl::AsciiStrToUpper(optimization_level);
   force_all_fp16_ = optimization_level == "UNSAFE_FORCE_ALL";
+  if (force_all_fp16_ && mode_ == AutoMixedPrecisionMode::MKL) {
+    // Many ops do not support bfloat16 on the CPU so we disallowing forcing to
+    // bfloat16.
+    return errors::InvalidArgument(
+        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_LEVEL cannot be set to "
+        "UNSAFE_FORCE_ALL when MKL is used");
+  }
 
-  fp16_whitelist_ =
-      AutoMixedPrecisionLists::WhiteList(cuda_version_, cudnn_version_);
-  fp16_blacklist_ = AutoMixedPrecisionLists::BlackList();
-  fp16_graylist_ = AutoMixedPrecisionLists::GrayList();
-  fp16_clearlist_ = AutoMixedPrecisionLists::ClearList();
-  TF_RETURN_IF_ERROR(ValidateLists(fp16_whitelist_, fp16_blacklist_,
-                                   fp16_graylist_, fp16_clearlist_));
+  std::unique_ptr mp_lists =
+      get_mixed_precision_lists();
+  f16_whitelist_ = mp_lists->WhiteList();
+  f16_blacklist_ = mp_lists->BlackList();
+  f16_graylist_ = mp_lists->GrayList();
+  f16_clearlist_ = mp_lists->ClearList();
+  TF_RETURN_IF_ERROR(ValidateLists(f16_whitelist_, f16_blacklist_,
+                                   f16_graylist_, f16_clearlist_));
 
   size_t timestamp = Env::Default()->NowMicros() / 1000;
   TF_RETURN_IF_ERROR(PrintDebugLogs(/* preop = */ true, timestamp));
 
   VLOG(2) << "Identifying nodes that should be processed";
   for (const NodeDef& node : graph_->node()) {
-    if (!MustPreserve(node) && IsOnGPU(node) &&
-        (ShouldIgnorePerformance() || IsOnSuitableGPUArch(node))) {
+    bool should_process;
+    switch (mode_) {
+      case AutoMixedPrecisionMode::CUDA:
+        should_process =
+            !MustPreserve(node) && IsOnDevice(node, DEVICE_GPU) &&
+            (ShouldIgnorePerformance() || IsOnSuitableGPUArch(node));
+        break;
+      case AutoMixedPrecisionMode::MKL:
+        should_process = !MustPreserve(node) && IsOnDevice(node, DEVICE_CPU);
+        break;
+    }
+    if (should_process) {
       should_process_nodes_.insert(&node);
     } else {
       LogSkippedNode(node);
@@ -1260,29 +1303,29 @@ Status AutoMixedPrecisionImpl::Optimize() {
   for (const auto& cluster : tensor_list_clusters) {
     VLOG(1) << "Found safe Tensor List cluster of size " << cluster.size();
     for (const NodeDef* node : cluster) {
-      VLOG(2) << "Cluster member: " << node->op() << " node " << node->name();
+      VLOG(2) << "  Cluster member: " << node->op() << " node " << node->name();
     }
     FindTensorListImplicitFloat32Edges(cluster, &ephemeral_edges);
   }
   TF_RETURN_IF_ERROR(graph_type_view_.AddEphemeralEdges(ephemeral_edges));
 
-  // The goal here is to change performance-critical ops to fp16, and to do so
-  // with the minimal number of casts, subject to the constraint that the
+  // The goal here is to change performance-critical ops to fp16 or bf16, and to
+  // do so with the minimal number of casts, subject to the constraint that the
   // model's convergence is not affected. This is achieved by first identifying
-  // which nodes should be changed to fp16 and then inserting casts at the
-  // boundaries between fp16/non-fp16 nodes.
+  // which nodes should be changed to f16 and then inserting casts at the
+  // boundaries between f16/non-f16 nodes.
 
-  // The algorithm for deciding which nodes to change to fp16 is as follows:
+  // The algorithm for deciding which nodes to change to f16 is as follows:
   // 1) Add all performance-critical ops (aka "whitelist" ops) to the white_set.
   //    This is done under the assumption that whitelist ops are always
-  //    numerically-safe in fp16 and that they are the most important ops for
+  //    numerically-safe in f16 and that they are the most important ops for
   //    improving performance.
   // 2) Add nodes to the black_set iff they are numerically-dangerous (aka
   //    "blacklist" ops) or they are on a forward path from a blacklist node to
   //    a black/gray node (including the node at the end of the path) through
   //    non-numerically-dangerous ops (aka "greylist" and "clearlist" ops).
   //    This is done to prevent numerically-dangerous ops and their downstream
-  //    effects from being changed to fp16, which would risk breaking the
+  //    effects from being changed to f16, which would risk breaking the
   //    numerical accuracy of the model.
   // 3) For all remaining nodes that are not considered dangerous (greylist
   //    and clearlist ops), find those that are between (i.e., both upstream
@@ -1480,7 +1523,7 @@ void AutoMixedPrecisionImpl::AddWhitelistOps(
     const NodeTypeId& root = *graph_type_view_.GetNode(root_idx);
     if (!ShouldProcess(*root.node)) continue;
     bool force_white = force_all_fp16_ && CanForceFP16(*root.node);
-    if (fp16_whitelist_.count(root.node->op()) || force_white) {
+    if (f16_whitelist_.count(root.node->op()) || force_white) {
       bool inserted = white_set->insert(root_idx).second;
       if (VLOG_IS_ON(2) && inserted) {
         VLOG(2) << "Painting type " << root.type_attr.DebugString()
@@ -1504,8 +1547,8 @@ void AutoMixedPrecisionImpl::PropagateBlackFwdThroughClearAndGray(
   absl::flat_hash_set upstream_of_black_or_gray_set;
   for (int root_idx = 0; root_idx < graph_type_view_.num_nodes(); ++root_idx) {
     const NodeTypeId& root = *graph_type_view_.GetNode(root_idx);
-    if (!(fp16_blacklist_.count(root.node->op()) ||
-          fp16_graylist_.count(root.node->op()))) {
+    if (!(f16_blacklist_.count(root.node->op()) ||
+          f16_graylist_.count(root.node->op()))) {
       continue;
     }
     DfsTypeTraversal(graph_type_view_, {&root},
@@ -1514,7 +1557,7 @@ void AutoMixedPrecisionImpl::PropagateBlackFwdThroughClearAndGray(
                        const NodeTypeId& item = *graph_type_view_.GetNode(idx);
                        return idx == root_idx ||
                               (!upstream_of_black_or_gray_set.count(idx) &&
-                               fp16_clearlist_.count(item.node->op()));
+                               f16_clearlist_.count(item.node->op()));
                      }),
                      DfsTypeCallbacks::PreOrder([&](int idx) {
                        upstream_of_black_or_gray_set.insert(idx);
@@ -1524,7 +1567,7 @@ void AutoMixedPrecisionImpl::PropagateBlackFwdThroughClearAndGray(
   // Propagate black forward through nodes in upstream_of_black_or_gray_set.
   for (int root_idx = 0; root_idx < graph_type_view_.num_nodes(); ++root_idx) {
     const NodeTypeId& root = *graph_type_view_.GetNode(root_idx);
-    if (black_set->count(root_idx) || !fp16_blacklist_.count(root.node->op())) {
+    if (black_set->count(root_idx) || !f16_blacklist_.count(root.node->op())) {
       continue;
     }
     DfsTypeTraversal(
@@ -1552,7 +1595,7 @@ void AutoMixedPrecisionImpl::AddClearAndGrayToWhiteIfBetweenWhite(
   absl::flat_hash_set downstream_of_white_set;
   for (int root_idx = 0; root_idx < graph_type_view_.num_nodes(); ++root_idx) {
     const NodeTypeId& root = *graph_type_view_.GetNode(root_idx);
-    if (!ShouldProcess(*root.node) || !fp16_whitelist_.count(root.node->op())) {
+    if (!ShouldProcess(*root.node) || !f16_whitelist_.count(root.node->op())) {
       continue;
     }
     DfsTypeTraversal(
@@ -1561,14 +1604,14 @@ void AutoMixedPrecisionImpl::AddClearAndGrayToWhiteIfBetweenWhite(
           const NodeTypeId& item = *graph_type_view_.GetNode(idx);
           return idx == root_idx ||
                  (!downstream_of_white_set.count(idx) &&
-                  !fp16_whitelist_.count(item.node->op()) &&
+                  !f16_whitelist_.count(item.node->op()) &&
                   !black_set.count(idx) && ShouldProcess(*item.node) &&
                   // TODO(benbarsdell): Consider allowing propagation through
                   // ops that are already float16 in order to reduce the number
                   // of casts.
-                  IsFloat32(item) && SupportsFloat16(item) &&
-                  (fp16_clearlist_.count(item.node->op()) ||
-                   fp16_graylist_.count(item.node->op())));
+                  IsFloat32(item) && SupportsF16(item) &&
+                  (f16_clearlist_.count(item.node->op()) ||
+                   f16_graylist_.count(item.node->op())));
         }),
         DfsTypeCallbacks::PreOrder(
             [&](int idx) { downstream_of_white_set.insert(idx); }));
@@ -1579,7 +1622,7 @@ void AutoMixedPrecisionImpl::AddClearAndGrayToWhiteIfBetweenWhite(
   for (int root_idx = 0; root_idx < graph_type_view_.num_nodes(); ++root_idx) {
     const NodeTypeId& root = *graph_type_view_.GetNode(root_idx);
     if (!ShouldProcess(*root.node) || upstream_of_white_set.count(root_idx) ||
-        !fp16_whitelist_.count(root.node->op())) {
+        !f16_whitelist_.count(root.node->op())) {
       continue;
     }
     DfsTypeTraversal(
@@ -1620,8 +1663,8 @@ void AutoMixedPrecisionImpl::PropagateWhiteThroughClear(
           return idx == root_idx ||
                  (!white_set->count(idx) && !black_set.count(idx) &&
                   ShouldProcess(*item.node) && IsFloat32(item) &&
-                  SupportsFloat16(item) &&
-                  (fp16_clearlist_.count(item.node->op())) &&
+                  SupportsF16(item) &&
+                  (f16_clearlist_.count(item.node->op())) &&
                   // We don't propagate (backwards) through nodes that read
                   // Variables because it can break the behavior of TensorBoard
                   // visualization and/or (in the case of Enter nodes) the model
@@ -1806,13 +1849,13 @@ void AutoMixedPrecisionImpl::MakeCastsWhiteIfAllOutputsWhite(
   }
 }
 
-// Changes all white-painted type attributes to DT_HALF, and inserts Cast nodes
-// at node outputs for all edges that connect white-painted <->
-// non-white-painted type attributes.
+// Changes all white-painted type attributes to DT_HALF or DT_BFLOAT16, and
+// inserts Cast nodes at node outputs for all edges that connect
+// white-painted <-> non-white-painted type attributes.
 Status AutoMixedPrecisionImpl::ChangeTypeAttrsAndAddCasts(
     const absl::flat_hash_set& white_set) {
   int num_nodes_changed = 0;
-  int num_nonvar_casts_to_fp16 = 0;
+  int num_nonvar_casts_to_f16 = 0;
   int num_nodes_preop = graph_->node_size();
   for (int node_idx = 0; node_idx < num_nodes_preop; ++node_idx) {
     NodeDef* node = graph_->mutable_node(node_idx);
@@ -1829,8 +1872,9 @@ Status AutoMixedPrecisionImpl::ChangeTypeAttrsAndAddCasts(
       bool src_is_white = white_set.count(node_type_idx);
       if (src_is_white) {
         VLOG(1) << "Changing type " << type_attr.DebugString() << " of "
-                << node->op() << " node " << node->name() << " to DT_HALF";
-        if (!SetDataType(node, type_attr, DT_HALF)) {
+                << node->op() << " node " << node->name() << " to "
+                << DataTypeString(target_dtype_);
+        if (!SetDataType(node, type_attr, target_dtype_)) {
           return errors::Internal("Failed to set type attribute");
         }
         ++num_nodes_changed;
@@ -1855,16 +1899,16 @@ Status AutoMixedPrecisionImpl::ChangeTypeAttrsAndAddCasts(
           bool dst_is_white = white_set.count(dst_type_idx);
           if (src_is_white != dst_is_white) {
             if (!added_cast_node) {
-              bool to_fp16 = dst_is_white;
+              bool to_f16 = dst_is_white;
               VLOG(1) << "Inserting cast to "
-                      << (to_fp16 ? "DT_HALF" : "DT_FLOAT") << " at "
-                      << src.node->op() << " " << src.node->name() << ":"
-                      << src.port_id;
+                      << (to_f16 ? DataTypeString(target_dtype_) : "DT_FLOAT")
+                      << " at " << src.node->op() << " " << src.node->name()
+                      << ":" << src.port_id;
               added_cast_node = graph_view_.AddNode(
-                  BuildCastNode(src, to_fp16, src.node->device()));
-              if (to_fp16 && !IsConstant(*node) && !IsVariable(*node) &&
+                  BuildCastNode(src, to_f16, src.node->device()));
+              if (to_f16 && !IsConstant(*node) && !IsVariable(*node) &&
                   !NodeImplicitlyReadsNonResourceVariable(*node)) {
-                ++num_nonvar_casts_to_fp16;
+                ++num_nonvar_casts_to_f16;
               }
             }
             TF_RETURN_IF_ERROR(graph_view_.UpdateRegularFaninByPort(
@@ -1874,9 +1918,13 @@ Status AutoMixedPrecisionImpl::ChangeTypeAttrsAndAddCasts(
       }
     }
   }
+  // Use Python type names (e.g. float16) instead of C++ type names (e.g. half)
+  // since many Python users will see this message.
+  const char* type_str = target_dtype_ == DT_HALF ? "float16" : "bfloat16";
   LOG(INFO) << "Converted " << num_nodes_changed << "/" << num_nodes_preop
-            << " nodes to float16 precision using " << num_nonvar_casts_to_fp16
-            << " cast(s) to float16 (excluding Const and Variable casts)";
+            << " nodes to " << type_str << " precision using "
+            << num_nonvar_casts_to_f16 << " cast(s) to " << type_str
+            << " (excluding Const and Variable casts)";
   return Status::OK();
 }
 
@@ -1902,12 +1950,23 @@ Status AutoMixedPrecision::Optimize(Cluster* cluster, const GrapplerItem& item,
     return errors::InvalidArgument("cluster == nullptr");
   }
 
+#if !defined(INTEL_MKL) || !defined(ENABLE_INTEL_MKL_BFLOAT16)
+  if (mode_ == AutoMixedPrecisionMode::MKL) {
+    return errors::Unimplemented(
+        "The auto_mixed_precision_mkl optimizer cannot be used since "
+        "this build of TensorFlow is not compiled with MKL support for bfloat16. "
+        "For information on MKL builds, see: "
+        "https://software.intel.com/en-us/articles/intel-optimization-for-"
+        "tensorflow-installation-guide");
+  }
+#endif
+
   // Start by copying input graph to output.
   *output = item.graph;
 
   int num_gpus = ShouldIgnorePerformance() ? GetNumGPUs(*cluster)
                                            : GetNumGPUs(*cluster, kMinGPUArch);
-  if (num_gpus < 1) {
+  if (num_gpus < 1 && mode_ == AutoMixedPrecisionMode::CUDA) {
     // AutoMixedPrecision is currently only tuned for GPU.
     LOG(WARNING) << "No (suitable) GPUs detected, skipping " << name()
                  << " graph optimizer";
@@ -1916,7 +1975,7 @@ Status AutoMixedPrecision::Optimize(Cluster* cluster, const GrapplerItem& item,
 
   // Optimize the output graph in-place.
   AutoMixedPrecisionImpl optimizer(cluster, item.NodesToPreserve(), output,
-                                   item.id);
+                                   item.id, mode_);
   if (item.id == "tf_graph") {
     LOG(INFO) << "Running " << name() << " graph optimizer";
   } else {
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
index 163d1f6923f..c41ba7d2821 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
@@ -22,16 +22,25 @@ limitations under the License.
 namespace tensorflow {
 namespace grappler {
 
-// Convert data types to float16 where appropriate to improve performance on
-// GPUs.
+enum class AutoMixedPrecisionMode { CUDA, MKL };
+
+// Convert data types to float16 or bfloat16 where appropriate to improve
+// performance on GPUs or CPUs.
 class AutoMixedPrecision : public GraphOptimizer {
  public:
+  // If 'mode' is CUDA, converts nodes to float16 on Nvidia GPUs. If MKL,
+  // converts nodes to bfloat16 on CPUs in order to take advantage of MKL
+  // performance improvements with bfloat16.
   explicit AutoMixedPrecision(
-      RewriterConfig::Toggle opt_level = RewriterConfig::ON) {}
+      AutoMixedPrecisionMode mode = AutoMixedPrecisionMode::CUDA)
+      : mode_(mode) {}
 
   ~AutoMixedPrecision() override {}
 
-  string name() const override { return "auto_mixed_precision"; };
+  string name() const override {
+    return mode_ == AutoMixedPrecisionMode::CUDA ? "auto_mixed_precision_cuda"
+                                                 : "auto_mixed_precision_mkl";
+  };
 
   bool UsesFunctionLibrary() const override { return false; }
 
@@ -40,6 +49,9 @@ class AutoMixedPrecision : public GraphOptimizer {
 
   void Feedback(Cluster* cluster, const GrapplerItem& item,
                 const GraphDef& optimize_output, double result) override;
+
+ private:
+  const AutoMixedPrecisionMode mode_;
 };
 
 }  // end namespace grappler
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
index d3d13e2edc0..c6016548117 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
@@ -23,10 +23,44 @@ limitations under the License.
 namespace tensorflow {
 namespace grappler {
 
+// Represents the four lists of ops: the white list, gray list, black list, and
+// clear list. These lists determine which ops are converted to fp16/bf16
+// (referred to as 'f16' for short) and which ops stay as fp32.
 class AutoMixedPrecisionLists {
- private:
-  static void UpdateList(gtl::FlatSet* list, const string& to_add,
-                         const string& to_remove) {
+ public:
+
+  virtual ~AutoMixedPrecisionLists() {}
+
+  // Returns the set of ops that are considered numerically-safe (for execution
+  // in f16), performance-critical, and can run in f16. These ops are always
+  // converted to f16.
+  virtual gtl::FlatSet WhiteList() = 0;
+  // Returns the set of ops that can run in f16 and are considered numerically-
+  // safe (for execution in f16), but which may be made unsafe by an upstream
+  // blacklist op.
+  virtual gtl::FlatSet GrayList() = 0;
+  // Returns the set of ops that are considered numerically-dangerous (i.e.,
+  // unsafe for execution in f16) and whose effects may also be observed in
+  // downstream nodes (e.g. for f16, in Exp -> Add, the Add is unsafe due to
+  // the Exp).
+  virtual gtl::FlatSet BlackList() = 0;
+  // Returns the set of ops that do not have numerically-significant effects
+  // (i.e., they are always considered safe for execution in f16 precision), and
+  // can run in f16.
+  virtual gtl::FlatSet ClearList() = 0;
+
+ protected:
+  // Adds or removes ops from list if certain environmental variables are set.
+  static void UpdateList(const string& list_name, gtl::FlatSet* list) {
+    CHECK(list_name == "WHITELIST" || list_name == "GRAYLIST" ||  // Crash OK.
+          list_name == "BLACKLIST" || list_name == "CLEARLIST");
+    string add_env_var =
+        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_" + list_name + "_ADD";
+    string remove_env_var =
+        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_" + list_name + "_REMOVE";
+    string to_add, to_remove;
+    TF_CHECK_OK(ReadStringFromEnvVar(add_env_var, "", &to_add));
+    TF_CHECK_OK(ReadStringFromEnvVar(remove_env_var, "", &to_remove));
     for (const auto& x : str_util::Split(to_add, ",")) {
       list->insert(x);
     }
@@ -35,6 +69,35 @@ class AutoMixedPrecisionLists {
     }
   }
 
+  // Subclasses should include these on the ClearList.
+  static void AddTensorListOps(gtl::FlatSet* list) {
+    // Note: if a data structure op (such as TensorListPopBack) is added here,
+    // IsTensorListReaderOp or IsTensorListWriterOp may need to be modified
+    constexpr char* tensor_list_ops[] = {
+        "TensorListConcat",
+        "TensorListConcatLists",
+        "TensorListConcatV2",
+        "TensorListGather",
+        "TensorListGetItem",
+        "TensorListPopBack",
+        "TensorListPushBack",
+        "TensorListPushBackBatch",
+        "TensorListFromTensor",
+        "TensorListScatter",
+        "TensorListScatterV2",
+        "TensorListScatterIntoExistingList",
+        "TensorListSetItem",
+        "TensorListSplit",
+        "TensorListStack"
+    };
+    for (auto op : tensor_list_ops) {
+      list->insert(op);
+    }
+  }
+};
+
+class AutoMixedPrecisionListsCuda : public AutoMixedPrecisionLists {
+ private:
   static bool IsPseudoFastMath() {
     string optimization_level;
     TF_CHECK_OK(
@@ -45,16 +108,10 @@ class AutoMixedPrecisionLists {
   }
 
  public:
-  // Returns the set of ops that are considered numerically-safe (for execution
-  // in fp16) and performance-critical. These ops are always converted to fp16.
-  static gtl::FlatSet WhiteList(int cuda_version, int cudnn_version) {
-    string to_add, to_remove;
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_WHITELIST_ADD", "", &to_add));
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_WHITELIST_REMOVE", "",
-        &to_remove));
+  AutoMixedPrecisionListsCuda(int cuda_version, int cudnn_version)
+      : cuda_version_(cuda_version), cudnn_version_(cudnn_version) {}
 
+  gtl::FlatSet WhiteList() override {
     auto list = gtl::FlatSet{
         "BlockLSTM",
         "BlockLSTMV2",
@@ -81,12 +138,12 @@ class AutoMixedPrecisionLists {
         // "DepthwiseConv2dNativeBackpropInput",
         "MatMul",
     };
-    if (cuda_version >= 9010) {
+    if (cuda_version_ >= 9010) {
       // Fp16 BatchMatMul is slow before CUDA 9.1.
       list.insert("BatchMatMul");
       list.insert("BatchMatMulV2");
     }
-    if (cudnn_version >= 7602) {
+    if (cudnn_version_ >= 7602) {
       // Fp16 3D conv is slow before CUDNN 7.6.2.
       list.insert("Conv3D");
       list.insert("Conv3DBackpropFilter");
@@ -94,22 +151,14 @@ class AutoMixedPrecisionLists {
       list.insert("Conv3DBackpropInput");
       list.insert("Conv3DBackpropInputV2");
     }
-    UpdateList(&list, to_add, to_remove);
+    UpdateList("WHITELIST", &list);
     return list;
   }
 
-  // Returns the set of ops that are considered numerically-safe (for execution
-  // in fp16), but which may be made unsafe by an upstream blacklist op.
-  static gtl::FlatSet GrayList() {
+  gtl::FlatSet GrayList() override {
     if (IsPseudoFastMath()) {
       return gtl::FlatSet{};
     }
-    string to_add, to_remove;
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_GRAYLIST_ADD", "", &to_add));
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_GRAYLIST_REMOVE", "",
-        &to_remove));
 
     auto list = gtl::FlatSet{
         "Add",
@@ -156,23 +205,14 @@ class AutoMixedPrecisionLists {
         "Tanh",
         "TanhGrad",
     };
-    UpdateList(&list, to_add, to_remove);
+    UpdateList("GRAYLIST", &list);
     return list;
   }
 
-  // Returns the set of ops that are considered numerically-dangerous (i.e.,
-  // unsafe for execution in fp16) and whose effects may also be observed in
-  // downstream nodes (e.g., in Exp -> Add, the Add is unsafe due to the Exp).
-  static gtl::FlatSet BlackList() {
+  gtl::FlatSet BlackList() override {
     if (IsPseudoFastMath()) {
       return gtl::FlatSet{};
     }
-    string to_add, to_remove;
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_BLACKLIST_ADD", "", &to_add));
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_BLACKLIST_REMOVE", "",
-        &to_remove));
 
     auto list = gtl::FlatSet{
         "Exp",
@@ -185,22 +225,14 @@ class AutoMixedPrecisionLists {
         "SparseSoftmaxCrossEntropyWithLogits",
         "Sum",
     };
-    UpdateList(&list, to_add, to_remove);
+    UpdateList("BLACKLIST", &list);
     return list;
   }
 
-  // Returns the set of ops that do not have numerically-significant effects
-  // (i.e., they are always considered safe for execution in fp16 precision).
-  static gtl::FlatSet ClearList() {
+  gtl::FlatSet ClearList() override {
     if (IsPseudoFastMath()) {
       return gtl::FlatSet{};
     }
-    string to_add, to_remove;
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_CLEARLIST_ADD", "", &to_add));
-    TF_CHECK_OK(ReadStringFromEnvVar(
-        "TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_CLEARLIST_REMOVE", "",
-        &to_remove));
 
     auto list = gtl::FlatSet{
         "Abs",
@@ -291,21 +323,6 @@ class AutoMixedPrecisionLists {
         "StridedSlice",
         "StridedSliceGrad",
         "Switch",
-        "TensorListConcat",
-        "TensorListConcatLists",
-        "TensorListConcatV2",
-        "TensorListGather",
-        "TensorListGetItem",
-        "TensorListPopBack",
-        "TensorListPushBack",
-        "TensorListPushBackBatch",
-        "TensorListFromTensor",
-        "TensorListScatter",
-        "TensorListScatterV2",
-        "TensorListScatterIntoExistingList",
-        "TensorListSetItem",
-        "TensorListSplit",
-        "TensorListStack",
         "Tile",
         "TopK",
         "TopKV2",
@@ -313,7 +330,125 @@ class AutoMixedPrecisionLists {
         "Where",
         "ZerosLike",
     };
-    UpdateList(&list, to_add, to_remove);
+    AddTensorListOps(&list);
+    UpdateList("CLEARLIST", &list);
+    return list;
+  }
+
+ private:
+  int cuda_version_;
+  int cudnn_version_;
+};
+
+class AutoMixedPrecisionListsMkl : public AutoMixedPrecisionLists {
+ private:
+
+ public:
+  AutoMixedPrecisionListsMkl() {}
+
+  // Only ops which are supported by MKL in bfloat16 should be added to the
+  // white list, gray list, or clear list.
+  gtl::FlatSet WhiteList() override {
+    auto list = gtl::FlatSet{
+        "Conv2D",
+        "Conv2DBackpropFilter",
+        "Conv2DBackpropInput",
+        "Conv3D",
+        "Conv3DBackpropFilterV2",
+        "Conv3DBackpropInputV2",
+        "DepthwiseConv2dNative",
+        "DepthwiseConv2dNativeBackpropFilter",
+        "DepthwiseConv2dNativeBackpropInput",
+        "MatMul",
+        "BatchMatMul",
+        "BatchMatMulV2"
+    };
+
+    UpdateList("WHITELIST", &list);
+    return list;
+  }
+
+  gtl::FlatSet GrayList() override {
+    auto list = gtl::FlatSet{
+        "Add",
+        "AddN",
+        "AddV2",
+        "AvgPool",
+        "AvgPool3D",
+        "AvgPool3DGrad",
+        "AvgPoolGrad",
+        "BiasAdd",
+        "BiasAddGrad",
+        "BiasAddV1",
+        "FusedBatchNormV2",
+        "FusedBatchNormGradV2",
+        "FusedBatchNormV3",
+        "FusedBatchNormGradV3",
+        "LeakyRelu",
+        "LeakyReluGrad",
+        "Mul",
+        "Sub",
+    };
+    UpdateList("GRAYLIST", &list);
+    return list;
+  }
+
+  gtl::FlatSet BlackList() override {
+    auto list = gtl::FlatSet{
+        "Exp",
+        "Expm1",
+        "L2Loss",
+        "Mean",
+        "Pow",
+        "SaveV2",
+        "Softmax",
+        "SoftmaxCrossEntropyWithLogits",
+        "SparseSoftmaxCrossEntropyWithLogits",
+        "Sum",
+    };
+    UpdateList("BLACKLIST", &list);
+    return list;
+  }
+
+  gtl::FlatSet ClearList() override {
+    auto list = gtl::FlatSet{
+        "Concat",
+        "ConcatV2",
+        "Enter",
+        "EnsureShape",
+        "Equal",
+        "Exit",
+        "ExpandDims",
+        "Identity",
+        "MaxPool",
+        "MaxPool3D",
+        "MaxPool3DGrad",
+        "MaxPoolGrad",
+        "MaxPoolV2",
+        "Maximum",
+        "Merge",
+        "NextIteration",
+        "PreventGradient",
+        "Relu",
+        "Relu6",
+        "Relu6Grad",
+        "ReluGrad",
+        "Reshape",
+        "Select",
+        "SelectV2",
+        "Shape",
+        "ShapeN",
+        "Slice",
+        "Split",
+        "SplitV",
+        "Squeeze",
+        "StopGradient",
+        "Switch",
+        "Transpose",
+        "ZerosLike",
+    };
+    AddTensorListOps(&list);
+    UpdateList("CLEARLIST", &list);
     return list;
   }
 };
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc b/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
index 951279d37cd..248d8dd4266 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_test.cc
@@ -13,12 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// Currently, this test only passes when TensorFlow passes with CUDA, because
-// otherwise the optimizer will not turn clearlist nodes to float16. When
-// looking at clearlist nodes, this optimizer checks if the nodes have a float16
-// GPU OpKernel, but without CUDA there are no GPU OpKernels at all.
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-
 #include "tensorflow/core/grappler/optimizers/auto_mixed_precision.h"
 
 #include 
@@ -70,6 +64,31 @@ Tensor GenerateRandomTensorInRange(const TensorShape& shape, double minval,
   return tensor;
 }
 
+void VerifyGraphsEquivalent(const GraphDef& original_graph,
+                            const GraphDef& optimized_graph,
+                            const string& func) {
+  EXPECT_EQ(original_graph.node_size(), optimized_graph.node_size()) << func;
+  GraphView optimized_view(&optimized_graph);
+  for (int i = 0; i < original_graph.node_size(); ++i) {
+    const NodeDef& original = original_graph.node(i);
+    const NodeDef& optimized = *optimized_view.GetNode(original.name());
+    EXPECT_EQ(original.name(), optimized.name()) << func;
+    EXPECT_EQ(original.op(), optimized.op()) << func;
+    EXPECT_EQ(original.input_size(), optimized.input_size()) << func;
+    if (original.input_size() == optimized.input_size()) {
+      for (int j = 0; j < original.input_size(); ++j) {
+        EXPECT_EQ(original.input(j), optimized.input(j)) << func;
+      }
+    }
+  }
+}
+
+// Currently, this test suite only passes when TensorFlow passes with CUDA,
+// because otherwise the optimizer will not turn clearlist nodes to float16.
+// When looking at clearlist nodes, this optimizer checks if the nodes have a
+// float16 GPU OpKernel, but without CUDA there are no GPU OpKernels at all.
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
 const std::pair kMinGPUArch = {7, 0};
 
 class AutoMixedPrecisionTest : public GrapplerTest {
@@ -184,25 +203,6 @@ class AutoMixedPrecisionTest : public GrapplerTest {
   bool gpu_available_;
 };
 
-void VerifyGraphsEquivalent(const GraphDef& original_graph,
-                            const GraphDef& optimized_graph,
-                            const string& func) {
-  EXPECT_EQ(original_graph.node_size(), optimized_graph.node_size()) << func;
-  GraphView optimized_view(&optimized_graph);
-  for (int i = 0; i < original_graph.node_size(); ++i) {
-    const NodeDef& original = original_graph.node(i);
-    const NodeDef& optimized = *optimized_view.GetNode(original.name());
-    EXPECT_EQ(original.name(), optimized.name()) << func;
-    EXPECT_EQ(original.op(), optimized.op()) << func;
-    EXPECT_EQ(original.input_size(), optimized.input_size()) << func;
-    if (original.input_size() == optimized.input_size()) {
-      for (int j = 0; j < original.input_size(); ++j) {
-        EXPECT_EQ(original.input(j), optimized.input(j)) << func;
-      }
-    }
-  }
-}
-
 TEST_F(AutoMixedPrecisionTest, NoOp) {
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
   Output input = ops::Const(s.WithOpName("input"), 1.234f, {32});
@@ -1164,8 +1164,188 @@ TEST_F(AutoMixedPrecisionTest, TanhOp) {
       });
 }
 
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#if INTEL_MKL
+#ifdef ENABLE_INTEL_MKL_BFLOAT16
+
+class AutoMixedPrecisionMklTest : public GrapplerTest {
+ protected:
+  void SetUp() override {
+    virtual_cluster_.reset(new SingleMachine(/* timeout_s = */ 10, 1, 0));
+    TF_CHECK_OK(virtual_cluster_->Provision());
+  }
+  void TearDown() override { TF_CHECK_OK(virtual_cluster_->Shutdown()); }
+
+  std::unique_ptr virtual_cluster_;
+};
+
+TEST_F(AutoMixedPrecisionMklTest, AlreadyBf16) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output input = ops::Const(s.WithOpName("input"), 1.f, {32, 32});
+  Output cst1 = ops::Cast(s.WithOpName("cst1"), input, DT_BFLOAT16);
+  Output wht1 = ops::MatMul(s.WithOpName("wht1"), cst1, cst1);
+  Output clr1 = ops::Relu(s.WithOpName("clr1"), wht1);
+  Output cst2 = ops::Cast(s.WithOpName("cst2"), clr1, DT_FLOAT);
+  Output clr2 = ops::Relu(s.WithOpName("clr2"), cst2);
+  Output fetch = ops::Identity(s.WithOpName("fetch"), clr2);
+
+  GrapplerItem item;
+  item.fetch = {"fetch"};
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
+  AutoMixedPrecision optimizer{AutoMixedPrecisionMode::MKL};
+  GraphDef output;
+  TF_ASSERT_OK(optimizer.Optimize(virtual_cluster_.get(), item, &output));
+  VLOG(1) << output.DebugString();
+
+  VerifyGraphsEquivalent(item.graph, output, __FUNCTION__);
+  GraphView output_view(&output);
+  EXPECT_EQ(output_view.GetNode("input")->attr().at("dtype").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("cst1")->attr().at("DstT").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("wht1")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("clr1")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("cst2")->attr().at("SrcT").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("cst2")->attr().at("DstT").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("clr2")->attr().at("T").type(), DT_FLOAT);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectTensorNear(tensors_expected[i], tensors[i], 1e-6);
+  }
+}
+
+TEST_F(AutoMixedPrecisionMklTest, Simple) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output input = ops::Const(s.WithOpName("input"), 1.f / 32, {32, 32});
+  Output blk1 = ops::Exp(s.WithOpName("blk1"), input);
+  Output clr1 = ops::Relu(s.WithOpName("clr1"), blk1);
+  Output gry1 = ops::Sqrt(s.WithOpName("gry1"), clr1);
+  Output clr2 = ops::Relu(s.WithOpName("clr2"), gry1);
+  Output wht1 = ops::MatMul(s.WithOpName("wht1"), clr2, clr2);
+  Output clr3 = ops::Relu(s.WithOpName("clr3"), wht1);
+  Output blk2 = ops::Log(s.WithOpName("blk2"), clr3);
+  Output clr4 = ops::Relu(s.WithOpName("clr4"), blk2);
+  Output blk3 = ops::SparseMatMul(s.WithOpName("blk3"), clr4, clr4);
+  Output clr5 = ops::Relu(s.WithOpName("clr5"), blk3);
+  Output fetch = ops::Identity(s.WithOpName("fetch"), clr5);
+
+  GrapplerItem item;
+  item.fetch = {"fetch"};
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
+  AutoMixedPrecision optimizer{AutoMixedPrecisionMode::MKL};
+  GraphDef output;
+  TF_ASSERT_OK(optimizer.Optimize(virtual_cluster_.get(), item, &output));
+
+  VLOG(1) << output.DebugString();
+
+  GraphView output_view(&output);
+  EXPECT_EQ(output.node_size(), item.graph.node_size() + 2);
+  EXPECT_EQ(output_view.GetNode("input")->attr().at("dtype").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("blk1")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("clr1")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("gry1")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("clr2")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("wht1")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("clr3")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("blk2")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("clr4")->attr().at("T").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("blk3")->attr().at("Ta").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("blk3")->attr().at("Tb").type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("clr5")->attr().at("T").type(), DT_FLOAT);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectClose(tensors_expected[i], tensors[i], -1, 5e-4);
+  }
+}
+
+TEST_F(AutoMixedPrecisionMklTest, TensorListSetGet) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  tensorflow::Input shape = {32, 32};
+  auto tl1 = ops::TensorListReserve(s.WithOpName("tl1"), {32, 32}, 8, DT_FLOAT);
+  Output input = ops::Const(s.WithOpName("input"), 1.f / 32, {32, 32});
+  Output idx1 = ops::Const(s.WithOpName("idx1"), 1);
+  Output idx2 = ops::Const(s.WithOpName("idx2"), 2);
+  Output idx3 = ops::Const(s.WithOpName("idx3"), 3);
+  auto tl1w1 =
+      ops::TensorListSetItem(s.WithOpName("tl1w1"), tl1.handle, idx1, input);
+  Output wht1 = ops::MatMul(s.WithOpName("wht1"), input, input);
+  auto tl1w2 =
+      ops::TensorListSetItem(s.WithOpName("tl1w2"), tl1.handle, idx2, wht1);
+  // Ensure that TensorListResize doesn't cause any problems.
+  Output tl1rs =
+      ops::TensorListResize(s.WithOpName("tl1rs"), tl1w2.output_handle, 6);
+  Output tl1r1 = ops::TensorListGetItem(s.WithOpName("tl1r1"), tl1rs, idx2,
+                                        shape, DT_FLOAT)
+                     .item;
+  Output gry1 = ops::Mul(s.WithOpName("gry1"), tl1r1, tl1r1);
+  Output wht2 = ops::MatMul(s.WithOpName("wht2"), gry1, gry1);
+  auto tl1w3 =
+      ops::TensorListSetItem(s.WithOpName("tl1w3"), tl1.handle, idx3, wht2);
+  Output tl1r2 =
+      ops::TensorListGetItem(s.WithOpName("tl1r2"), tl1w3.output_handle, idx3,
+                             shape, DT_FLOAT)
+          .item;
+  auto tl2 = ops::TensorListReserve(s.WithOpName("tl2"), shape, 8, DT_FLOAT);
+  auto tl2w1 =
+      ops::TensorListSetItem(s.WithOpName("tl2w1"), tl2.handle, idx1, input);
+  Output tl2r1 =
+      ops::TensorListGetItem(s.WithOpName("tl2r1"), tl2w1.output_handle, idx1,
+                             shape, DT_FLOAT)
+          .item;
+  Output fetch1 = ops::Identity(s.WithOpName("fetch1"), tl1r2);
+  Output fetch2 = ops::Identity(s.WithOpName("fetch2"), tl2r1);
+
+  GrapplerItem item;
+  item.fetch = {"fetch1", "fetch2"};
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
+  AutoMixedPrecision optimizer{AutoMixedPrecisionMode::MKL};
+  GraphDef output;
+  TF_ASSERT_OK(optimizer.Optimize(virtual_cluster_.get(), item, &output));
+
+  VLOG(1) << output.DebugString();
+
+  GraphView output_view(&output);
+  EXPECT_EQ(output.node_size(), item.graph.node_size() + 2);
+  const char* type_key = "element_dtype";
+  EXPECT_EQ(output_view.GetNode("tl1")->attr().at(type_key).type(),
+            DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("tl1w1")->attr().at(type_key).type(),
+            DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("wht1")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("tl1w2")->attr().at(type_key).type(),
+            DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("tl1r1")->attr().at(type_key).type(),
+            DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("gry1")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("wht2")->attr().at("T").type(), DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("tl1w3")->attr().at(type_key).type(),
+            DT_BFLOAT16);
+  EXPECT_EQ(output_view.GetNode("tl2")->attr().at(type_key).type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("tl2w1")->attr().at(type_key).type(), DT_FLOAT);
+  EXPECT_EQ(output_view.GetNode("tl2r1")->attr().at(type_key).type(), DT_FLOAT);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectClose(tensors_expected[i], tensors[i], -1, 1e-2);
+  }
+}
+
+#endif  // ENABLE_INTEL_MKL_BFLOAT16
+#endif  // INTEL_MKL
+
 }  // namespace
 }  // namespace grappler
 }  // namespace tensorflow
-
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index cd0d44e8e12..2f1c869965d 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -188,7 +188,9 @@ std::unique_ptr MetaOptimizer::MakeNewOptimizer(
   MK_OPT("remap", new Remapper(cfg_.remapping()));
   MK_OPT("layout", new GenericLayoutOptimizer());
   MK_OPT("auto_mixed_precision",
-         new AutoMixedPrecision(cfg_.auto_mixed_precision()));
+         new AutoMixedPrecision(AutoMixedPrecisionMode::CUDA));
+  MK_OPT("auto_mixed_precision_mkl",
+         new AutoMixedPrecision(AutoMixedPrecisionMode::MKL));
   MK_OPT("memory", new MemoryOptimizer(RewriterConfig::MANUAL));
   MK_OPT("common_subgraph_elimination",
          new CommonSubgraphElimination(cfg_.common_subgraph_elimination()));
@@ -249,7 +251,11 @@ Status MetaOptimizer::InitializeOptimizers(
   }
   if (AutoMixedPrecisionEnabled(cfg_.auto_mixed_precision())) {
     optimizers->push_back(
-        MakeUnique(cfg_.auto_mixed_precision()));
+        MakeUnique(AutoMixedPrecisionMode::CUDA));
+  }
+  if (AutoMixedPrecisionEnabled(cfg_.auto_mixed_precision_mkl())) {
+    optimizers->push_back(
+        MakeUnique(AutoMixedPrecisionMode::MKL));
   }
   if (cfg_.pin_to_host_optimization() == RewriterConfig::ON) {
     optimizers->push_back(MakeUnique());
@@ -835,6 +841,7 @@ bool MetaOptimizerEnabled(const ConfigProto& cfg) {
          rewrite_cfg.scoped_allocator_optimization() == RewriterConfig::ON ||
          rewrite_cfg.pin_to_host_optimization() == RewriterConfig::ON ||
          AutoMixedPrecisionEnabled(rewrite_cfg.auto_mixed_precision()) ||
+         AutoMixedPrecisionEnabled(rewrite_cfg.auto_mixed_precision_mkl()) ||
          !rewrite_cfg.optimizers().empty() ||
          !rewrite_cfg.custom_optimizers().empty();
 }
diff --git a/tensorflow/core/protobuf/rewriter_config.proto b/tensorflow/core/protobuf/rewriter_config.proto
index 38c3ad7ae57..9520db92742 100644
--- a/tensorflow/core/protobuf/rewriter_config.proto
+++ b/tensorflow/core/protobuf/rewriter_config.proto
@@ -85,11 +85,15 @@ message RewriterConfig {
   // Enable the swap of kernel implementations based on the device placement
   // (default is ON).
   Toggle implementation_selector = 22;
-  // Optimize data types (default is OFF).
-  // e.g., This will try to use float16 on GPU which is faster.
+  // Optimize data types for CUDA (default is OFF).
+  // This will try to use float16 on GPU which is faster.
   // Note that this can change the numerical stability of the graph and may
   // require the use of loss scaling to maintain model convergence.
   Toggle auto_mixed_precision = 23;
+  // Optimize data types for MKL (default is OFF).
+  // This will try to use bfloat16 on CPUs, which is faster.
+  // Note that this can change the numerical stability of the graph.
+  Toggle auto_mixed_precision_mkl = 25;
   // Disable the entire meta optimizer (off by default).
   bool disable_meta_optimizer = 19;
 
diff --git a/tensorflow/python/grappler/auto_mixed_precision_test.py b/tensorflow/python/grappler/auto_mixed_precision_test.py
index 494f6fc78fc..0de12b9eca8 100644
--- a/tensorflow/python/grappler/auto_mixed_precision_test.py
+++ b/tensorflow/python/grappler/auto_mixed_precision_test.py
@@ -19,8 +19,8 @@ from __future__ import division
 from __future__ import print_function
 
 import os
-import unittest
 
+from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.core.framework import types_pb2
@@ -209,7 +209,7 @@ def _make_node_with_color(color, input_tensor, name=None):
   if color == 'c':  # Clear node
     return nn.relu(input_tensor, name=name)
   if color == 'b':  # Black node
-    return math_ops.sqrt(math_ops.pow(input_tensor, 2.), name=name)
+    return math_ops.pow(math_ops.pow(input_tensor, 2.), 0.5, name=name)
   raise ValueError('Invalid node color: ' + str(color))
 
 
@@ -231,18 +231,21 @@ def _build_simple_loop_graph(inp_colors, body_colors, out_colors):
   return a
 
 
-def _get_config(auto_mixed_precision=True):
+def _get_config(auto_mixed_precision_mode):
   """Returns a ConfigProto with auto mixed precision enabled if appropriate."""
-  if auto_mixed_precision:
-    rewrite_config = rewriter_config_pb2.RewriterConfig(
-        auto_mixed_precision=rewriter_config_pb2.RewriterConfig.ON,
-        # do not remove duplicated nodes
-        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
+  rewrite_config = rewriter_config_pb2.RewriterConfig(
+      # do not remove duplicated nodes
+      arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
+      # do not turn Conv2D and other nodes into _FusedConv2D
+      remapping=rewriter_config_pb2.RewriterConfig.OFF,
+  )
+  if auto_mixed_precision_mode == 'cuda':
+    rewrite_config.auto_mixed_precision = rewriter_config_pb2.RewriterConfig.ON
+  elif auto_mixed_precision_mode == 'mkl':
+    rewrite_config.auto_mixed_precision_mkl = (
+        rewriter_config_pb2.RewriterConfig.ON)
   else:
-    rewrite_config = rewriter_config_pb2.RewriterConfig(
-        auto_mixed_precision=rewriter_config_pb2.RewriterConfig.OFF,
-        # do not remove duplicated nodes
-        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
+    assert auto_mixed_precision_mode is None
   rewrite_config.min_graph_nodes = -1
   graph_options = config_pb2.GraphOptions(
       rewrite_options=rewrite_config, build_cost_model=1)
@@ -255,19 +258,33 @@ def _is_cast_to_fp16(node_name):
   return node_name.endswith('-CastToFp16-AutoMixedPrecision')
 
 
+def _is_cast_to_bf16(node_name):
+  return node_name.endswith('-CastToBf16-AutoMixedPrecision')
+
+
 def _is_cast_to_fp32(node_name):
   return node_name.endswith('-CastToFp32-AutoMixedPrecision')
 
 
-def _count_casts(nodes):
+def _count_casts(mode, nodes):
+  """Counts the number of casts to f16 and fp32."""
   num_to_fp16 = 0
+  num_to_bf16 = 0
   num_to_fp32 = 0
   for node in nodes:
     if _is_cast_to_fp16(node.name):
       num_to_fp16 += 1
+    if _is_cast_to_bf16(node.name):
+      num_to_bf16 += 1
     elif _is_cast_to_fp32(node.name):
       num_to_fp32 += 1
-  return num_to_fp16, num_to_fp32
+  if mode == 'cuda':
+    assert num_to_bf16 == 0
+    return num_to_fp16, num_to_fp32
+  else:
+    assert mode == 'mkl'
+    assert num_to_fp16 == 0
+    return num_to_bf16, num_to_fp32
 
 
 def _build_node_map(nodes):
@@ -303,7 +320,7 @@ def _example_noninlined_funcdef(features):
   return features * math_ops.sigmoid(features)
 
 
-class AutoMixedPrecisionTest(test.TestCase):
+class AutoMixedPrecisionTest(test.TestCase, parameterized.TestCase):
   """Tests the Grappler auto mixed precision optimizer."""
   IGNORE_PERF_VAR = 'TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'
 
@@ -311,8 +328,8 @@ class AutoMixedPrecisionTest(test.TestCase):
 
   def setUp(self):
     super(AutoMixedPrecisionTest, self).setUp()
-    # Enable the tests to be run on pre-Volta GPUs by telling the grappler pass
-    # to ignore performance and always transform the graph.
+    # Enable the CUDA tests to be run on pre-Volta GPUs by telling the grappler
+    # pass to ignore performance and always transform the graph.
     self._original_ignore_perf_value = os.getenv(self.IGNORE_PERF_VAR)
     os.environ[self.IGNORE_PERF_VAR] = '1'
 
@@ -323,24 +340,33 @@ class AutoMixedPrecisionTest(test.TestCase):
       del os.environ[self.IGNORE_PERF_VAR]
     super(AutoMixedPrecisionTest, self).tearDown()
 
-  def _assert_output_fp16(self, node_map, node_name, output_port=0):
-    self.assertEqual(node_map[node_name].output_info[output_port].dtype,
-                     types_pb2.DT_HALF)
+  def _lower_precision_dtype(self, mode):
+    return dtypes.float16 if mode == 'cuda' else dtypes.bfloat16
 
-  def _run(self, fetches):
+  def _assert_output_f16(self, mode, node_map, node_name, output_port=0):
+    self.assertEqual(node_map[node_name].output_info[output_port].dtype,
+                     self._lower_precision_dtype(mode).as_datatype_enum)
+
+  def _run(self, mode, fetches):
     """Runs the graph and returns the evaluation of the fetches."""
-    with session.Session(config=_get_config(False)) as sess:
+    with session.Session(config=_get_config(None)) as sess:
       sess.run(variables.global_variables_initializer())
       output_val_ref = self.evaluate(fetches)
 
-    with session.Session(config=_get_config()) as sess:
+    with session.Session(config=_get_config(mode)) as sess:
       sess.run(variables.global_variables_initializer())
       metadata = config_pb2.RunMetadata()
       output_val = sess.run(fetches, run_metadata=metadata)
 
     return output_val_ref, output_val, metadata.cost_graph
 
-  def _run_simple_loop_test(self, inp, body, out):
+  def _maybe_skip(self, mode):
+    if mode == 'cuda' and not test.is_gpu_available(cuda_only=True):
+      self.skipTest('No GPU is available')
+    if mode == 'mkl' and not test_util.IsMklEnabled():
+      self.skipTest('MKL is not enabled')
+
+  def _run_simple_loop_test(self, mode, inp, body, out):
     """Runs a test of a simple loop.
 
     The loop has different node colors in different sections of the graph. The
@@ -359,398 +385,441 @@ class AutoMixedPrecisionTest(test.TestCase):
       out: A string of letters indicating the colors and expected dtypes of the
         output nodes.
     """
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      expected_types = []
-      for section in [inp, body, out]:
-        section_expected_types = []
-        for color in section:
-          if color.isupper():
-            expected_type = types_pb2.DT_HALF
-          else:
-            expected_type = types_pb2.DT_FLOAT
-          section_expected_types.append(expected_type)
-        expected_types.append(section_expected_types)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    expected_types = []
+    for section in [inp, body, out]:
+      section_expected_types = []
+      for color in section:
+        if color.isupper():
+          expected_type = self._lower_precision_dtype(mode).as_datatype_enum
+        else:
+          expected_type = types_pb2.DT_FLOAT
+        section_expected_types.append(expected_type)
+      expected_types.append(section_expected_types)
 
-      a = _build_simple_loop_graph(inp, body, out)
-      output_val_ref, output_val, cost_graph = self._run(a)
-      node_map = _build_node_map(cost_graph.node)
+    a = _build_simple_loop_graph(inp, body, out)
+    output_val_ref, output_val, cost_graph = self._run(mode, a)
+    node_map = _build_node_map(cost_graph.node)
 
-      section_names = ['input', 'while/body', 'output']
-      all_types_correct = True
-      for section_name, expected_types in zip(section_names, expected_types):
-        for i, expected_type in enumerate(expected_types):
-          node_name = section_name + '_%i' % i
-          output_port = 0
-          optimized_type = node_map[node_name].output_info[output_port].dtype
-          if optimized_type != expected_type:
-            print('Expected node %s to have type %s but got type %s' %
-                  (node_name, expected_type, optimized_type))
-            all_types_correct = False
-      self.assertTrue(all_types_correct)
+    section_names = ['input', 'while/body', 'output']
+    all_types_correct = True
+    for section_name, expected_types in zip(section_names, expected_types):
+      for i, expected_type in enumerate(expected_types):
+        node_name = section_name + '_%i' % i
+        output_port = 0
+        optimized_type = node_map[node_name].output_info[output_port].dtype
+        if optimized_type != expected_type:
+          print('Expected node %s to have type %s but got type %s' %
+                (node_name, expected_type, optimized_type))
+          all_types_correct = False
+    self.assertTrue(all_types_correct)
+    if mode == 'mkl':
+      self.assertAllClose(output_val_ref, output_val, atol=2e-2, rtol=2e-2)
+    else:
       self.assertAllClose(output_val_ref, output_val, atol=2e-3, rtol=1e-3)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_conv_bn(self):
+  def test_conv_bn(self, mode):
     """Test graph with convolution followed by batch norm."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 1])
-      x = _conv_bn(x)
-      output = _conv_bn(x)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 1])
+    x = _conv_bn(x)
+    output = _conv_bn(x)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
-      num_to_fp16, num_to_fp32 = _count_casts(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
+    num_to_f16, num_to_fp32 = _count_casts(mode, cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'Conv2D')
-      self._assert_output_fp16(node_map, 'FusedBatchNormV3')
-      self._assert_output_fp16(node_map, 'Conv2D_1')
-      self.assertEqual(num_to_fp16,
-                       3)  # Before Conv2D:0, Conv2D:1, Conv2D_1:1
-      self.assertEqual(num_to_fp32, 1)  # After FusedBatchNormV3:0
+    self._assert_output_f16(mode, node_map, 'Conv2D')
+    self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
+    self._assert_output_f16(mode, node_map, 'Conv2D_1')
+    self.assertEqual(num_to_f16, 3)  # Before Conv2D:0, Conv2D:1, Conv2D_1:1
+    self.assertEqual(num_to_fp32, 1)  # After FusedBatchNormV3:0
+    if mode == 'mkl':
+      tol = 1e-2
+    elif test.is_built_with_rocm():
       # Bump up the tolerance for the ROCm platform
       # The default tolerance (1e-3) results in a tiny fraction (<1%) of
       # miscompares on ROCm platform, and hence the tolerance bump
-      tol = 2e-3 if test.is_built_with_rocm else 1e-3
-      self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
+      tol = 2e-3
+    else:
+      tol = 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
-  # TODO: enable these tests when cuDNN is upgraded to >= 7.6.2. Same with the
-  # test_conv3d() below.
-  @unittest.skip('Test case should be skipped when cuDNN < 7.6.2')
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_conv3d_bn(self):
+  def test_conv3d_bn(self, mode):
     """Test graph with convolution followed by batch norm."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 8, 1])
-      x = _conv3d_bn(x)
-      output = _conv3d_bn(x)
+    self._maybe_skip(mode)
+    if mode == 'cuda':
+      # TODO: enable these tests when cuDNN is upgraded to >= 7.6.2.
+      self.skipTest('Test case should be skipped when cuDNN < 7.6.2')
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 8, 1])
+    x = _conv3d_bn(x)
+    output = _conv3d_bn(x)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
-      num_to_fp16, num_to_fp32 = _count_casts(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
+    num_to_fp16, num_to_fp32 = _count_casts(mode, cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'Conv3D')
-      self._assert_output_fp16(node_map, 'FusedBatchNormV3')
-      self._assert_output_fp16(node_map, 'Conv3D_1')
-      self.assertEqual(num_to_fp16, 3)  # Before Conv3D:0, Conv3D:1, Conv3D_1:1
-      self.assertEqual(num_to_fp32, 1)  # After FusedBatchNormV3:0
-      self.assertAllClose(output_val_ref, output_val, atol=1e-2, rtol=1e-2)
+    self._assert_output_f16(mode, node_map, 'Conv3D')
+    self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
+    self._assert_output_f16(mode, node_map, 'Conv3D_1')
+    self.assertEqual(num_to_fp16, 3)  # Before Conv3D:0, Conv3D:1, Conv3D_1:1
+    self.assertEqual(num_to_fp32, 1)  # After FusedBatchNormV3:0
+    self.assertAllClose(output_val_ref, output_val, atol=1e-2, rtol=1e-2)
 
-  @unittest.skip('Test case should be skipped when cuDNN < 7.6.2')
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_conv3d(self):
+  def test_conv3d(self, mode):
     """Test grad ops with convolution3d graph."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 8, 1])
-      f = _weight([3, 3, 3, 1, 6])
-      y = _conv3d(x, f)
-      y = array_ops.identity(y)
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x, f])
-      output = (y, g)
+    self._maybe_skip(mode)
+    if mode == 'cuda':
+      # TODO: enable these tests when cuDNN is upgraded to >= 7.6.2.
+      self.skipTest('Test case should be skipped when cuDNN < 7.6.2')
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 8, 1])
+    f = _weight([3, 3, 3, 1, 6])
+    y = _conv3d(x, f)
+    y = array_ops.identity(y)
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x, f])
+    output = (y, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
-      self._assert_output_fp16(node_map, 'Conv3D')
-      self._assert_output_fp16(node_map,
-                               'gradients/Conv3D_grad/Conv3DBackpropInputV2')
-      self._assert_output_fp16(node_map,
-                               'gradients/Conv3D_grad/Conv3DBackpropFilterV2')
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
+    self._assert_output_f16(mode, node_map, 'Conv3D')
+    self._assert_output_f16(mode, node_map,
+                            'gradients/Conv3D_grad/Conv3DBackpropInputV2')
+    self._assert_output_f16(mode, node_map,
+                            'gradients/Conv3D_grad/Conv3DBackpropFilterV2')
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    tol = 5e-2 if mode == 'mkl' else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  # TODO(reedwm): Fix and enable this test with MKL. Currently this crashes with
+  # MKL
+  @parameterized.parameters(['cuda'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_conv_bn_dropout(self):
+  def test_conv_bn_dropout(self, mode):
     """Test dropout precision of convolution batch norm graph."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 1])
-      y = _conv_bn(x)
-      y = nn.dropout(y, rate=0.5)
-      y = math_ops.add(y, 1, name='addition')
-      y = _conv_bn(y)
-      y = array_ops.identity(y)
-      optimizer = gradient_descent.GradientDescentOptimizer(
-          learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x])
-      output = (y, g)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 1])
+    y = _conv_bn(x)
+    y = nn.dropout(y, rate=0.5)
+    y = math_ops.add(y, 1, name='addition')
+    y = _conv_bn(y)
+    y = array_ops.identity(y)
+    optimizer = gradient_descent.GradientDescentOptimizer(
+        learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x])
+    output = (y, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
-      self._assert_output_fp16(node_map, 'Conv2D')
-      self._assert_output_fp16(node_map, 'FusedBatchNormV3')
-      # We do not assert dropout's dtype because we do not want to rely on the
-      # node names of dropout's internal implementation.
-      self._assert_output_fp16(node_map, 'addition')
-      self._assert_output_fp16(node_map, 'Conv2D_1')
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
+    self._assert_output_f16(mode, node_map, 'Conv2D')
+    self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
+    # We do not assert dropout's dtype because we do not want to rely on the
+    # node names of dropout's internal implementation.
+    self._assert_output_f16(mode, node_map, 'addition')
+    self._assert_output_f16(mode, node_map, 'Conv2D_1')
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      # Bump up the tolerance for the ROCm platform
-      # The default tolerance (1e-3) results in a tiny fraction (<1%) of
-      # miscompares on ROCm platform, and hence the tolerance bump
-      tol = 2e-3 if test.is_built_with_rocm else 1e-3
-      self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    # Bump up the tolerance for the ROCm platform
+    # The default tolerance (1e-3) results in a tiny fraction (<1%) of
+    # miscompares on ROCm platform, and hence the tolerance bump
+    tol = 2e-3 if test.is_built_with_rocm else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  # TODO(reedwm): Fix and enable this test with MKL. Currently this crashes with
+  # MKL
+  @parameterized.parameters(['cuda'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_conv_pool(self):
+  def test_conv_pool(self, mode):
     """Test graph with convolution followed by pooling."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 1])
-      output = _conv_pool(x)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 1])
+    output = _conv_pool(x)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
-      num_to_fp16, num_to_fp32 = _count_casts(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
+    num_to_f16, num_to_fp32 = _count_casts(mode, cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'Conv2D')
-      self._assert_output_fp16(node_map, 'Relu')
-      self._assert_output_fp16(node_map, 'MaxPool')
-      self._assert_output_fp16(node_map, 'Conv2D_1')
-      self.assertEqual(num_to_fp16, 4)
-      self.assertEqual(num_to_fp32, 1)
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'Conv2D')
+    self._assert_output_f16(mode, node_map, 'Relu')
+    self._assert_output_f16(mode, node_map, 'MaxPool')
+    self._assert_output_f16(mode, node_map, 'Conv2D_1')
+    self.assertEqual(num_to_f16, 4)
+    self.assertEqual(num_to_fp32, 1)
+    tol = 5e-3 if mode == 'mkl' else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_simple_loop(self):
+  def test_simple_loop(self, mode):
     """Test graph with while loop."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([8, 8])
-      y = _simple_loop(x, _matmul_act)[1]
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x])
-      output = (y, g)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([8, 8])
+    y = _simple_loop(x, _matmul_act)[1]
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x])
+    output = (y, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'while/MatMul')
-      self._assert_output_fp16(node_map, 'while/Relu')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'while/MatMul')
+    self._assert_output_f16(mode, node_map, 'while/Relu')
+    tol = 1e-2 if mode == 'mkl' else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_loop_with_vars_intertwined(self):
+  def test_loop_with_vars_intertwined(self, mode):
     """Test graph with intertwined while loops."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([8, 8])
-      _, _, k, l = _loop_vars_intertwined(
-          array_ops.ones(array_ops.shape(x)), x, _matmul_act, _matmul_act)
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(k, [x])
-      output = (k, l, g)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([8, 8])
+    _, _, k, l = _loop_vars_intertwined(
+        array_ops.ones(array_ops.shape(x)), x, _matmul_act, _matmul_act)
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(k, [x])
+    output = (k, l, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'while/MatMul')
-      self._assert_output_fp16(node_map, 'while/Relu')
-      self._assert_output_fp16(node_map, 'while/MatMul_1')
-      self._assert_output_fp16(node_map, 'while/Relu_1')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'while/MatMul')
+    self._assert_output_f16(mode, node_map, 'while/Relu')
+    self._assert_output_f16(mode, node_map, 'while/MatMul_1')
+    self._assert_output_f16(mode, node_map, 'while/Relu_1')
+    tol = 5e-3 if mode == 'mkl' else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  @parameterized.parameters(['cuda'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_multi_paths(self):
+  def test_multi_paths(self, mode):
     """Test graph with multiple paths."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([2, 8, 8, 3])
-      x1, x2, x3 = array_ops.split(x, num_or_size_splits=3, axis=3)
-      y1 = _conv_pool(x1)
-      y2 = _conv_pool(x2)
-      y3 = _conv_pool(x3)
-      y = array_ops.concat([y1, y2, y3], axis=3)
-      y = array_ops.identity(y)
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x])
-      output = (y, g)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([2, 8, 8, 3])
+    x1, x2, x3 = array_ops.split(x, num_or_size_splits=3, axis=3)
+    y1 = _conv_pool(x1)
+    y2 = _conv_pool(x2)
+    y3 = _conv_pool(x3)
+    y = array_ops.concat([y1, y2, y3], axis=3)
+    y = array_ops.identity(y)
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x])
+    output = (y, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'split')
-      for suffix in [''] + ['_%i' % i for i in range(1, 6)]:
-        self._assert_output_fp16(node_map, 'Conv2D' + suffix)
-        self._assert_output_fp16(node_map, 'Relu' + suffix)
-        self._assert_output_fp16(node_map, 'MaxPool' + suffix)
-      self._assert_output_fp16(node_map, 'concat')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'split')
+    for suffix in [''] + ['_%i' % i for i in range(1, 6)]:
+      self._assert_output_f16(mode, node_map, 'Conv2D' + suffix)
+      self._assert_output_f16(mode, node_map, 'Relu' + suffix)
+      self._assert_output_f16(mode, node_map, 'MaxPool' + suffix)
+    self._assert_output_f16(mode, node_map, 'concat')
+    self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_multi_paths_2(self):
+  def test_multi_paths_2(self, mode):
     """Test graph with multiple paths."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([8, 8])
-      y1 = _matmul_act(x)
-      y2 = _matmul_act(x)
-      y = y1 + y2 + x
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x])
-      output = (g, y)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([8, 8])
+    y1 = _matmul_act(x)
+    y2 = _matmul_act(x)
+    y = y1 + y2 + x
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x])
+    output = (g, y)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'MatMul')
-      self._assert_output_fp16(node_map, 'Relu')
-      self._assert_output_fp16(node_map, 'MatMul_1')
-      self._assert_output_fp16(node_map, 'Relu_1')
+    self._assert_output_f16(mode, node_map, 'MatMul')
+    self._assert_output_f16(mode, node_map, 'Relu')
+    self._assert_output_f16(mode, node_map, 'MatMul_1')
+    self._assert_output_f16(mode, node_map, 'Relu_1')
+    if mode == 'mkl':
+      tol = 2e-2
+    elif test.is_built_with_rocm():
       # Bump up the tolerance for the ROCm platform
       # The default tolerance (1e-3) results in a tiny fraction (<1%) of
       # miscompares on ROCm platform, and hence the tolerance bump
-      tol = 2e-3 if test.is_built_with_rocm else 1e-3
-      self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
+      tol = 2e-3
+    else:
+      tol = 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  @parameterized.parameters(['cuda'])  # MKL doesn't support bf16 Sigmoid
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_recurrent_lstm(self):
+  def test_recurrent_lstm(self, mode):
     """Test graph with recurrent lstm."""
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      init_c = _input([8, 4])
-      init_h = _input([8, 4])
-      _, _, h, _ = _recurrent_lstm(init_c, init_h)
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(h, [init_c, init_h])
-      output = (h, g)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    init_c = _input([8, 4])
+    init_h = _input([8, 4])
+    _, _, h, _ = _recurrent_lstm(init_c, init_h)
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(h, [init_c, init_h])
+    output = (h, g)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'while/concat')
-      self._assert_output_fp16(node_map, 'while/MatMul')
-      self._assert_output_fp16(node_map, 'while/split')
-      self._assert_output_fp16(node_map, 'while/Sigmoid')
-      self._assert_output_fp16(node_map, 'while/Sigmoid_1')
-      self._assert_output_fp16(node_map, 'while/Sigmoid_2')
-      self._assert_output_fp16(node_map, 'while/Tanh')
-      self._assert_output_fp16(node_map, 'while/Tanh_1')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'while/concat')
+    self._assert_output_f16(mode, node_map, 'while/MatMul')
+    self._assert_output_f16(mode, node_map, 'while/split')
+    self._assert_output_f16(mode, node_map, 'while/Sigmoid')
+    self._assert_output_f16(mode, node_map, 'while/Sigmoid_1')
+    self._assert_output_f16(mode, node_map, 'while/Sigmoid_2')
+    self._assert_output_f16(mode, node_map, 'while/Tanh')
+    self._assert_output_f16(mode, node_map, 'while/Tanh_1')
+    self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('v1 loop test')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_1(self):
-    self._run_simple_loop_test('W', 'C', 'C')
+  def test_propagation_through_simple_loop_1(self, mode):
+    self._run_simple_loop_test(mode, 'W', 'C', 'C')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('v1 loop test')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_2(self):
-    self._run_simple_loop_test('C', 'C', 'W')
+  def test_propagation_through_simple_loop_2(self, mode):
+    self._run_simple_loop_test(mode, 'C', 'C', 'W')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('v1 loop test')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_3(self):
-    self._run_simple_loop_test('W', 'G', 'W')
+  def test_propagation_through_simple_loop_3(self, mode):
+    self._run_simple_loop_test(mode, 'W', 'G', 'W')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('v1 loop test')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_4(self):
-    self._run_simple_loop_test('W', 'gbg', 'W')
+  def test_propagation_through_simple_loop_4(self, mode):
+    self._run_simple_loop_test(mode, 'W', 'gbg', 'W')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_5(self):
-    self._run_simple_loop_test('b', 'gWC', 'c')
+  def test_propagation_through_simple_loop_5(self, mode):
+    self._run_simple_loop_test(mode, 'b', 'gWC', 'c')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_6(self):
-    self._run_simple_loop_test('b', 'CWCG', 'C')
+  def test_propagation_through_simple_loop_6(self, mode):
+    self._run_simple_loop_test(mode, 'b', 'CWCG', 'C')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_7(self):
-    self._run_simple_loop_test('C', 'GWCG', 'C')
+  def test_propagation_through_simple_loop_7(self, mode):
+    self._run_simple_loop_test(mode, 'C', 'GWCG', 'C')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_v1_only('b/138749235')
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_propagation_through_simple_loop_8(self):
-    self._run_simple_loop_test('C', 'CgbgWC', 'g')
+  def test_propagation_through_simple_loop_8(self, mode):
+    self._run_simple_loop_test(mode, 'C', 'CgbgWC', 'g')
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_noninlined_funcdef(self):
+  def test_noninlined_funcdef(self, mode):
     """Test graph with non-inlined function subgraph.
 
     This requires the grappler pass to handle an OpDef that only appears in the
     graph's function registry instead of the global op registry.
     """
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(0)
-      x = _input([8, 8])
-      y = _matmul_act(x)
-      y = _example_noninlined_funcdef(y)
-      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
-      g = optimizer.compute_gradients(y, [x])
-      output = (g, y)
+    self._maybe_skip(mode)
+    random_seed.set_random_seed(0)
+    x = _input([8, 8])
+    y = _matmul_act(x)
+    y = _example_noninlined_funcdef(y)
+    optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+    g = optimizer.compute_gradients(y, [x])
+    output = (g, y)
 
-      output_val_ref, output_val, cost_graph = self._run(output)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, output)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'MatMul')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'MatMul')
+    tol = 1e-2 if mode == 'mkl' else 1e-3
+    self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
 
+  @parameterized.parameters(['cuda', 'mkl'])
   @test_util.run_deprecated_v1
   @test_util.disable_xla('This test does not pass with XLA')
-  def test_ingraph_train_loop(self):
+  def test_ingraph_train_loop(self, mode):
     """Tests a graph containing a while loop around a training update.
 
     This requires the grappler pass to take special care with its handling of
     Enter ops that appear in front of reads from non-resource variables. See
     the use of NodeImplicitlyReadsVariable in auto_mixed_precision.cc.
     """
+    self._maybe_skip(mode)
     if tf2.enabled():
       # This test tests non-resource variables, which are only used in TF1.
       self.skipTest('TensorFlow 1 required')
-    if test.is_gpu_available(cuda_only=True):
-      random_seed.set_random_seed(1234)
-      np.random.seed(1234)
-      num_iter, bs, nchan, nclass = 100, 64, 32, 100
+    random_seed.set_random_seed(1234)
+    np.random.seed(1234)
+    num_iter, bs, nchan, nclass = 100, 64, 32, 100
 
-      data = np.random.normal(size=(bs * num_iter, nchan)).astype(np.float32)
-      labels = np.random.randint(nclass, size=(bs * num_iter,))
-      ds = dataset_ops.Dataset.from_tensor_slices((data, labels))
-      ds = ds.batch(bs).prefetch(3)
-      it = ds.make_one_shot_iterator()
+    data = np.random.normal(size=(bs * num_iter, nchan)).astype(np.float32)
+    labels = np.random.randint(nclass, size=(bs * num_iter,))
+    ds = dataset_ops.Dataset.from_tensor_slices((data, labels))
+    ds = ds.batch(bs).prefetch(3)
+    it = ds.make_one_shot_iterator()
 
-      def body(_, i):
-        i += 1
-        x, yt = it.get_next()
-        dense = layers.Dense(nclass)
-        y = dense(x)
-        loss = losses.sparse_softmax_cross_entropy(yt, y)
-        opt = adam.AdamOptimizer()
-        train_op = opt.minimize(loss, var_list=dense.trainable_weights)
-        with ops.control_dependencies([train_op]):
-          loss = array_ops.identity(loss)
-        return loss, i
+    def body(_, i):
+      i += 1
+      x, yt = it.get_next()
+      dense = layers.Dense(nclass)
+      y = dense(x)
+      loss = losses.sparse_softmax_cross_entropy(yt, y)
+      opt = adam.AdamOptimizer()
+      train_op = opt.minimize(loss, var_list=dense.trainable_weights)
+      with ops.control_dependencies([train_op]):
+        loss = array_ops.identity(loss)
+      return loss, i
 
-      begin, end = constant_op.constant(0), constant_op.constant(num_iter)
-      loss, _ = control_flow_ops.while_loop(
-          lambda loss, i: math_ops.less(i, end), body, [0.0, begin])
+    begin, end = constant_op.constant(0), constant_op.constant(num_iter)
+    loss, _ = control_flow_ops.while_loop(
+        lambda loss, i: math_ops.less(i, end), body, [0.0, begin])
 
-      output_val_ref, output_val, cost_graph = self._run(loss)
-      node_map = _build_node_map(cost_graph.node)
+    output_val_ref, output_val, cost_graph = self._run(mode, loss)
+    node_map = _build_node_map(cost_graph.node)
 
-      self._assert_output_fp16(node_map, 'while/dense/MatMul')
-      self._assert_output_fp16(
-          node_map, 'while/gradients/while/dense/MatMul_grad/MatMul_1')
-      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+    self._assert_output_f16(mode, node_map, 'while/dense/MatMul')
+    self._assert_output_f16(
+        mode, node_map, 'while/gradients/while/dense/MatMul_grad/MatMul_1')
+    self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
 
   # TODO(benbarsdell): Add tests for list ops (TensorList*) that pass through
   # graph source/sink nodes, similar to the TensorListThroughFunction C++ test.

From ebe063eb74ab1ee80bfb0d4447fb7c842e0ca27b Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Thu, 18 Jun 2020 12:38:59 -0700
Subject: [PATCH 0521/1390] Create an empty test for python_op_gen.

PiperOrigin-RevId: 317159004
Change-Id: Iabc4810f9c51c257d62dc3e7bd20d96131939d5d
---
 tensorflow/python/BUILD                       | 15 ++++++-
 tensorflow/python/framework/python_op_gen.cc  | 16 ++++---
 tensorflow/python/framework/python_op_gen.h   | 14 ++++++-
 .../python/framework/python_op_gen_test.cc    | 42 +++++++++++++++++++
 4 files changed, 80 insertions(+), 7 deletions(-)
 create mode 100644 tensorflow/python/framework/python_op_gen_test.cc

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index a4e72bf2460..de9cf9a24c7 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -3,7 +3,7 @@
 #  ":platform" - Low-level and platform-specific Python code.
 
 load("//tensorflow:tensorflow.bzl", "py_strict_library")
-load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cuda_library", "tf_gen_op_wrapper_py")
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cc_test", "tf_cuda_library", "tf_gen_op_wrapper_py")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_monitoring_python_deps")
@@ -1236,6 +1236,19 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_cc_test(
+    name = "python_op_gen_test",
+    srcs = ["framework/python_op_gen_test.cc"],
+    deps = [
+        ":python_op_gen",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:op_gen_lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
 py_library(
     name = "framework_for_generated_wrappers",
     srcs_version = "PY2AND3",
diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index ca0c5d9ef1a..0f84c6a063d 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -981,9 +981,9 @@ void GenEagerPythonOp::AddRawOpExport(const string& parameters) {
                      function_name_, "))\n");
 }
 
-string GetPythonOps(const OpList& ops, const ApiDefMap& api_defs,
-                    const std::vector& hidden_ops,
-                    const string& source_file_name = "") {
+string GetPythonOpsImpl(const OpList& ops, const ApiDefMap& api_defs,
+                        const std::vector& hidden_ops,
+                        const string& source_file_name = "") {
   string result;
   // Header
   // TODO(josh11b): Mention the library for which wrappers are being generated.
@@ -1069,11 +1069,17 @@ from tensorflow.python.util.tf_export import tf_export
 
 }  // namespace
 
+string GetPythonOps(const OpList& ops, const ApiDefMap& api_defs,
+                    const std::vector& hidden_ops,
+                    const string& source_file_name) {
+  return GetPythonOpsImpl(ops, api_defs, hidden_ops, source_file_name);
+}
+
 void PrintPythonOps(const OpList& ops, const ApiDefMap& api_defs,
                     const std::vector& hidden_ops,
                     const string& source_file_name) {
   printf("%s",
-         GetPythonOps(ops, api_defs, hidden_ops, source_file_name).c_str());
+         GetPythonOpsImpl(ops, api_defs, hidden_ops, source_file_name).c_str());
 }
 
 string GetPythonWrappers(const char* op_list_buf, size_t op_list_len) {
@@ -1081,7 +1087,7 @@ string GetPythonWrappers(const char* op_list_buf, size_t op_list_len) {
   ops.ParseFromArray(op_list_buf, op_list_len);
 
   ApiDefMap api_def_map(ops);
-  return GetPythonOps(ops, api_def_map, {});
+  return GetPythonOpsImpl(ops, api_def_map, {});
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/python/framework/python_op_gen.h b/tensorflow/python/framework/python_op_gen.h
index 22fcc452fbb..f1cd6e49013 100644
--- a/tensorflow/python/framework/python_op_gen.h
+++ b/tensorflow/python/framework/python_op_gen.h
@@ -23,8 +23,20 @@ limitations under the License.
 
 namespace tensorflow {
 
+// Returns a string containing the generated Python code for the given Ops.
+// ops is a protobuff, typically generated using OpRegistry::Global()->Export.
+// api_defs is typically constructed directly from ops.
 // hidden_ops should be a list of Op names that should get a leading _
-// in the output. Prints the output to stdout.
+// in the output.
+// source_file_name is optional and contains the name of the original C++ source
+// file where the ops' REGISTER_OP() calls reside.
+string GetPythonOps(const OpList& ops, const ApiDefMap& api_defs,
+                    const std::vector& hidden_ops,
+                    const string& source_file_name);
+
+// Prints the output of GetPrintOps to stdout.
+// hidden_ops should be a list of Op names that should get a leading _
+// in the output.
 // Optional fourth argument is the name of the original C++ source file
 // where the ops' REGISTER_OP() calls reside.
 void PrintPythonOps(const OpList& ops, const ApiDefMap& api_defs,
diff --git a/tensorflow/python/framework/python_op_gen_test.cc b/tensorflow/python/framework/python_op_gen_test.cc
new file mode 100644
index 00000000000..5185086fdd3
--- /dev/null
+++ b/tensorflow/python/framework/python_op_gen_test.cc
@@ -0,0 +1,42 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/python/framework/python_op_gen.h"
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(PythonOpGen, Basic) {
+  OpList ops;
+  OpRegistry::Global()->Export(false, &ops);
+
+  ApiDefMap api_def_map(ops);
+
+  string code = GetPythonOps(ops, api_def_map, {}, "");
+
+  EXPECT_TRUE(absl::StrContains(code, "def case"));
+
+  // TODO(mdan): Add tests to verify type annotations are correctly added.
+}
+
+// TODO(mdan): Include more tests with synhtetic ops and api defs.
+
+}  // namespace
+}  // namespace tensorflow

From b0418130b440fa63b588390f37d0ab21b1c4731c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 12:40:15 -0700
Subject: [PATCH 0522/1390] Go: Update generated wrapper functions for
 TensorFlow ops.

PiperOrigin-RevId: 317159284
Change-Id: I438a0e11d3e678fd3d5bc153a5b0ff353e0fe86a
---
 tensorflow/go/op/wrappers.go | 794 +++++++++++++++++------------------
 1 file changed, 397 insertions(+), 397 deletions(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 188bb3b78bb..ace8e58fdcd 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -158,65 +158,6 @@ func FakeQuantWithMinMaxArgsGradient(scope *Scope, gradients tf.Output, inputs t
 	return op.Output(0)
 }
 
-// Applies sparse addition to `input` using individual values or slices
-//
-// from `updates` according to indices `indices`.  The updates are non-aliasing:
-// `input` is only modified in-place if no other operations will use it.
-// Otherwise, a copy of `input` is made.  This operation has a gradient with
-// respect to both `input` and `updates`.
-//
-// `input` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
-//
-// `indices` must be integer tensor, containing indices into `input`.
-// It must be shape \\([d_0, ..., d_{Q-2}, K]\\) where `0 < K <= P`.
-//
-// The innermost dimension of `indices` (with length `K`) corresponds to
-// indices into elements (if `K = P`) or `(P-K)`-dimensional slices
-// (if `K < P`) along the `K`th dimension of `input`.
-//
-// `updates` is `Tensor` of rank `Q-1+P-K` with shape:
-//
-// $$[d_0, ..., d_{Q-2}, input.shape[K], ..., input.shape[P-1]].$$
-//
-// For example, say we want to add 4 scattered elements to a rank-1 tensor to 8
-// elements. In Python, that addition would look like this:
-//
-//     input = tf.constant([1, 2, 3, 4, 5, 6, 7, 8])
-//     indices = tf.constant([[4], [3], [1], [7]])
-//     updates = tf.constant([9, 10, 11, 12])
-//     output = tf.scatter_nd_non_aliasing_add(input, indices, updates)
-//     with tf.Session() as sess:
-//       print(sess.run(output))
-//
-// The resulting value `output` would look like this:
-//
-//     [1, 13, 3, 14, 14, 6, 7, 20]
-//
-// See `tf.scatter_nd` for more details about how to make updates to slices.
-//
-// Arguments:
-//	input: A Tensor.
-//	indices: A Tensor. Must be one of the following types: `int32`, `int64`.
-// A tensor of indices into `input`.
-//	updates: A Tensor. Must have the same type as ref. A tensor of updated values
-// to add to `input`.
-//
-// Returns A `Tensor` with the same shape as `input`, containing values of `input`
-// updated with `updates`.
-func ScatterNdNonAliasingAdd(scope *Scope, input tf.Output, indices tf.Output, updates tf.Output) (output tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	opspec := tf.OpSpec{
-		Type: "ScatterNdNonAliasingAdd",
-		Input: []tf.Input{
-			input, indices, updates,
-		},
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // Subtracts sparse `updates` from an existing tensor according to `indices`.
 //
 // This operation creates a new tensor by subtracting sparse `updates` from the
@@ -13210,115 +13151,6 @@ func ResizeNearestNeighborGrad(scope *Scope, grads tf.Output, size tf.Output, op
 	return op.Output(0)
 }
 
-// A placeholder op for a value that will be fed into the computation.
-//
-// Arguments:
-//	dtype: The type of elements in the tensor.
-//	shape: The shape of the tensor.
-//
-// Returns A tensor that will be provided using the infeed mechanism.
-func InfeedDequeue(scope *Scope, dtype tf.DataType, shape tf.Shape) (output tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	attrs := map[string]interface{}{"dtype": dtype, "shape": shape}
-	opspec := tf.OpSpec{
-		Type: "InfeedDequeue",
-
-		Attrs: attrs,
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
-// Encodes a `RaggedTensor` into a `variant` Tensor.
-//
-//
-// Encodes the given `RaggedTensor` and returns a `variant` Tensor. If
-// `batched_input` is True, then input `RaggedTensor` is unbatched along the
-// zero-th dimension, each component `RaggedTensor` is encoded into a scalar
-// `variant` Tensor, and these are stacked to return a 1-D `variant` Tensor.
-// If `batched_input` is False, then the input `RaggedTensor` is encoded as is and
-// a scalar `variant` Tensor is returned. A `RaggedTensor` is encoded by first
-// creating a 1-D `variant` Tensor with `ragged_rank + 1` elements, containing the
-// splits and values Tensors of the `RaggedTensor`. Then the 1-D `variant` Tensor
-// is wrapped in a scalar `variant` Tensor. See `RaggedTensorFromVariant` for the
-// corresponding decoding logic.
-//
-//
-// Arguments:
-//	rt_nested_splits: A list of one or more Tensors representing the splits of the input
-// `RaggedTensor`.
-//	rt_dense_values: A Tensor representing the values of the input `RaggedTensor`.
-//	batched_input: A `bool` denoting whether the input is a batched `RaggedTensor`.
-//
-// Returns A `variant` Tensor that containing encoded `RaggedTensor`.
-func RaggedTensorToVariant(scope *Scope, rt_nested_splits []tf.Output, rt_dense_values tf.Output, batched_input bool) (encoded_ragged tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	attrs := map[string]interface{}{"batched_input": batched_input}
-	opspec := tf.OpSpec{
-		Type: "RaggedTensorToVariant",
-		Input: []tf.Input{
-			tf.OutputList(rt_nested_splits), rt_dense_values,
-		},
-		Attrs: attrs,
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
-// ResizeNearestNeighborAttr is an optional argument to ResizeNearestNeighbor.
-type ResizeNearestNeighborAttr func(optionalAttr)
-
-// ResizeNearestNeighborAlignCorners sets the optional align_corners attribute to value.
-//
-// value: If true, the centers of the 4 corner pixels of the input and output tensors are
-// aligned, preserving the values at the corner pixels. Defaults to false.
-// If not specified, defaults to false
-func ResizeNearestNeighborAlignCorners(value bool) ResizeNearestNeighborAttr {
-	return func(m optionalAttr) {
-		m["align_corners"] = value
-	}
-}
-
-// ResizeNearestNeighborHalfPixelCenters sets the optional half_pixel_centers attribute to value.
-// If not specified, defaults to false
-func ResizeNearestNeighborHalfPixelCenters(value bool) ResizeNearestNeighborAttr {
-	return func(m optionalAttr) {
-		m["half_pixel_centers"] = value
-	}
-}
-
-// Resize `images` to `size` using nearest neighbor interpolation.
-//
-// Arguments:
-//	images: 4-D with shape `[batch, height, width, channels]`.
-//	size: = A 1-D int32 Tensor of 2 elements: `new_height, new_width`.  The
-// new size for the images.
-//
-// Returns 4-D with shape
-// `[batch, new_height, new_width, channels]`.
-func ResizeNearestNeighbor(scope *Scope, images tf.Output, size tf.Output, optional ...ResizeNearestNeighborAttr) (resized_images tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	attrs := map[string]interface{}{}
-	for _, a := range optional {
-		a(attrs)
-	}
-	opspec := tf.OpSpec{
-		Type: "ResizeNearestNeighbor",
-		Input: []tf.Input{
-			images, size,
-		},
-		Attrs: attrs,
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // Runs multiple additive regression ensemble predictors on input instances and
 //
 // computes the logits. It is designed to be used during prediction.
@@ -15836,6 +15668,65 @@ func SqrtGrad(scope *Scope, y tf.Output, dy tf.Output) (z tf.Output) {
 	return op.Output(0)
 }
 
+// Applies sparse addition to `input` using individual values or slices
+//
+// from `updates` according to indices `indices`.  The updates are non-aliasing:
+// `input` is only modified in-place if no other operations will use it.
+// Otherwise, a copy of `input` is made.  This operation has a gradient with
+// respect to both `input` and `updates`.
+//
+// `input` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.
+//
+// `indices` must be integer tensor, containing indices into `input`.
+// It must be shape \\([d_0, ..., d_{Q-2}, K]\\) where `0 < K <= P`.
+//
+// The innermost dimension of `indices` (with length `K`) corresponds to
+// indices into elements (if `K = P`) or `(P-K)`-dimensional slices
+// (if `K < P`) along the `K`th dimension of `input`.
+//
+// `updates` is `Tensor` of rank `Q-1+P-K` with shape:
+//
+// $$[d_0, ..., d_{Q-2}, input.shape[K], ..., input.shape[P-1]].$$
+//
+// For example, say we want to add 4 scattered elements to a rank-1 tensor to 8
+// elements. In Python, that addition would look like this:
+//
+//     input = tf.constant([1, 2, 3, 4, 5, 6, 7, 8])
+//     indices = tf.constant([[4], [3], [1], [7]])
+//     updates = tf.constant([9, 10, 11, 12])
+//     output = tf.scatter_nd_non_aliasing_add(input, indices, updates)
+//     with tf.Session() as sess:
+//       print(sess.run(output))
+//
+// The resulting value `output` would look like this:
+//
+//     [1, 13, 3, 14, 14, 6, 7, 20]
+//
+// See `tf.scatter_nd` for more details about how to make updates to slices.
+//
+// Arguments:
+//	input: A Tensor.
+//	indices: A Tensor. Must be one of the following types: `int32`, `int64`.
+// A tensor of indices into `input`.
+//	updates: A Tensor. Must have the same type as ref. A tensor of updated values
+// to add to `input`.
+//
+// Returns A `Tensor` with the same shape as `input`, containing values of `input`
+// updated with `updates`.
+func ScatterNdNonAliasingAdd(scope *Scope, input tf.Output, indices tf.Output, updates tf.Output) (output tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "ScatterNdNonAliasingAdd",
+		Input: []tf.Input{
+			input, indices, updates,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // MutableHashTableOfTensorsV2Attr is an optional argument to MutableHashTableOfTensorsV2.
 type MutableHashTableOfTensorsV2Attr func(optionalAttr)
 
@@ -37168,235 +37059,6 @@ func SdcaFprint(scope *Scope, input tf.Output) (output tf.Output) {
 	return op.Output(0)
 }
 
-// ParseSequenceExampleAttr is an optional argument to ParseSequenceExample.
-type ParseSequenceExampleAttr func(optionalAttr)
-
-// ParseSequenceExampleNcontextSparse sets the optional Ncontext_sparse attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func ParseSequenceExampleNcontextSparse(value int64) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["Ncontext_sparse"] = value
-	}
-}
-
-// ParseSequenceExampleNcontextDense sets the optional Ncontext_dense attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func ParseSequenceExampleNcontextDense(value int64) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["Ncontext_dense"] = value
-	}
-}
-
-// ParseSequenceExampleNfeatureListSparse sets the optional Nfeature_list_sparse attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func ParseSequenceExampleNfeatureListSparse(value int64) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["Nfeature_list_sparse"] = value
-	}
-}
-
-// ParseSequenceExampleNfeatureListDense sets the optional Nfeature_list_dense attribute to value.
-// If not specified, defaults to 0
-//
-// REQUIRES: value >= 0
-func ParseSequenceExampleNfeatureListDense(value int64) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["Nfeature_list_dense"] = value
-	}
-}
-
-// ParseSequenceExampleContextSparseTypes sets the optional context_sparse_types attribute to value.
-//
-// value: A list of Ncontext_sparse types; the data types of data in
-// each context Feature given in context_sparse_keys.
-// Currently the ParseSingleSequenceExample supports DT_FLOAT (FloatList),
-// DT_INT64 (Int64List), and DT_STRING (BytesList).
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func ParseSequenceExampleContextSparseTypes(value []tf.DataType) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["context_sparse_types"] = value
-	}
-}
-
-// ParseSequenceExampleFeatureListDenseTypes sets the optional feature_list_dense_types attribute to value.
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func ParseSequenceExampleFeatureListDenseTypes(value []tf.DataType) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["feature_list_dense_types"] = value
-	}
-}
-
-// ParseSequenceExampleContextDenseShapes sets the optional context_dense_shapes attribute to value.
-//
-// value: A list of Ncontext_dense shapes; the shapes of data in
-// each context Feature given in context_dense_keys.
-// The number of elements in the Feature corresponding to context_dense_key[j]
-// must always equal context_dense_shapes[j].NumEntries().
-// The shape of context_dense_values[j] will match context_dense_shapes[j].
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func ParseSequenceExampleContextDenseShapes(value []tf.Shape) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["context_dense_shapes"] = value
-	}
-}
-
-// ParseSequenceExampleFeatureListSparseTypes sets the optional feature_list_sparse_types attribute to value.
-//
-// value: A list of Nfeature_list_sparse types; the data types
-// of data in each FeatureList given in feature_list_sparse_keys.
-// Currently the ParseSingleSequenceExample supports DT_FLOAT (FloatList),
-// DT_INT64 (Int64List), and DT_STRING (BytesList).
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func ParseSequenceExampleFeatureListSparseTypes(value []tf.DataType) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["feature_list_sparse_types"] = value
-	}
-}
-
-// ParseSequenceExampleFeatureListDenseShapes sets the optional feature_list_dense_shapes attribute to value.
-//
-// value: A list of Nfeature_list_dense shapes; the shapes of
-// data in each FeatureList given in feature_list_dense_keys.
-// The shape of each Feature in the FeatureList corresponding to
-// feature_list_dense_key[j] must always equal
-// feature_list_dense_shapes[j].NumEntries().
-// If not specified, defaults to <>
-//
-// REQUIRES: len(value) >= 0
-func ParseSequenceExampleFeatureListDenseShapes(value []tf.Shape) ParseSequenceExampleAttr {
-	return func(m optionalAttr) {
-		m["feature_list_dense_shapes"] = value
-	}
-}
-
-// Transforms a vector of brain.SequenceExample protos (as strings) into typed tensors.
-//
-// Arguments:
-//	serialized: A vector containing binary serialized SequenceExample protos.
-//	debug_name: A vector containing the names of the serialized protos.
-// May contain, for example, table key (descriptive) name for the
-// corresponding serialized proto.  This is purely useful for debugging
-// purposes, and the presence of values here has no effect on the output.
-// May also be an empty vector if no name is available.
-//	context_dense_defaults: A list of Ncontext_dense Tensors (some may be empty).
-// context_dense_defaults[j] provides default values
-// when the SequenceExample's context map lacks context_dense_key[j].
-// If an empty Tensor is provided for context_dense_defaults[j],
-// then the Feature context_dense_keys[j] is required.
-// The input type is inferred from context_dense_defaults[j], even when it's
-// empty.  If context_dense_defaults[j] is not empty, its shape must match
-// context_dense_shapes[j].
-//	feature_list_dense_missing_assumed_empty: A vector listing the
-// FeatureList keys which may be missing from the SequenceExamples.  If the
-// associated FeatureList is missing, it is treated as empty.  By default,
-// any FeatureList not listed in this vector must exist in the SequenceExamples.
-//	context_sparse_keys: A list of Ncontext_sparse string Tensors (scalars).
-// The keys expected in the Examples' features associated with context_sparse
-// values.
-//	context_dense_keys: A list of Ncontext_dense string Tensors (scalars).
-// The keys expected in the SequenceExamples' context features associated with
-// dense values.
-//	feature_list_sparse_keys: A list of Nfeature_list_sparse string Tensors
-// (scalars).  The keys expected in the FeatureLists associated with sparse
-// values.
-//	feature_list_dense_keys: A list of Nfeature_list_dense string Tensors (scalars).
-// The keys expected in the SequenceExamples' feature_lists associated
-// with lists of dense values.
-func ParseSequenceExample(scope *Scope, serialized tf.Output, debug_name tf.Output, context_dense_defaults []tf.Output, feature_list_dense_missing_assumed_empty []string, context_sparse_keys []string, context_dense_keys []string, feature_list_sparse_keys []string, feature_list_dense_keys []string, optional ...ParseSequenceExampleAttr) (context_sparse_indices []tf.Output, context_sparse_values []tf.Output, context_sparse_shapes []tf.Output, context_dense_values []tf.Output, feature_list_sparse_indices []tf.Output, feature_list_sparse_values []tf.Output, feature_list_sparse_shapes []tf.Output, feature_list_dense_values []tf.Output, feature_list_dense_lengths []tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	attrs := map[string]interface{}{"feature_list_dense_missing_assumed_empty": feature_list_dense_missing_assumed_empty, "context_sparse_keys": context_sparse_keys, "context_dense_keys": context_dense_keys, "feature_list_sparse_keys": feature_list_sparse_keys, "feature_list_dense_keys": feature_list_dense_keys}
-	for _, a := range optional {
-		a(attrs)
-	}
-	opspec := tf.OpSpec{
-		Type: "ParseSequenceExample",
-		Input: []tf.Input{
-			serialized, debug_name, tf.OutputList(context_dense_defaults),
-		},
-		Attrs: attrs,
-	}
-	op := scope.AddOperation(opspec)
-	if scope.Err() != nil {
-		return
-	}
-	var idx int
-	var err error
-	if context_sparse_indices, idx, err = makeOutputList(op, idx, "context_sparse_indices"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if context_sparse_values, idx, err = makeOutputList(op, idx, "context_sparse_values"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if context_sparse_shapes, idx, err = makeOutputList(op, idx, "context_sparse_shapes"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if context_dense_values, idx, err = makeOutputList(op, idx, "context_dense_values"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if feature_list_sparse_indices, idx, err = makeOutputList(op, idx, "feature_list_sparse_indices"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if feature_list_sparse_values, idx, err = makeOutputList(op, idx, "feature_list_sparse_values"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if feature_list_sparse_shapes, idx, err = makeOutputList(op, idx, "feature_list_sparse_shapes"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if feature_list_dense_values, idx, err = makeOutputList(op, idx, "feature_list_dense_values"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	if feature_list_dense_lengths, idx, err = makeOutputList(op, idx, "feature_list_dense_lengths"); err != nil {
-		scope.UpdateErr("ParseSequenceExample", err)
-		return
-	}
-	return context_sparse_indices, context_sparse_values, context_sparse_shapes, context_dense_values, feature_list_sparse_indices, feature_list_sparse_values, feature_list_sparse_shapes, feature_list_dense_values, feature_list_dense_lengths
-}
-
-// Returns true if queue is closed.
-//
-// This operation returns true if the queue is closed and false if the queue
-// is open.
-//
-// Arguments:
-//	handle: The handle to a queue.
-func QueueIsClosedV2(scope *Scope, handle tf.Output) (is_closed tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	opspec := tf.OpSpec{
-		Type: "QueueIsClosedV2",
-		Input: []tf.Input{
-			handle,
-		},
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // Inverse 3D fast Fourier transform.
 //
 // Computes the inverse 3-dimensional discrete Fourier transform over the
@@ -38804,6 +38466,27 @@ func StatefulTruncatedNormal(scope *Scope, resource tf.Output, algorithm tf.Outp
 	return op.Output(0)
 }
 
+// Returns true if queue is closed.
+//
+// This operation returns true if the queue is closed and false if the queue
+// is open.
+//
+// Arguments:
+//	handle: The handle to a queue.
+func QueueIsClosedV2(scope *Scope, handle tf.Output) (is_closed tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "QueueIsClosedV2",
+		Input: []tf.Input{
+			handle,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // Checks whether a quantile stream has been initialized.
 //
 // An Op that checks if quantile stream resource is initialized.
@@ -38826,6 +38509,214 @@ func IsBoostedTreesQuantileStreamResourceInitialized(scope *Scope, quantile_stre
 	return op.Output(0)
 }
 
+// ParseSequenceExampleAttr is an optional argument to ParseSequenceExample.
+type ParseSequenceExampleAttr func(optionalAttr)
+
+// ParseSequenceExampleNcontextSparse sets the optional Ncontext_sparse attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func ParseSequenceExampleNcontextSparse(value int64) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["Ncontext_sparse"] = value
+	}
+}
+
+// ParseSequenceExampleNcontextDense sets the optional Ncontext_dense attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func ParseSequenceExampleNcontextDense(value int64) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["Ncontext_dense"] = value
+	}
+}
+
+// ParseSequenceExampleNfeatureListSparse sets the optional Nfeature_list_sparse attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func ParseSequenceExampleNfeatureListSparse(value int64) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["Nfeature_list_sparse"] = value
+	}
+}
+
+// ParseSequenceExampleNfeatureListDense sets the optional Nfeature_list_dense attribute to value.
+// If not specified, defaults to 0
+//
+// REQUIRES: value >= 0
+func ParseSequenceExampleNfeatureListDense(value int64) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["Nfeature_list_dense"] = value
+	}
+}
+
+// ParseSequenceExampleContextSparseTypes sets the optional context_sparse_types attribute to value.
+//
+// value: A list of Ncontext_sparse types; the data types of data in
+// each context Feature given in context_sparse_keys.
+// Currently the ParseSingleSequenceExample supports DT_FLOAT (FloatList),
+// DT_INT64 (Int64List), and DT_STRING (BytesList).
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func ParseSequenceExampleContextSparseTypes(value []tf.DataType) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["context_sparse_types"] = value
+	}
+}
+
+// ParseSequenceExampleFeatureListDenseTypes sets the optional feature_list_dense_types attribute to value.
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func ParseSequenceExampleFeatureListDenseTypes(value []tf.DataType) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["feature_list_dense_types"] = value
+	}
+}
+
+// ParseSequenceExampleContextDenseShapes sets the optional context_dense_shapes attribute to value.
+//
+// value: A list of Ncontext_dense shapes; the shapes of data in
+// each context Feature given in context_dense_keys.
+// The number of elements in the Feature corresponding to context_dense_key[j]
+// must always equal context_dense_shapes[j].NumEntries().
+// The shape of context_dense_values[j] will match context_dense_shapes[j].
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func ParseSequenceExampleContextDenseShapes(value []tf.Shape) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["context_dense_shapes"] = value
+	}
+}
+
+// ParseSequenceExampleFeatureListSparseTypes sets the optional feature_list_sparse_types attribute to value.
+//
+// value: A list of Nfeature_list_sparse types; the data types
+// of data in each FeatureList given in feature_list_sparse_keys.
+// Currently the ParseSingleSequenceExample supports DT_FLOAT (FloatList),
+// DT_INT64 (Int64List), and DT_STRING (BytesList).
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func ParseSequenceExampleFeatureListSparseTypes(value []tf.DataType) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["feature_list_sparse_types"] = value
+	}
+}
+
+// ParseSequenceExampleFeatureListDenseShapes sets the optional feature_list_dense_shapes attribute to value.
+//
+// value: A list of Nfeature_list_dense shapes; the shapes of
+// data in each FeatureList given in feature_list_dense_keys.
+// The shape of each Feature in the FeatureList corresponding to
+// feature_list_dense_key[j] must always equal
+// feature_list_dense_shapes[j].NumEntries().
+// If not specified, defaults to <>
+//
+// REQUIRES: len(value) >= 0
+func ParseSequenceExampleFeatureListDenseShapes(value []tf.Shape) ParseSequenceExampleAttr {
+	return func(m optionalAttr) {
+		m["feature_list_dense_shapes"] = value
+	}
+}
+
+// Transforms a vector of brain.SequenceExample protos (as strings) into typed tensors.
+//
+// Arguments:
+//	serialized: A vector containing binary serialized SequenceExample protos.
+//	debug_name: A vector containing the names of the serialized protos.
+// May contain, for example, table key (descriptive) name for the
+// corresponding serialized proto.  This is purely useful for debugging
+// purposes, and the presence of values here has no effect on the output.
+// May also be an empty vector if no name is available.
+//	context_dense_defaults: A list of Ncontext_dense Tensors (some may be empty).
+// context_dense_defaults[j] provides default values
+// when the SequenceExample's context map lacks context_dense_key[j].
+// If an empty Tensor is provided for context_dense_defaults[j],
+// then the Feature context_dense_keys[j] is required.
+// The input type is inferred from context_dense_defaults[j], even when it's
+// empty.  If context_dense_defaults[j] is not empty, its shape must match
+// context_dense_shapes[j].
+//	feature_list_dense_missing_assumed_empty: A vector listing the
+// FeatureList keys which may be missing from the SequenceExamples.  If the
+// associated FeatureList is missing, it is treated as empty.  By default,
+// any FeatureList not listed in this vector must exist in the SequenceExamples.
+//	context_sparse_keys: A list of Ncontext_sparse string Tensors (scalars).
+// The keys expected in the Examples' features associated with context_sparse
+// values.
+//	context_dense_keys: A list of Ncontext_dense string Tensors (scalars).
+// The keys expected in the SequenceExamples' context features associated with
+// dense values.
+//	feature_list_sparse_keys: A list of Nfeature_list_sparse string Tensors
+// (scalars).  The keys expected in the FeatureLists associated with sparse
+// values.
+//	feature_list_dense_keys: A list of Nfeature_list_dense string Tensors (scalars).
+// The keys expected in the SequenceExamples' feature_lists associated
+// with lists of dense values.
+func ParseSequenceExample(scope *Scope, serialized tf.Output, debug_name tf.Output, context_dense_defaults []tf.Output, feature_list_dense_missing_assumed_empty []string, context_sparse_keys []string, context_dense_keys []string, feature_list_sparse_keys []string, feature_list_dense_keys []string, optional ...ParseSequenceExampleAttr) (context_sparse_indices []tf.Output, context_sparse_values []tf.Output, context_sparse_shapes []tf.Output, context_dense_values []tf.Output, feature_list_sparse_indices []tf.Output, feature_list_sparse_values []tf.Output, feature_list_sparse_shapes []tf.Output, feature_list_dense_values []tf.Output, feature_list_dense_lengths []tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{"feature_list_dense_missing_assumed_empty": feature_list_dense_missing_assumed_empty, "context_sparse_keys": context_sparse_keys, "context_dense_keys": context_dense_keys, "feature_list_sparse_keys": feature_list_sparse_keys, "feature_list_dense_keys": feature_list_dense_keys}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "ParseSequenceExample",
+		Input: []tf.Input{
+			serialized, debug_name, tf.OutputList(context_dense_defaults),
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	if scope.Err() != nil {
+		return
+	}
+	var idx int
+	var err error
+	if context_sparse_indices, idx, err = makeOutputList(op, idx, "context_sparse_indices"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if context_sparse_values, idx, err = makeOutputList(op, idx, "context_sparse_values"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if context_sparse_shapes, idx, err = makeOutputList(op, idx, "context_sparse_shapes"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if context_dense_values, idx, err = makeOutputList(op, idx, "context_dense_values"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if feature_list_sparse_indices, idx, err = makeOutputList(op, idx, "feature_list_sparse_indices"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if feature_list_sparse_values, idx, err = makeOutputList(op, idx, "feature_list_sparse_values"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if feature_list_sparse_shapes, idx, err = makeOutputList(op, idx, "feature_list_sparse_shapes"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if feature_list_dense_values, idx, err = makeOutputList(op, idx, "feature_list_dense_values"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	if feature_list_dense_lengths, idx, err = makeOutputList(op, idx, "feature_list_dense_lengths"); err != nil {
+		scope.UpdateErr("ParseSequenceExample", err)
+		return
+	}
+	return context_sparse_indices, context_sparse_values, context_sparse_shapes, context_dense_values, feature_list_sparse_indices, feature_list_sparse_values, feature_list_sparse_shapes, feature_list_dense_values, feature_list_dense_lengths
+}
+
 // Fast Fourier transform.
 //
 // Computes the 1-dimensional discrete Fourier transform over the inner-most
@@ -41742,6 +41633,115 @@ func ResourceApplyAdam(scope *Scope, var_ tf.Output, m tf.Output, v tf.Output, b
 	return scope.AddOperation(opspec)
 }
 
+// ResizeNearestNeighborAttr is an optional argument to ResizeNearestNeighbor.
+type ResizeNearestNeighborAttr func(optionalAttr)
+
+// ResizeNearestNeighborAlignCorners sets the optional align_corners attribute to value.
+//
+// value: If true, the centers of the 4 corner pixels of the input and output tensors are
+// aligned, preserving the values at the corner pixels. Defaults to false.
+// If not specified, defaults to false
+func ResizeNearestNeighborAlignCorners(value bool) ResizeNearestNeighborAttr {
+	return func(m optionalAttr) {
+		m["align_corners"] = value
+	}
+}
+
+// ResizeNearestNeighborHalfPixelCenters sets the optional half_pixel_centers attribute to value.
+// If not specified, defaults to false
+func ResizeNearestNeighborHalfPixelCenters(value bool) ResizeNearestNeighborAttr {
+	return func(m optionalAttr) {
+		m["half_pixel_centers"] = value
+	}
+}
+
+// Resize `images` to `size` using nearest neighbor interpolation.
+//
+// Arguments:
+//	images: 4-D with shape `[batch, height, width, channels]`.
+//	size: = A 1-D int32 Tensor of 2 elements: `new_height, new_width`.  The
+// new size for the images.
+//
+// Returns 4-D with shape
+// `[batch, new_height, new_width, channels]`.
+func ResizeNearestNeighbor(scope *Scope, images tf.Output, size tf.Output, optional ...ResizeNearestNeighborAttr) (resized_images tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "ResizeNearestNeighbor",
+		Input: []tf.Input{
+			images, size,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
+// A placeholder op for a value that will be fed into the computation.
+//
+// Arguments:
+//	dtype: The type of elements in the tensor.
+//	shape: The shape of the tensor.
+//
+// Returns A tensor that will be provided using the infeed mechanism.
+func InfeedDequeue(scope *Scope, dtype tf.DataType, shape tf.Shape) (output tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{"dtype": dtype, "shape": shape}
+	opspec := tf.OpSpec{
+		Type: "InfeedDequeue",
+
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
+// Encodes a `RaggedTensor` into a `variant` Tensor.
+//
+//
+// Encodes the given `RaggedTensor` and returns a `variant` Tensor. If
+// `batched_input` is True, then input `RaggedTensor` is unbatched along the
+// zero-th dimension, each component `RaggedTensor` is encoded into a scalar
+// `variant` Tensor, and these are stacked to return a 1-D `variant` Tensor.
+// If `batched_input` is False, then the input `RaggedTensor` is encoded as is and
+// a scalar `variant` Tensor is returned. A `RaggedTensor` is encoded by first
+// creating a 1-D `variant` Tensor with `ragged_rank + 1` elements, containing the
+// splits and values Tensors of the `RaggedTensor`. Then the 1-D `variant` Tensor
+// is wrapped in a scalar `variant` Tensor. See `RaggedTensorFromVariant` for the
+// corresponding decoding logic.
+//
+//
+// Arguments:
+//	rt_nested_splits: A list of one or more Tensors representing the splits of the input
+// `RaggedTensor`.
+//	rt_dense_values: A Tensor representing the values of the input `RaggedTensor`.
+//	batched_input: A `bool` denoting whether the input is a batched `RaggedTensor`.
+//
+// Returns A `variant` Tensor that containing encoded `RaggedTensor`.
+func RaggedTensorToVariant(scope *Scope, rt_nested_splits []tf.Output, rt_dense_values tf.Output, batched_input bool) (encoded_ragged tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{"batched_input": batched_input}
+	opspec := tf.OpSpec{
+		Type: "RaggedTensorToVariant",
+		Input: []tf.Input{
+			tf.OutputList(rt_nested_splits), rt_dense_values,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // ResourceApplyKerasMomentumAttr is an optional argument to ResourceApplyKerasMomentum.
 type ResourceApplyKerasMomentumAttr func(optionalAttr)
 

From 07d5aa230954276d9de201217b18e86815328ab8 Mon Sep 17 00:00:00 2001
From: Sachin Joglekar 
Date: Thu, 18 Jun 2020 12:58:11 -0700
Subject: [PATCH 0523/1390] Use Requantize after per-channel Conv only if
 tensor min/max is different from RELU min/max beyond a limit. Also rectify
 RELU1 bound to (-1, 1)

PiperOrigin-RevId: 317162909
Change-Id: Ice90226436cccf49507bd17877222da755d22644
---
 .../hexagon/builders/conv_2d_builder.cc       | 11 ++-
 .../hexagon/builders/tests/conv_test.cc       | 84 +++++++++++++++++++
 2 files changed, 92 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc b/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
index 97db6bf8fd0..a366522e35c 100644
--- a/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 #include 
 
+#include 
 #include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
@@ -197,7 +198,7 @@ TfLiteStatus Conv2dOpBuilder::PopulateSubGraph(const TfLiteIntArray* inputs,
     conv_output_min = 0;
     conv_output_max = 6;
   } else if (activation == kTfLiteActRelu1) {
-    conv_output_min = 0;
+    conv_output_min = -1;
     conv_output_max = 1;
   } else if (activation == kTfLiteActRelu) {
     conv_output_min = 0;
@@ -351,8 +352,12 @@ TfLiteStatus Conv2dOpBuilder::PopulateSubGraph(const TfLiteIntArray* inputs,
     output_max_tensor = AddOutput(sizeof(float), 4, kScalarShape);
   }
 
-  // Requantize if activation was not None.
-  if (activation != kTfLiteActNone) {
+  // Requantize if activation was not None & the TFLite tensor's min/max is
+  // different (diff > 1e-2) from the RELU bounds.
+  const float min_bound_diff = std::abs(conv_output_min - output_min);
+  const float max_bound_diff = std::abs(conv_output_max - output_max);
+  if (activation != kTfLiteActNone &&
+      (min_bound_diff > 0.01 || max_bound_diff > 0.01)) {
     auto* requantized_min_const = graph_builder_->AddConstNodeWithData(
         kScalarShape, reinterpret_cast(&output_min), sizeof(output_min));
     auto* requantized_max_const = graph_builder_->AddConstNodeWithData(
diff --git a/tensorflow/lite/delegates/hexagon/builders/tests/conv_test.cc b/tensorflow/lite/delegates/hexagon/builders/tests/conv_test.cc
index 13fd768fded..eed1bf29aae 100644
--- a/tensorflow/lite/delegates/hexagon/builders/tests/conv_test.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/tests/conv_test.cc
@@ -207,6 +207,43 @@ TEST(QuantizedConvolutionOpModel, SimpleConvTestReLU6Activation) {
                   1e-5)));
 }
 
+// Same as above, but the output min/max matches the RELU bounds.
+// Therefore, a Requantize node will not get added after Supernode.
+TEST(QuantizedConvolutionOpModel,
+     SimpleConvTestReLU6Activation_NoRequantizeRequired) {
+  QuantizedConvolutionOpModel m(
+      BuiltinOperator_CONV_2D, {TensorType_UINT8, {2, 2, 4, 1}, -63.5, 64},
+      {TensorType_UINT8, {3, 2, 2, 1}, -63.5, 64}, {TensorType_UINT8, {}, 0, 6},
+      Padding_VALID, /**dilation_factor**/ 1,
+      /**stride**/ 2, ActivationFunctionType_RELU6);
+  m.SetInput({
+      // First batch
+      1, 1, 1, 1,  // row = 1
+      2, 2, 2, 2,  // row = 2
+      // Second batch
+      1, 2, 3, 4,  // row = 1
+      1, 2, 3, 4,  // row = 2
+  });
+  m.SetFilter({
+      1, 2, 3, 4,    // first 2x2 filter
+      -1, 1, -1, 1,  // second 2x2 filter
+      -1, -1, 1, 1,  // third 2x2 filter
+  });
+  m.SetBias({1, 2, 3});
+
+  m.ApplyDelegateAndInvoke();
+
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(ArrayFloatNear(
+                  {
+                      6, 2, 5,  // first batch, left
+                      6, 2, 5,  // first batch, right
+                      6, 4, 3,  // second batch, left
+                      6, 4, 3,  // second batch, right
+                  },
+                  2e-2)));
+}
+
 TEST(QuantizedConvolutionOpModel, SimplePerTensor_Int8) {
   QuantizedConvolutionOpModel m(
       BuiltinOperator_CONV_2D,
@@ -512,6 +549,53 @@ TEST(QuantizedConvolutionOpModel, DepthwiseConvSimplePerTensor_Int8) {
       ElementsAreArray(ArrayFloatNear({43, 48, 40, 52, 3, -4, 4, 4}, 0.6f)));
 }
 
+TEST(QuantizedConvolutionOpModel, DepthwiseConvSimplePerTensor_Int8_RELU1) {
+  QuantizedConvolutionOpModel m(
+      BuiltinOperator_DEPTHWISE_CONV_2D,
+      {TensorType_INT8, {1, 2, 3, 1}, -63.5, 64, 0.5, -1},
+      {TensorType_INT8,
+       // [1 * 2 * 2 * 4] as [input_channel, y, x, output_channel]
+       {1, 2, 2, 4},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{0.1, 2, 3, 0.4},
+       /*per_channel_quantization_offsets=*/{0, 0, 0, 0},
+       /*channel_index=*/3},
+      {TensorType_INT8, {}, -63.5, 64, 0.5, -1}, Padding_VALID,
+      /**dilation_factor**/ 1,
+      /**stride**/ 1, ActivationFunctionType_RELU_N1_TO_1);
+  m.SetInt8Input({
+      // [1 * 2 * 3 * 1] as [batch, y, x, input_channel]
+      3,   // batch = 0, y = 0, x = 0
+      1,   // batch = 0, y = 0, x = 1
+      -2,  // batch = 0, y = 0, x = 2
+      4,   // batch = 0, y = 1, x = 0
+      2,   // batch = 0, y = 1, x = 1
+      -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetPerChannelQuantizedFilter({
+      // [1 * 2 * 2 * 4] as [input_channel, y, x, output_channel]
+      // depth multiplier = 2
+      1, 2, 3, 4,  // y = 0, x = 0
+      3, 4, 5, 6,  // y = 0, x = 1
+      7, 8, 5, 6,  // y = 1, x = 0
+      3, 4, 1, 2,  // y = 1, x = 1
+  });
+  m.SetPerChannelQuantizedBias({3, -2, 4, 6});
+
+  // Reference output.
+  m.Invoke();
+  auto reference_output = m.GetDequantizedOutput();
+
+  m.ApplyDelegateAndInvoke();
+
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(ArrayFloatNear(reference_output, 1e-2)));
+}
+
 TEST(QuantizedConvolutionOpModel, DepthwiseConvSimplePerAxis_Int8) {
   QuantizedConvolutionOpModel m(
       BuiltinOperator_DEPTHWISE_CONV_2D,

From 41e7392f58391c18aa872c638ff2e1ac72326bcd Mon Sep 17 00:00:00 2001
From: Rahul Joshi 
Date: Thu, 18 Jun 2020 13:09:04 -0700
Subject: [PATCH 0524/1390] - Eliminate all uses of passes that mark function
 visibility since the visibility   is now set correctly when importing. -
 Update tf_saved_model dialect verification to verify that exported functions
 are   marked public. - Eliminate function_visibility.mlir test. This test
 fails after the   tf_saved_model verification changes since its run
 tf.entry_function based   visibility on a tf_saved_model MLIR module. Also,
 these passes will be removed. - Fix TPURewritePass to mark the appropriate
 visibility on the serialized MLIR   attached to tf._TPUCompileMlir op.

PiperOrigin-RevId: 317165278
Change-Id: I8e8f6de4b56e89c303815edc3b34bcf0a4e82d2d
---
 .../compiler/mlir/lite/tf_tfl_passes.cc       |   3 -
 tensorflow/compiler/mlir/tensorflow/BUILD     |   1 -
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  14 +-
 .../tensorflow/tests/function_visibility.mlir |  47 -----
 .../tf_saved_model_delete_unused_funcs.mlir   |  96 ----------
 .../tf_saved_model_freeze_global_tensors.mlir |   2 +-
 .../tensorflow/tests/tf_saved_model_ops.mlir  |   2 +-
 .../tests/tf_saved_model_ops_invalid.mlir     |  34 +++-
 ...timize_global_tensors_interprocedural.mlir |  22 +--
 .../transforms/mark_function_visibility.cc    | 165 ------------------
 .../mlir/tensorflow/transforms/passes.h       |  22 ---
 .../tensorflow/transforms/tpu_rewrite_pass.cc |   3 +
 .../tensorflow/utils/compile_mlir_util.cc     |   3 -
 tensorflow/compiler/tf2xla/mlir_tf2xla.cc     |   5 -
 14 files changed, 59 insertions(+), 360 deletions(-)
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/tests/function_visibility.mlir
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_delete_unused_funcs.mlir
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/mark_function_visibility.cc

diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
index 589515d6246..008098f62ba 100644
--- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
+++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
@@ -212,9 +212,6 @@ void CreateTFLStandardPipeline(OpPassManager& pm,
 
   // Saved model pass to mark global tensors immutable.
   pm.addPass(mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass());
-  // Used to mark non-exported functions in saved model private.
-  pm.addPass(mlir::tf_saved_model::
-                 CreateMarkFunctionVisibilityUsingSavedModelLinkagePass());
   // Op fusion pass.
   pm.addPass(mlir::TFL::CreatePrepareCompositeFunctionsPass());
 
diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 17ed0e36a28..54e57512c32 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -491,7 +491,6 @@ cc_library(
         "transforms/graph_pruning.cc",
         "transforms/launch_to_device_attribute.cc",
         "transforms/layout_optimization.cc",
-        "transforms/mark_function_visibility.cc",
         "transforms/materialize_mlir_passthrough_op.cc",
         "transforms/optimize.cc",
         "transforms/optimize_global_tensors.cc",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index 6af70158e14..d59532fef65 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -229,8 +229,20 @@ static LogicalResult VerifySavedModelModule(
     }
   }
   for (auto func : module.getOps()) {
+    const bool is_exported = IsExported(func);
+
+    if (is_exported && func.getVisibility() != FuncOp::Visibility::Public) {
+      return func.emitError()
+             << "exported function @" << func.getName() << " should be public";
+    }
+
+    if (!is_exported && func.getVisibility() == FuncOp::Visibility::Public) {
+      return func.emitError() << "non-exported function @" << func.getName()
+                              << " should be private";
+    }
+
     if (HasAnyTfSavedModelArgAttr(func)) {
-      if (!IsExported(func)) {
+      if (!is_exported) {
         return func.emitError()
                << "can only apply 'tf_saved_model' argument attributes "
                   "to exported functions";
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/function_visibility.mlir b/tensorflow/compiler/mlir/tensorflow/tests/function_visibility.mlir
deleted file mode 100644
index 55af3cffde3..00000000000
--- a/tensorflow/compiler/mlir/tensorflow/tests/function_visibility.mlir
+++ /dev/null
@@ -1,47 +0,0 @@
-// RUN: tf-opt -tf-saved-model-mark-func-visibility -split-input-file %s | FileCheck --check-prefix=SAVEDMODEL %s
-// RUN: tf-opt -tf-mark-func-visibility -split-input-file -verify-diagnostics %s | FileCheck %s
-
-
-module attributes {tf_saved_model.semantics} {
-  // SAVEDMODEL: func @func_exported_1() attributes {tf_saved_model.exported_names = ["func_exported_1"]}
-  func @func_exported_1() attributes {tf_saved_model.exported_names = ["func_exported_1"]} {
-    "tf.some_call"() {callee = {callee = {callee = @child}}} : () -> ()
-    return
-  }
-
-  // SAVEDMODEL: func @func_exported_2() attributes {tf_saved_model.exported_names = ["func_exported_2"]}
-  func @func_exported_2() attributes {tf_saved_model.exported_names = ["func_exported_2"]} {
-    "tf.some_call"() {callee = {callee = {callee = @child}}} : () -> ()
-    return
-  }
-
-  // SAVEDMODEL: func @func_not_exported() attributes {sym_visibility = "private"}
-  func @func_not_exported() {
-    return
-  }
-
-}
-
-// -----
-
-module {
-  // CHECK: func @func_with_entry_spec(%arg0: tensor<1xi32>) -> tensor<1xi32> attributes {tf.entry_function = {inputs = "x", outputs = "y"}}
-  func @func_with_entry_spec(%arg0: tensor<1xi32>) -> tensor<1xi32> attributes {tf.entry_function = {inputs = "x", outputs = "y"}} {
-    return %arg0 : tensor<1xi32>
-  }
-
-  // CHECK: func @func_without_entry_spec(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi32> attributes {sym_visibility = "private"}
-  func @func_without_entry_spec(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> tensor<*xi32> {
-    %0 = "tf.AddV2"(%arg0, %arg1) {T = i32, device = ""} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
-    return %0 : tensor<*xi32>
-  }
-}
-
-// -----
-
-module {
-  // expected-error @+1 {{can't overwrite the visibility of function private_func_with_entry_spec with private visibility}}
-  func @private_func_with_entry_spec(%arg0: tensor<1xi32>) -> tensor<1xi32> attributes {tf.entry_function = {inputs = "x", outputs = "y"}, sym_visibility = "private"} {
-    return %arg0 : tensor<1xi32>
-  }
-}
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_delete_unused_funcs.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_delete_unused_funcs.mlir
deleted file mode 100644
index 6f2c47a935f..00000000000
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_delete_unused_funcs.mlir
+++ /dev/null
@@ -1,96 +0,0 @@
-// RUN: tf-opt -tf-saved-model-mark-func-visibility -symbol-dce -split-input-file %s | FileCheck %s
-
-module attributes {tf_saved_model.semantics} {
-
-  // Test case: Unused function should be deleted.
-
-  // CHECK-NOT: func @unused
-  func @unused() {
-    return
-  }
-
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // Test case: Root calls child. Child should not be deleted.
-
-  // CHECK: func @root
-  func @root() attributes {tf_saved_model.exported_names = ["root"]} {
-    "tf.some_call"() { callee = @child } : () -> ()
-    return
-  }
-
-  // CHECK: func @child
-  func @child() {
-    return
-  }
-
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // Test case: Don't crash if attribute that doesn't reference a func.
-
-  "tf.some_opaque_global_variable"() { sym_name = "some_global" } : () -> ()
-
-  func @root2() attributes {tf_saved_model.exported_names = ["root2"]} {
-    "tf.do_something_with_a_global"() { global = @some_global } : () -> ()
-    return
-  }
-
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // Test case: Delete recursively dead cycle.
-
-  // CHECK-NOT: func @recursively_dead0
-  func @recursively_dead0() {
-    "tf.some_call"() { callee = @recursively_dead1 } : () -> ()
-    return
-  }
-  // CHECK-NOT: func @recursively_dead1
-  func @recursively_dead1() {
-    "tf.some_call"() { callee = @recursively_dead0 } : () -> ()
-    return
-  }
-
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // Test case: Root calls child with a deeply nested symbol reference.
-  // Child should not be deleted.
-
-  // CHECK: func @root
-  func @root() attributes {tf_saved_model.exported_names = ["root"]} {
-    "tf.some_call"() {callee = {callee = {callee = @child}}} : () -> ()
-    return
-  }
-
-  // CHECK: func @child
-  func @child() {
-    return
-  }
-
-}
-
-// -----
-
-// Test case: If the module doesn't have tf_saved_model semantics, then this
-// pass shouldn't do anything.
-module {
-  // CHECK: func @not_dead()
-  func @not_dead() {
-    return
-  }
-}
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir
index 38627b41b68..6c32a3bc4d6 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir
@@ -64,7 +64,7 @@ module attributes {tf_saved_model.semantics} {
     return
   }
 
-  func @f_callee(%arg0: tensor>>) {
+  func @f_callee(%arg0: tensor>>) attributes {sym_visibility = "private"} {
     return
   }
 }
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
index aa1f996da07..05e7638645f 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
@@ -40,7 +40,7 @@ module attributes {tf_saved_model.semantics} {
     return %arg0 : tensor
   }
 
-  func @f() {
+  func @f() attributes {sym_visibility = "private"} {
     return
   }
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index 544600cf6b8..f04e1a60b36 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -3,7 +3,7 @@
 module attributes {tf_saved_model.semantics} {
 
   // expected-error@+1 {{unknown tf_saved_model dialect arg attribute 'tf_saved_model.not_a_real_arg_attr'}}
-  func @f(%arg0: tensor {tf_saved_model.not_a_real_arg_attr = 1 : i32}) {
+  func @f(%arg0: tensor {tf_saved_model.not_a_real_arg_attr = 1 : i32}) attributes {sym_visibility = "private"} {
     return
   }
 
@@ -233,7 +233,7 @@ module attributes {tf_saved_model.semantics} {
   "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor, value = dense<1.> : tensor<1xf32> } : () -> ()
   // expected-error@+1 {{can only apply 'tf_saved_model' argument attributes to exported functions}}
   func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v})
-  -> (tensor {tf_saved_model.index_path = []}) {
+  -> (tensor {tf_saved_model.index_path = []}) attributes {sym_visibility = "private"} {
     %0 = "tf.ReadVariableOp"(%arg0) : (tensor>>) -> tensor
     return %0 : tensor
   }
@@ -273,7 +273,7 @@ module attributes {tf_saved_model.semantics} {
 
   // expected-error@+1 {{the initializer function should have no output}}
   "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> {
+  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
     %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
     return %0 : tensor<1xf32>
   }
@@ -286,8 +286,34 @@ module attributes {tf_saved_model.semantics} {
   "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
   // expected-error@+1 {{there must be no more than one session_initializer op}}
   "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> {
+  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
     %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
     return %0 : tensor<1xf32>
   }
 }
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{exported function @f should be public}}
+  func @f(
+    %arg0: tensor {tf.resource_name = "resource"}
+  ) attributes { sym_visibility = "private", tf_saved_model.exported_names = ["foo.some_func"] } {
+    return
+  }
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{non-exported function @f should be private}}
+  func @f(
+    %arg0: tensor {tf.resource_name = "resource"}
+  ) {
+    return
+  }
+
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir
index 91e8c9c4b66..14a0006cd3b 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors_interprocedural.mlir
@@ -20,12 +20,12 @@ module attributes {tf_saved_model.semantics} {
     return %val : tensor
   }
 
-  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee_callee} : (tensor<*x!tf.resource>) -> (tensor)
     return %val : tensor
   }
 
-  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf.resource>) -> tensor
     return %val : tensor
   }
@@ -59,7 +59,7 @@ module attributes {tf_saved_model.semantics} {
     return %val : tensor
   }
 
-  func @f_common(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_common(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf.resource>) -> tensor
     return %val : tensor
   }
@@ -85,7 +85,7 @@ module attributes {tf_saved_model.semantics} {
     return %val_2 : tensor
   }
 
-  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %cst_1 = constant dense<2.0> : tensor
     return %cst_1 : tensor
   }
@@ -112,13 +112,13 @@ module attributes {tf_saved_model.semantics} {
   }
 
   // CHECK: func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee_callee} : (tensor<*x!tf.resource>) -> (tensor)
     return %val : tensor
   }
 
   // CHECK: func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %c0 = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor
     "tf.AssignVariableOp"(%arg0, %c0) : (tensor<*x!tf.resource>, tensor) -> ()
     return %c0 : tensor
@@ -146,13 +146,13 @@ module attributes {tf_saved_model.semantics} {
   }
 
   // CHECK: func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee_callee} : (tensor<*x!tf.resource>) -> (tensor)
     return %val : tensor
   }
 
   // CHECK: func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f_callee_callee(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %c0 = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor
     "tf.AssignVariableOp"(%arg0, %c0) : (tensor<*x!tf.resource>, tensor) -> ()
     return %c0 : tensor
@@ -179,13 +179,13 @@ module attributes {tf_saved_model.semantics} {
 
 
   // CHECK: func @f(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @g} : (tensor<*x!tf.resource>) -> (tensor)
     return %val : tensor
   }
 
   // CHECK: func @g(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @g(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @g(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %val = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f} : (tensor<*x!tf.resource>) -> (tensor)
     return %val : tensor
   }
@@ -212,7 +212,7 @@ module attributes {tf_saved_model.semantics} {
 
 
   // CHECK: func @f(%arg0: tensor<*x!tf.resource>) -> tensor
-  func @f(%arg0: tensor<*x!tf.resource>) -> tensor {
+  func @f(%arg0: tensor<*x!tf.resource>) -> tensor attributes {sym_visibility = "private"} {
     %c0 = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor
     "tf.AssignAddVariableOp"(%arg0, %c0) : (tensor<*x!tf.resource>, tensor) -> ()
     return %c0 : tensor
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/mark_function_visibility.cc b/tensorflow/compiler/mlir/tensorflow/transforms/mark_function_visibility.cc
deleted file mode 100644
index 31a80a4ecdb..00000000000
--- a/tensorflow/compiler/mlir/tensorflow/transforms/mark_function_visibility.cc
+++ /dev/null
@@ -1,165 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "llvm/ADT/STLExtras.h"
-#include "mlir/IR/Module.h"  // from @llvm-project
-#include "mlir/Pass/Pass.h"  // from @llvm-project
-#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
-
-#define DEBUG_TYPE "tf-shape-inference"
-
-namespace mlir {
-
-namespace {
-
-LogicalResult MarkFunctionVisibility(
-    ModuleOp module, llvm::function_ref IsExternalVisible) {
-  LogicalResult result = success();
-
-  for (auto func : module.getOps()) {
-    FuncOp::Visibility old_visibility = func.getVisibility();
-
-    FuncOp::Visibility visibility = IsExternalVisible(func)
-                                        ? FuncOp::Visibility::Public
-                                        : FuncOp::Visibility::Private;
-
-    auto get_visibility_name = [](FuncOp::Visibility v) {
-      return v == FuncOp::Visibility::Public
-                 ? "public"
-                 : v == FuncOp::Visibility::Private ? "private" : "nested";
-    };
-
-    if (old_visibility != SymbolTable::Visibility::Public &&
-        old_visibility != visibility) {
-      result = func.emitError()
-               << "can't overwrite the visibility of function "
-               << func.getName() << " with "
-               << get_visibility_name(old_visibility) << " visibility";
-    }
-
-    LLVM_DEBUG(llvm::dbgs()
-               << "function " << func.getName() << " has "
-               << get_visibility_name(visibility) << " visibility \n");
-
-    func.setVisibility(visibility);
-  }
-
-  return result;
-}
-
-}  // anonymous namespace
-
-namespace TF {
-
-LogicalResult MarkFunctionVisibilityUsingEntryFunctionSpecification(
-    ModuleOp module) {
-  auto HasEntryFunctionSpecification = [](FuncOp func) -> bool {
-    auto attrs = func.getAttrOfType("tf.entry_function");
-    return attrs && !attrs.empty();
-  };
-  return MarkFunctionVisibility(module, HasEntryFunctionSpecification);
-}
-
-namespace {
-struct MarkFunctionVisibilityUsingEntryFunctionSpecificationPass
-    : public PassWrapper<
-          MarkFunctionVisibilityUsingEntryFunctionSpecificationPass,
-          OperationPass> {
-  void runOnOperation() override {
-    if (failed(MarkFunctionVisibilityUsingEntryFunctionSpecification(
-            getOperation()))) {
-      signalPassFailure();
-    }
-  }
-};
-}  // namespace
-
-static PassRegistration<
-    MarkFunctionVisibilityUsingEntryFunctionSpecificationPass>
-    pass("tf-mark-func-visibility",
-         "Use tf.entry_function to mark function visibility.");
-
-std::unique_ptr>
-CreateMarkFunctionVisibilityUsingEntryFunctionSpecificationPass() {
-  return std::make_unique<
-      MarkFunctionVisibilityUsingEntryFunctionSpecificationPass>();
-}
-
-// Marks the main function with public visibility, while other functions are
-// marked with private visibility.
-LogicalResult MarkOnlyMainFunctionWithPublicVisibility(ModuleOp module) {
-  for (auto func : module.getOps()) {
-    if (func.getName() == "main") {
-      func.setVisibility(FuncOp::Visibility::Public);
-    } else {
-      func.setVisibility(FuncOp::Visibility::Private);
-    }
-  }
-  return success();
-}
-
-namespace {
-struct MarkOnlyMainFunctionWithPublicVisibilityPass
-    : public PassWrapper> {
-  void runOnOperation() override {
-    if (failed(MarkOnlyMainFunctionWithPublicVisibility(getOperation()))) {
-      signalPassFailure();
-    }
-  }
-};
-}  // namespace
-
-std::unique_ptr>
-CreateMarkOnlyMainFunctionWithPublicVisibilityPass() {
-  return std::make_unique();
-}
-
-}  // namespace TF
-
-namespace tf_saved_model {
-
-static LogicalResult MarkFunctionVisibilityUsingSavedModelLinkage(
-    ModuleOp module) {
-  if (!tf_saved_model::HasTfSavedModelSemantics(module)) {
-    return success();
-  }
-  return MarkFunctionVisibility(module, tf_saved_model::IsExported);
-}
-
-namespace {
-struct MarkFunctionVisibilityUsingSavedModelLinkagePass
-    : public PassWrapper> {
-  void runOnOperation() override {
-    if (failed(MarkFunctionVisibilityUsingSavedModelLinkage(getOperation()))) {
-      signalPassFailure();
-    }
-  }
-};
-}  // namespace
-
-static PassRegistration pass(
-    "tf-saved-model-mark-func-visibility",
-    "Use tf_saved_model linkage information to mark function visibility.");
-
-std::unique_ptr>
-CreateMarkFunctionVisibilityUsingSavedModelLinkagePass() {
-  return std::make_unique();
-}
-
-}  // namespace tf_saved_model
-
-}  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
index 7158d0f6be0..5cb15027fc5 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
@@ -117,21 +117,6 @@ std::unique_ptr> CreatePromoteVarHandlesToArgsPass();
 std::unique_ptr>
 CreateConvertReadonlyReferenceVariablesToResourceVariablesPass();
 
-// Marks function visibility using tf.entry_function specification. That is,
-// functions with tf.entry_function attributes are marked with public
-// visibility while the other functions are marked with private visibility.
-LogicalResult MarkFunctionVisibilityUsingEntryFunctionSpecification(
-    ModuleOp module);
-// Creates a pass that uses tf.entry_function specification to mark function
-// visibility.
-std::unique_ptr>
-CreateMarkFunctionVisibilityUsingEntryFunctionSpecificationPass();
-
-// Creates a pass that marks the main function with public visibility, while
-// other functions are marked with private visibility.
-std::unique_ptr>
-CreateMarkOnlyMainFunctionWithPublicVisibilityPass();
-
 // Creates a simple device assignment pass on TF dialect for CoreRT use case.
 std::unique_ptr> CreateSimpleTFDeviceAssignmentPass(
     llvm::StringRef default_device);
@@ -315,13 +300,6 @@ std::unique_ptr> CreateOptimizeGlobalTensorsPass();
 // Creates a pass that freezes tf_saved_model.global_tensor ops.
 std::unique_ptr> CreateFreezeGlobalTensorsPass();
 
-// Creates a pass that uses tf_saved_model dialect linkage information
-// to mark function visibility. That is, exported functions are marked with
-// public visibility while the other functions are marked with private
-// visibility.
-std::unique_ptr>
-CreateMarkFunctionVisibilityUsingSavedModelLinkagePass();
-
 }  // namespace tf_saved_model
 
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_rewrite_pass.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_rewrite_pass.cc
index 696882cd105..ec9b3df525f 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_rewrite_pass.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_rewrite_pass.cc
@@ -146,6 +146,9 @@ LogicalResult EncapsulateFuncAndSerialize(FuncOp entry_func,
       // We can simply change name of TPU program's main function because there
       // should be no other reference to it.
       clone.setName("main");
+      clone.setVisibility(FuncOp::Visibility::Public);
+    } else {
+      clone.setVisibility(FuncOp::Visibility::Private);
     }
     symbol_table.insert(clone);
   }
diff --git a/tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc b/tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
index fd1ba3b1901..dac2fea87e2 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
@@ -267,9 +267,6 @@ Status ConvertMLIRToXlaComputation(
     const XlaCompiler::ShapeRepresentationFn shape_representation_fn,
     std::vector> custom_legalization_passes) {
   mlir::PassManager tf2xla(module_op.getContext());
-  // Mark main function as public, and other functions as private.
-  tf2xla.addPass(
-      mlir::TF::CreateMarkOnlyMainFunctionWithPublicVisibilityPass());
   tf2xla.addNestedPass(mlir::createCanonicalizerPass());
   tf2xla.addPass(mlir::TF::CreateTensorListOpsDecompositionPass());
   tf2xla.addPass(mlir::TF::CreateStackOpsDecompositionPass());
diff --git a/tensorflow/compiler/tf2xla/mlir_tf2xla.cc b/tensorflow/compiler/tf2xla/mlir_tf2xla.cc
index 43793be56a7..60d1f3da0c5 100644
--- a/tensorflow/compiler/tf2xla/mlir_tf2xla.cc
+++ b/tensorflow/compiler/tf2xla/mlir_tf2xla.cc
@@ -165,11 +165,6 @@ Status ConvertGraphDefToXlaViaMlir(
   device_set.AddDevice(&device);
   AddDevicesToOp(*module, &device_set);
 
-  if (failed(mlir::TF::MarkFunctionVisibilityUsingEntryFunctionSpecification(
-          *module))) {
-    return errors::Internal("Problem with mark function visibility");
-  }
-
   TF_RETURN_IF_ERROR(mlir::TF::RunBridgeWithStandardPipeline(
       *module, /*enable_logging=*/VLOG_IS_ON(1), /*enable_inliner=*/true));
 

From ededf6f4b9d1488c1d27df58b047fd5da6ad6c73 Mon Sep 17 00:00:00 2001
From: Sharada Shiddibhavi 
Date: Thu, 18 Jun 2020 13:14:37 -0700
Subject: [PATCH 0525/1390] Update tensorflow/core/util/mkl_util.h

Addressing review comments

Co-authored-by: Penporn Koanantakool <38085909+penpornk@users.noreply.github.com>
---
 tensorflow/core/util/mkl_util.h | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index 996984eebc0..854d6e349cb 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -1538,13 +1538,8 @@ class MklDnnData {
   /// Set function for data buffer of user memory primitive.
   inline void SetUsrMemDataHandle(const Tensor* tensor,
                                   std::shared_ptr t_stream = nullptr) {
-    CHECK_NOTNULL(user_memory_);
-    CHECK_NOTNULL(tensor);
-#ifdef ENABLE_MKLDNN_THREADPOOL
-    user_memory_->set_data_handle(GetTensorBuffer(tensor), *t_stream);
-#else
-    user_memory_->set_data_handle(GetTensorBuffer(tensor));
-#endif  // ENABLE_MKLDNN_THREADPOOL
+    SetUsrMemDataHandle(GetTensorBuffer(tensor), t_stream);
+  }
   }
 
   /// allocate function for data buffer

From 51f3da9ca884951be412bccb766c3700ba2255f3 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Thu, 18 Jun 2020 13:15:43 -0700
Subject: [PATCH 0526/1390] [Resubmit] Remove dynamic dimension of strided
 slice grad if input to strided slice is static.

If we slice a dynamic shaped tensor from a static tensor, the output
of the gradient should still be static. Unfortunately this cannot be
deduced alone by xla, so extra information is needed from the tf2xla
bridge.

PiperOrigin-RevId: 317166566
Change-Id: Ic3a16826242947a29cafd51b2f5c19e65d531fb9
---
 .../tf2xla/kernels/strided_slice_op.cc        | 24 +++++++++++++++++
 tensorflow/compiler/xla/client/xla_builder.cc | 23 ++++++++++++++++
 tensorflow/compiler/xla/client/xla_builder.h  |  6 +++++
 .../compiler/xla/client/xla_builder_test.cc   | 26 +++++++++++++++++++
 4 files changed, 79 insertions(+)

diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
index 2684c982600..784b790767c 100644
--- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
@@ -350,6 +350,30 @@ class StridedSliceGradOp : public XlaOpKernel {
       grad = xla::Rev(grad, dimensions_to_reverse);
     }
     grad = xla::Pad(grad, zero, padding_config);
+
+    xla::XlaOp dynamic_shape = ctx->Input(0);
+    xla::Shape grad_shape = ctx->builder()->GetShape(grad).ValueOrDie();
+    ctx->set_dynamic_dimension_is_minus_one(true);
+    std::vector dynamic_size;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(0, &dynamic_size));
+    // Input of strided_slice_op has to have the same shape as output.
+    DCHECK_EQ(grad_shape.rank(), input_shape.dims());
+    for (int64 dim = 0; dim < input_shape.dims(); ++dim) {
+      DCHECK_EQ(grad_shape.dimensions(dim), input_shape.dim_size(dim));
+      if (dynamic_size[dim] == -1) {
+        // Input is a dynamic dimension, set the same dynamic dimension size in
+        // the output.
+        auto dim_size = xla::Slice(dynamic_shape, {dim}, {dim + 1}, {1});
+        auto dim_size_scalar =
+            xla::Reshape(xla::ShapeUtil::MakeScalarShape(xla::S32), dim_size);
+        grad = xla::SetDimensionSize(grad, dim_size_scalar, dim);
+      } else if (grad_shape.is_dynamic_dimension(dim)) {
+        // Input is static but output is dynamic, respect input and remove any
+        // dynamic dim in the output.
+        grad = xla::RemoveDynamicDimension(grad, dim);
+      }
+    }
+
     ctx->SetOutput(0, grad);
   }
 
diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index bfba48862f6..c7b6a7f9491 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -2727,6 +2727,25 @@ XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64 dimension) {
   });
 }
 
+XlaOp XlaBuilder::RemoveDynamicDimension(XlaOp operand, int64 dimension) {
+  return ReportErrorOrReturn([&]() -> StatusOr {
+    HloInstructionProto instr;
+    TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+
+    Shape shape = *operand_shape;
+    shape.set_dynamic_dimension(dimension, false);
+    // Setting an op's dynamic dimension to its static size removes the dynamic
+    // dimension.
+    XlaOp static_size =
+        ConstantR0(this, operand_shape->dimensions(dimension));
+
+    *instr.mutable_shape() = shape.ToProto();
+    instr.add_dimensions(dimension);
+    return AddInstruction(std::move(instr), HloOpcode::kSetDimensionSize,
+                          {operand, static_size});
+  });
+}
+
 XlaOp XlaBuilder::SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension) {
   return ReportErrorOrReturn([&]() -> StatusOr {
     HloInstructionProto instr;
@@ -3827,4 +3846,8 @@ XlaOp SetDimensionSize(const XlaOp operand, const XlaOp val, int64 dimension) {
   return operand.builder()->SetDimensionSize(operand, val, dimension);
 }
 
+XlaOp RemoveDynamicDimension(const XlaOp operand, int64 dimension) {
+  return operand.builder()->RemoveDynamicDimension(operand, dimension);
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h
index ffa6a7c3439..b8af180b83e 100644
--- a/tensorflow/compiler/xla/client/xla_builder.h
+++ b/tensorflow/compiler/xla/client/xla_builder.h
@@ -704,6 +704,8 @@ class XlaBuilder {
 
   XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
+  XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
+
   StatusOr AddInstruction(HloInstructionProto&& instr, HloOpcode opcode,
                                  absl::Span operands = {});
 
@@ -1151,6 +1153,7 @@ class XlaBuilder {
 
   friend XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
   friend XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
+  friend XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
 
  protected:
   // Returns OK status if the given op was built using this builder. Otherwise,
@@ -2149,6 +2152,9 @@ XlaOp GetDimensionSize(XlaOp operand, int64 dimension);
 
 XlaOp SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension);
 
+// Returns the same op but with dynamic dimension removed.
+XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
+
 // Implementation details below this point.
 //
 
diff --git a/tensorflow/compiler/xla/client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_builder_test.cc
index 4fa47077fca..7011c946203 100644
--- a/tensorflow/compiler/xla/client/xla_builder_test.cc
+++ b/tensorflow/compiler/xla/client/xla_builder_test.cc
@@ -556,6 +556,32 @@ TEST_F(XlaBuilderTest, DynamicParameter) {
   EXPECT_TRUE(param_shape.is_dynamic_dimension(0));
 }
 
+TEST_F(XlaBuilderTest, SetDimensionSize) {
+  XlaBuilder b(TestName());
+  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
+  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
+  auto set_dim_size = SetDimensionSize(p0, p1, 0);
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          BuildHloModule(&b, /*root=*/set_dim_size));
+  const Shape& root_shape =
+      module->entry_computation()->root_instruction()->shape();
+  EXPECT_TRUE(root_shape.is_dynamic_dimension(0));
+}
+
+TEST_F(XlaBuilderTest, RemoveDimensionSize) {
+  XlaBuilder b(TestName());
+  auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {10}), "p0");
+  auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(S32, {}), "p1");
+  auto set_dim_size = SetDimensionSize(p0, p1, 0);
+  auto remove_dim_size = RemoveDynamicDimension(set_dim_size, 0);
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          BuildHloModule(&b, /*root=*/remove_dim_size));
+  const Shape& root_shape =
+      module->entry_computation()->root_instruction()->shape();
+  // Dynamic dimension has been removed.
+  EXPECT_FALSE(root_shape.is_dynamic_dimension(0));
+}
+
 TEST_F(XlaBuilderTest, DynamicUnary) {
   XlaBuilder b(TestName());
   Shape tuple_param_shape = ShapeUtil::MakeTupleShape(

From 8944a3eeb18c2374f02759324cbeded018c5868b Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Thu, 18 Jun 2020 13:23:05 -0700
Subject: [PATCH 0527/1390] Enable type annotations for python/autograph.

PiperOrigin-RevId: 317168016
Change-Id: I57341cc0347ab8eb008949e191da1415069203b7
---
 tensorflow/python/autograph/pyct/parser.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/autograph/pyct/parser.py b/tensorflow/python/autograph/pyct/parser.py
index 747d56e401d..9ac7c2ef2a6 100644
--- a/tensorflow/python/autograph/pyct/parser.py
+++ b/tensorflow/python/autograph/pyct/parser.py
@@ -22,6 +22,7 @@ from __future__ import division
 from __future__ import print_function
 
 import re
+import sys
 import textwrap
 import tokenize
 
@@ -33,11 +34,18 @@ from tensorflow.python.autograph.pyct import errors
 from tensorflow.python.autograph.pyct import inspect_utils
 
 
-STANDARD_PREAMBLE = textwrap.dedent("""
-    from __future__ import division
-    from __future__ import print_function
+PY2_PREAMBLE = textwrap.dedent("""
+from __future__ import division
+from __future__ import print_function
 """)
-STANDARD_PREAMBLE_LEN = 2
+PY3_PREAMBLE = ''
+
+if sys.version_info >= (3,):
+  STANDARD_PREAMBLE = PY3_PREAMBLE
+else:
+  STANDARD_PREAMBLE = PY2_PREAMBLE
+
+STANDARD_PREAMBLE_LEN = STANDARD_PREAMBLE.count('__future__')
 
 
 _LEADING_WHITESPACE = re.compile(r'\s*')

From f8431d0c293a34c5bfb91cf0c57384eaa47a9911 Mon Sep 17 00:00:00 2001
From: Sharada Shiddibhavi 
Date: Thu, 18 Jun 2020 13:29:48 -0700
Subject: [PATCH 0528/1390] Update mkl_util.h

---
 tensorflow/core/util/mkl_util.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index 854d6e349cb..eb1a105e07c 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -1540,7 +1540,6 @@ class MklDnnData {
                                   std::shared_ptr t_stream = nullptr) {
     SetUsrMemDataHandle(GetTensorBuffer(tensor), t_stream);
   }
-  }
 
   /// allocate function for data buffer
   inline void AllocateBuffer(size_t size) {

From 2663edb6691b50de8ca9445c311d02b72faa1bf5 Mon Sep 17 00:00:00 2001
From: Ken Franko 
Date: Thu, 18 Jun 2020 13:29:28 -0700
Subject: [PATCH 0529/1390] Change TPUExtractOutsideCompilation pass for a
 Module pass.

This is needed for getting the devices from the module for assigning host device for outside compilation launch op.

PiperOrigin-RevId: 317169244
Change-Id: I734e7eeef3fdb037045d070ffd736be4ef8edee1
---
 .../compiler/mlir/tensorflow/transforms/passes.h      |  3 ++-
 .../transforms/tpu_extract_outside_compilation.cc     | 11 ++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
index 5cb15027fc5..a34be28c809 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
@@ -281,7 +281,8 @@ std::unique_ptr> CreateTPUHostComputationExpansionPass();
 
 // Creates a pass that extract outside compilation (CPU ops inside TPU cluster)
 // ops to a separate parallel_execute region to run on CPU.
-std::unique_ptr> CreateTPUExtractOutsideCompilationPass();
+std::unique_ptr>
+CreateTPUExtractOutsideCompilationPass();
 
 // Populates the supplied passmanager with the passes required to run the
 void CreateTPUBridgePipeline(OpPassManager& pm);
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
index a2a19108326..503c9869557 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tpu_extract_outside_compilation.cc
@@ -49,8 +49,9 @@ using OutsideClusterMap =
 // TODO(b/154363171): Add example tranformations.
 
 struct TPUExtractOutsideCompilation
-    : public PassWrapper {
-  void runOnFunction() override;
+    : public PassWrapper> {
+  void runOnOperation() override;
 };
 
 // Collects and clusters ops in `block` with the same `_xla_outside_compilation`
@@ -305,9 +306,9 @@ void CreateParallelExecuteFromOutsideClusters(
   }
 }
 
-void TPUExtractOutsideCompilation::runOnFunction() {
+void TPUExtractOutsideCompilation::runOnOperation() {
   auto extract_result =
-      getFunction().walk([&](tf_device::ClusterOp tpu_cluster) {
+      getOperation().walk([&](tf_device::ClusterOp tpu_cluster) {
         OutsideClusterMap clusters;
         if (failed(CollectAndGroupOutsideClusterOps(&tpu_cluster.GetBody(),
                                                     &clusters)))
@@ -325,7 +326,7 @@ void TPUExtractOutsideCompilation::runOnFunction() {
 
 }  // namespace
 
-std::unique_ptr>
+std::unique_ptr>
 CreateTPUExtractOutsideCompilationPass() {
   return std::make_unique();
 }

From 6558da5a66ad6863e47abfe596eee2290524b1b7 Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Thu, 18 Jun 2020 13:30:08 -0700
Subject: [PATCH 0530/1390] Apply new TraceMe APIs.

PiperOrigin-RevId: 317169381
Change-Id: I2259895a8dde21e25661a239b9d4f5911a454adb
---
 tensorflow/compiler/xla/pjrt/BUILD               |  2 ++
 tensorflow/compiler/xla/pjrt/pjrt_client.cc      | 16 ++++++++--------
 tensorflow/core/profiler/lib/connected_traceme.h |  1 +
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/BUILD b/tensorflow/compiler/xla/pjrt/BUILD
index dd50d0577d4..e401a798d68 100644
--- a/tensorflow/compiler/xla/pjrt/BUILD
+++ b/tensorflow/compiler/xla/pjrt/BUILD
@@ -141,7 +141,9 @@ cc_library(
         "//tensorflow/compiler/xla/service/gpu:gpu_executable_run_options",
         "//tensorflow/core:allocator",
         "//tensorflow/core:lib",
+        "//tensorflow/core/profiler/lib:connected_traceme",
         "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/core/profiler/lib:traceme_encode",
         "//tensorflow/stream_executor:event",
         "//tensorflow/stream_executor:stream",
         "//tensorflow/stream_executor/host:host_platform_id",
diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.cc b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
index ccb72b7ce30..ef259cf1cfd 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.cc
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
@@ -98,7 +98,9 @@ limitations under the License.
 #include "tensorflow/core/platform/mem.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/profiler/lib/connected_traceme.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
+#include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/stream_executor/device_memory.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
 #include "tensorflow/stream_executor/event.h"
@@ -1429,10 +1431,9 @@ StatusOr PjRtExecutable::EnqueueExecution(
     int executable_idx, const RunId& run_id, const ExecuteOptions& options,
     Device* device, std::vector* device_buffers) const {
   int device_ordinal = device->local_device_state()->device_ordinal();
-  tensorflow::profiler::TraceMe traceme([&] {
-    return absl::StrCat("LocalExecutable::Execute#run_id=", run_id.ToInt(),
-                        "#");
-  });
+  tensorflow::profiler::TraceMeConsumer activity(
+      "LocalExecutable::Execute", tensorflow::profiler::ContextType::kPjRt,
+      run_id.ToInt());
   VLOG(3) << "Replica " << replica << ", partition " << partition
           << " mapped to device ordinal for execution: " << device_ordinal;
 
@@ -1721,10 +1722,9 @@ PjRtExecutable::ExecuteOnLocalDevices(
     absl::Span> argument_handles,
     const ExecuteOptions& options) const {
   RunId run_id;
-  tensorflow::profiler::TraceMe traceme([&] {
-    return absl::StrCat(
-        "LocalExecutable::ExecuteOnLocalDevices#run_id=", run_id.ToInt(), "#");
-  });
+  tensorflow::profiler::TraceMeProducer activity(
+      "LocalExecutable::ExecuteOnLocalDevices",
+      tensorflow::profiler::ContextType::kPjRt, run_id.ToInt());
 
   const int num_local_devices = local_devices_.size();
 
diff --git a/tensorflow/core/profiler/lib/connected_traceme.h b/tensorflow/core/profiler/lib/connected_traceme.h
index ed8b4ac1ad2..b55c4407fe6 100644
--- a/tensorflow/core/profiler/lib/connected_traceme.h
+++ b/tensorflow/core/profiler/lib/connected_traceme.h
@@ -29,6 +29,7 @@ enum class ContextType : int {
   kGeneric,
   kTfExecutor,
   kSharedBatchScheduler,
+  kPjRt,
 };
 
 /*

From 61d83075ad34fd33d7d44ea3341a055e26e775a6 Mon Sep 17 00:00:00 2001
From: Niranjan Hasabnis 
Date: Wed, 17 Jun 2020 14:13:29 -0700
Subject: [PATCH 0531/1390] Fixing MklTanh compilation error with DNNL0

---
 .../core/common_runtime/mkl_layout_pass.cc    |  3 +++
 .../common_runtime/mkl_layout_pass_test.cc    |  3 +++
 tensorflow/core/kernels/mkl_relu_op.cc        | 23 +++++++++++++++++--
 tensorflow/core/kernels/mkl_relu_op_test.cc   |  5 +++-
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/common_runtime/mkl_layout_pass.cc b/tensorflow/core/common_runtime/mkl_layout_pass.cc
index f6e42fc7e8c..778d5445cb2 100644
--- a/tensorflow/core/common_runtime/mkl_layout_pass.cc
+++ b/tensorflow/core/common_runtime/mkl_layout_pass.cc
@@ -682,12 +682,15 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     rinfo_.push_back(
         {csinfo_.requantize, mkl_op_registry::GetMklOpName(csinfo_.requantize),
          CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation});
+#ifdef ENABLE_MKLDNN_V1
+    // Optimized TanhGrad support exists only in DNNL 1.x.
     rinfo_.push_back({csinfo_.tanh, mkl_op_registry::GetMklOpName(csinfo_.tanh),
                       CopyAttrsAll, AlwaysRewrite,
                       kRewriteForLayoutPropagation});
     rinfo_.push_back(
         {csinfo_.tanh_grad, mkl_op_registry::GetMklOpName(csinfo_.tanh_grad),
          CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation});
+#endif  // ENABLE_MKLDNN_V1
     rinfo_.push_back(
         {csinfo_.reshape, mkl_op_registry::GetMklOpName(csinfo_.reshape),
          CopyAttrsAll, AlwaysRewrite, kRewriteForLayoutPropagation});
diff --git a/tensorflow/core/common_runtime/mkl_layout_pass_test.cc b/tensorflow/core/common_runtime/mkl_layout_pass_test.cc
index 9971f6c5d7e..d480c0a49ce 100644
--- a/tensorflow/core/common_runtime/mkl_layout_pass_test.cc
+++ b/tensorflow/core/common_runtime/mkl_layout_pass_test.cc
@@ -3024,6 +3024,8 @@ REGISTER_TEST_ALL_TYPES(NodeRewrite_LeakyReluLeakyReluGrad_Positive);
 // clang-format on
 
 // clang-format off
+#ifdef ENABLE_MKLDNN_V1
+
 #define REGISTER_TEST(NAME, T, INPUT)                                        \
   TEST_F(MklLayoutPassTest, NAME##_##T) {                                    \
     DCHECK_EQ(kTensorOrdering, MklTfTensorOrdering::TENSORS_CONTIGUOUS);     \
@@ -3081,6 +3083,7 @@ REGISTER_TEST_ALL_TYPES(NodeRewrite_TanhGrad_Positive);
 }
 REGISTER_TEST_ALL_TYPES(NodeRewrite_TanhTanhGrad_Positive);
 #undef REGISTER_TEST
+#endif  // ENABLE_MKLDNN_V1
 // clang-format on
 
 TEST_F(MklLayoutPassTest, NodeRewrite_AvgPool_Positive) {
diff --git a/tensorflow/core/kernels/mkl_relu_op.cc b/tensorflow/core/kernels/mkl_relu_op.cc
index 6d79b8f3282..70aa1e937d3 100644
--- a/tensorflow/core/kernels/mkl_relu_op.cc
+++ b/tensorflow/core/kernels/mkl_relu_op.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include 
 
 #include "mkldnn.hpp"
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -27,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/util/mkl_types.h"
 #include "tensorflow/core/util/mkl_util.h"
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
 using mkldnn::algorithm;
 using mkldnn::eltwise_forward;
@@ -269,7 +269,7 @@ class MklEltwiseBwdParams {
 
   MklEltwiseBwdParams(const memory::dims& src_dims,
                       const memory::desc& common_md, algorithm alg_kind,
-                      float alpha, float beta, int forward_input_type)
+                      float alpha, float beta, int forward_input_type = -1)
       : src_dims(src_dims),
         common_md(common_md),
         alg_kind(alg_kind),
@@ -644,7 +644,10 @@ class MklReluGradOpBase : public OpKernel {
   virtual int GetDiffSrcIndex() const { return 0; }
   // What is the type of input tensor that grad op receives from forward op --
   // is it 'x' (SRC) or 'y' (DST). For Relu-family, it is 'x', so fwd op SRC.
+
+#ifdef ENABLE_MKLDNN_V1
   virtual int GetTypeOfInputTensorFromFwdOp() const { return MKLDNN_ARG_SRC; }
+#endif
 
   void Compute(OpKernelContext* context) {
     try {
@@ -736,8 +739,16 @@ class MklReluGradOpBase : public OpKernel {
         common_md = src_md;
       }
 
+#ifdef ENABLE_MKLDNN_V1
       MklEltwiseBwdParams bwdParams(src_dims, common_md, alg_kind, alpha_,
                                        beta_, GetTypeOfInputTensorFromFwdOp());
+#else
+      // MKLDNN V0 does not support reusing output of forward op in backward.
+      // So this optimization works only in MKLDNN v1.
+      MklEltwiseBwdParams bwdParams(src_dims, common_md, alg_kind, alpha_,
+                                       beta_);
+#endif  // ENABLE_MKLDNN_V1
+
       MklEltwiseBwdPrimitive* eltwise_bwd =
           MklEltwiseBwdPrimitiveFactory::Get(bwdParams);
 
@@ -962,6 +973,11 @@ class MklEluGradOp
   }
 };
 
+#ifdef ENABLE_MKLDNN_V1
+// Optimized TanhGrad support exists in DNNL1.x only
+// (eltwise_tanh_use_dst_for_bwd). We can still support it with DNNL0.x, but
+// it will not be optimized. So we disable it for DNNL0.x.
+
 template 
 class MklTanhOp : public MklReluOpBase {
  public:
@@ -1043,6 +1059,7 @@ class MklTanhGradOp
         (static_cast(user_g))[0] * (static_cast(1) - tanh * tanh);
   }
 };
+#endif  // ENABLE_MKLDNN_V1
 
 #define RELU6_UPPER_BOUND 6.0f
 template 
@@ -1227,6 +1244,7 @@ TF_CALL_bfloat16(REGISTER_RELU_MKL_SUPPORTED_KERNELS_TYPES);
 TF_CALL_float(REGISTER_ELU_MKL_SUPPORTED_KERNELS_TYPES);
 TF_CALL_bfloat16(REGISTER_ELU_MKL_SUPPORTED_KERNELS_TYPES);
 
+#ifdef ENABLE_MKLDNN_V1
 #define REGISTER_TANH_MKL_SUPPORTED_KERNELS_TYPES(type)        \
   REGISTER_KERNEL_BUILDER(                                     \
       Name("_MklTanh")                                         \
@@ -1242,6 +1260,7 @@ TF_CALL_bfloat16(REGISTER_ELU_MKL_SUPPORTED_KERNELS_TYPES);
       MklTanhGradOp);
 TF_CALL_float(REGISTER_TANH_MKL_SUPPORTED_KERNELS_TYPES);
 TF_CALL_bfloat16(REGISTER_TANH_MKL_SUPPORTED_KERNELS_TYPES);
+#endif
 
 #define REGISTER_RELU6_MKL_SUPPORTED_KERNELS_TYPES(type)       \
   REGISTER_KERNEL_BUILDER(                                     \
diff --git a/tensorflow/core/kernels/mkl_relu_op_test.cc b/tensorflow/core/kernels/mkl_relu_op_test.cc
index d1fdf7ab4ae..86d7f979c1f 100644
--- a/tensorflow/core/kernels/mkl_relu_op_test.cc
+++ b/tensorflow/core/kernels/mkl_relu_op_test.cc
@@ -15,8 +15,8 @@ limitations under the License.
 
 #ifdef INTEL_MKL
 
-#include "mkldnn.hpp"
 #include "absl/strings/match.h"
+#include "mkldnn.hpp"
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/cc/ops/nn_ops.h"
 #include "tensorflow/cc/ops/standard_ops.h"
@@ -121,8 +121,11 @@ static Graph* Activation(const string& op_name, const string& kind,
   BM(OP, 32, 64, 128, 256, cpu); \
   BM(OP, 33, 65, 129, 257, cpu);
 
+#ifdef ENABLE_MKLDNN_V1
+// Optimized MKLDNN TanhGrad support exists in DNNL1.x only.
 TEST_ALL_SIZES(Tanh)
 TEST_ALL_SIZES(TanhGrad)
+#endif  // ENABLE_MKLDNN_V1
 TEST_ALL_SIZES(Relu)
 TEST_ALL_SIZES(ReluGrad)
 TEST_ALL_SIZES(Elu)

From e0962f4c374f4cbf78ad27fd2391f976c1a2050d Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Thu, 18 Jun 2020 13:35:31 -0700
Subject: [PATCH 0532/1390] Remove the old grouping rule for PJRT.

PiperOrigin-RevId: 317170383
Change-Id: I96973a2d2cd0ca1fc786bc7500deb6b4fedd0534
---
 tensorflow/core/profiler/utils/group_events.cc  | 5 +----
 tensorflow/core/profiler/utils/xplane_schema.cc | 4 ----
 tensorflow/core/profiler/utils/xplane_schema.h  | 3 ---
 3 files changed, 1 insertion(+), 11 deletions(-)

diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index be8dd506b0c..0772cff7b97 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -635,10 +635,7 @@ std::vector CreateInterThreadConnectInfoList() {
        {StatType::kStepId, StatType::kIterNum}},
       {HostEventType::kKernelLaunch,
        HostEventType::kKernelExecute,
-       {StatType::kCorrelationId}},
-      {HostEventType::kLocalExecutableExecuteOnLocalDevice,
-       HostEventType::kLocalExecutableExecute,
-       {StatType::kRunId}}};
+       {StatType::kCorrelationId}}};
   return connect_info_list;
 }
 
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index be53dcbdc01..5ca8326d72c 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -91,10 +91,6 @@ const HostEventTypeMap& GetHostEventTypeMap() {
       {"WhileOp-StartBody", kWhileOpStartBody},
       {"ForOp", kForOp},
       {"PartitionedCallOp", kPartitionedCallOp},
-      // XLA related.
-      {"LocalExecutable::ExecuteOnLocalDevices",
-       kLocalExecutableExecuteOnLocalDevice},
-      {"LocalExecutable::Execute", kLocalExecutableExecute},
       // tf.data related.
       {"IteratorGetNextOp::DoCompute", kIteratorGetNextOp},
       {"IteratorGetNextAsOptionalOp::DoCompute", kIteratorGetNextAsOptionalOp},
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index a31814cef06..41774deaa59 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -81,9 +81,6 @@ enum HostEventType {
   kWhileOpStartBody,
   kForOp,
   kPartitionedCallOp,
-  // XLA related.
-  kLocalExecutableExecuteOnLocalDevice,
-  kLocalExecutableExecute,
   // tf.data related.
   kIteratorGetNextOp,
   kIteratorGetNextAsOptionalOp,

From 8f700fb2e0da382f1e2e9630f56a7922a8799a59 Mon Sep 17 00:00:00 2001
From: Berkin Ilbeyi 
Date: Thu, 18 Jun 2020 13:39:10 -0700
Subject: [PATCH 0533/1390] [XLA] Propagate memory spaces recursively inside
 nested fusions.

PiperOrigin-RevId: 317171110
Change-Id: I65004edb7498acb2f3b4238d9afbbb5d3930aab5
---
 .../xla/service/memory_space_propagation.cc   |  80 +++++++---
 .../xla/service/memory_space_propagation.h    |  11 +-
 .../service/memory_space_propagation_test.cc  | 148 ++++++++++++++++++
 3 files changed, 214 insertions(+), 25 deletions(-)

diff --git a/tensorflow/compiler/xla/service/memory_space_propagation.cc b/tensorflow/compiler/xla/service/memory_space_propagation.cc
index 80eb4017477..2eb15b14eaf 100644
--- a/tensorflow/compiler/xla/service/memory_space_propagation.cc
+++ b/tensorflow/compiler/xla/service/memory_space_propagation.cc
@@ -29,36 +29,78 @@ StatusOr MemorySpacePropagation::Run(HloModule* module) {
         // Propagate the operand subshapes.
         for (int operand_idx = 0; operand_idx < instruction->operand_count();
              ++operand_idx) {
-          modified |=
-              PropagateSubshapes(instruction->operand(operand_idx)->shape(),
-                                 instruction->fused_parameter(operand_idx));
+          for (const ShapeUtil::IndexedShape& indexed_shape :
+               ShapeUtil::GetLeafShapes(
+                   instruction->operand(operand_idx)->shape())) {
+            int64 memory_space = indexed_shape.shape.layout().memory_space();
+            modified |= Propagate(indexed_shape.index,
+                                  instruction->fused_parameter(operand_idx),
+                                  memory_space);
+          }
         }
 
         // Propagate output subshapes.
-        modified |= PropagateSubshapes(instruction->shape(),
-                                       instruction->fused_expression_root());
+        for (const ShapeUtil::IndexedShape& indexed_shape :
+             ShapeUtil::GetLeafShapes(instruction->shape())) {
+          int64 memory_space = indexed_shape.shape.layout().memory_space();
+          modified |=
+              Propagate(indexed_shape.index,
+                        instruction->fused_expression_root(), memory_space);
+        }
       }
     }
   }
   return modified;
 }
 
-bool MemorySpacePropagation::PropagateSubshapes(
-    const Shape& caller_shape, const HloInstruction* callee_instruction) const {
+bool MemorySpacePropagation::Propagate(ShapeIndexView index,
+                                       const HloInstruction* callee_instruction,
+                                       int64 memory_space) const {
   bool modified = false;
-  for (const ShapeUtil::IndexedShape& indexed_shape :
-       ShapeUtil::GetLeafShapes(caller_shape)) {
-    int64 memory_space = indexed_shape.shape.layout().memory_space();
-    const HloValue& value = dataflow_analysis_->GetUniqueValueAt(
-        callee_instruction, indexed_shape.index);
+  const HloValue& value = dataflow_analysis_->GetUniqueValueAt(
+      callee_instruction, index.ToShapeIndex());
 
-    for (const HloPosition& position : value.positions()) {
-      Shape* shape = ShapeUtil::GetMutableSubshape(
-          position.instruction->mutable_shape(), position.index);
-      if (shape->layout().memory_space() != memory_space) {
-        shape->mutable_layout()->set_memory_space(memory_space);
-        modified = true;
-      }
+  for (const HloPosition& position : value.positions()) {
+    HloInstruction* instruction = position.instruction;
+    Shape* shape = ShapeUtil::GetMutableSubshape(instruction->mutable_shape(),
+                                                 position.index);
+    if (shape->layout().memory_space() == memory_space) {
+      continue;
+    }
+    shape->mutable_layout()->set_memory_space(memory_space);
+    modified = true;
+
+    // For fusion outputs, propagate the memory space to the fusion root.
+    if (instruction->opcode() == HloOpcode::kFusion) {
+      Propagate(position.index, instruction->fused_expression_root(),
+                memory_space);
+    }
+
+    const HloInstruction* parent_fusion =
+        instruction->parent()->FusionInstruction();
+    // For nested fusion roots, pop one level up and propagate the memory space
+    // to the output of the calling fusion instruction.
+    if (instruction == instruction->parent()->root_instruction() &&
+        parent_fusion->parent()->IsFusionComputation()) {
+      Propagate(position.index, parent_fusion, memory_space);
+    }
+
+    // For nested fusion parameters, pop one level up and propagate the memory
+    // space to the operand of the calling fusion instruction.
+    if (instruction->opcode() == HloOpcode::kParameter &&
+        parent_fusion->parent()->IsFusionComputation()) {
+      const HloInstruction* fusion_operand =
+          parent_fusion->operand(instruction->parameter_number());
+      Propagate(position.index, fusion_operand, memory_space);
+    }
+  }
+
+  for (const HloUse& use : value.uses()) {
+    // For fusion uses, propagate the memory space to the fusion parameter.
+    if (use.instruction->opcode() == HloOpcode::kFusion) {
+      modified |= Propagate(
+          use.operand_index,
+          use.instruction->fused_parameter(use.operand_number), memory_space);
     }
   }
   return modified;
diff --git a/tensorflow/compiler/xla/service/memory_space_propagation.h b/tensorflow/compiler/xla/service/memory_space_propagation.h
index 65a1dfd14a6..510e9e69f79 100644
--- a/tensorflow/compiler/xla/service/memory_space_propagation.h
+++ b/tensorflow/compiler/xla/service/memory_space_propagation.h
@@ -31,12 +31,11 @@ class MemorySpacePropagation : public HloModulePass {
   StatusOr Run(HloModule* module) override;
 
  private:
-  // Given the caller shape (operand or output) and its corresponding
-  // insturction in the fused computation (parameter or root), propagates the
-  // memory space to all the subshapes in the callee side. Returns true if the
-  // module is modified.
-  bool PropagateSubshapes(const Shape& caller_shape,
-                          const HloInstruction* callee_instruction) const;
+  // Given the shape index (operand or output) and its corresponding instruction
+  // in the fused computation (parameter or root), propagates the memory space
+  // in the callee side. Returns true if the module is modified.
+  bool Propagate(ShapeIndexView index, const HloInstruction* callee_instruction,
+                 int64 memory_space) const;
 
   std::unique_ptr dataflow_analysis_;
 };
diff --git a/tensorflow/compiler/xla/service/memory_space_propagation_test.cc b/tensorflow/compiler/xla/service/memory_space_propagation_test.cc
index 8d74958f6aa..de45af5a190 100644
--- a/tensorflow/compiler/xla/service/memory_space_propagation_test.cc
+++ b/tensorflow/compiler/xla/service/memory_space_propagation_test.cc
@@ -199,5 +199,153 @@ TEST_F(MemorySpacePropagationTest, TupleOutput) {
   EXPECT_EQ(module->Hash(), ref->Hash());
 }
 
+TEST_F(MemorySpacePropagationTest, NestedInputFusion) {
+  // Tests propagating the memory space to nested fusions on the input side.
+  absl::string_view hlo_string = R"(
+  HloModule NestedFusion
+
+  %bitcast_fusion {
+    %bf_param = s32[3,2]{0,1:T(128)} parameter(0)
+    ROOT %bitcast = s32[6]{0:T(128)} bitcast(%bf_param)
+  }
+
+  %fused_computation {
+    %param_1.3 = s32[1]{0:T(128)} parameter(1)
+    %constant.2 = s32[]{:T(128)} constant(-2147483648)
+    %pad.2 = s32[6]{0:T(128)} pad(s32[1]{0:T(128)} %param_1.3, s32[]{:T(128)} %constant.2), padding=0_5
+    %param_2.3 = s32[5]{0:T(128)} parameter(2)
+    %pad.3 = s32[6]{0:T(128)} pad(s32[5]{0:T(128)} %param_2.3, s32[]{:T(128)} %constant.2), padding=1_0
+    %maximum.1 = s32[6]{0:T(128)} maximum(s32[6]{0:T(128)} %pad.2, s32[6]{0:T(128)} %pad.3)
+    %param_0.1 = s32[3,2]{0,1:T(128)} parameter(0)
+    %fusion.1 = s32[6]{0:T(128)} fusion(%param_0.1), kind=kLoop, calls=bitcast_fusion
+    ROOT %add.0 = s32[6]{0:T(128)} add(s32[6]{0:T(128)} %maximum.1, s32[6]{0:T(128)} %fusion.1)
+  }
+
+  ENTRY %entry {
+    %param0 = s32[3,2]{0,1:T(128)} parameter(0)
+    %param1 = s32[1]{0:T(128)} parameter(1)
+    %param2 = s32[5]{0:T(128)} parameter(2)
+    %arg0 = s32[3,2]{0,1:T(128)S(1)} copy(%param0)
+    %arg1 = s32[1]{0:T(128)} copy(%param1)
+    %arg2 = s32[5]{0:T(128)S(1)} copy(%param2)
+    %fusion = s32[6]{0:T(128)S(1)} fusion(s32[3,2]{0,1:T(128)S(1)} %arg0, s32[1]{0:T(128)} %arg1, s32[5]{0:T(128)S(1)} %arg2), kind=kLoop, calls=%fused_computation
+    ROOT %root = s32[6]{0:T(128)} copy(%fusion)
+  }
+  )";
+  absl::string_view expected_hlo_string = R"(
+  HloModule NestedFusion
+
+  %bitcast_fusion {
+    %bf_param = s32[3,2]{0,1:T(128)S(1)} parameter(0)
+    ROOT %bitcast = s32[6]{0:T(128)S(1)} bitcast(%bf_param)
+  }
+
+  %fused_computation {
+    %param_1.3 = s32[1]{0:T(128)} parameter(1)
+    %constant.2 = s32[]{:T(128)} constant(-2147483648)
+    %pad.2 = s32[6]{0:T(128)} pad(s32[1]{0:T(128)} %param_1.3, s32[]{:T(128)} %constant.2), padding=0_5
+    %param_2.3 = s32[5]{0:T(128)S(1)} parameter(2)
+    %pad.3 = s32[6]{0:T(128)} pad(s32[5]{0:T(128)} %param_2.3, s32[]{:T(128)} %constant.2), padding=1_0
+    %maximum.1 = s32[6]{0:T(128)} maximum(s32[6]{0:T(128)} %pad.2, s32[6]{0:T(128)} %pad.3)
+    %param_0.1 = s32[3,2]{0,1:T(128)S(1)} parameter(0)
+    %fusion.1 = s32[6]{0:T(128)S(1)} fusion(%param_0.1), kind=kLoop, calls=bitcast_fusion
+    ROOT %add.0 = s32[6]{0:T(128)S(1)} add(s32[6]{0:T(128)} %maximum.1, s32[6]{0:T(128)S(1)} %fusion.1)
+  }
+
+  ENTRY %entry {
+    %param0 = s32[3,2]{0,1:T(128)} parameter(0)
+    %param1 = s32[1]{0:T(128)} parameter(1)
+    %param2 = s32[5]{0:T(128)} parameter(2)
+    %arg0 = s32[3,2]{0,1:T(128)S(1)} copy(%param0)
+    %arg1 = s32[1]{0:T(128)} copy(%param1)
+    %arg2 = s32[5]{0:T(128)S(1)} copy(%param2)
+    %fusion = s32[6]{0:T(128)S(1)} fusion(s32[3,2]{0,1:T(128)S(1)} %arg0, s32[1]{0:T(128)} %arg1, s32[5]{0:T(128)S(1)} %arg2), kind=kLoop, calls=%fused_computation
+    ROOT %root = s32[6]{0:T(128)} copy(%fusion)
+  }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnUnverifiedModule(hlo_string));
+  MemorySpacePropagation memory_space_propagation;
+  EXPECT_TRUE(memory_space_propagation.Run(module.get()).ValueOrDie());
+  TF_EXPECT_OK(Verify(module.get()));
+  TF_ASSERT_OK_AND_ASSIGN(auto ref,
+                          ParseAndReturnVerifiedModule(expected_hlo_string));
+  EXPECT_EQ(module->Hash(), ref->Hash());
+}
+
+TEST_F(MemorySpacePropagationTest, NestedOutputFusion) {
+  // Tests propagating the memory space to nested fusions on the output side.
+  absl::string_view hlo_string = R"(
+  HloModule NestedFusion
+
+  %bitcast_fusion {
+    %bf_param = s32[6]{0:T(128)} parameter(0)
+    ROOT %bitcast = s32[3,2]{0,1:T(128)} bitcast(%bf_param)
+  }
+
+  %fused_computation {
+    %param_1.3 = s32[1]{0:T(128)} parameter(1)
+    %constant.2 = s32[]{:T(128)} constant(-2147483648)
+    %pad.2 = s32[6]{0:T(128)} pad(s32[1]{0:T(128)} %param_1.3, s32[]{:T(128)} %constant.2), padding=0_5
+    %param_2.3 = s32[5]{0:T(128)} parameter(2)
+    %pad.3 = s32[6]{0:T(128)} pad(s32[5]{0:T(128)} %param_2.3, s32[]{:T(128)} %constant.2), padding=1_0
+    %maximum.1 = s32[6]{0:T(128)} maximum(s32[6]{0:T(128)} %pad.2, s32[6]{0:T(128)} %pad.3)
+    %param_0.1 = s32[6]{0:T(128)} parameter(0)
+    %add.0 = s32[6]{0:T(128)} add(s32[6]{0:T(128)} %maximum.1, s32[6]{0:T(128)} %param_0.1)
+    ROOT %fusion.1 = s32[3,2]{0,1:T(128)} fusion(%add.0), kind=kLoop, calls=bitcast_fusion
+  }
+
+  ENTRY %entry {
+    %param0 = s32[6]{0:T(128)} parameter(0)
+    %param1 = s32[1]{0:T(128)} parameter(1)
+    %param2 = s32[5]{0:T(128)} parameter(2)
+    %arg0 = s32[6]{0:T(128)S(1)} copy(%param0)
+    %arg1 = s32[1]{0:T(128)} copy(%param1)
+    %arg2 = s32[5]{0:T(128)S(1)} copy(%param2)
+    %fusion = s32[3,2]{0,1:T(128)S(1)} fusion(s32[6]{0:T(128)S(1)} %arg0, s32[1]{0:T(128)} %arg1, s32[5]{0:T(128)S(1)} %arg2), kind=kLoop, calls=%fused_computation
+    ROOT %root = s32[3,2]{0,1:T(128)} copy(%fusion)
+  }
+  )";
+  absl::string_view expected_hlo_string = R"(
+  HloModule NestedFusion
+
+  %bitcast_fusion {
+    %bf_param = s32[6]{0:T(128)S(1)} parameter(0)
+    ROOT %bitcast = s32[3,2]{0,1:T(128)S(1)} bitcast(%bf_param)
+  }
+
+  %fused_computation {
+    %param_1.3 = s32[1]{0:T(128)} parameter(1)
+    %constant.2 = s32[]{:T(128)} constant(-2147483648)
+    %pad.2 = s32[6]{0:T(128)} pad(s32[1]{0:T(128)} %param_1.3, s32[]{:T(128)} %constant.2), padding=0_5
+    %param_2.3 = s32[5]{0:T(128)S(1)} parameter(2)
+    %pad.3 = s32[6]{0:T(128)} pad(s32[5]{0:T(128)} %param_2.3, s32[]{:T(128)} %constant.2), padding=1_0
+    %maximum.1 = s32[6]{0:T(128)} maximum(s32[6]{0:T(128)} %pad.2, s32[6]{0:T(128)} %pad.3)
+    %param_0.1 = s32[6]{0:T(128)S(1)} parameter(0)
+    %add.0 = s32[6]{0:T(128)S(1)} add(s32[6]{0:T(128)} %maximum.1, s32[6]{0:T(128)S(1)} %param_0.1)
+    ROOT %fusion.1 = s32[3,2]{0,1:T(128)S(1)} fusion(%add.0), kind=kLoop, calls=bitcast_fusion
+  }
+
+  ENTRY %entry {
+    %param0 = s32[6]{0:T(128)} parameter(0)
+    %param1 = s32[1]{0:T(128)} parameter(1)
+    %param2 = s32[5]{0:T(128)} parameter(2)
+    %arg0 = s32[6]{0:T(128)S(1)} copy(%param0)
+    %arg1 = s32[1]{0:T(128)} copy(%param1)
+    %arg2 = s32[5]{0:T(128)S(1)} copy(%param2)
+    %fusion = s32[3,2]{0,1:T(128)S(1)} fusion(s32[6]{0:T(128)S(1)} %arg0, s32[1]{0:T(128)} %arg1, s32[5]{0:T(128)S(1)} %arg2), kind=kLoop, calls=%fused_computation
+    ROOT %root = s32[3,2]{0,1:T(128)} copy(%fusion)
+  }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnUnverifiedModule(hlo_string));
+  MemorySpacePropagation memory_space_propagation;
+  EXPECT_TRUE(memory_space_propagation.Run(module.get()).ValueOrDie());
+  TF_EXPECT_OK(Verify(module.get()));
+  TF_ASSERT_OK_AND_ASSIGN(auto ref,
+                          ParseAndReturnVerifiedModule(expected_hlo_string));
+  EXPECT_EQ(module->Hash(), ref->Hash());
+}
+
 }  // namespace
 }  // namespace xla

From fc5151130813140eb4189f77dd3a759c4077836d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 13:41:07 -0700
Subject: [PATCH 0534/1390] Delete orphaned comment

The comment for `loader_spec_` stuck around even after it was removed.

PiperOrigin-RevId: 317171521
Change-Id: Iddb6029fdad9cd5ef33bc4f4ea2653caed305658
---
 tensorflow/compiler/xla/service/gpu/kernel_thunk.h | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h
index 88351881f3a..25acabb239b 100644
--- a/tensorflow/compiler/xla/service/gpu/kernel_thunk.h
+++ b/tensorflow/compiler/xla/service/gpu/kernel_thunk.h
@@ -77,8 +77,6 @@ class KernelThunk : public Thunk {
   // Will be set by IrEmitterUnnested.
   LaunchDimensions launch_dimensions_;
 
-  // Describes how to load this kernel. ExecuteOnStream reuses this loader
-  // specification for all executions.
   mutable tensorflow::mutex mutex_;
 
   // Loaded kernels for each `StreamExecutor`.  Requires pointer stability of

From 8452c9f80ee02cb71fb72f638d3bdef754f15297 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Thu, 18 Jun 2020 13:58:25 -0700
Subject: [PATCH 0535/1390] Added status_group fuzzer

---
 .../security/fuzzing/status_group_fuzz.cc     | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)
 create mode 100644 tensorflow/security/fuzzing/status_group_fuzz.cc

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
new file mode 100644
index 00000000000..979fd444b48
--- /dev/null
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -0,0 +1,83 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+#include 
+
+#include "tensorflow/core/platform/status.h"
+
+#include 
+
+// This is a fuzzer for `tensorflow::StatusGroup`. Since `Status` is used almost
+// everywhere, we need to ensure that the common functionality is safe. We don't
+// expect many crashes from this fuzzer
+
+namespace {
+
+tensorflow::error::Code BuildRandomErrorCode(uint32_t code){
+
+  // We cannot build a `Status` with error_code of 0 and a message, so force
+  // error code to be non-zero.
+  if (code == 0) {
+    return tensorflow::error::UNKNOWN;
+  }
+
+  return static_cast(code);
+}
+
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
+  tensorflow::error::Code error_code;
+
+  std::string error_message = "ERROR";
+
+  tensorflow::Status s, derived_s;
+
+  tensorflow::StatusGroup sg;
+
+  bool is_derived;
+
+  uint32_t code;
+
+  FuzzedDataProvider fuzzed_data(data, size);
+
+  while(fuzzed_data.remaining_bytes() > 0) {
+    code = fuzzed_data.ConsumeIntegral();
+
+    error_code = BuildRandomErrorCode(code);
+
+    is_derived = fuzzed_data.ConsumeBool();
+
+    s = tensorflow::Status(error_code, error_message);
+
+    if(is_derived) {
+      derived_s = tensorflow::StatusGroup::MakeDerived(s);
+
+      sg.Update(derived_s);
+
+    } else {
+      sg.Update(s);
+
+    }
+  }
+
+  sg.as_summary_status();
+
+  sg.as_concatenated_status();
+
+  sg.AttachLogMessages();
+
+  return 0;
+}
+
+}  // namespace

From c41f4652b45bf70f20686e612b41574b4b8139d7 Mon Sep 17 00:00:00 2001
From: Marissa Ikonomidis 
Date: Thu, 18 Jun 2020 13:52:21 -0700
Subject: [PATCH 0536/1390] Add an option to enable MLIR bridge for tpu_py_test
 rule

If enable_mlir_bridge is True, a new test will be generated that runs with the MLIR bridge enabled.
This option is off by default.

PiperOrigin-RevId: 317173675
Change-Id: I332e1ae24cf82fceea20fd0aff2cec7c9b236a24
---
 tensorflow/core/platform/default/distribute.bzl | 3 +++
 tensorflow/python/framework/test_util.py        | 3 +++
 tensorflow/python/tpu/tpu.bzl                   | 2 ++
 3 files changed, 8 insertions(+)

diff --git a/tensorflow/core/platform/default/distribute.bzl b/tensorflow/core/platform/default/distribute.bzl
index 46a5d826a79..b16d5e8cff7 100644
--- a/tensorflow/core/platform/default/distribute.bzl
+++ b/tensorflow/core/platform/default/distribute.bzl
@@ -22,6 +22,7 @@ def distribute_py_test(
         full_precision = False,
         disable_v2 = False,
         disable_v3 = False,
+        disable_mlir_bridge = True,
         **kwargs):
     """Generates py_test targets for CPU and GPU.
 
@@ -40,6 +41,7 @@ def distribute_py_test(
         full_precision: unused.
         disable_v2: whether tests for TPU version 2 should be generated.
         disable_v3: whether tests for TPU version 3 should be generated.
+        disable_mlir_bridge: whether to also run this with the mlir bridge enabled.
         **kwargs: extra keyword arguments to the non-tpu test.
     """
 
@@ -77,6 +79,7 @@ def distribute_py_test(
             tags = tpu_tags,
             disable_v2 = disable_v2,
             disable_v3 = disable_v3,
+            disable_mlir_bridge = disable_mlir_bridge,
         )
 
 register_extension_info(
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index a46bb7c9bda..8ddbcf34f3b 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -1933,6 +1933,9 @@ class TensorFlowTestCase(googletest.TestCase):
       # disable it here.
       pywrap_tf_session.TF_SetXlaConstantFoldingDisabled(True)
 
+    if is_mlir_bridge_enabled():
+      context.context().enable_mlir_bridge = True
+
     self._threads = []
     self._tempdir = None
     self._cached_session = None
diff --git a/tensorflow/python/tpu/tpu.bzl b/tensorflow/python/tpu/tpu.bzl
index 5453702d64d..3c26d9b49bf 100644
--- a/tensorflow/python/tpu/tpu.bzl
+++ b/tensorflow/python/tpu/tpu.bzl
@@ -25,6 +25,7 @@ def tpu_py_test(
         disable_v2 = False,
         disable_v3 = False,
         disable_experimental = False,
+        disable_mlir_bridge = True,
         args = [],
         **kwargs):
     """Generates identical unit test variants for various Cloud TPU versions.
@@ -37,6 +38,7 @@ def tpu_py_test(
         disable_v2: If true, don't generate TPU v2 tests.
         disable_v3: If true, don't generate TPU v3 tests.
         disable_experimental: Unused.
+        disable_mlir_bridge: Unused.
         args: Arguments to apply to tests.
         **kwargs: Additional named arguments to apply to tests.
     """

From dda51e1c94160b8252c51dd0ddca445d821ba8b9 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Thu, 18 Jun 2020 13:59:11 -0700
Subject: [PATCH 0537/1390] Added status group fuzzer build rules

---
 tensorflow/security/fuzzing/BUILD | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/security/fuzzing/BUILD b/tensorflow/security/fuzzing/BUILD
index 9b5aeec2d36..871baa0055b 100644
--- a/tensorflow/security/fuzzing/BUILD
+++ b/tensorflow/security/fuzzing/BUILD
@@ -19,6 +19,14 @@ tf_fuzz_target(
     ],
 )
 
+tf_fuzz_target(
+    name = "status_group_fuzz",
+    srcs = ["status_group_fuzz.cc"],
+    deps = [
+        "//tensorflow/core/platform:status",
+    ],
+)
+
 # A trivial fuzzer with no pre-specified corpus.
 # TODO(mihaimaruseac): Move fuzz_session and the op fuzzers to a subdirectory
 tf_fuzz_target(

From ef52b4e0886b7212471462643e92e98bea0253be Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Thu, 18 Jun 2020 13:53:10 -0700
Subject: [PATCH 0538/1390] [XLA/Client] Implement LocalClient::Run which
 supports buffer donation

PiperOrigin-RevId: 317173848
Change-Id: If92955ac5051376fbf0932b773b675459497c0c4
---
 .../compiler/xla/client/local_client.cc       | 47 +++++++++++++++----
 tensorflow/compiler/xla/client/local_client.h |  9 ++++
 .../tests/multiple_devices_on_host_test.cc    |  3 +-
 tensorflow/compiler/xla/tests/while_test.cc   |  6 ++-
 4 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index afe115deda8..f71e8a2d56d 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -168,6 +168,26 @@ LocalExecutable::RunHelper(const absl::Span argument_shapes,
   return std::make_pair(service_options, std::move(stream));
 }
 
+StatusOr LocalExecutable::GetExecutableRunOptions(
+    absl::Span argument_shapes,
+    const ExecutableRunOptions& run_options) {
+  TF_ASSIGN_OR_RETURN(auto options_and_stream,
+                      RunHelper(argument_shapes, run_options));
+  ExecutableRunOptions options = options_and_stream.first.run_options();
+  options.set_device_ordinal(-1);
+  return options;
+}
+
+template 
+static StatusOr BlockHostUntilDoneAfterAsyncCall(
+    se::Stream* stream, std::function()> async_callback) {
+  StatusOr result = async_callback();
+  Status block_status = stream->BlockHostUntilDone();
+  TF_RETURN_IF_ERROR(result.status());
+  TF_RETURN_IF_ERROR(block_status);
+  return result;
+}
+
 StatusOr LocalExecutable::Run(
     const absl::Span arguments,
     ExecutableRunOptions run_options) {
@@ -176,15 +196,24 @@ StatusOr LocalExecutable::Run(
   for (const ShapedBuffer* const arg : arguments) {
     argument_shapes.push_back(&arg->on_host_shape());
   }
-  TF_ASSIGN_OR_RETURN(auto options_and_stream,
-                      RunHelper(argument_shapes, run_options));
-  ExecutableRunOptions options = options_and_stream.first.run_options();
-  options.set_device_ordinal(-1);
-  auto result = RunAsync(arguments, options);
-  Status block_status = options.stream()->BlockHostUntilDone();
-  TF_RETURN_IF_ERROR(result.status());
-  TF_RETURN_IF_ERROR(block_status);
-  return result;
+  TF_ASSIGN_OR_RETURN(ExecutableRunOptions options,
+                      GetExecutableRunOptions(argument_shapes, run_options));
+  return BlockHostUntilDoneAfterAsyncCall(
+      options.stream(), [&] { return RunAsync(arguments, options); });
+}
+
+StatusOr LocalExecutable::Run(
+    std::vector arguments, ExecutableRunOptions run_options) {
+  std::vector argument_shapes;
+  argument_shapes.reserve(arguments.size());
+  for (const ExecutionInput& arg : arguments) {
+    argument_shapes.push_back(&arg.shape());
+  }
+  TF_ASSIGN_OR_RETURN(ExecutableRunOptions options,
+                      GetExecutableRunOptions(argument_shapes, run_options));
+  return BlockHostUntilDoneAfterAsyncCall(
+      options.stream(),
+      [&] { return RunAsync(argument_shapes, std::move(arguments), options); });
 }
 
 static std::shared_ptr DumpArguments(
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index 7cdeb9dcbf6..b00f5cc6801 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -51,6 +51,11 @@ class LocalExecutable {
       const absl::Span arguments,
       ExecutableRunOptions run_options);
 
+  // Similar to Run(), but allows for donating argument buffers to the
+  // executable.
+  StatusOr Run(std::vector arguments,
+                                ExecutableRunOptions run_options);
+
   // Similar to Run(), but need not block the host waiting for the computation
   // to complete before returning.
   StatusOr RunAsync(
@@ -85,6 +90,10 @@ class LocalExecutable {
       const absl::Span argument_shapes,
       ExecutableRunOptions run_options);
 
+  StatusOr GetExecutableRunOptions(
+      absl::Span argument_shapes,
+      const ExecutableRunOptions& run_options);
+
   // The ordinal of the device which this executable was compiled for. The
   // executable can run on all equivalent devices (as determined by
   // Backend::devices_equivalent).
diff --git a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
index 2b19aaded9c..2231fc6feab 100644
--- a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
+++ b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
@@ -45,7 +45,8 @@ void CompileAndExecute(
       xla::ClientLibrary::GetXlaService(client->platform())
           ->backend()
           .memory_allocator());
-  StatusOr result = executable->Run({}, execute_options);
+  StatusOr result =
+      executable->Run(absl::Span(), execute_options);
   {
     absl::MutexLock lock(results_mutex);
     results->emplace_back(device_ordinal, std::move(result));
diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc
index d575bbb1f3e..8e8c3605cc7 100644
--- a/tensorflow/compiler/xla/tests/while_test.cc
+++ b/tensorflow/compiler/xla/tests/while_test.cc
@@ -1324,14 +1324,16 @@ void BM_WhileLoop(int num_iters) {
   options.set_allocator(&allocator);
   const int kWarmups = 2;
   for (int i = 0; i < kWarmups; ++i) {
-    auto result = executable->Run({}, options);
+    auto result =
+        executable->Run(absl::Span(), options);
     ASSERT_TRUE(result.ok());
   }
 
   // Run benchmark.
   tensorflow::testing::StartTiming();
   for (int i = 0; i < num_iters; ++i) {
-    auto result = executable->Run({}, options);
+    auto result =
+        executable->Run(absl::Span(), options);
     ASSERT_TRUE(result.ok());
   }
 }

From 3833402726d72d04dd2821e89c642f613e80a531 Mon Sep 17 00:00:00 2001
From: jonah-kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Thu, 18 Jun 2020 14:28:47 -0700
Subject: [PATCH 0539/1390] Cast optimizer parameters as python floats during
 serialization.

Accounted only for the case in which the hyper parameter is neither a callable nor a tensor, to avoid any troublesome casts.
---
 tensorflow/python/keras/optimizer_v2/optimizer_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
index c55b332bfc0..d8992bbe3e0 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
@@ -910,7 +910,7 @@ class OptimizerV2(trackable.Trackable):
       return value()
     if tensor_util.is_tensor(value):
       return backend.get_value(value)
-    return value
+    return float(value)
 
   def variables(self):
     """Returns variables of this Optimizer based on the order created."""

From bc1c0e86a677d9b1e5d3e3f0da85c445c2a7efe2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 14:24:06 -0700
Subject: [PATCH 0540/1390] Wrap save/restore logic in tf.function when in
 eager mode. This allows parallel saving and restoring when using multiple
 devices.

PiperOrigin-RevId: 317180143
Change-Id: Icdc740d02beb7c2d3236191add3b72fa103fc134
---
 .../grappler/optimizers/function_optimizer.cc |   8 +-
 .../parallel_device/parallel_device_test.py   |   4 -
 .../python/framework/auto_control_deps.py     |   2 +-
 tensorflow/python/training/saving/BUILD       |   1 -
 .../training/saving/functional_saver.py       | 111 ++++++------------
 .../training/saving/functional_saver_test.py  |  17 +--
 6 files changed, 42 insertions(+), 101 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index 0e156aaa84c..a66e645e04b 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -837,6 +837,7 @@ const bool IsExemptFromSideEffectsExecutionValidation(const string& op) {
        "ParameterizedTruncatedNormal", "TruncatedNormal", "RandomShuffle",
        "Multinomial", "RandomGamma", "RandomGammaGrad", "RandomPoisson",
        "RandomPoissonV2",
+       // LINT.ThenChange(//tensorflow/python/framework/auto_control_deps.py)
 
        // ReadVariableOp marked as stateful because it consumes DT_RESOURCE,
        // but it can't generate any observable side-effect.
@@ -850,12 +851,7 @@ const bool IsExemptFromSideEffectsExecutionValidation(const string& op) {
        // the same device_ordinal on the same host.
        "EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
        "EnqueueTPUEmbeddingSparseTensorBatch",
-       "EnqueueTPUEmbeddingRaggedTensorBatch",
-
-       // SaveV2 and RestoreV2 should be allowed to operate in parallel on
-       // multiple hosts.
-       "SaveV2", "RestoreV2"});
-  // LINT.ThenChange(//tensorflow/python/framework/auto_control_deps.py)
+       "EnqueueTPUEmbeddingRaggedTensorBatch"});
   return exemption->contains(op);
 }
 
diff --git a/tensorflow/python/distribute/parallel_device/parallel_device_test.py b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
index 1429c522aba..8fc3dcb5816 100644
--- a/tensorflow/python/distribute/parallel_device/parallel_device_test.py
+++ b/tensorflow/python/distribute/parallel_device/parallel_device_test.py
@@ -172,8 +172,6 @@ class ParallelDeviceTests(_VirtualDeviceTestCase):
       config.set_synchronous_execution(previous)
 
   def test_checkpointing(self):
-    self.skipTest(
-        "Disable saving until SaveableObject's methods are traceable.")
     prefix = os.path.join(self.get_temp_dir(), "ckpt")
     with self.device.scope():
       different_values = self.device.pack(
@@ -265,8 +263,6 @@ class LayerTests(_VirtualDeviceTestCase):
     self.assertIn(self.device.components[1], final_kernels[1].backing_device)
 
   def test_training_loop(self):
-    self.skipTest(
-        "Disable saving until SaveableObject's methods are traceable.")
     for _ in range(5):
       layer = _Dense(5)
       checkpoint = tracking.Checkpoint(layer=layer)
diff --git a/tensorflow/python/framework/auto_control_deps.py b/tensorflow/python/framework/auto_control_deps.py
index 4b47735e0bf..51dcb248b11 100644
--- a/tensorflow/python/framework/auto_control_deps.py
+++ b/tensorflow/python/framework/auto_control_deps.py
@@ -100,7 +100,7 @@ _ORDER_INSENSITIVE_STATEFUL_OPS = [
     "CudnnRNNV2", "CudnnRNNV3", "CudnnRNNBackpropV2", "CudnnRNNBackpropV3",
     "EnqueueTPUEmbeddingSparseBatch", "EnqueueTPUEmbeddingIntegerBatch",
     "EnqueueTPUEmbeddingSparseTensorBatch",
-    "EnqueueTPUEmbeddingRaggedTensorBatch", "RestoreV2", "SaveV2"
+    "EnqueueTPUEmbeddingRaggedTensorBatch"
 ]
 # LINT.ThenChange(//tensorflow/core/grappler/optimizers/function_optimizer.cc)
 
diff --git a/tensorflow/python/training/saving/BUILD b/tensorflow/python/training/saving/BUILD
index 12940840309..670a4c35c6f 100644
--- a/tensorflow/python/training/saving/BUILD
+++ b/tensorflow/python/training/saving/BUILD
@@ -43,7 +43,6 @@ cuda_py_test(
         ":checkpoint_options",
         ":functional_saver",
         ":saveable_hook",
-        "//tensorflow/python/eager:remote",
         "//tensorflow/python/eager:test",
     ],
 )
diff --git a/tensorflow/python/training/saving/functional_saver.py b/tensorflow/python/training/saving/functional_saver.py
index 3a9b565470d..c4334e096df 100644
--- a/tensorflow/python/training/saving/functional_saver.py
+++ b/tensorflow/python/training/saving/functional_saver.py
@@ -21,7 +21,6 @@ from __future__ import print_function
 import uuid
 
 from tensorflow.core.protobuf import saver_pb2
-from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -162,8 +161,7 @@ class MultiDeviceSaver(object):
         self._after_restore_callbacks.append(saveable.after_restore)
 
       if is_saveable:
-        host_device = saveable_object_util.set_cpu0(saveable.device)
-        saveables_by_device.setdefault(host_device, []).append(saveable)
+        saveables_by_device.setdefault(saveable.device, []).append(saveable)
 
     self._single_device_savers = {
         device: _SingleDeviceSaver(saveables)
@@ -249,50 +247,33 @@ class MultiDeviceSaver(object):
       tmp_checkpoint_prefix = string_ops.string_join(
           [file_prefix, sharded_suffix])
 
-    def save_fn():
-      num_shards = len(self._single_device_savers)
-      sharded_saves = []
-      sharded_prefixes = []
-      num_shards_tensor = constant_op.constant(num_shards, name="num_shards")
-      last_device = None
-      for shard, (device, saver) in enumerate(
-          sorted(self._single_device_savers.items())):
-        last_device = device
-        with ops.device(saveable_object_util.set_cpu0(device)):
-          shard_prefix = sharded_filename(tmp_checkpoint_prefix, shard,
-                                          num_shards_tensor)
-        sharded_prefixes.append(shard_prefix)
-        with ops.device(device):
-          # _SingleDeviceSaver will use the CPU device when necessary, but
-          # initial read operations should be placed on the SaveableObject's
-          # device.
-          sharded_saves.append(saver.save(shard_prefix, options))
+    num_shards = len(self._single_device_savers)
+    sharded_saves = []
+    sharded_prefixes = []
+    num_shards_tensor = constant_op.constant(num_shards, name="num_shards")
+    last_device = None
+    for shard, (device, saver) in enumerate(
+        sorted(self._single_device_savers.items())):
+      last_device = device
+      with ops.device(saveable_object_util.set_cpu0(device)):
+        shard_prefix = sharded_filename(tmp_checkpoint_prefix, shard,
+                                        num_shards_tensor)
+      sharded_prefixes.append(shard_prefix)
+      with ops.device(device):
+        # _SingleDeviceSaver will use the CPU device when necessary, but initial
+        # read operations should be placed on the SaveableObject's device.
+        sharded_saves.append(saver.save(shard_prefix, options))
 
-      with ops.control_dependencies(sharded_saves):
-        # Merge on the io_device if specified, otherwise co-locates the merge op
-        # with the last device used.
-        merge_device = (
-            options.experimental_io_device or
-            saveable_object_util.set_cpu0(last_device))
-        with ops.device(merge_device):
-          # V2 format write path consists of a metadata merge step.  Once
-          # merged, attempts to delete the temporary directory,
-          # "_temp".
-          return gen_io_ops.merge_v2_checkpoints(
-              sharded_prefixes, file_prefix, delete_old_dirs=True)
-
-    # Since this will causes a function re-trace on each save, limit this to the
-    # cases where it is needed: eager and when there are multiple tasks/single
-    # device savers. Note that the retrace is needed to ensure we pickup the
-    # latest values of options like experimental_io_device.
-    if context.executing_eagerly() and len(self._single_device_savers) > 1:
-      # Explicitly place the identity op on the first device.
-      @def_function.function(experimental_compile=False)
-      def tf_function_save():
-        save_fn()
-      tf_function_save()
-    else:
-      return save_fn()
+    with ops.control_dependencies(sharded_saves):
+      # Merge on the io_device if specified, otherwise co-locates the merge op
+      # with the last device used.
+      merge_device = (options.experimental_io_device or
+                      saveable_object_util.set_cpu0(last_device))
+      with ops.device(merge_device):
+        # V2 format write path consists of a metadata merge step.  Once merged,
+        # attempts to delete the temporary directory, "_temp".
+        return gen_io_ops.merge_v2_checkpoints(
+            sharded_prefixes, file_prefix, delete_old_dirs=True)
 
   def restore(self, file_prefix, options=None):
     """Restore the saveable objects from a checkpoint with `file_prefix`.
@@ -306,38 +287,12 @@ class MultiDeviceSaver(object):
       A dictionary mapping from SaveableObject names to restore operations.
     """
     options = options or checkpoint_options.CheckpointOptions()
-
-    def restore_fn():
-      restore_ops = {}
-      # Sort by device name to avoid propagating non-deterministic dictionary
-      # ordering in some Python versions.
-      for device, saver in sorted(self._single_device_savers.items()):
-        with ops.device(device):
-          restore_ops.update(saver.restore(file_prefix, options))
-
-      return restore_ops
-
-    # Since this will causes a function re-trace on each save, limit this to the
-    # cases where it is needed: eager and when there are multiple tasks/single
-    # device savers. Note that the retrace is needed to ensure we pickup the
-    # latest values of options like experimental_io_device.
-    if context.executing_eagerly() and len(self._single_device_savers) > 1:
-      first_device, _ = list(self._single_device_savers.items())[0]
-      @def_function.function(experimental_compile=False)
-      def tf_function_restore():
-        restore_ops = restore_fn()
-        restore_tensors = {}
-        # tf.functions must return tensors, thus we use control dependencies so
-        # that we can return a tensor which depends on the given op.
-        with ops.device(saveable_object_util.set_cpu0(first_device)):
-          for name, op in restore_ops.items():
-            with ops.control_dependencies([op]):
-              restore_tensors[name] = array_ops.identity(file_prefix)
-        return restore_tensors
-
-      restore_ops = tf_function_restore()
-    else:
-      restore_ops = restore_fn()
+    restore_ops = {}
+    # Sort by device name to avoid propagating non-deterministic dictionary
+    # ordering in some Python versions.
+    for device, saver in sorted(self._single_device_savers.items()):
+      with ops.device(device):
+        restore_ops.update(saver.restore(file_prefix, options))
 
     for callback in self._after_restore_callbacks:
       callback()
diff --git a/tensorflow/python/training/saving/functional_saver_test.py b/tensorflow/python/training/saving/functional_saver_test.py
index 8f3eef4fb9c..7db32ff72d7 100644
--- a/tensorflow/python/training/saving/functional_saver_test.py
+++ b/tensorflow/python/training/saving/functional_saver_test.py
@@ -21,7 +21,6 @@ from __future__ import print_function
 import os
 
 from tensorflow.python.eager import context
-from tensorflow.python.eager import remote
 from tensorflow.python.eager import test
 from tensorflow.python.eager import wrap_function
 from tensorflow.python.framework import config
@@ -30,7 +29,6 @@ from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import gfile
-from tensorflow.python.training import server_lib
 from tensorflow.python.training.saving import checkpoint_options
 from tensorflow.python.training.saving import functional_saver
 from tensorflow.python.training.saving import saveable_hook
@@ -128,16 +126,13 @@ class SaverTest(test.TestCase):
     second_saver.restore(save_path)
     self.assertEqual(2., self.evaluate(v2))
 
-  def test_checkpoint_is_sharded_by_task(self):
-    servers = [server_lib.Server.create_local_server() for _ in range(3)]
-    cluster_spec = server_lib.ClusterSpec({
-        "worker": [s.target[len("grpc://"):] for s in servers]})
-    remote.connect_to_cluster(cluster_spec)
-    with ops.device("/job:worker/task:0/cpu:0"):
+  @test_util.run_in_graph_and_eager_modes
+  def test_checkpoint_is_sharded_by_device(self):
+    with ops.device("cpu:0"):
       v0 = resource_variable_ops.ResourceVariable(0.)
-    with ops.device("/job:worker/task:1/cpu:0"):
+    with ops.device("cpu:1"):
       v1 = resource_variable_ops.ResourceVariable(1.)
-    with ops.device("/job:worker/task:2/cpu:0"):
+    with ops.device("cpu:2"):
       v2 = resource_variable_ops.ResourceVariable(2.)
 
     self.evaluate([v0.initializer, v1.initializer, v2.initializer])
@@ -172,7 +167,7 @@ class SaverTest(test.TestCase):
         list(saveable_object_util.saveable_objects_for_op(v2, "v2")))
     prefix = os.path.join(self.get_temp_dir(), "ckpt")
     self.evaluate(saver.save(constant_op.constant(prefix), self.local_options))
-    self.assertEqual(2, len(gfile.Glob(prefix + "*")))
+    self.assertEqual(4, len(gfile.Glob(prefix + "*")))
     self.evaluate(v0.assign(-1.))
     self.evaluate(v1.assign(-1.))
     self.evaluate(v2.assign(-1.))

From cb6e1ed5d8a406861398c428ca5fd6b84b439357 Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Thu, 18 Jun 2020 14:28:42 -0700
Subject: [PATCH 0541/1390] Return `debug_string` when creating
 CompilationCacheKey.

PiperOrigin-RevId: 317181056
Change-Id: I02198244c1c3749ff1ecf4e0647b8daa80dd868c
---
 tensorflow/core/tpu/kernels/BUILD             |  16 ++
 .../kernels/tpu_compilation_cache_external.cc | 127 ---------------
 .../kernels/tpu_compilation_cache_external.h  |   8 -
 .../core/tpu/kernels/tpu_compile_c_api.h      |  19 ++-
 tensorflow/core/tpu/kernels/tpu_op_util.cc    | 151 ++++++++++++++++++
 tensorflow/core/tpu/kernels/tpu_op_util.h     |  40 +++++
 6 files changed, 223 insertions(+), 138 deletions(-)
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_util.cc
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_util.h

diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index 94d3c8edf2b..9d38eb71f3c 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -405,6 +405,22 @@ cc_library(
     alwayslink = True,
 )
 
+cc_library(
+    name = "tpu_op_util",
+    srcs = ["tpu_op_util.cc"],
+    hdrs = ["tpu_op_util.h"],
+    deps = [
+        ":tpu_compilation_cache_key",
+        ":tpu_compile_c_api_hdrs",
+        ":tpu_mesh_state_interface",
+        "//tensorflow/compiler/xla:xla_data_proto_cc",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
 cc_library(
     name = "tpu_util",
     srcs = ["tpu_util.cc"],
diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc
index 8cee90e8e55..c4442fc95d5 100644
--- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc
@@ -49,70 +49,6 @@ void PopulateEntry(const std::string& key, CompiledSubgraph* entry,
       absl::make_unique(std::move(tpu_program_group));
   entry->initialized = true;
 }
-
-// Return fingerprint_in_metadata if it's not empty; otherwise read input tensor
-// data to compute the fingerprint.
-std::string GuaranteedConstFingerprint(
-    const string& fingerprint_in_metadata,
-    const OpInputList& guaranteed_constants) {
-  if (fingerprint_in_metadata.empty()) {
-    uint64_t fingerprint = 0;
-    for (const auto& constant : guaranteed_constants) {
-      fingerprint = TpuCompile_CreateGuaranteedConstFingerprint(
-          fingerprint, constant.tensor_data().data(),
-          constant.tensor_data().size());
-    }
-    return std::to_string(fingerprint);
-  } else {
-    return fingerprint_in_metadata;
-  }
-}
-
-std::string CreateShapePrefix(
-    const std::vector& dynamic_shapes) {
-  std::string shapes_prefix;
-  for (const TensorShape& shape : dynamic_shapes) {
-    for (int64 size : shape.dim_sizes()) {
-      absl::StrAppend(&shapes_prefix, size, ",");
-    }
-    absl::StrAppend(&shapes_prefix, ";");
-  }
-  return shapes_prefix;
-}
-
-// Include compilation configurations of the arguments that are not captured
-// by the called graph.
-std::string CreateConfigPrefix(const TPUCompileMetadataProto& metadata) {
-  std::string config_prefix;
-  for (const auto& arg : metadata.args()) {
-    if (arg.is_same_data_across_replicas()) {
-      absl::StrAppend(&config_prefix, ":s");
-      // Same.
-    } else {
-      // Different.
-      absl::StrAppend(&config_prefix, ":");
-    }
-    if (arg.enable_xla_sharding() ==
-        tpu::TPUCompileMetadataProto::Arg::ALLOWED) {
-      // Enabled.
-      absl::StrAppend(&config_prefix, "e");
-    }
-    if (arg.unrestricted_layout()) {
-      // Unrestricted.
-      absl::StrAppend(&config_prefix, ":u");
-    }
-    absl::StrAppend(&config_prefix, ",type(", arg.dtype(), ")");
-    if (arg.has_shape()) {
-      absl::StrAppend(&config_prefix, ",shape(");
-      for (const auto& dim : arg.shape().dim()) {
-        absl::StrAppend(&config_prefix, dim.size(), ",");
-      }
-      absl::StrAppend(&config_prefix, ")");
-    }
-  }
-  return config_prefix;
-}
-
 }  // namespace
 
 TpuCompilationCacheExternal::EntryRefImpl::EntryRefImpl(
@@ -196,68 +132,5 @@ CompiledSubgraph* TpuCompilationCacheExternal::InitializeEntry(
   marked_for_eviction_size_ += main_entry->total_size;
   return main_entry;
 }
-
-/*static*/ TpuCompilationCacheKey
-TpuCompilationCacheExternal::CreateCompilationCacheKey(
-    absl::string_view function_name, uint64 function_library_fingerprint,
-    absl::string_view mlir_module,
-    const tensorflow::OpInputList& guaranteed_constants,
-    const std::vector& dynamic_shapes,
-    const tensorflow::tpu::TPUCompileMetadataProto& metadata,
-    const TpuMeshStateInterface& mesh_state) {
-  VLOG(1) << "FunctionLibraryFingerprint:" << function_library_fingerprint;
-  std::string shapes_prefix = CreateShapePrefix(dynamic_shapes);
-  VLOG(1) << "shapes_prefix = " << shapes_prefix;
-  std::string config_prefix = CreateConfigPrefix(metadata);
-  VLOG(1) << "config_prefix = " << config_prefix;
-  std::vector flattened_device_ids;
-  if (metadata.has_device_assignment()) {
-    for (const auto& device :
-         metadata.device_assignment().computation_devices()) {
-      flattened_device_ids.insert(flattened_device_ids.end(),
-                                  device.replica_device_ids().begin(),
-                                  device.replica_device_ids().end());
-    }
-  }
-  // TODO(henrytan): return the debug_string.
-  const char* prefix =
-      TpuCompile_CreateCompilationCacheKey(CompilationCacheKeyProperty{
-          config_prefix.data(),
-          shapes_prefix.data(),
-          function_name.data(),
-          mlir_module.data(),
-          flattened_device_ids.data(),
-          flattened_device_ids.size(),
-          guaranteed_constants.size(),
-          function_library_fingerprint,
-          metadata.num_cores_per_replica(),
-          metadata.num_replicas(),
-          mesh_state.data(),
-      });
-  auto buffer_cleanup = gtl::MakeCleanup([prefix]() { delete[] prefix; });
-  TpuCompilationCacheKey key;
-  key.prefix = prefix;
-
-  // Guaranteed constants can be different across sessions. Use session_handle
-  // and guaranteed_const fingerprint to guarantee no collision.
-  if (guaranteed_constants.size() > 0) {
-    key.has_guaranteed_const = true;
-    key.session_handle = metadata.session_handle();
-    // Both `metadata` and `guaranteed_constants` lifetime are captured by
-    // reference based on the assumption that these variables lifetime is
-    // managed through the `TPUCompileOpKernelImpl` that outlives the
-    // lifetime of the compilation cache lookups.
-    string fingerprint;
-    key.guaranteed_const_fingerprint = [&metadata, &guaranteed_constants,
-                                        fingerprint]() mutable {
-      if (fingerprint.empty()) {
-        fingerprint = GuaranteedConstFingerprint(
-            metadata.guaranteed_const_fingerprint(), guaranteed_constants);
-      }
-      return fingerprint;
-    };
-  }
-  return key;
-}
 }  // namespace tpu
 }  // namespace tensorflow
diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
index 2c75cb4d053..fe251326a43 100644
--- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
@@ -63,14 +63,6 @@ class TpuCompilationCacheExternal : public TpuCompilationCacheInterface {
   explicit TpuCompilationCacheExternal(int64 max_cache_size)
       : TpuCompilationCacheInterface(max_cache_size) {}
 
-  static TpuCompilationCacheKey CreateCompilationCacheKey(
-      absl::string_view function_name, uint64 function_library_fingerprint,
-      absl::string_view mlir_module,
-      const tensorflow::OpInputList& guaranteed_constants,
-      const std::vector& dynamic_shapes,
-      const tensorflow::tpu::TPUCompileMetadataProto& metadata,
-      const TpuMeshStateInterface& mesh_state);
-
   string DebugString() const override { return "TpuCompilationCacheExternal"; }
 
  private:
diff --git a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
index d1546ed9610..c101e489d56 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
@@ -42,6 +42,13 @@ struct CompilationCacheKeyProperty {
   const XLA_TpuMeshState* mesh_state;
 };
 
+// Compilation cache key result returning both the key and a more verbose debug
+// version.
+struct CompilationCacheKeyResult {
+  const char* key;
+  const char* debug_string;
+};
+
 extern "C" {
 
 // Returns the number of available TPU core count.
@@ -49,9 +56,14 @@ TFTPU_CAPI_EXPORT int TpuTopology_AvailableCoreCount(
     const XLA_TpuMeshState* mesh_state, TpuCoreTypeEnum tpu_core_type);
 
 // Creates a unique compilation cache `key` used for `put` and `get` operations.
-// Returned buffer is heap-allocated and must be owned.
-TFTPU_CAPI_EXPORT const char* TpuCompile_CreateCompilationCacheKey(
-    CompilationCacheKeyProperty property);
+// Returned buffers are heap-allocated and must be owned.
+TFTPU_CAPI_EXPORT CompilationCacheKeyResult
+TpuCompile_CreateCompilationCacheKey(CompilationCacheKeyProperty property);
+
+// Destroys the CompilationCacheKeyResult returned by calling the
+// `TpuCompile_CreateCompilationCacheKey` API.
+TFTPU_CAPI_EXPORT void TpuCompile_DestroyCompilationCacheKey(
+    CompilationCacheKeyResult result);
 
 // Creates a guaranteed const fingerprint. Guarantee const is normally used in
 // TPU inference to avoid re-copying unchanged variables onto the TPU device.
@@ -75,6 +87,7 @@ TFTPU_CAPI_EXPORT void TpuCompile_BuildXLADeviceAssignment(
 struct TfTpu_CompileApiFn {
   TFTPU_ADD_FN_IN_STRUCT(TpuTopology_AvailableCoreCount);
   TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CreateCompilationCacheKey);
+  TFTPU_ADD_FN_IN_STRUCT(TpuCompile_DestroyCompilationCacheKey);
   TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CreateGuaranteedConstFingerprint);
   TFTPU_ADD_FN_IN_STRUCT(TpuCompile_CompileAheadOfTime);
   TFTPU_ADD_FN_IN_STRUCT(TpuCompile_BuildXLADeviceAssignment);
diff --git a/tensorflow/core/tpu/kernels/tpu_op_util.cc b/tensorflow/core/tpu/kernels/tpu_op_util.cc
new file mode 100644
index 00000000000..e2f717fea8b
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_op_util.cc
@@ -0,0 +1,151 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/core/tpu/kernels/tpu_op_util.h"
+
+#include 
+
+#include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
+
+namespace tensorflow {
+namespace tpu {
+namespace {
+// Return fingerprint_in_metadata if it's not empty; otherwise read input tensor
+// data to compute the fingerprint.
+std::string GuaranteedConstFingerprint(
+    const string& fingerprint_in_metadata,
+    const OpInputList& guaranteed_constants) {
+  if (fingerprint_in_metadata.empty()) {
+    uint64_t fingerprint = 0;
+    for (const auto& constant : guaranteed_constants) {
+      fingerprint = TpuCompile_CreateGuaranteedConstFingerprint(
+          fingerprint, constant.tensor_data().data(),
+          constant.tensor_data().size());
+    }
+    return std::to_string(fingerprint);
+  } else {
+    return fingerprint_in_metadata;
+  }
+}
+
+std::string CreateShapePrefix(
+    const std::vector& dynamic_shapes) {
+  std::string shapes_prefix;
+  for (const TensorShape& shape : dynamic_shapes) {
+    for (int64 size : shape.dim_sizes()) {
+      absl::StrAppend(&shapes_prefix, size, ",");
+    }
+    absl::StrAppend(&shapes_prefix, ";");
+  }
+  return shapes_prefix;
+}
+
+// Include compilation configurations of the arguments that are not captured
+// by the called graph.
+std::string CreateConfigPrefix(const TPUCompileMetadataProto& metadata) {
+  std::string config_prefix;
+  for (const auto& arg : metadata.args()) {
+    if (arg.is_same_data_across_replicas()) {
+      absl::StrAppend(&config_prefix, ":s");
+      // Same.
+    } else {
+      // Different.
+      absl::StrAppend(&config_prefix, ":");
+    }
+    if (arg.enable_xla_sharding() ==
+        tpu::TPUCompileMetadataProto::Arg::ALLOWED) {
+      // Enabled.
+      absl::StrAppend(&config_prefix, "e");
+    }
+    if (arg.unrestricted_layout()) {
+      // Unrestricted.
+      absl::StrAppend(&config_prefix, ":u");
+    }
+    absl::StrAppend(&config_prefix, ",type(", arg.dtype(), ")");
+    if (arg.has_shape()) {
+      absl::StrAppend(&config_prefix, ",shape(");
+      for (const auto& dim : arg.shape().dim()) {
+        absl::StrAppend(&config_prefix, dim.size(), ",");
+      }
+      absl::StrAppend(&config_prefix, ")");
+    }
+  }
+  return config_prefix;
+}
+}  // namespace
+
+TpuCompilationCacheKey CreateCompilationCacheKey(
+    absl::string_view function_name, uint64 function_library_fingerprint,
+    absl::string_view mlir_module, const OpInputList& guaranteed_constants,
+    const std::vector& dynamic_shapes,
+    const TPUCompileMetadataProto& metadata,
+    const TpuMeshStateInterface& mesh_state) {
+  VLOG(1) << "FunctionLibraryFingerprint:" << function_library_fingerprint;
+  std::string shapes_prefix = CreateShapePrefix(dynamic_shapes);
+  VLOG(1) << "shapes_prefix = " << shapes_prefix;
+  std::string config_prefix = CreateConfigPrefix(metadata);
+  VLOG(1) << "config_prefix = " << config_prefix;
+  std::vector flattened_device_ids;
+  if (metadata.has_device_assignment()) {
+    for (const auto& device :
+         metadata.device_assignment().computation_devices()) {
+      flattened_device_ids.insert(flattened_device_ids.end(),
+                                  device.replica_device_ids().begin(),
+                                  device.replica_device_ids().end());
+    }
+  }
+  CompilationCacheKeyResult result =
+      TpuCompile_CreateCompilationCacheKey(CompilationCacheKeyProperty{
+          config_prefix.data(),
+          shapes_prefix.data(),
+          function_name.data(),
+          mlir_module.data(),
+          flattened_device_ids.data(),
+          flattened_device_ids.size(),
+          guaranteed_constants.size(),
+          function_library_fingerprint,
+          metadata.num_cores_per_replica(),
+          metadata.num_replicas(),
+          mesh_state.data(),
+      });
+  auto buffer_cleanup = gtl::MakeCleanup(
+      [result]() { TpuCompile_DestroyCompilationCacheKey(result); });
+  TpuCompilationCacheKey key;
+  key.prefix = result.key;
+  key.debug_string = result.debug_string;
+
+  // Guaranteed constants can be different across sessions. Use session_handle
+  // and guaranteed_const fingerprint to guarantee no collision.
+  if (guaranteed_constants.size() > 0) {
+    key.has_guaranteed_const = true;
+    key.session_handle = metadata.session_handle();
+    // Both `metadata` and `guaranteed_constants` lifetime are captured by
+    // reference based on the assumption that these variables lifetime is
+    // managed through the `TPUCompileOpKernelImpl` that outlives the
+    // lifetime of the compilation cache lookups.
+    string fingerprint;
+    key.guaranteed_const_fingerprint = [&metadata, &guaranteed_constants,
+                                        fingerprint]() mutable {
+      if (fingerprint.empty()) {
+        fingerprint = GuaranteedConstFingerprint(
+            metadata.guaranteed_const_fingerprint(), guaranteed_constants);
+      }
+      return fingerprint;
+    };
+  }
+  return key;
+}
+}  // namespace tpu
+}  // namespace tensorflow
diff --git a/tensorflow/core/tpu/kernels/tpu_op_util.h b/tensorflow/core/tpu/kernels/tpu_op_util.h
new file mode 100644
index 00000000000..0a9657ca05e
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_op_util.h
@@ -0,0 +1,40 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_UTIL_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_UTIL_H_
+
+#include 
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tpu/kernels/tpu_compilation_cache_key.h"
+#include "tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h"
+
+namespace tensorflow {
+namespace tpu {
+// Creates a unique compilation cache `key`.
+TpuCompilationCacheKey CreateCompilationCacheKey(
+    absl::string_view function_name, uint64 function_library_fingerprint,
+    absl::string_view mlir_module, const OpInputList& guaranteed_constants,
+    const std::vector& dynamic_shapes,
+    const TPUCompileMetadataProto& metadata,
+    const TpuMeshStateInterface& mesh_state);
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_UTIL_H_

From 9d33f296d1edc2f656e253cf2a015d36daedd5c1 Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Thu, 18 Jun 2020 15:00:50 -0700
Subject: [PATCH 0542/1390] Prep change for publishing TPU Ops.

PiperOrigin-RevId: 317188030
Change-Id: I29f9236c0ade6bf586c8a52ead977b5d31aec357
---
 tensorflow/core/tpu/kernels/BUILD             | 11 +++
 .../tpu_compilation_cache_entry_unloader.h    | 69 +++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h

diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index 9d38eb71f3c..a41747ee8c5 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -450,3 +450,14 @@ cc_library(
     hdrs = ["tpu_compile_op.h"],
     deps = ["//tensorflow/core:framework"],
 )
+
+cc_library(
+    name = "tpu_compilation_cache_entry_unloader",
+    hdrs = ["tpu_compilation_cache_entry_unloader.h"],
+    deps = [
+        ":tpu_compilation_cache_interface",
+        "//tensorflow/core:framework",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h b/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h
new file mode 100644
index 00000000000..c298d8fcc12
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h
@@ -0,0 +1,69 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.h"
+
+namespace tensorflow {
+namespace tpu {
+
+class TpuCompilationCacheEntryUnloader : public ResourceBase {
+ public:
+  explicit TpuCompilationCacheEntryUnloader(TpuCompilationCacheInterface* cache)
+      : cache_(cache) {
+    // Hold a reference to the cache until the unloader is destroyed.
+    cache_->Ref();
+    VLOG(1) << "Will unload compilation cache entries when session closes.";
+  }
+
+  ~TpuCompilationCacheEntryUnloader() override {
+    absl::MutexLock lock(&mu_);
+    for (int64 uid : cache_entry_uids_) {
+      Status s = cache_->MarkEntryForEviction(uid);
+      if (!s.ok()) {
+        LOG(WARNING) << "MarkEntryForEviction in "
+                        "~CompilationCacheEntryUnloader fails with error "
+                     << s;
+      }
+    }
+    // Release our reference to the cache.
+    cache_->Unref();
+  }
+
+  // Add cache entry uid to be unloaded in destructor.
+  void AddCacheEntryUid(int64 uid) {
+    absl::MutexLock lock(&mu_);
+    cache_entry_uids_.insert(uid);
+  }
+
+  std::string DebugString() const override {
+    return "CompilationCacheEntryUnloader";
+  }
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(TpuCompilationCacheEntryUnloader);
+  mutable absl::Mutex mu_;
+  TpuCompilationCacheInterface* cache_;  // Not owned.
+  absl::flat_hash_set cache_entry_uids_ ABSL_GUARDED_BY(mu_);
+};
+
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TPU_KERNELS_COMPILATION_CACHE_ENTRY_UNLOADER_H_

From 35b978db57eaa87f32e7c9c3e9a7c323e595c978 Mon Sep 17 00:00:00 2001
From: Sean Silva 
Date: Thu, 18 Jun 2020 15:00:53 -0700
Subject: [PATCH 0543/1390] Move tfl-device-index-selector to TF directory.

There's nothing lite-specific about this pass.

PiperOrigin-RevId: 317188038
Change-Id: Iac9799e296e043aabf7aeabec2e8f72d07c77178
---
 tensorflow/compiler/mlir/lite/BUILD                  |  1 -
 tensorflow/compiler/mlir/lite/tf_tfl_passes.cc       |  2 +-
 tensorflow/compiler/mlir/lite/transforms/passes.h    |  3 ---
 tensorflow/compiler/mlir/tensorflow/BUILD            |  1 +
 .../transforms/device_index_selector.cc              | 12 ++++++------
 .../compiler/mlir/tensorflow/transforms/passes.h     |  3 +++
 .../tests/tf_device_index_selector.mlir              |  2 +-
 7 files changed, 12 insertions(+), 12 deletions(-)
 rename tensorflow/compiler/mlir/{lite => tensorflow}/transforms/device_index_selector.cc (92%)
 rename tensorflow/compiler/{mlir/lite => tensorflow}/tests/tf_device_index_selector.mlir (94%)

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 8e9d615053c..8d4efeb3d60 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -314,7 +314,6 @@ tf_cc_test(
 cc_library(
     name = "tensorflow_lite_legalize_tf",
     srcs = [
-        "transforms/device_index_selector.cc",
         "transforms/dilated_conv.cc",
         "transforms/generated_legalize_tf.inc",
         "transforms/generated_lower_static_tensor_list.inc",
diff --git a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
index 008098f62ba..fed2896035b 100644
--- a/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
+++ b/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc
@@ -63,7 +63,7 @@ void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config,
   standard_pipeline_options.enable_inliner = false;
   standard_pipeline_options.form_clusters = pass_config.form_clusters;
   mlir::TF::CreateTFStandardPipeline(*pass_manager, standard_pipeline_options);
-  pass_manager->addPass(mlir::TFL::CreateDeviceIndexSelectorPass());
+  pass_manager->addPass(mlir::TF::CreateDeviceIndexSelectorPass());
 
   if (pass_config.shape_inference) {
     pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h
index 01e5eb1cb68..105c9394fb4 100644
--- a/tensorflow/compiler/mlir/lite/transforms/passes.h
+++ b/tensorflow/compiler/mlir/lite/transforms/passes.h
@@ -91,9 +91,6 @@ std::unique_ptr> CreateWhileOutlinePass();
 // Verifies runtime constraints.
 std::unique_ptr> CreateRuntimeVerifyPass();
 
-// Creates function pass to select device index/fold tf.DeviceIndex.
-std::unique_ptr> CreateDeviceIndexSelectorPass();
-
 }  // namespace TFL
 
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 54e57512c32..7c0d427e87b 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -475,6 +475,7 @@ cc_library(
         "transforms/cluster_outlining.cc",
         "transforms/collection_ops_util.cc",
         "transforms/decompose_resource_ops_pass.cc",
+        "transforms/device_index_selector.cc",
         "transforms/einsum.cc",
         "transforms/executor_island_coarsening.cc",
         "transforms/executor_tpuv1_inline_tpu_island.cc",
diff --git a/tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
similarity index 92%
rename from tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc
rename to tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
index d4aed750dc8..550647a915a 100644
--- a/tensorflow/compiler/mlir/lite/transforms/device_index_selector.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
@@ -21,11 +21,11 @@ limitations under the License.
 #include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
-#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
 
 namespace mlir {
-namespace TFL {
+namespace TF {
 namespace {
 
 // Folds the DeviceIndex op to a constant value. The DeviceIndex return the
@@ -55,8 +55,8 @@ void DeviceIndexSelector::runOnOperation() {
   // Convert all the DeviceIndex ops to constant values.
   func.getBody().walk([](TF::DeviceIndexOp op) {
     // This just selects the default in all cases where DeviceIndex feeds into
-    // tf.Case. This could be enhanced based on explicit TFLite specification or
-    // TAC in future.
+    // tf.Case. This could be enhanced to have some sort of policy in the
+    // future.
     OpBuilder b(op);
     RankedTensorType type = RankedTensorType::get({}, b.getIntegerType(32));
     int index = op.device_names().size();
@@ -79,7 +79,7 @@ std::unique_ptr> CreateDeviceIndexSelectorPass() {
 }
 
 static PassRegistration pass(
-    "tfl-device-index-selector", "Fold tf.DeviceIndex to constant");
+    "tf-device-index-selector", "Fold tf.DeviceIndex to constant");
 
-}  // namespace TFL
+}  // namespace TF
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
index a34be28c809..168b317641d 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h
@@ -147,6 +147,9 @@ std::unique_ptr> CreateLegalizeHloToTfPass();
 // generally used beyond exporting to runtimes that supports these ops. In the
 // future these fusions may be codegen'd automatically.
 std::unique_ptr> CreateFusedKernelMatcherPass();
+
+// Creates function pass to select device index/fold tf.DeviceIndex.
+std::unique_ptr> CreateDeviceIndexSelectorPass();
 }  // namespace TF
 
 namespace tf_executor {
diff --git a/tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir b/tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
similarity index 94%
rename from tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir
rename to tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
index 1ac7f30d644..7fc2b210f91 100644
--- a/tensorflow/compiler/mlir/lite/tests/tf_device_index_selector.mlir
+++ b/tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
@@ -1,6 +1,6 @@
 // Test DeviceIndex selector.
 
-// RUN: tf-opt --tfl-device-index-selector %s | FileCheck %s
+// RUN: tf-opt --tf-device-index-selector %s | FileCheck %s
 
 // CHECK-LABEL: func @select
 func @select(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) {

From 834fe68f365e1d7f082b596fe87471ce84c2c8ec Mon Sep 17 00:00:00 2001
From: Pete Warden 
Date: Thu, 18 Jun 2020 15:16:08 -0700
Subject: [PATCH 0544/1390] Optimized Arduino library by enabling
 precompilation Precompilation allows Arduino users to build their sketches
 much faster, but requires some support from the library properties to enable.
 This has recently been upgraded to suppor the 'full' mode, as shown in
 https://github.com/arduino/arduino-cli/pull/611, so we want to take advantage
 of this.

PiperOrigin-RevId: 317191283
Change-Id: Ie44a31ba45105f65fdad0da487290aff5fa2a179
---
 tensorflow/lite/micro/tools/make/templates/library.properties | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/micro/tools/make/templates/library.properties b/tensorflow/lite/micro/tools/make/templates/library.properties
index e41fd8d8fbe..6e02748a0b4 100644
--- a/tensorflow/lite/micro/tools/make/templates/library.properties
+++ b/tensorflow/lite/micro/tools/make/templates/library.properties
@@ -7,4 +7,5 @@ paragraph=This library runs TensorFlow machine learning models on microcontrolle
 category=Data Processing
 url=https://www.tensorflow.org/lite/microcontrollers/overview
 ldflags=-lm
-includes=TensorFlowLite.h
\ No newline at end of file
+includes=TensorFlowLite.h
+precompiled=full

From 852cde437fdd062f52c42e47344029897ee67afd Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Thu, 18 Jun 2020 15:27:56 -0700
Subject: [PATCH 0545/1390] Make error_message constant

Co-authored-by: Mihai Maruseac 
---
 tensorflow/security/fuzzing/status_group_fuzz.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index 979fd444b48..5e2b7eec403 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -39,7 +39,7 @@ tensorflow::error::Code BuildRandomErrorCode(uint32_t code){
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   tensorflow::error::Code error_code;
 
-  std::string error_message = "ERROR";
+  const std::string error_message = "ERROR";
 
   tensorflow::Status s, derived_s;
 

From 14e942faaa56b749c81c52595d04a9bb7f26fa02 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Thu, 18 Jun 2020 15:31:26 -0700
Subject: [PATCH 0546/1390] Update spacing and variable declaration

---
 .../security/fuzzing/status_group_fuzz.cc     | 27 ++++---------------
 1 file changed, 5 insertions(+), 22 deletions(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index 5e2b7eec403..52d83d00866 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -14,9 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include 
 #include 
-
 #include "tensorflow/core/platform/status.h"
-
 #include 
 
 // This is a fuzzer for `tensorflow::StatusGroup`. Since `Status` is used almost
@@ -37,37 +35,22 @@ tensorflow::error::Code BuildRandomErrorCode(uint32_t code){
 }
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
-  tensorflow::error::Code error_code;
-
   const std::string error_message = "ERROR";
-
-  tensorflow::Status s, derived_s;
-
   tensorflow::StatusGroup sg;
-
-  bool is_derived;
-
-  uint32_t code;
-
   FuzzedDataProvider fuzzed_data(data, size);
 
   while(fuzzed_data.remaining_bytes() > 0) {
-    code = fuzzed_data.ConsumeIntegral();
+    uint32_t code = fuzzed_data.ConsumeIntegral();
+    tensorflow::error::Code error_code = BuildRandomErrorCode(code);
+    bool is_derived = fuzzed_data.ConsumeBool();
 
-    error_code = BuildRandomErrorCode(code);
-
-    is_derived = fuzzed_data.ConsumeBool();
-
-    s = tensorflow::Status(error_code, error_message);
+    tensorflow::Status s = tensorflow::Status(error_code, error_message);
 
     if(is_derived) {
-      derived_s = tensorflow::StatusGroup::MakeDerived(s);
-
+      tensorflow::Status derived_s = tensorflow::StatusGroup::MakeDerived(s);
       sg.Update(derived_s);
-
     } else {
       sg.Update(s);
-
     }
   }
 

From 71bbebbf4d04c1bcb6ed44e2156087c9fec06e9e Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Thu, 18 Jun 2020 15:32:15 -0700
Subject: [PATCH 0547/1390] Moved final StatusGroup method calls

---
 tensorflow/security/fuzzing/status_group_fuzz.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index 52d83d00866..bc80cd72bc9 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -55,9 +55,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   }
 
   sg.as_summary_status();
-
   sg.as_concatenated_status();
-
   sg.AttachLogMessages();
 
   return 0;

From a82b75c82b63c4397b3d6a215e439ca77e687a84 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 15:35:42 -0700
Subject: [PATCH 0548/1390] [XLA/Client] Implement LocalClient::Run which
 supports buffer donation

PiperOrigin-RevId: 317195199
Change-Id: If4d35d0627fa068a0c2b522fdae52466abd21f51
---
 .../compiler/xla/client/local_client.cc       | 47 ++++---------------
 tensorflow/compiler/xla/client/local_client.h |  9 ----
 .../tests/multiple_devices_on_host_test.cc    |  3 +-
 tensorflow/compiler/xla/tests/while_test.cc   |  6 +--
 4 files changed, 12 insertions(+), 53 deletions(-)

diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index f71e8a2d56d..afe115deda8 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -168,26 +168,6 @@ LocalExecutable::RunHelper(const absl::Span argument_shapes,
   return std::make_pair(service_options, std::move(stream));
 }
 
-StatusOr LocalExecutable::GetExecutableRunOptions(
-    absl::Span argument_shapes,
-    const ExecutableRunOptions& run_options) {
-  TF_ASSIGN_OR_RETURN(auto options_and_stream,
-                      RunHelper(argument_shapes, run_options));
-  ExecutableRunOptions options = options_and_stream.first.run_options();
-  options.set_device_ordinal(-1);
-  return options;
-}
-
-template 
-static StatusOr BlockHostUntilDoneAfterAsyncCall(
-    se::Stream* stream, std::function()> async_callback) {
-  StatusOr result = async_callback();
-  Status block_status = stream->BlockHostUntilDone();
-  TF_RETURN_IF_ERROR(result.status());
-  TF_RETURN_IF_ERROR(block_status);
-  return result;
-}
-
 StatusOr LocalExecutable::Run(
     const absl::Span arguments,
     ExecutableRunOptions run_options) {
@@ -196,24 +176,15 @@ StatusOr LocalExecutable::Run(
   for (const ShapedBuffer* const arg : arguments) {
     argument_shapes.push_back(&arg->on_host_shape());
   }
-  TF_ASSIGN_OR_RETURN(ExecutableRunOptions options,
-                      GetExecutableRunOptions(argument_shapes, run_options));
-  return BlockHostUntilDoneAfterAsyncCall(
-      options.stream(), [&] { return RunAsync(arguments, options); });
-}
-
-StatusOr LocalExecutable::Run(
-    std::vector arguments, ExecutableRunOptions run_options) {
-  std::vector argument_shapes;
-  argument_shapes.reserve(arguments.size());
-  for (const ExecutionInput& arg : arguments) {
-    argument_shapes.push_back(&arg.shape());
-  }
-  TF_ASSIGN_OR_RETURN(ExecutableRunOptions options,
-                      GetExecutableRunOptions(argument_shapes, run_options));
-  return BlockHostUntilDoneAfterAsyncCall(
-      options.stream(),
-      [&] { return RunAsync(argument_shapes, std::move(arguments), options); });
+  TF_ASSIGN_OR_RETURN(auto options_and_stream,
+                      RunHelper(argument_shapes, run_options));
+  ExecutableRunOptions options = options_and_stream.first.run_options();
+  options.set_device_ordinal(-1);
+  auto result = RunAsync(arguments, options);
+  Status block_status = options.stream()->BlockHostUntilDone();
+  TF_RETURN_IF_ERROR(result.status());
+  TF_RETURN_IF_ERROR(block_status);
+  return result;
 }
 
 static std::shared_ptr DumpArguments(
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index b00f5cc6801..7cdeb9dcbf6 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -51,11 +51,6 @@ class LocalExecutable {
       const absl::Span arguments,
       ExecutableRunOptions run_options);
 
-  // Similar to Run(), but allows for donating argument buffers to the
-  // executable.
-  StatusOr Run(std::vector arguments,
-                                ExecutableRunOptions run_options);
-
   // Similar to Run(), but need not block the host waiting for the computation
   // to complete before returning.
   StatusOr RunAsync(
@@ -90,10 +85,6 @@ class LocalExecutable {
       const absl::Span argument_shapes,
       ExecutableRunOptions run_options);
 
-  StatusOr GetExecutableRunOptions(
-      absl::Span argument_shapes,
-      const ExecutableRunOptions& run_options);
-
   // The ordinal of the device which this executable was compiled for. The
   // executable can run on all equivalent devices (as determined by
   // Backend::devices_equivalent).
diff --git a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
index 2231fc6feab..2b19aaded9c 100644
--- a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
+++ b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
@@ -45,8 +45,7 @@ void CompileAndExecute(
       xla::ClientLibrary::GetXlaService(client->platform())
           ->backend()
           .memory_allocator());
-  StatusOr result =
-      executable->Run(absl::Span(), execute_options);
+  StatusOr result = executable->Run({}, execute_options);
   {
     absl::MutexLock lock(results_mutex);
     results->emplace_back(device_ordinal, std::move(result));
diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc
index 8e8c3605cc7..d575bbb1f3e 100644
--- a/tensorflow/compiler/xla/tests/while_test.cc
+++ b/tensorflow/compiler/xla/tests/while_test.cc
@@ -1324,16 +1324,14 @@ void BM_WhileLoop(int num_iters) {
   options.set_allocator(&allocator);
   const int kWarmups = 2;
   for (int i = 0; i < kWarmups; ++i) {
-    auto result =
-        executable->Run(absl::Span(), options);
+    auto result = executable->Run({}, options);
     ASSERT_TRUE(result.ok());
   }
 
   // Run benchmark.
   tensorflow::testing::StartTiming();
   for (int i = 0; i < num_iters; ++i) {
-    auto result =
-        executable->Run(absl::Span(), options);
+    auto result = executable->Run({}, options);
     ASSERT_TRUE(result.ok());
   }
 }

From 40ef6a7ad67a973f971eae13b59b5f25777d037e Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Thu, 18 Jun 2020 22:43:22 +0000
Subject: [PATCH 0549/1390] Map args & attrs to types and use map to annotate

---
 tensorflow/python/framework/python_op_gen.cc | 149 ++++++++-----------
 1 file changed, 59 insertions(+), 90 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 8c4d0f5b753..79c8800418c 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -47,6 +47,7 @@ constexpr char kEagerFallbackSuffix[] = "_eager_fallback";
 
 std::unordered_map dtypes_map {
       {"_dtypes.float16", "_dtypes.Float16"},
+      {"_dtypes.half", "_dtypes.Half"},
       {"_dtypes.float32", "_dtypes.Float32"},
       {"_dtypes.float64", "_dtypes.Float64"},
       {"_dtypes.bfloat16", "_dtypes.BFloat16"},
@@ -162,7 +163,8 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
 
   bool AddEagerFastPathAndGraphCode(const string& parameters,
                                     const std::vector& output_sizes,
-                                    const string& eager_not_allowed_error);
+                                    const string& eager_not_allowed_error,
+                                    std::unordered_map& type_map);
   bool AddEagerFallbackCode(const string& parameters,
                             const std::vector& output_sizes,
                             const string& num_outputs_expr,
@@ -179,8 +181,11 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
   void AddRawOpExport(const string& parameters);
 
   void GenerateTypeVars();
-  string GetTypeAnnotatedParams();
-  void AddReturnTypeAnnotation();
+
+  std::unordered_map GetTypeAnnotationMap();
+
+  void AddReturnTypeAnnotation(std::unordered_map& type_map);
+
   void AddAttrForArg(const string& attr, int arg_index) {
     gtl::InsertIfNotPresent(&inferred_attrs_, attr,
                             op_def_.input_arg(arg_index).name());
@@ -343,13 +348,22 @@ string GenEagerPythonOp::Code() {
     param_names_.push_back(param_and_default.first);
   }
 
-  string parameters;
+  std::unordered_map type_map;
+  // Only populate map for whitelisted ops
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    strings::StrAppend(¶meters, GetTypeAnnotatedParams());
-  } else {
-    for (const auto& param : params_no_default_) {
-      if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
-      strings::StrAppend(¶meters, param.GetRenameTo());
+    type_map = GetTypeAnnotationMap();
+  }
+
+  string parameters;
+  for (const auto& param : params_no_default_) {
+    if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
+    strings::StrAppend(¶meters, param.GetRenameTo());
+
+    // Add type annotations to param
+    if (type_map.find(param.GetName()) != type_map.end()) {
+      if(!type_map[param.GetName()].empty()) {
+        strings::StrAppend(¶meters, ": ", type_map[param.GetName()]);
+      }
     }
   }
 
@@ -393,7 +407,7 @@ string GenEagerPythonOp::Code() {
   string eager_not_allowed_error = GetEagerNotAllowedError();
 
   if (!AddEagerFastPathAndGraphCode(parameters_with_defaults, output_sizes,
-                                    eager_not_allowed_error)) {
+                                    eager_not_allowed_error, type_map)) {
     return result_;
   }
 
@@ -405,61 +419,56 @@ string GenEagerPythonOp::Code() {
   return prelude_ + result_;
 }
 
-string GenEagerPythonOp::GetTypeAnnotatedParams() {
-  // holds mappings from param name to its type annotation
-  std::unordered_map param_type_map;
+std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
+  std::unordered_map type_map;
+  // Mapping attrs to TypeVars
   for (int i = 0; i allowed_types;
-      bool has_dtype_half = false;
       for (int t : attr.allowed_values().list().type()) {
-        if (t == 19) { // DT_HALF = 19
-          has_dtype_half = true;
-          break;
-        }
         DataType dtype = static_cast(t);
         const string py_dtype = python_op_gen_internal::DataTypeToPython(dtype, "_dtypes.");
         if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
@@ -482,9 +486,6 @@ void GenEagerPythonOp::GenerateTypeVars() {
         }
       }
 
-      // Do not create a type variable that includes the dtype half
-      if (has_dtype_half) continue;
-
       // If all dtypes are allowed, add them all
       if (allowed_types.empty()) {
         for (std::pair map_dtype : dtypes_map) {
@@ -509,48 +510,16 @@ void GenEagerPythonOp::GenerateTypeVars() {
   if(added_typevar) strings::StrAppend(&result_, "\n");
 }
 
-void GenEagerPythonOp::AddReturnTypeAnnotation() {
-  string return_type = "";
+void GenEagerPythonOp::AddReturnTypeAnnotation(std::unordered_map& type_map) {
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
-    // If the "type" field is set, the return Tensor has a single DataType
-    if (arg.type() != 0) {
-      const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
-      if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
-        strings::StrAppend(&return_type, "_ops.Tensor[", dtypes_map[py_dtype], "]");
+    // Add type annotations to param
+    if (type_map.find(arg.name()) != type_map.end()) {
+      if (!type_map[arg.name()].empty()) {
+        result_.erase(result_.length() - 2);
+        strings::StrAppend(&result_, " -> ", type_map[arg.name()], ":\n");
       }
     }
-    else {
-      for (int i = 0; i allowed_types;
-          for (int t : attr.allowed_values().list().type()) {
-            // Do not add type annotations when return type can be half
-            if (t == 19) return; // DT_HALF = 19
-            DataType dtype = static_cast(t);
-            const string py_dtype = python_op_gen_internal::DataTypeToPython(dtype, "_dtypes.");
-            allowed_types.emplace_back(py_dtype);
-          }
-
-          std::sort(allowed_types.begin(), allowed_types.end());
-
-          string typevar_dtypes;
-          for (std::vector::iterator it = allowed_types.begin(); it != allowed_types.end(); ++it) {
-            if (!typevar_dtypes.empty()) strings::StrAppend(&typevar_dtypes, ", ");
-            strings::StrAppend(&typevar_dtypes, *it);
-          }
-
-          const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
-          strings::StrAppend(&return_type, "_ops.Tensor[", type_var_name, "]");
-        }
-      }
-    }
-
-    if (!return_type.empty()) {
-      result_.erase(result_.length() - 2);
-      strings::StrAppend(&result_, " -> ", return_type, ":\n");
-    }
   }
 }
 
@@ -876,7 +845,7 @@ void GenEagerPythonOp::AddEagerFunctionTeardown(
 
 bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
     const string& parameters, const std::vector& output_sizes,
-    const string& eager_not_allowed_error) {
+    const string& eager_not_allowed_error, std::unordered_map& type_map) {
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
     GenerateTypeVars();
   }
@@ -887,7 +856,7 @@ bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
   AddExport();
   AddDefLine(function_name_, parameters);
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    AddReturnTypeAnnotation();
+    AddReturnTypeAnnotation(type_map);
   }
   AddDocStringDescription();
   AddDocStringArgs();

From 0deffad6acbc2f5848022bf8ae360c9adbdf1ef8 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Thu, 18 Jun 2020 15:50:38 -0700
Subject: [PATCH 0550/1390] Make `return_state` as explicit kwarg in the
 Conv2DLSTM layer.

It was previously hide in the **kwargs, and we are also missing documentation for it.

The existing test case should already cover the functionality of it.

PiperOrigin-RevId: 317197835
Change-Id: Icfae1e177eeb886b41345078f6b93f282a94df5b
---
 .../keras/layers/convolutional_recurrent.py   | 43 +++++++++++--------
 ...orflow.keras.layers.-conv-l-s-t-m2-d.pbtxt |  2 +-
 ...orflow.keras.layers.-conv-l-s-t-m2-d.pbtxt |  2 +-
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/tensorflow/python/keras/layers/convolutional_recurrent.py b/tensorflow/python/keras/layers/convolutional_recurrent.py
index 19831429b73..6c812204cba 100644
--- a/tensorflow/python/keras/layers/convolutional_recurrent.py
+++ b/tensorflow/python/keras/layers/convolutional_recurrent.py
@@ -753,7 +753,9 @@ class ConvLSTM2D(ConvRNN2D):
       the `recurrent_kernel` weights matrix.
     bias_constraint: Constraint function applied to the bias vector.
     return_sequences: Boolean. Whether to return the last output
-      in the output sequence, or the full sequence.
+      in the output sequence, or the full sequence. (default False)
+    return_state: Boolean Whether to return the last state
+      in addition to the output. (default False)
     go_backwards: Boolean (default False).
       If True, process the input sequence backwards.
     stateful: Boolean (default False). If True, the last state
@@ -786,22 +788,27 @@ class ConvLSTM2D(ConvRNN2D):
         `(samples, time, rows, cols, channels)`
 
   Output shape:
-    - If `return_sequences`
-       - If data_format='channels_first'
-          5D tensor with shape:
-          `(samples, time, filters, output_row, output_col)`
-       - If data_format='channels_last'
-          5D tensor with shape:
-          `(samples, time, output_row, output_col, filters)`
-    - Else
-      - If data_format ='channels_first'
-          4D tensor with shape:
-          `(samples, filters, output_row, output_col)`
-      - If data_format='channels_last'
-          4D tensor with shape:
-          `(samples, output_row, output_col, filters)`
-      where `o_row` and `o_col` depend on the shape of the filter and
-      the padding
+    - If `return_state`: a list of tensors. The first tensor is
+      the output. The remaining tensors are the last states,
+      each 4D tensor with shape:
+      `(samples, filters, new_rows, new_cols)`
+      if data_format='channels_first'
+      or 4D tensor with shape:
+      `(samples, new_rows, new_cols, filters)`
+      if data_format='channels_last'.
+      `rows` and `cols` values might have changed due to padding.
+    - If `return_sequences`: 5D tensor with shape:
+      `(samples, timesteps, filters, new_rows, new_cols)`
+      if data_format='channels_first'
+      or 5D tensor with shape:
+      `(samples, timesteps, new_rows, new_cols, filters)`
+      if data_format='channels_last'.
+    - Else, 4D tensor with shape:
+      `(samples, filters, new_rows, new_cols)`
+      if data_format='channels_first'
+      or 4D tensor with shape:
+      `(samples, new_rows, new_cols, filters)`
+      if data_format='channels_last'.
 
   Raises:
     ValueError: in case of invalid constructor arguments.
@@ -834,6 +841,7 @@ class ConvLSTM2D(ConvRNN2D):
                recurrent_constraint=None,
                bias_constraint=None,
                return_sequences=False,
+               return_state=False,
                go_backwards=False,
                stateful=False,
                dropout=0.,
@@ -863,6 +871,7 @@ class ConvLSTM2D(ConvRNN2D):
                           dtype=kwargs.get('dtype'))
     super(ConvLSTM2D, self).__init__(cell,
                                      return_sequences=return_sequences,
+                                     return_state=return_state,
                                      go_backwards=go_backwards,
                                      stateful=stateful,
                                      **kwargs)
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
index f77d613e354..958d06a0d0f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
@@ -207,7 +207,7 @@ tf_class {
   }
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'recurrent_activation\', \'use_bias\', \'kernel_initializer\', \'recurrent_initializer\', \'bias_initializer\', \'unit_forget_bias\', \'kernel_regularizer\', \'recurrent_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'recurrent_constraint\', \'bias_constraint\', \'return_sequences\', \'go_backwards\', \'stateful\', \'dropout\', \'recurrent_dropout\'], varargs=None, keywords=kwargs, defaults=[\'(1, 1)\', \'valid\', \'None\', \'(1, 1)\', \'tanh\', \'hard_sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'False\', \'False\', \'False\', \'0.0\', \'0.0\'], "
+    argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'recurrent_activation\', \'use_bias\', \'kernel_initializer\', \'recurrent_initializer\', \'bias_initializer\', \'unit_forget_bias\', \'kernel_regularizer\', \'recurrent_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'recurrent_constraint\', \'bias_constraint\', \'return_sequences\', \'return_state\', \'go_backwards\', \'stateful\', \'dropout\', \'recurrent_dropout\'], varargs=None, keywords=kwargs, defaults=[\'(1, 1)\', \'valid\', \'None\', \'(1, 1)\', \'tanh\', \'hard_sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'False\', \'False\', \'False\', \'False\', \'0.0\', \'0.0\'], "
   }
   member_method {
     name: "add_loss"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
index f77d613e354..958d06a0d0f 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.keras.layers.-conv-l-s-t-m2-d.pbtxt
@@ -207,7 +207,7 @@ tf_class {
   }
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'recurrent_activation\', \'use_bias\', \'kernel_initializer\', \'recurrent_initializer\', \'bias_initializer\', \'unit_forget_bias\', \'kernel_regularizer\', \'recurrent_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'recurrent_constraint\', \'bias_constraint\', \'return_sequences\', \'go_backwards\', \'stateful\', \'dropout\', \'recurrent_dropout\'], varargs=None, keywords=kwargs, defaults=[\'(1, 1)\', \'valid\', \'None\', \'(1, 1)\', \'tanh\', \'hard_sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'False\', \'False\', \'False\', \'0.0\', \'0.0\'], "
+    argspec: "args=[\'self\', \'filters\', \'kernel_size\', \'strides\', \'padding\', \'data_format\', \'dilation_rate\', \'activation\', \'recurrent_activation\', \'use_bias\', \'kernel_initializer\', \'recurrent_initializer\', \'bias_initializer\', \'unit_forget_bias\', \'kernel_regularizer\', \'recurrent_regularizer\', \'bias_regularizer\', \'activity_regularizer\', \'kernel_constraint\', \'recurrent_constraint\', \'bias_constraint\', \'return_sequences\', \'return_state\', \'go_backwards\', \'stateful\', \'dropout\', \'recurrent_dropout\'], varargs=None, keywords=kwargs, defaults=[\'(1, 1)\', \'valid\', \'None\', \'(1, 1)\', \'tanh\', \'hard_sigmoid\', \'True\', \'glorot_uniform\', \'orthogonal\', \'zeros\', \'True\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'False\', \'False\', \'False\', \'False\', \'0.0\', \'0.0\'], "
   }
   member_method {
     name: "add_loss"

From e08382691bfb897d584c5d5a8e8a0abe0472373d Mon Sep 17 00:00:00 2001
From: Chenkai Kuang 
Date: Thu, 18 Jun 2020 15:54:03 -0700
Subject: [PATCH 0551/1390] Make "map_resources" overridable by subclass of
 `Trackable`.

This allows moving the implementation of map_resources from `tf.saved_model.save` to subclass of `Trackable`, e.g, Variable, DistributedVariable.

This is a non-functional change.

PiperOrigin-RevId: 317198449
Change-Id: I4aa48d4974b6547b5de8ac0f5c38f3da29d364bc
---
 tensorflow/python/distribute/BUILD            |  4 +--
 tensorflow/python/distribute/values.py        | 12 +++++++
 .../experimental/autocast_variable.py         |  7 ++++
 .../python/ops/resource_variable_ops.py       |  7 ++++
 tensorflow/python/saved_model/save.py         | 36 +++++--------------
 tensorflow/python/training/tracking/base.py   | 18 ++++++++++
 .../python/training/tracking/tracking.py      | 13 +++++++
 7 files changed, 66 insertions(+), 31 deletions(-)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 96559a9a740..7208807a18c 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -744,14 +744,12 @@ py_library(
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:framework_ops",
         "//tensorflow/python:math_ops",
-        "//tensorflow/python:tensor_util",
+        "//tensorflow/python:resource_variable_ops",
         "//tensorflow/python:tf_export",
         "//tensorflow/python:type_spec",
-        "//tensorflow/python:util",
         "//tensorflow/python:variable_scope",
         "//tensorflow/python:variables",
         "//tensorflow/python/eager:context",
-        "//tensorflow/python/eager:tape",
         "//tensorflow/python/training/saving:saveable_object",
         "//tensorflow/python/training/saving:saveable_object_util",
         "//tensorflow/python/training/tracking:base",
diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index d0ed27c69de..60b2ea4fe31 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -32,6 +32,7 @@ from tensorflow.python.framework import type_spec
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.training.saving import saveable_object
@@ -793,6 +794,17 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       return ops.convert_to_tensor(
           self._get(), dtype=dtype, name=name, as_ref=as_ref)
 
+  def _map_resources(self):
+    """For implementing `Trackable`."""
+    new_obj = resource_variable_ops.copy_to_graph_uninitialized(self._primary)
+    obj_map, resource_map = {}, {}
+    for v in self._values:
+      obj_map[v] = new_obj
+      resource_map[v.handle] = new_obj.handle
+    obj_map[self] = new_obj
+    resource_map[self] = new_obj.handle
+    return obj_map, resource_map
+
 
 class _DistributedVariableSaveable(saveable_object.SaveableObject):
   """Class for defining how to restore a DistributedVariable."""
diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
index 7d0abe30581..57e8ced65a0 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
@@ -285,6 +285,13 @@ class AutoCastVariable(variables.Variable, core.Tensor):
     # models with normal variables, and vice versa.
     return self._variable._gather_saveables_for_checkpoint()  # pylint:disable=protected-access
 
+  def _map_resources(self):
+    # By delegating this method to the wrapped variable, SavedModel with
+    # AutoCastVariables are identical to SavedModel with normal variables.
+    obj_map, resource_map = self._variable._map_resources()  # pylint:disable=protected-access
+    obj_map[self] = obj_map[self._variable]
+    return obj_map, resource_map
+
   # TODO(reedwm): Maybe encode the fact the variable is an AutoCastVariable in
   # to_proto().
   def to_proto(self, export_scope=None):
diff --git a/tensorflow/python/ops/resource_variable_ops.py b/tensorflow/python/ops/resource_variable_ops.py
index 25f6347f034..cb235fcbe2d 100644
--- a/tensorflow/python/ops/resource_variable_ops.py
+++ b/tensorflow/python/ops/resource_variable_ops.py
@@ -633,6 +633,13 @@ class BaseResourceVariable(variables.VariableV1, core.Tensor):
     return gen_state_ops.resource_count_up_to(self.handle, limit=limit,
                                               T=self.dtype)
 
+  def _map_resources(self):
+    """For implementing `Trackable`."""
+    new_variable = copy_to_graph_uninitialized(self)
+    obj_map = {self: new_variable}
+    resource_map = {self._handle: new_variable.handle}
+    return obj_map, resource_map
+
   def _read_variable_op(self):
     variable_accessed(self)
 
diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py
index 5844c80995f..802ce1d61b7 100644
--- a/tensorflow/python/saved_model/save.py
+++ b/tensorflow/python/saved_model/save.py
@@ -19,14 +19,12 @@ from __future__ import division
 from __future__ import print_function
 
 import collections
-import copy
 import os
 
 from tensorflow.core.framework import versions_pb2
 from tensorflow.core.protobuf import meta_graph_pb2
 from tensorflow.core.protobuf import saved_model_pb2
 from tensorflow.core.protobuf import saved_object_graph_pb2
-from tensorflow.python.distribute import distribute_utils as ds_utils
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.eager import function as defun
@@ -241,7 +239,7 @@ class _SaveableView(object):
     Creates resource handle ops in the current default graph, whereas
     `accessible_objects` will be from an eager context. Resource mapping adds
     resource handle ops to the main GraphDef of a SavedModel, which allows the
-    C++ loader API to interact with variables.
+    C++ loader API to interact with resources.
 
     Returns:
       A tuple of (object_map, resource_map, asset_info):
@@ -265,33 +263,15 @@ class _SaveableView(object):
         asset_index={})
 
     for node_id, obj in enumerate(self.nodes):
-      if isinstance(obj, tracking.CapturableResource):
-        new_obj = object_map[obj] = copy.copy(obj)
-        # pylint: disable=protected-access
-        with ops.device(obj._resource_device):
-          new_resource = new_obj._create_resource()
-        new_obj._resource_handle = new_resource
-        # pylint: enable=protected-access
-        resource_map[obj.resource_handle] = new_resource
-        self.captured_tensor_node_ids[obj.resource_handle] = node_id
-      elif (ds_utils.is_distributed_variable(obj) or
-            resource_variable_ops.is_resource_variable(obj)):
-        obj_to_copy = obj._primary if ds_utils.is_distributed_variable(  # pylint: disable=protected-access
-            obj) else obj
-        new_variable = resource_variable_ops.copy_to_graph_uninitialized(
-            obj_to_copy)
-        if ds_utils.is_distributed_variable(obj):
-          self.captured_tensor_node_ids[obj] = node_id
-          for v in obj.values:
-            object_map[v] = new_variable
-            resource_map[v.handle] = new_variable.handle
-            self.captured_tensor_node_ids[v.handle] = node_id
-        object_map[obj] = new_variable
-        resource_map[obj.handle] = new_variable.handle
-        self.captured_tensor_node_ids[obj.handle] = node_id
-      elif isinstance(obj, tracking.Asset):
+      if isinstance(obj, tracking.Asset):
         _process_asset(obj, asset_info, resource_map)
         self.captured_tensor_node_ids[obj.asset_path] = node_id
+      elif isinstance(obj, base.Trackable):
+        node_object_map, node_resource_map = obj._map_resources()  # pylint: disable=protected-access
+        for capturable in node_resource_map.keys():
+          self.captured_tensor_node_ids[capturable] = node_id
+        object_map.update(node_object_map)
+        resource_map.update(node_resource_map)
 
     # Note: some concrete functions can have been realized when tracing other
     # functions, and might closure-capture tensors from their parent functions.
diff --git a/tensorflow/python/training/tracking/base.py b/tensorflow/python/training/tracking/base.py
index e3cd9828724..ea76ad8db47 100644
--- a/tensorflow/python/training/tracking/base.py
+++ b/tensorflow/python/training/tracking/base.py
@@ -1021,3 +1021,21 @@ class Trackable(object):
     """
     del serialization_cache
     return dict()
+
+  def _map_resources(self):
+    """Makes new resource handle ops corresponding to existing resource tensors.
+
+    Internal sub-classes can override this to inform model saving how to add new
+    resource handle ops to the main GraphDef of a SavedModel (TF 1.x style
+    graph), which allows session based APIs (e.g, C++ loader API) to interact
+    with resources owned by this object.
+
+    Returns:
+      A tuple of (object_map, resource_map):
+        object_map: A dictionary mapping from objects that hold existing
+          resource tensors to replacement objects created to hold the new
+          resource tensors.
+        resource_map: A dictionary mapping from existing resource tensors to
+          newly created resource tensors.
+    """
+    return {}, {}
diff --git a/tensorflow/python/training/tracking/tracking.py b/tensorflow/python/training/tracking/tracking.py
index 553f0ec73bf..fb2735e6445 100644
--- a/tensorflow/python/training/tracking/tracking.py
+++ b/tensorflow/python/training/tracking/tracking.py
@@ -17,6 +17,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import copy
 import functools
 import weakref
 
@@ -243,6 +244,18 @@ class CapturableResource(base.Trackable):
         self._resource_handle = self._create_resource()
     return self._resource_handle
 
+  def _map_resources(self):
+    """For implementing `Trackable`."""
+    new_obj = copy.copy(self)
+    # pylint: disable=protected-access
+    with ops.device(self._resource_device):
+      new_resource = new_obj._create_resource()
+    new_obj._resource_handle = new_resource
+    # pylint: enable=protected-access
+    obj_map = {self: new_obj}
+    resource_map = {self.resource_handle: new_resource}
+    return obj_map, resource_map
+
   def _list_functions_for_serialization(self, unused_functions):
     @def_function.function(input_signature=[], autograph=False)
     def _creator():

From 39504c25d9de697d3568bc4d370722d0f48376cf Mon Sep 17 00:00:00 2001
From: Smit Hinsu 
Date: Thu, 18 Jun 2020 15:55:24 -0700
Subject: [PATCH 0552/1390] Fix bug in xla-legalize-tf-with-tf2xla pass by
 handling non-tensor operands

Currently, it only expects tensor operands but that is not applicable for non tensorflow dialect ops.

PiperOrigin-RevId: 317198672
Change-Id: I1387e664de740d044ef535f6903e07d63fa02f6d
---
 .../mlir/xla/tests/legalize-tf-with-tf2xla.mlir      | 12 ++++++++++--
 .../mlir/xla/transforms/legalize_tf_with_tf2xla.cc   |  6 +++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/tests/legalize-tf-with-tf2xla.mlir b/tensorflow/compiler/mlir/xla/tests/legalize-tf-with-tf2xla.mlir
index b8a6df54519..86a7f2b9e09 100644
--- a/tensorflow/compiler/mlir/xla/tests/legalize-tf-with-tf2xla.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/legalize-tf-with-tf2xla.mlir
@@ -35,7 +35,7 @@ func @not_whitelisted_op(%arg0: tensor<3xi32>, %arg1: tensor, %arg2: tensor
 // CHECK-LABEL: unranked_operand
 func @unranked_operand(%arg0: tensor<*xf32>) -> tensor<*xf32> {
   // CHECK: tf.Abs
-  // expected-remark@+1 {{lowering requires static shaped operands}}
+  // expected-remark@+1 {{lowering requires static shaped tensor operands}}
   %0 = "tf.Abs"(%arg0) : (tensor<*xf32>) -> tensor<*xf32>
 
   return %0 : tensor<*xf32>
@@ -44,12 +44,20 @@ func @unranked_operand(%arg0: tensor<*xf32>) -> tensor<*xf32> {
 // CHECK-LABEL: dynamic_operand
 func @dynamic_operand(%arg0: tensor) -> tensor {
   // CHECK: tf.Abs
-  // expected-remark@+1 {{lowering requires static shaped operands}}
+  // expected-remark@+1 {{lowering requires static shaped tensor operands}}
   %0 = "tf.Abs"(%arg0) : (tensor) -> tensor
 
   return %0 : tensor
 }
 
+// CHECK-LABEL: tuple_type
+func @tuple_type(%arg0: tuple, tensor>) -> tensor {
+  // Verifies that the pass can handle operands of non-tensor type like tuple
+  // from non TensorFlow ops.
+  %0 = "xla_hlo.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple, tensor>) -> tensor
+  return %0 : tensor
+}
+
 // CHECK-LABEL: unsupported_dtype
 func @unsupported_dtype(%arg0: tensor<2x!tf.variant>) -> tensor<2x!tf.variant> {
   // CHECK: tf.AddN
diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
index e57d6938efb..ef79c8868bb 100644
--- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
@@ -337,9 +337,9 @@ LogicalResult FuncLegalizer::LegalizeOp(Operation* op) {
 
   // Only static shaped operands are supported in XLA builders for now.
   for (Type ty : op->getOperandTypes()) {
-    auto ranked_ty = ty.cast();
-    if (!ranked_ty.hasStaticShape()) {
-      op->emitRemark() << "lowering requires static shaped operands";
+    auto ranked_ty = ty.dyn_cast();
+    if (!ranked_ty || !ranked_ty.hasStaticShape()) {
+      op->emitRemark() << "lowering requires static shaped tensor operands";
       return success();
     }
   }

From 8d34408863b650564076f148edad9f91508abf04 Mon Sep 17 00:00:00 2001
From: Smit Hinsu 
Date: Thu, 18 Jun 2020 16:02:23 -0700
Subject: [PATCH 0553/1390] Auto-generate following TensorFlow ops related to
 image

ResizeBilinearGrad ResizeBilinear AdjustContrastv2 ResizeNearestNeighbor
AdjustSaturation AdjustHue RGBToHSV HSVToRGB

PiperOrigin-RevId: 317199967
Change-Id: I1953acf599f2f7de686bda73b654e4c7b98dffd5
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    | 153 ++++++++++++++++++
 1 file changed, 153 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index 3a5deb9c569..dcd083fc398 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -164,6 +164,81 @@ def TF_AddV2Op : TF_Op<"AddV2", [Commutative, NoSideEffect, ResultsBroadcastable
   let hasFolder = 1;
 }
 
+def TF_AdjustContrastv2Op : TF_Op<"AdjustContrastv2", [NoSideEffect]> {
+  let summary = "Adjust the contrast of one or more images.";
+
+  let description = [{
+`images` is a tensor of at least 3 dimensions.  The last 3 dimensions are
+interpreted as `[height, width, channels]`.  The other dimensions only
+represent a collection of images, such as `[batch, height, width, channels].`
+
+Contrast is adjusted independently for each channel of each image.
+
+For each channel, the Op first computes the mean of the image pixels in the
+channel and then adjusts each component of each pixel to
+`(x - mean) * contrast_factor + mean`.
+  }];
+
+  let arguments = (ins
+    TensorOf<[F16, F32]>:$images,
+    F32Tensor:$contrast_factor
+  );
+
+  let results = (outs
+    TensorOf<[F16, F32]>:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_AdjustHueOp : TF_Op<"AdjustHue", [NoSideEffect]> {
+  let summary = "Adjust the hue of one or more images.";
+
+  let description = [{
+`images` is a tensor of at least 3 dimensions.  The last dimension is
+interpreted as channels, and must be three.
+
+The input image is considered in the RGB colorspace. Conceptually, the RGB
+colors are first mapped into HSV. A delta is then applied all the hue values,
+and then remapped back to RGB colorspace.
+  }];
+
+  let arguments = (ins
+    TensorOf<[F16, F32]>:$images,
+    F32Tensor:$delta
+  );
+
+  let results = (outs
+    TensorOf<[F16, F32]>:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
+def TF_AdjustSaturationOp : TF_Op<"AdjustSaturation", [NoSideEffect]> {
+  let summary = "Adjust the saturation of one or more images.";
+
+  let description = [{
+`images` is a tensor of at least 3 dimensions.  The last dimension is
+interpreted as channels, and must be three.
+
+The input image is considered in the RGB colorspace. Conceptually, the RGB
+colors are first mapped into HSV. A scale is then applied all the saturation
+values, and then remapped back to RGB colorspace.
+  }];
+
+  let arguments = (ins
+    TensorOf<[F16, F32]>:$images,
+    F32Tensor:$scale
+  );
+
+  let results = (outs
+    TensorOf<[F16, F32]>:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
 def TF_AllOp : TF_Op<"All", [NoSideEffect]> {
   let summary = [{
 Computes the "logical and" of elements across dimensions of a tensor.
@@ -3866,6 +3941,28 @@ tf.math.greater_equal(x, y) ==> [True, False, True, True]
   TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
 }
 
+def TF_HSVToRGBOp : TF_Op<"HSVToRGB", [NoSideEffect]> {
+  let summary = "Convert one or more images from HSV to RGB.";
+
+  let description = [{
+Outputs a tensor of the same shape as the `images` tensor, containing the RGB
+value of the pixels. The output is only well defined if the value in `images`
+are in `[0,1]`.
+
+See `rgb_to_hsv` for a description of the HSV encoding.
+  }];
+
+  let arguments = (ins
+    TF_FpTensor:$images
+  );
+
+  let results = (outs
+    TF_FpTensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
 def TF_HashTableV2Op : TF_Op<"HashTableV2", []> {
   let summary = "Creates a non-initialized hash table.";
 
@@ -6733,6 +6830,41 @@ the dimension is padded with zeros.
   TF_DerivedResultTypeAttr Tcomplex = TF_DerivedResultTypeAttr<0>;
 }
 
+def TF_RGBToHSVOp : TF_Op<"RGBToHSV", [NoSideEffect]> {
+  let summary = "Converts one or more images from RGB to HSV.";
+
+  let description = [{
+Outputs a tensor of the same shape as the `images` tensor, containing the HSV
+value of the pixels. The output is only well defined if the value in `images`
+are in `[0,1]`.
+
+`output[..., 0]` contains hue, `output[..., 1]` contains saturation, and
+`output[..., 2]` contains value. All HSV values are in `[0,1]`. A hue of 0
+corresponds to pure red, hue 1/3 is pure green, and 2/3 is pure blue.
+
+Usage Example:
+
+>>> blue_image = tf.stack([
+...    tf.zeros([5,5]),
+...    tf.zeros([5,5]),
+...    tf.ones([5,5])],
+...    axis=-1)
+>>> blue_hsv_image = tf.image.rgb_to_hsv(blue_image)
+>>> blue_hsv_image[0,0].numpy()
+array([0.6666667, 1. , 1. ], dtype=float32)
+  }];
+
+  let arguments = (ins
+    TF_FpTensor:$images
+  );
+
+  let results = (outs
+    TF_FpTensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+}
+
 def TF_RandomGammaGradOp : TF_Op<"RandomGammaGrad", [NoSideEffect, ResultsBroadcastableShape]>,
                            WithBroadcastableBinOpBuilder {
   let summary = [{
@@ -7230,6 +7362,27 @@ Input images can be of different types but output images are always float.
   TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
 }
 
+def TF_ResizeBilinearGradOp : TF_Op<"ResizeBilinearGrad", [NoSideEffect]> {
+  let summary = "Computes the gradient of bilinear interpolation.";
+
+  let description = [{
+  }];
+
+  let arguments = (ins
+    F32Tensor:$grads,
+    TF_FpTensor:$original_image,
+
+    DefaultValuedAttr:$align_corners,
+    DefaultValuedAttr:$half_pixel_centers
+  );
+
+  let results = (outs
+    TF_FpTensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<1>;
+}
+
 def TF_ResizeNearestNeighborOp : TF_Op<"ResizeNearestNeighbor", [NoSideEffect]> {
   let summary = [{
 Resize `images` to `size` using nearest neighbor interpolation.

From 67544cd4bbdf8070adebbb077439cac300f479ca Mon Sep 17 00:00:00 2001
From: Ken Franko 
Date: Thu, 18 Jun 2020 16:02:51 -0700
Subject: [PATCH 0554/1390] Add more outside compilation tests including
 multiple clusters and more variety of inputs/outputs.

PiperOrigin-RevId: 317200078
Change-Id: Id26e99059097073299ef5f681fae053b082ec149
---
 .../tpu/tpu_outside_compilation_test.py       | 95 ++++++++++++++++++-
 1 file changed, 91 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/tpu/tpu_outside_compilation_test.py b/tensorflow/python/tpu/tpu_outside_compilation_test.py
index f7ecb294c44..54c2598324c 100644
--- a/tensorflow/python/tpu/tpu_outside_compilation_test.py
+++ b/tensorflow/python/tpu/tpu_outside_compilation_test.py
@@ -18,13 +18,18 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
+import numpy as np
+
 from tensorflow.python.distribute import tpu_strategy as tpu_lib
 from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
 from tensorflow.python.eager import def_function
 from tensorflow.python.eager import remote
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import flags
 from tensorflow.python.tpu import tpu
@@ -52,7 +57,7 @@ def get_tpu_strategy():
   return tpu_lib.TPUStrategy(resolver)
 
 
-class TpuOutsideCompilationTest(test.TestCase):
+class TpuOutsideCompilationTest(test.TestCase, parameterized.TestCase):
 
   def testResourceVariableAssignOnHost(self):
     strategy = get_tpu_strategy()
@@ -79,6 +84,26 @@ class TpuOutsideCompilationTest(test.TestCase):
     self.assertAllEqual(4.0 * strategy.num_replicas_in_sync, v2.numpy())
     self.assertAllEqual(5.0, v.numpy())
 
+  def testHostNoInput(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn():
+      logging_ops.print_v2("Outside compiled")
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        tpu.outside_compilation(outside_fn)
+        return x2 + 5.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(35., shape=(strategy.num_replicas_in_sync)))
+
   def testHostInputOnly(self):
     strategy = get_tpu_strategy()
 
@@ -120,13 +145,71 @@ class TpuOutsideCompilationTest(test.TestCase):
         strategy.experimental_local_results(train_step()),
         constant_op.constant(36., shape=(strategy.num_replicas_in_sync)))
 
-  def testOutsideCompilationControlFlowIf(self):
+  def testHostMultipleInputs(self):
+    strategy = get_tpu_strategy()
+    val0 = np.arange(6).reshape((2, 3)).astype(np.float32)
+    val1 = np.arange(6).reshape((3, 2)).astype(np.float32)
+
+    def outside_fn(arg0, arg1):
+      tmp = array_ops.reshape(arg1, array_ops.shape(arg0))
+      ret0 = arg0 + tmp
+      ret1 = math_ops.matmul(arg0, arg1)
+      ret2 = array_ops.concat([arg0, tmp], 0)
+      return ret0, ret1, ret2
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x, y):
+        a = x + 7.0
+        b = y * 2.0
+        c, d, e = tpu.outside_compilation(outside_fn, a, b)
+        return (math_ops.reduce_max(c) + math_ops.reduce_min(d) +
+                math_ops.reduce_sum(e))
+
+      return strategy.run(tpu_fn, args=(val0, val1))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(213., shape=(strategy.num_replicas_in_sync)))
+
+  def testMultipleClusters(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn1(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x + 6.0
+
+    def outside_fn2(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x - 18.0
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        output1 = tpu.outside_compilation(outside_fn1, x2)
+        x3 = output1 + 3.0
+        output2 = tpu.outside_compilation(outside_fn2, x3)
+        return output2
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(21., shape=(strategy.num_replicas_in_sync)))
+
+  @parameterized.parameters((True), (False))
+  def testOutsideCompilationControlFlowIf(self, take_true_branch):
     strategy = get_tpu_strategy()
 
     def outside_fn(x):
       logging_ops.print_v2("Outside compiled", x)
       return x + 6.0
 
+    input_value = 51.0 if take_true_branch else 25.0
+
     @def_function.function
     def train_step():
 
@@ -137,11 +220,15 @@ class TpuOutsideCompilationTest(test.TestCase):
         else:
           return x2
 
-      return strategy.run(tpu_fn, args=(25.0,))
+      return strategy.run(tpu_fn, args=(input_value,))
 
+    output_value = 36.0
+    if take_true_branch:
+      output_value = 56.0
     self.assertAllEqual(
         strategy.experimental_local_results(train_step()),
-        constant_op.constant(36., shape=(strategy.num_replicas_in_sync)))
+        constant_op.constant(
+            output_value, shape=(strategy.num_replicas_in_sync)))
 
   def testOutsideCompilationControlFlowWhile(self):
     strategy = get_tpu_strategy()

From 4aea552e064cf92330e07e83a3b5a1ca2a7034d0 Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Thu, 18 Jun 2020 16:15:22 -0700
Subject: [PATCH 0555/1390] Publishing tpu_op_consts to tpu kernels library.

PiperOrigin-RevId: 317202394
Change-Id: Ib6a1f350af7384513a3744084a9959ed86278d1f
---
 tensorflow/core/tpu/kernels/BUILD             | 11 +++++-
 .../kernels/tpu_compilation_cache_external.h  |  5 +--
 tensorflow/core/tpu/kernels/tpu_op_consts.cc  | 24 ++++++++++++
 tensorflow/core/tpu/kernels/tpu_op_consts.h   | 39 +++++++++++++++++++
 4 files changed, 74 insertions(+), 5 deletions(-)
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_consts.cc
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_consts.h

diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index a41747ee8c5..a9f2202cd45 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -321,6 +321,7 @@ cc_library(
         ":tpu_compile_c_api_hdrs",
         ":tpu_compile_op_support",
         ":tpu_mesh_state_interface",
+        ":tpu_op_consts",
         ":tpu_program_group",
         ":tpu_util",
         ":trace_util_hdrs",
@@ -433,7 +434,6 @@ cc_library(
         "//tensorflow/compiler/xla/client:compile_only_client",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
     alwayslink = 1,
@@ -461,3 +461,12 @@ cc_library(
         "@com_google_absl//absl/synchronization",
     ],
 )
+
+cc_library(
+    name = "tpu_op_consts",
+    srcs = ["tpu_op_consts.cc"],
+    hdrs = ["tpu_op_consts.h"],
+    deps = [
+        "@com_google_absl//absl/base:core_headers",
+    ],
+)
diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
index fe251326a43..86615b15d4c 100644
--- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_external.h
@@ -38,15 +38,12 @@ limitations under the License.
 #include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h"
+#include "tensorflow/core/tpu/kernels/tpu_op_consts.h"
 #include "tensorflow/core/tpu/kernels/tpu_program_group.h"
 
 namespace tensorflow {
 namespace tpu {
 
-constexpr char kCompilationCacheResourceName[] = "tpu_compilation_cache";
-constexpr char kCompilationCacheUnloaderResourceName[] =
-    "tpu_compilation_cache_unloader";
-
 class TpuCompilationCacheExternal : public TpuCompilationCacheInterface {
  public:
   using Status = ::stream_executor::port::Status;
diff --git a/tensorflow/core/tpu/kernels/tpu_op_consts.cc b/tensorflow/core/tpu/kernels/tpu_op_consts.cc
new file mode 100644
index 00000000000..e5e1aacb3cc
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_op_consts.cc
@@ -0,0 +1,24 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/core/tpu/kernels/tpu_op_consts.h"
+
+namespace tensorflow {
+namespace tpu {
+const char kCompilationCacheResourceName[] = "tpu_compilation_cache";
+const char kCompiledProtoCacheResourceName[] = "tpu_proto_cache";
+const char kCompilationCacheUnloaderResourceName[] =
+    "tpu_compilation_cache_unloader";
+}  // namespace tpu
+}  // namespace tensorflow
diff --git a/tensorflow/core/tpu/kernels/tpu_op_consts.h b/tensorflow/core/tpu/kernels/tpu_op_consts.h
new file mode 100644
index 00000000000..25223b7e429
--- /dev/null
+++ b/tensorflow/core/tpu/kernels/tpu_op_consts.h
@@ -0,0 +1,39 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_CONSTS_H_
+#define TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_CONSTS_H_
+
+#include "absl/base/attributes.h"
+
+namespace tensorflow {
+namespace tpu {
+
+// Resource names in the ResourceMgr.
+//
+// Name of cache for compiled TPU ISA protos. CompilationCache is created by
+// ConfigureDistributedTpuOp, so only the master has a CompilationCache.
+ABSL_CONST_INIT extern const char kCompilationCacheResourceName[];
+// Name of base class allowing Execute Ops to look up ISA protos.
+// CompiledProtoCache is created by InitializeHostForDistributedTpuOp, so each
+// tpu_worker has a CompiledProtoCache.
+ABSL_CONST_INIT extern const char kCompiledProtoCacheResourceName[];
+// Name of cache unloader for compiled TPU ISA protos. Cache unloader should be
+// put into TPU_SYSTEM device resource manager. Inference may use it to unload
+// cache entries created during lifetime of a DirectSession.
+ABSL_CONST_INIT extern const char kCompilationCacheUnloaderResourceName[];
+
+}  // namespace tpu
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_TPU_KERNELS_TPU_OP_CONSTS_H_

From aab151356d2334a9d6cec71ce5165e6e6c45c793 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 16:24:14 -0700
Subject: [PATCH 0556/1390] [pjrt] Refresh stream error status in strategic
 places to flush out silent failures.

PiperOrigin-RevId: 317204018
Change-Id: If75a3ad9ec846ce1621cdba92a2dc738b65b7001
---
 tensorflow/compiler/xla/pjrt/local_device_state.cc |  4 ++++
 tensorflow/compiler/xla/pjrt/pjrt_client.cc        | 12 +++++++++---
 tensorflow/stream_executor/stream.cc               |  7 ++++++-
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/local_device_state.cc b/tensorflow/compiler/xla/pjrt/local_device_state.cc
index d173c891c95..a229e56001e 100644
--- a/tensorflow/compiler/xla/pjrt/local_device_state.cc
+++ b/tensorflow/compiler/xla/pjrt/local_device_state.cc
@@ -127,11 +127,15 @@ std::unique_ptr LocalDeviceState::BorrowStreamFromPool() {
   } else {
     std::unique_ptr stream = std::move(usage_stream_pool_.top());
     usage_stream_pool_.pop();
+    stream->RefreshStatus().IgnoreError();  // Can return error::Unimplemented
+    QCHECK(stream->ok());
     return stream;
   }
 }
 
 void LocalDeviceState::ReturnStreamToPool(std::unique_ptr stream) {
+  stream->RefreshStatus().IgnoreError();  // Can return error::Unimplemented
+  QCHECK(stream->ok());
   absl::MutexLock lock(&mu_);
   usage_stream_pool_.push(std::move(stream));
 }
diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.cc b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
index ef259cf1cfd..46f592100c9 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.cc
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
@@ -751,16 +751,22 @@ StatusOr> PjRtBuffer::FromHostLiteral(
     // memory that has already been allocated, and a possible Event
     // allocation.
 
+    se::Stream* h2d_stream = local_device->host_to_device_stream();
     ShapedBuffer buffer = device_buffer->AsShapedBuffer(
         compact_shape, on_device_shape, client->client()->platform());
     TF_CHECK_OK(transfer_manager->TransferLiteralToDeviceAsync(
-        local_device->host_to_device_stream(), literal, buffer));
+        h2d_stream, literal, buffer));
 
     std::shared_ptr event =
         device_buffer->definition_events()[0];
     TF_CHECK_OK(AddDestinationBufferSynchronization(
-        local_device, std::move(device_buffer), event,
-        local_device->host_to_device_stream()));
+        local_device, std::move(device_buffer), event, h2d_stream));
+
+    // This can sometimes catch the case where the literal memory has been
+    // freed before the H2D transfer was issued.
+    h2d_stream->RefreshStatus()
+        .IgnoreError();  // Can return error::Unimplemented
+    QCHECK(h2d_stream->ok());
   };
   client->h2d_transfer_pool()->Schedule(transfer_h2d);
   return py_buffer;
diff --git a/tensorflow/stream_executor/stream.cc b/tensorflow/stream_executor/stream.cc
index c63565c65a8..da418122375 100644
--- a/tensorflow/stream_executor/stream.cc
+++ b/tensorflow/stream_executor/stream.cc
@@ -285,7 +285,12 @@ Stream::~Stream() {
 
 port::Status Stream::RefreshStatus() {
   port::Status status = parent_->GetStatus(this);
-  CheckStatus(status);
+  // We should not put the stream in an error state, just because the GetStatus
+  // method is unimplemented.
+  if (status != port::Status(port::error::UNIMPLEMENTED,
+                             "GetStatus is not supported on this executor.")) {
+    CheckStatus(status);
+  }
   return status;
 }
 

From 1363f0f6e89cdc73ca46d43475bdb35750ed2e50 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 19:28:34 -0400
Subject: [PATCH 0557/1390] add save load test for pathlib path

---
 tensorflow/python/keras/saving/save_test.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 5c5846fe738..98d78735ad1 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -71,6 +71,12 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
     save.save_model(self.model, path)
     self.assert_saved_model(path)
 
+  @test_util.run_v2_only
+  def test_save_format_defaults_pathlib(self):
+    path = pathlib.Path(self.get_temp_dir()) / 'model_path'
+    save.save_model(self.model, path)
+    self.assert_saved_model(path)
+
   @test_util.run_v2_only
   def test_save_hdf5(self):
     path = os.path.join(self.get_temp_dir(), 'model')
@@ -81,6 +87,13 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
         'requires the model to be a Functional model or a Sequential model.'):
       save.save_model(self.subclassed_model, path, save_format='h5')
 
+  @test_util.run_v2_only
+  def test_save_load_hdf5_pathlib(self):
+    if sys.version_info >= (3, 6):
+      path = pathlib.Path(self.get_temp_dir()) / 'model'
+      save.save_model(self.model, path, save_format='h5')
+      save.load_model(path)
+
   @test_util.run_v2_only
   def test_save_tf(self):
     path = os.path.join(self.get_temp_dir(), 'model')

From 94e37f84f19384e685420ef7f90382fcfe719498 Mon Sep 17 00:00:00 2001
From: Dero Gharibian 
Date: Thu, 18 Jun 2020 16:25:11 -0700
Subject: [PATCH 0558/1390] Remove unnecessary string copy.

PiperOrigin-RevId: 317204219
Change-Id: I85fab345945b6ea4f428f8aedc861eb79e5fd7e0
---
 tensorflow/core/kernels/summary_op.cc | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/summary_op.cc b/tensorflow/core/kernels/summary_op.cc
index 386a8964dba..f4c91fc9ff1 100644
--- a/tensorflow/core/kernels/summary_op.cc
+++ b/tensorflow/core/kernels/summary_op.cc
@@ -52,7 +52,8 @@ class SummaryScalarOp : public OpKernel {
     Summary s;
     for (int i = 0; i < Ttags.size(); i++) {
       Summary::Value* v = s.add_value();
-      v->set_tag(string(Ttags(i)));  // NOLINT
+      const tstring& Ttags_i = Ttags(i);
+      v->set_tag(Ttags_i.data(), Ttags_i.size());
       v->set_simple_value(float(Tvalues(i)));
     }
 
@@ -102,7 +103,8 @@ class SummaryHistoOp : public OpKernel {
 
     Summary s;
     Summary::Value* v = s.add_value();
-    v->set_tag(string(tags.scalar()()));  // NOLINT
+    const tstring& tags0 = tags.scalar()();
+    v->set_tag(tags0.data(), tags0.size());
     histo.EncodeToProto(v->mutable_histo(), false /* Drop zero buckets */);
 
     Tensor* summary_tensor = nullptr;

From e7733c8dd4cae1a1a03053b452acbbded7ab6014 Mon Sep 17 00:00:00 2001
From: Reed 
Date: Thu, 18 Jun 2020 16:52:51 -0700
Subject: [PATCH 0559/1390] Add LINT.IfChange

---
 .../core/grappler/optimizers/auto_mixed_precision_lists.h       | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
index c6016548117..9041d934017 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
@@ -73,6 +73,7 @@ class AutoMixedPrecisionLists {
   static void AddTensorListOps(gtl::FlatSet* list) {
     // Note: if a data structure op (such as TensorListPopBack) is added here,
     // IsTensorListReaderOp or IsTensorListWriterOp may need to be modified
+    // LINT.IfChange
     constexpr char* tensor_list_ops[] = {
         "TensorListConcat",
         "TensorListConcatLists",
@@ -90,6 +91,7 @@ class AutoMixedPrecisionLists {
         "TensorListSplit",
         "TensorListStack"
     };
+    // LINT.ThenChange(//tensorflow/core/grappler/optimizers/auto_mixed_precision.cc)
     for (auto op : tensor_list_ops) {
       list->insert(op);
     }

From 1c72a6c65e1733b55286c6361142a39d699732dc Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 19:53:22 -0400
Subject: [PATCH 0560/1390] convert export_dir as pathlike object to str in
 parse_saved_model

---
 tensorflow/python/saved_model/loader_impl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index 2df2bea428e..71f6ed16c9d 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -83,11 +83,11 @@ def parse_saved_model(export_dir):
   """
   # Build the path to the SavedModel in pbtxt format.
   path_to_pbtxt = os.path.join(
-      compat.as_bytes(export_dir),
+      compat.as_bytes(compat.path_to_str(export_dir)),
       compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
   # Build the path to the SavedModel in pb format.
   path_to_pb = os.path.join(
-      compat.as_bytes(export_dir),
+      compat.as_bytes(compat.path_to_str(export_dir)),
       compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
 
   # Parse the SavedModel protocol buffer.

From 83b09270dc34308ef60f2f68de540a4cb213e1e4 Mon Sep 17 00:00:00 2001
From: Lucy Fox 
Date: Thu, 18 Jun 2020 16:52:46 -0700
Subject: [PATCH 0561/1390] Update FusedKernelMatcher pass to use upstream util
 to get stripped op name.

Added this util upstream in D81435, so now using that instead and deleting the unneeded code here.

PiperOrigin-RevId: 317209256
Change-Id: Id2d8a1fca34ca85e59a05a85bf7f6f59b425c7c1
---
 .../tensorflow/transforms/fused_kernel_matcher.cc     | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
index 4b10550df7b..d10f5e26e8f 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/fused_kernel_matcher.cc
@@ -52,11 +52,6 @@ struct FusedKernelMatcherPass
   void runOnFunction() override;
 };
 
-// Returns an op's name with the dialect prefix stripped off.
-StringRef GetOpNameWithoutDialect(Operation *op) {
-  return op->getName().getStringRef().split(".").second;
-}
-
 bool IsActivationFunction(Operation *op) {
   return isa(op) || isa(op) || isa(op);
 }
@@ -128,8 +123,8 @@ class FuseContractionWithBiasAdd : public OpRewritePattern {
     }
 
     SmallVector locations{contraction.getLoc(), bias_add.getLoc()};
-    SmallVector fused_ops{
-        StringAttr::get(GetOpNameWithoutDialect(bias_add), context)};
+    SmallVector fused_ops{StringAttr::get(
+        bias_add.getOperation()->getName().stripDialect(), context)};
 
     // BiasAdd may or may not feed into an activation function.
     auto activation = GetActivation(bias_add);
@@ -143,7 +138,7 @@ class FuseContractionWithBiasAdd : public OpRewritePattern {
     if (fuse_activation) {
       locations.push_back(activation->getLoc());
       fused_ops.push_back(
-          StringAttr::get(GetOpNameWithoutDialect(activation), context));
+          StringAttr::get(activation->getName().stripDialect(), context));
       result_type = activation->getResultTypes().front();
     } else {
       result_type = bias_add.getResult().getType();

From 5f2e0240ee7977042e41d9c29c349a7b14301290 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 19:57:23 -0400
Subject: [PATCH 0562/1390] modify docstring

---
 tensorflow/python/saved_model/loader_impl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index 71f6ed16c9d..06cd988130d 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -73,7 +73,7 @@ def parse_saved_model(export_dir):
   """Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`.
 
   Args:
-    export_dir: Directory containing the SavedModel file.
+    export_dir: String or Pathlike, path to the directory containing the SavedModel file.
 
   Returns:
     A `SavedModel` protocol buffer.

From 86770b177e88cfa25844188a655f8399f5c3526d Mon Sep 17 00:00:00 2001
From: jonah-kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Thu, 18 Jun 2020 17:00:04 -0700
Subject: [PATCH 0563/1390] Add show_dtype param to plot_model and
 model_to_dot.

---
 tensorflow/python/keras/utils/vis_utils.py | 29 +++++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/keras/utils/vis_utils.py b/tensorflow/python/keras/utils/vis_utils.py
index 3720708543f..21de8014c2a 100644
--- a/tensorflow/python/keras/utils/vis_utils.py
+++ b/tensorflow/python/keras/utils/vis_utils.py
@@ -69,6 +69,7 @@ def add_edge(dot, src, dst):
 @keras_export('keras.utils.model_to_dot')
 def model_to_dot(model,
                  show_shapes=False,
+                 show_dtype=False,
                  show_layer_names=True,
                  rankdir='TB',
                  expand_nested=False,
@@ -79,6 +80,7 @@ def model_to_dot(model,
   Arguments:
     model: A Keras model instance.
     show_shapes: whether to display shape information.
+    show_dtype: whether to display layer dtypes.
     show_layer_names: whether to display layer names.
     rankdir: `rankdir` argument passed to PyDot,
         a string specifying the format of the plot:
@@ -150,8 +152,11 @@ def model_to_dot(model,
     if isinstance(layer, wrappers.Wrapper):
       if expand_nested and isinstance(layer.layer,
                                       functional.Functional):
-        submodel_wrapper = model_to_dot(layer.layer, show_shapes,
-                                        show_layer_names, rankdir,
+        submodel_wrapper = model_to_dot(layer.layer, 
+                                        show_shapes,
+                                        show_dtype,
+                                        show_layer_names, 
+                                        rankdir,
                                         expand_nested,
                                         subgraph=True)
         # sub_w : submodel_wrapper
@@ -165,8 +170,11 @@ def model_to_dot(model,
         class_name = '{}({})'.format(class_name, child_class_name)
 
     if expand_nested and isinstance(layer, functional.Functional):
-      submodel_not_wrapper = model_to_dot(layer, show_shapes,
-                                          show_layer_names, rankdir,
+      submodel_not_wrapper = model_to_dot(layer, 
+                                          show_shapes, 
+                                          show_dtype,
+                                          show_layer_names, 
+                                          rankdir,
                                           expand_nested,
                                           subgraph=True)
       # sub_n : submodel_not_wrapper
@@ -180,6 +188,16 @@ def model_to_dot(model,
       label = '{}: {}'.format(layer_name, class_name)
     else:
       label = class_name
+  
+    # Rebuild the label as a table including the layer's dtype.
+    if show_dtype: 
+      def format_dtype(dtype):
+          if dtype is None:
+            return '?'
+          else:
+            return str(dtype)
+          
+      label = '%s|%s' % (label, dtype)
 
     # Rebuild the label as a table including input/output shapes.
     if show_shapes:
@@ -260,6 +278,7 @@ def model_to_dot(model,
 def plot_model(model,
                to_file='model.png',
                show_shapes=False,
+               show_dtype=False,
                show_layer_names=True,
                rankdir='TB',
                expand_nested=False,
@@ -286,6 +305,7 @@ def plot_model(model,
     model: A Keras model instance
     to_file: File name of the plot image.
     show_shapes: whether to display shape information.
+    show_dtype: whether to display layer dtypes.
     show_layer_names: whether to display layer names.
     rankdir: `rankdir` argument passed to PyDot,
         a string specifying the format of the plot:
@@ -300,6 +320,7 @@ def plot_model(model,
   """
   dot = model_to_dot(model,
                      show_shapes=show_shapes,
+                     show_dtype=show_dtype,
                      show_layer_names=show_layer_names,
                      rankdir=rankdir,
                      expand_nested=expand_nested,

From c3ca4da9b46e3148bdb913d08fec3fd2727158e1 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:00:06 -0400
Subject: [PATCH 0564/1390] only test pathlib on python v>3.6

---
 tensorflow/python/keras/saving/save_test.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 98d78735ad1..66364666841 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -73,9 +73,10 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
 
   @test_util.run_v2_only
   def test_save_format_defaults_pathlib(self):
-    path = pathlib.Path(self.get_temp_dir()) / 'model_path'
-    save.save_model(self.model, path)
-    self.assert_saved_model(path)
+    if sys.version_info >= (3, 6):
+      path = pathlib.Path(self.get_temp_dir()) / 'model_path'
+      save.save_model(self.model, path)
+      self.assert_saved_model(path)
 
   @test_util.run_v2_only
   def test_save_hdf5(self):

From e96543f6fbba9fa112d7ca1d731b64e3654e1629 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Thu, 18 Jun 2020 16:55:26 -0700
Subject: [PATCH 0565/1390] Add MultiProcessPoolRunner

Tensorflow initialization can take a long time when GPUs are present. We cannot afford starting a new group of workers for every single test. MultiProcessPoolRunner uses a pool of workers so that we can avoid the initialization cost. Compared to MultiProcessRunner, it doesn't support terminating workers.

Note that implementation wise we could build MultiProcessPoolRunner on top of MultiProcessRunner or vice-versa if there's no need to support termination. Since it's easier for MultiProcessPoolRunner not to support termination, we choose MultiProcessPoolRunner on top of MultiProcessRunner.

PiperOrigin-RevId: 317209754
Change-Id: Ia439028c81c5a9f87b16d631a170158724ce47d4
---
 .../python/distribute/multi_process_runner.py | 205 ++++++++++++++++--
 .../distribute/multi_process_runner_test.py   |  69 +++++-
 2 files changed, 249 insertions(+), 25 deletions(-)

diff --git a/tensorflow/python/distribute/multi_process_runner.py b/tensorflow/python/distribute/multi_process_runner.py
index 8699e59b410..ce36287a9da 100644
--- a/tensorflow/python/distribute/multi_process_runner.py
+++ b/tensorflow/python/distribute/multi_process_runner.py
@@ -67,8 +67,7 @@ except ImportError:
 # exception stack trace info is stored in exc_info to pass on to parent process
 # to be re-raised.
 _ProcessStatusInfo = collections.namedtuple(
-    '_ProcessStatusInfo',
-    ['task_type', 'is_successful', 'exc_info', 'return_value'])
+    '_ProcessStatusInfo', ['is_successful', 'exc_info', 'return_value'])
 
 # Information returned from a successful MultiProcessRunner run.
 MultiProcessRunnerResult = collections.namedtuple('MultiProcessRunnerResult',
@@ -124,6 +123,7 @@ class MultiProcessRunner(object):
                stream_stdout=True,
                list_stdout=False,
                use_dill_for_args=True,
+               daemon=False,
                args=None,
                kwargs=None):
     """Creates a multi-process runner.
@@ -157,6 +157,7 @@ class MultiProcessRunner(object):
       use_dill_for_args: Whether to use dill to pickle `args` and `kwargs`. dill
         can pickle more objects, but doesn't work with types in
         `multiprocessing` library like `Mutex`.
+      daemon: Whether to start processes as daemons.
       args: Positional arguments to be sent to functions run on processes.
       kwargs: Keyword arguments to be sent to functions run on processes.
 
@@ -188,6 +189,7 @@ class MultiProcessRunner(object):
     self._list_stdout = list_stdout
     self._dependence_on_chief = True
     self._use_dill_for_args = use_dill_for_args
+    self._daemon = daemon
     self._args = args or ()
     self._kwargs = kwargs or {}
 
@@ -268,7 +270,8 @@ class MultiProcessRunner(object):
         test_env=test_env,
         target=_ProcFunc(),
         args=(resources, test_env, proc_func, args, kwargs,
-              self._use_dill_for_args))
+              self._use_dill_for_args),
+        daemon=self._daemon)
     p.start()
     self._processes[(task_type, task_id)] = p
     self._outstanding_subprocess_count += 1
@@ -568,7 +571,6 @@ class _ProcFunc(object):
         time.sleep(0.1)
     self._resources.process_status_queue.put(
         _ProcessStatusInfo(
-            task_type=task_type,
             is_successful=True,
             exc_info=None,
             return_value=None))
@@ -628,17 +630,9 @@ class _ProcFunc(object):
     if test_env.v2_enabled:
       v2_compat.enable_v2_behavior()
 
-    try:
-      with self._runtime_mode(test_env.executing_eagerly):
-        return_value = proc_func(*args, **kwargs)
-        is_successful = True
-        exc_info = None
-
-    except Exception:  # pylint: disable=broad-except
-      # Capture all exceptions to be reported to parent process.
-      return_value = None
-      is_successful = False
-      exc_info = sys.exc_info()
+    with self._runtime_mode(test_env.executing_eagerly):
+      info = _run_contained(proc_func, args, kwargs)
+      self._resources.process_status_queue.put(info)
 
       # Re-raise the exception in addition to reporting it to the parent
       # process, so that even if `--test_timeout` flag is set and the
@@ -647,18 +641,183 @@ class _ProcFunc(object):
       # instead of silently suppressing the error due to early bazel
       # timeout. Raising an error in the subprocess produces stack trace in
       # the log, but the program continues running.
-      raise
+      if not info.is_successful:
+        six.reraise(*info.exc_info)
 
-    finally:
-      info = _ProcessStatusInfo(
-          task_type=test_env.task_type,
-          is_successful=is_successful,
-          exc_info=exc_info,
-          return_value=return_value)
-      self._resources.process_status_queue.put(info)
       self._close_streaming()
 
 
+class MultiProcessPoolRunner(object):
+  """A utility class to start a process pool to simulate a cluster.
+
+  It's similar to MultiProcessRunner, but uses a pool of processes to avoid the
+  expensive initialization cost of Tensorflow.
+  """
+
+  def __init__(self, cluster_spec, initializer=None):
+    """Creates a multi-process pool runner.
+
+    Args:
+      cluster_spec: Dict for cluster spec. The following is an example of
+        cluster with three workers.
+        {"worker": ["worker0.example.com:2222",
+                    "worker1.example.com:2222",
+                    "worker2.example.com:2222"]}
+      initializer: a callable to called at the startup of worker processes.
+
+    Raises:
+      RuntimeError: if `multi_process_runner.test_main()` is not called.
+      ValueError: if there are more than one chief in the `cluster_spec`.
+    """
+    self._cluster_spec = cluster_spec
+    self._initializer = initializer
+    self._conn = {}
+    self._runner = None
+
+  def __del__(self):
+    self._reset()
+
+  def _reset(self):
+    for conn in self._conn.values():
+      conn.close()
+    self._conn = {}
+    if self._runner is not None:
+      self._runner.join()
+      self._runner = None
+
+  def _start(self):
+    """Starts the worker pool."""
+    # We need different arguments for different processes so we're passing a
+    # no-op proc_func here and use start_single_process instead.
+    #
+    # We also need to start the process pool as daemon, so that they don't block
+    # the program from exiting. Note that __del__ may not get called when
+    # there's an exception. The user may also store a pool runner in a global
+    # object to share across test cases
+    self._runner = MultiProcessRunner(
+        proc_func=lambda: None,
+        cluster_spec=self._cluster_spec,
+        use_dill_for_args=False,
+        daemon=True)
+    if self._initializer:
+      initializer = dill.dumps(self._initializer, dill.HIGHEST_PROTOCOL)
+    else:
+      initializer = None
+    for task_type, addresses in self._cluster_spec.items():
+      for task_id, _ in enumerate(addresses):
+        conn1, conn2 = multiprocessing.Pipe(duplex=True)
+        self._conn[(task_type, task_id)] = conn1
+        self._runner.start_single_process(
+            task_type,
+            task_id,
+            proc_func=_pool_runner_worker,
+            args=(initializer, conn2))
+
+  def run(self, proc_func, args=None, kwargs=None):
+    """Runs `proc_func` with `args` and `kwargs` on all jobs.
+
+    Args:
+      proc_func: The function to be run.
+      args: Optional positional arguments to be supplied in `proc_func`.
+      kwargs: Optional keyword arguments to be supplied in `proc_func`.
+
+    Returns:
+      A list of return values.
+    """
+    if self._runner is None:
+      self._start()
+
+    # Since we start the processes as daemon they're going to be killed by
+    # SIGTERM when the program exits. We only turn on streaming during run() to
+    # avoid printing the stacktrace caused by the SIGTERM.
+    self._runner._stream_stdout = True  # pylint: disable=protected-access
+
+    try:
+      proc_func = dill.dumps(proc_func, dill.HIGHEST_PROTOCOL)
+      for conn in self._conn.values():
+        conn.send((proc_func, args or [], kwargs or {}))
+
+      process_statuses = []
+      for (task_type, task_id), conn in self._conn.items():
+        logging.info('Waiting for the result from %s-%d', task_type, task_id)
+        try:
+          process_statuses.append(conn.recv())
+        except EOFError:
+          # This shouldn't happen due to exceptions in proc_func. This usually
+          # means bugs in the runner.
+          self._reset()
+          raise RuntimeError('Unexpected EOF. Worker process may have died. '
+                             'Please report a bug')
+
+      return_values = []
+      for process_status in process_statuses:
+        assert isinstance(process_status, _ProcessStatusInfo)
+        if not process_status.is_successful:
+          six.reraise(*process_status.exc_info)
+        if process_status.return_value is not None:
+          return_values.append(process_status.return_value)
+
+      return return_values
+    finally:
+      self._runner._stream_stdout = False  # pylint: disable=protected-access
+
+
+def _pool_runner_worker(initializer, conn):
+  """Function that runs on the workers in a pool.
+
+  It listens for callables to run and returns the result until `conn` is closed.
+  It captures the exceptions during executing the callable and return it through
+  `conn`.
+
+  Args:
+    initializer: A callable to execute during startup.
+    conn: A multiprocessing.Connection object to listen for tasks and send
+      results.
+  """
+  if initializer:
+    initializer = dill.loads(initializer)
+    initializer()
+  while True:
+    try:
+      proc_func, args, kwargs = conn.recv()
+    except EOFError:
+      break
+    proc_func = dill.loads(proc_func)
+    info = _run_contained(proc_func, args, kwargs)
+    sys.stdout.flush()
+    sys.stderr.flush()
+    conn.send(info)
+
+
+def _run_contained(proc_func, args, kwargs):
+  """Runs `proc_func` with `args` and `kwargs`.
+
+  The function returns _ProcessStatusInfo which captures the return value and
+  the exception.
+
+  Args:
+    proc_func: The function to be run.
+    args: Optional positional arguments to be supplied in `proc_func`.
+    kwargs: Optional keyword arguments to be supplied in `proc_func`.
+
+  Returns:
+    a _ProcessStatusInfo.
+  """
+  try:
+    return_value = proc_func(*args, **kwargs)
+    is_successful = True
+    exc_info = None
+  except Exception:  # pylint: disable=broad-except
+    return_value = None
+    is_successful = False
+    exc_info = sys.exc_info()
+  finally:
+    return _ProcessStatusInfo(  # pylint: disable=lost-exception
+        is_successful=is_successful,
+        exc_info=exc_info,
+        return_value=return_value)
+
+
 class SubprocessTimeoutError(RuntimeError):
   """An error that indicates there is at least one subprocess timing out.
 
diff --git a/tensorflow/python/distribute/multi_process_runner_test.py b/tensorflow/python/distribute/multi_process_runner_test.py
index aeba43b6b7c..d76ef5a5a3c 100644
--- a/tensorflow/python/distribute/multi_process_runner_test.py
+++ b/tensorflow/python/distribute/multi_process_runner_test.py
@@ -22,6 +22,8 @@ import json
 import os
 import threading
 import time
+import unittest
+
 from absl import logging
 
 from tensorflow.python.distribute import multi_process_runner
@@ -45,7 +47,7 @@ def proc_func_that_adds_simple_return_data():
   return 'dummy_data'
 
 
-def proc_func_that_return_args_and_kwargs(*args, **kwargs):
+def proc_func_that_returns_args_and_kwargs(*args, **kwargs):
   return list(args) + list(kwargs.items())
 
 
@@ -53,6 +55,20 @@ def proc_func_with_barrier():
   return multi_process_runner.barrier()
 
 
+def proc_func_that_returns_pid():
+  return os.getpid()
+
+
+V = None
+
+
+def proc_func_that_sets_global(val):
+  global V
+  old_val = V
+  V = val
+  return old_val
+
+
 class MultiProcessRunnerTest(test.TestCase):
 
   def _worker_idx(self):
@@ -95,7 +111,7 @@ class MultiProcessRunnerTest(test.TestCase):
 
   def test_multi_process_runner_args_passed_correctly(self):
     return_value = multi_process_runner.run(
-        proc_func_that_return_args_and_kwargs,
+        proc_func_that_returns_args_and_kwargs,
         multi_worker_test_base.create_cluster_spec(num_workers=1),
         args=('a', 'b'),
         kwargs={
@@ -325,5 +341,54 @@ class MultiProcessRunnerTest(test.TestCase):
                 for line in list_to_assert))
 
 
+class MultiProcessPoolRunnerTest(test.TestCase):
+
+  def test_same_process_across_runs(self):
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
+    pid = runner.run(proc_func_that_returns_pid)
+    for _ in range(3):
+      self.assertAllEqual(runner.run(proc_func_that_returns_pid), pid)
+
+  def test_exceptions_in_sub_process(self):
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
+    pid = runner.run(proc_func_that_returns_pid)
+    with self.assertRaisesRegexp(ValueError, 'This is an error.'):
+      runner.run(proc_func_that_errors)
+    self.assertAllEqual(runner.run(proc_func_that_returns_pid), pid)
+
+  def test_tf_config(self):
+    cluster_spec = multi_worker_test_base.create_cluster_spec(
+        has_chief=True, num_workers=2)
+    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
+    result = runner.run(proc_func_that_adds_task_type_in_return_data)
+
+    job_count_dict = {'worker': 2, 'chief': 1}
+    for data in result:
+      job_count_dict[data] -= 1
+
+    self.assertEqual(job_count_dict['worker'], 0)
+    self.assertEqual(job_count_dict['chief'], 0)
+
+  @unittest.expectedFailure
+  def test_exception_in_main_process(self):
+    # When there's an exception in the main process, __del__() is not called.
+    # This test is to verify MultiProcessPoolRunner can cope with __del__() not
+    # being called.
+    cluster_spec = multi_worker_test_base.create_cluster_spec(
+        has_chief=True, num_workers=2)
+    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
+    runner.run(proc_func_that_returns_pid)
+    raise ValueError('failure')
+
+  def test_initializer(self):
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    runner = multi_process_runner.MultiProcessPoolRunner(
+        cluster_spec, initializer=lambda: proc_func_that_sets_global(1))
+    result = runner.run(proc_func_that_sets_global, args=(2,))
+    self.assertAllEqual(result, [1, 1])
+
+
 if __name__ == '__main__':
   multi_process_runner.test_main()

From 965b93f2ee2ecf6dc7152adc31a76802b0ede85f Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:05:14 -0400
Subject: [PATCH 0566/1390] test pathlib path for save weights

---
 tensorflow/python/keras/saving/save_test.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 66364666841..5b5da8c5047 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -119,6 +119,20 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
       save.save_model(self.model, path, save_format='tf')
       save.load_model(path)
 
+  @test_util.run_v2_only
+  def test_save_load_weights_tf_pathlib(self):
+    if sys.version_info >= (3, 6):
+      path = pathlib.Path(self.get_temp_dir()) / 'model'
+      self.model.save_weights(path, save_format='tf')
+      self.model.load_weights(path)
+
+  @test_util.run_v2_only
+  def test_save_load_weights_hdf5_pathlib(self):
+    if sys.version_info >= (3, 6):
+      path = pathlib.Path(self.get_temp_dir()) / 'model'
+      self.model.save_weights(path, save_format='h5')
+      self.model.load_weights(path)
+
   @combinations.generate(combinations.combine(mode=['graph', 'eager']))
   def test_saving_with_dense_features(self):
     cols = [

From fd3c95a3b4eb1418d6f87bed916ec17fff44adf0 Mon Sep 17 00:00:00 2001
From: jonah-kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Thu, 18 Jun 2020 17:07:30 -0700
Subject: [PATCH 0567/1390] Include show_dtype param in plot_model tests.

---
 tensorflow/python/keras/utils/vis_utils_test.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/keras/utils/vis_utils_test.py b/tensorflow/python/keras/utils/vis_utils_test.py
index 984014216be..0eb38a1d48f 100644
--- a/tensorflow/python/keras/utils/vis_utils_test.py
+++ b/tensorflow/python/keras/utils/vis_utils_test.py
@@ -36,7 +36,8 @@ class ModelToDotFormatTest(test.TestCase):
     model.add(keras.layers.Dense(5, name='dense'))
     dot_img_file = 'model_1.png'
     try:
-      vis_utils.plot_model(model, to_file=dot_img_file, show_shapes=True)
+      vis_utils.plot_model(model, to_file=dot_img_file, 
+                           show_shapes=True, show_dtype=True)
       self.assertTrue(file_io.file_exists(dot_img_file))
       file_io.delete_file(dot_img_file)
     except ImportError:
@@ -62,7 +63,8 @@ class ModelToDotFormatTest(test.TestCase):
     dot_img_file = 'model_2.png'
     try:
       vis_utils.plot_model(
-          model, to_file=dot_img_file, show_shapes=True, expand_nested=True)
+          model, to_file=dot_img_file, show_shapes=True,
+          show_dtype=True expand_nested=True)
       self.assertTrue(file_io.file_exists(dot_img_file))
       file_io.delete_file(dot_img_file)
     except ImportError:
@@ -76,7 +78,8 @@ class ModelToDotFormatTest(test.TestCase):
     dot_img_file = 'model_3.png'
     try:
       vis_utils.plot_model(
-          model, to_file=dot_img_file, show_shapes=True, expand_nested=True)
+          model, to_file=dot_img_file, show_shapes=True,
+          show_dtype=True, expand_nested=True)
       self.assertTrue(file_io.file_exists(dot_img_file))
       file_io.delete_file(dot_img_file)
     except ImportError:
@@ -88,7 +91,8 @@ class ModelToDotFormatTest(test.TestCase):
     dot_img_file = 'model_4.png'
     try:
       vis_utils.plot_model(
-          model, to_file=dot_img_file, show_shapes=True, expand_nested=True)
+          model, to_file=dot_img_file, show_shapes=True,
+          show_dtype=True expand_nested=True)
       self.assertTrue(file_io.file_exists(dot_img_file))
       file_io.delete_file(dot_img_file)
     except ImportError:

From ae76bc79213d4559b113899f438cf54283ec11c2 Mon Sep 17 00:00:00 2001
From: Marat Dukhan 
Date: Thu, 18 Jun 2020 17:19:14 -0700
Subject: [PATCH 0568/1390] Update XNNPACK dependency and document sparse
 inference capability

PiperOrigin-RevId: 317213816
Change-Id: I35431b40fd63d836d4fe979f65a71a181c0c820d
---
 tensorflow/lite/delegates/xnnpack/README.md | 37 +++++++++++++++++++++
 tensorflow/workspace.bzl                    |  8 ++---
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/README.md b/tensorflow/lite/delegates/xnnpack/README.md
index 97d2d5565db..d94e92c7306 100644
--- a/tensorflow/lite/delegates/xnnpack/README.md
+++ b/tensorflow/lite/delegates/xnnpack/README.md
@@ -238,6 +238,43 @@ Below is the list of current operators and limitations:
 * Fused `NONE`, `RELU`, `RELU_N1_TO_1`, and `RELU6` activations are supported,
   but fused `TANH` and `SIGN_BIT` activations are not.
 
+### Sparse Inference (experimental)
+
+XNNPACK backend supports sparse inference for CNN models described in the
+[Fast Sparse ConvNets](https://arxiv.org/abs/1911.09723) paper. This
+functionality must be enabled at build-time via
+`--define xnn_enable_sparse=true` Bazel flag. Sparse inference is restricted
+to subgraphs with the following operators:
+
+* Sparse subgraph must start with a 3x3 stride-2 `CONV_2D` operator with
+  padding 1 on each side, no dilation, and 3 input channels.
+* Sparse subgraph must end with a `MEAN` operator that does reduction across
+  spatial axes.
+* Sparse subgraph may contain the following operators:
+  * `CONV_2D` with 1x1 kernel and no padding. It is important to have high
+    sparsity (at least 70%) in the filter of this operator to get speedup
+    over dense inference.
+  * `DEPTHWISE_CONV_2D` with 3x3 kernel, stride 1, no dilation, and padding 1
+    on each side.
+  * `DEPTHWISE_CONV_2D` with 3x3 kernel, stride 2, no dilation, and padding 1
+    on each side.
+  * `DEPTHWISE_CONV_2D` with 5x5 kernel, stride 1, no dilation, and padding 2
+    on each side.
+  * `DEPTHWISE_CONV_2D` with 5x5 kernel, stride 2, no dilation, and padding 2
+    on each side.
+  * `ADD` and `MUL` operators where both inputs are 4D tensors. If one of the
+    inputs to `ADD` or `MUL` is a constant tensor, it must be representable as
+    either a scalar, or a 1D vector.
+  * Unary elementwise operators `ABS`, `CEIL`, `FLOOR`, `HARD_SWISH`,
+    `LEAKY_RELU`, `LOGISTIC`, `NEG`, `RELU`, `RELU6`, `RELU_N1_TO_1`, `ROUND`,
+    and `SQUARE`.
+
+Pre-trained [Fast Sparse ConvNets models](https://github.com/google-research/google-research/tree/master/fastconvnets)
+provide examples that satisfy these constrains.
+
+In addition to acceleration, sparse models get the compression benefit by
+storing only non-zero values in the [TensorFlow Lite file format](https://github.com/tensorflow/tensorflow/blob/4aea552e064cf92330e07e83a3b5a1ca2a7034d0/tensorflow/lite/schema/schema.fbs#L84-L109).
+
 ### Other limitations
 
 * Dynamically allocated (with `kTfLiteDynamic` allocation type) inputs and
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index f5b0b7537dc..52c573628ac 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -164,11 +164,11 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
 
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "4af883fea0a6ada106867f29670a6c0b7af74bee85d74a2e04356a670814a3d4",
-        strip_prefix = "XNNPACK-69a6a7667d96a84c596b0f4e00632b2037c17723",
+        sha256 = "2527a30464b43bd03f137b2c455a0381e49eae63d09cfeee128a717dfbe962d5",
+        strip_prefix = "XNNPACK-8b283aa30a3186c6e640aed520543e9c067132d2",
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/69a6a7667d96a84c596b0f4e00632b2037c17723.zip",
-            "https://github.com/google/XNNPACK/archive/69a6a7667d96a84c596b0f4e00632b2037c17723.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google/XNNPACK/archive/8b283aa30a3186c6e640aed520543e9c067132d2.zip",
+            "https://github.com/google/XNNPACK/archive/8b283aa30a3186c6e640aed520543e9c067132d2.zip",
         ],
     )
 

From 55e33877b8f81b5ede6365456ed3aa89bbe16d8e Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:37:52 -0400
Subject: [PATCH 0569/1390] Revert "test pathlib path for save weights"

This reverts commit 965b93f2ee2ecf6dc7152adc31a76802b0ede85f.
---
 tensorflow/python/keras/saving/save_test.py | 14 --------------
 1 file changed, 14 deletions(-)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 5b5da8c5047..66364666841 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -119,20 +119,6 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
       save.save_model(self.model, path, save_format='tf')
       save.load_model(path)
 
-  @test_util.run_v2_only
-  def test_save_load_weights_tf_pathlib(self):
-    if sys.version_info >= (3, 6):
-      path = pathlib.Path(self.get_temp_dir()) / 'model'
-      self.model.save_weights(path, save_format='tf')
-      self.model.load_weights(path)
-
-  @test_util.run_v2_only
-  def test_save_load_weights_hdf5_pathlib(self):
-    if sys.version_info >= (3, 6):
-      path = pathlib.Path(self.get_temp_dir()) / 'model'
-      self.model.save_weights(path, save_format='h5')
-      self.model.load_weights(path)
-
   @combinations.generate(combinations.combine(mode=['graph', 'eager']))
   def test_saving_with_dense_features(self):
     cols = [

From edc445eb184633f0ca10270bca6e9c0b834ab454 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:38:02 -0400
Subject: [PATCH 0570/1390] Revert "only test pathlib on python v>3.6"

This reverts commit c3ca4da9b46e3148bdb913d08fec3fd2727158e1.
---
 tensorflow/python/keras/saving/save_test.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 66364666841..98d78735ad1 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -73,10 +73,9 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
 
   @test_util.run_v2_only
   def test_save_format_defaults_pathlib(self):
-    if sys.version_info >= (3, 6):
-      path = pathlib.Path(self.get_temp_dir()) / 'model_path'
-      save.save_model(self.model, path)
-      self.assert_saved_model(path)
+    path = pathlib.Path(self.get_temp_dir()) / 'model_path'
+    save.save_model(self.model, path)
+    self.assert_saved_model(path)
 
   @test_util.run_v2_only
   def test_save_hdf5(self):

From 10fb2155fb720f9e0e70d9e48a934383b4b42c91 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:38:08 -0400
Subject: [PATCH 0571/1390] Revert "modify docstring"

This reverts commit 5f2e0240ee7977042e41d9c29c349a7b14301290.
---
 tensorflow/python/saved_model/loader_impl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index 06cd988130d..71f6ed16c9d 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -73,7 +73,7 @@ def parse_saved_model(export_dir):
   """Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`.
 
   Args:
-    export_dir: String or Pathlike, path to the directory containing the SavedModel file.
+    export_dir: Directory containing the SavedModel file.
 
   Returns:
     A `SavedModel` protocol buffer.

From 9f0e739f5f1623b74c1835e9209c7c25b4c49380 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:38:15 -0400
Subject: [PATCH 0572/1390] Revert "convert export_dir as pathlike object to
 str in parse_saved_model"

This reverts commit 1c72a6c65e1733b55286c6361142a39d699732dc.
---
 tensorflow/python/saved_model/loader_impl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/saved_model/loader_impl.py b/tensorflow/python/saved_model/loader_impl.py
index 71f6ed16c9d..2df2bea428e 100644
--- a/tensorflow/python/saved_model/loader_impl.py
+++ b/tensorflow/python/saved_model/loader_impl.py
@@ -83,11 +83,11 @@ def parse_saved_model(export_dir):
   """
   # Build the path to the SavedModel in pbtxt format.
   path_to_pbtxt = os.path.join(
-      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(export_dir),
       compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
   # Build the path to the SavedModel in pb format.
   path_to_pb = os.path.join(
-      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(export_dir),
       compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
 
   # Parse the SavedModel protocol buffer.

From 5869f657624ece35544c3fe54f999b2e3f449305 Mon Sep 17 00:00:00 2001
From: Yixing Fu 
Date: Thu, 18 Jun 2020 20:38:24 -0400
Subject: [PATCH 0573/1390] Revert "add save load test for pathlib path"

This reverts commit 1363f0f6e89cdc73ca46d43475bdb35750ed2e50.
---
 tensorflow/python/keras/saving/save_test.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/tensorflow/python/keras/saving/save_test.py b/tensorflow/python/keras/saving/save_test.py
index 98d78735ad1..5c5846fe738 100644
--- a/tensorflow/python/keras/saving/save_test.py
+++ b/tensorflow/python/keras/saving/save_test.py
@@ -71,12 +71,6 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
     save.save_model(self.model, path)
     self.assert_saved_model(path)
 
-  @test_util.run_v2_only
-  def test_save_format_defaults_pathlib(self):
-    path = pathlib.Path(self.get_temp_dir()) / 'model_path'
-    save.save_model(self.model, path)
-    self.assert_saved_model(path)
-
   @test_util.run_v2_only
   def test_save_hdf5(self):
     path = os.path.join(self.get_temp_dir(), 'model')
@@ -87,13 +81,6 @@ class TestSaveModel(test.TestCase, parameterized.TestCase):
         'requires the model to be a Functional model or a Sequential model.'):
       save.save_model(self.subclassed_model, path, save_format='h5')
 
-  @test_util.run_v2_only
-  def test_save_load_hdf5_pathlib(self):
-    if sys.version_info >= (3, 6):
-      path = pathlib.Path(self.get_temp_dir()) / 'model'
-      save.save_model(self.model, path, save_format='h5')
-      save.load_model(path)
-
   @test_util.run_v2_only
   def test_save_tf(self):
     path = os.path.join(self.get_temp_dir(), 'model')

From 723751b20ef2aa0a4af39cad2581fd483ae78ad7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 18:15:12 -0700
Subject: [PATCH 0574/1390] Add set_outfeed_config in XLA HloInstruction.

PiperOrigin-RevId: 317222410
Change-Id: I5de8a5067f1002a9d656d4e26d145ffe3fe372ed
---
 tensorflow/compiler/xla/service/hlo_instruction.cc | 4 ++++
 tensorflow/compiler/xla/service/hlo_instruction.h  | 3 +++
 tensorflow/compiler/xla/service/hlo_instructions.h | 1 +
 3 files changed, 8 insertions(+)

diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index cfa21b95dd2..6de76c1cc63 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -3908,6 +3908,10 @@ const string& HloInstruction::outfeed_config() const {
   return Cast(this)->outfeed_config();
 }
 
+void HloInstruction::set_outfeed_config(const string& config) {
+  return Cast(this)->set_outfeed_config(config);
+}
+
 const std::vector& HloInstruction::replica_groups() const {
   return Cast(this)->replica_groups();
 }
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 7a5d506b681..f3bb59ff625 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -1755,6 +1755,9 @@ class HloInstruction {
   // Returns the config for the Outfeed instruction.
   const string& outfeed_config() const;
 
+  // Delegates to HloOutfeedInstruction::set_outfeed_config.
+  void set_outfeed_config(const string& config);
+
   // Returns the shape for the Outfeed instruction.
   const Shape& outfeed_shape() const;
 
diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h
index 6da01dc088e..f5a963ef063 100644
--- a/tensorflow/compiler/xla/service/hlo_instructions.h
+++ b/tensorflow/compiler/xla/service/hlo_instructions.h
@@ -1141,6 +1141,7 @@ class HloOutfeedInstruction : public HloInstruction {
   const Shape& outfeed_shape() const { return outfeed_shape_; }
   // Returns the config for the Outfeed instruction.
   const string& outfeed_config() const { return outfeed_config_; }
+  void set_outfeed_config(const string& config) { outfeed_config_ = config; }
   // Returns a serialized representation of this instruction.
   HloInstructionProto ToProto() const override;
 

From 50763f6e3db1685193c53f0e19405f87f1bcc3b4 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Fri, 19 Jun 2020 01:21:19 +0000
Subject: [PATCH 0575/1390] Extend typing to variable types, Annotate params
 with defaults

---
 tensorflow/python/framework/python_op_gen.cc | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 79c8800418c..1ef3f9c342b 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -361,7 +361,7 @@ string GenEagerPythonOp::Code() {
 
     // Add type annotations to param
     if (type_map.find(param.GetName()) != type_map.end()) {
-      if(!type_map[param.GetName()].empty()) {
+      if (!type_map[param.GetName()].empty()) {
         strings::StrAppend(¶meters, ": ", type_map[param.GetName()]);
       }
     }
@@ -372,6 +372,19 @@ string GenEagerPythonOp::Code() {
     if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
     if (!parameters_with_defaults.empty())
       strings::StrAppend(¶meters_with_defaults, ", ");
+
+    // Add type annotations to param_and_default
+    if (type_map.find(param_and_default.first.GetName()) != type_map.end()) {
+      if (!type_map[param_and_default.first.GetName()].empty()) {
+        strings::StrAppend(¶meters, ": ", type_map[param_and_default.first.GetName()]);
+        strings::StrAppend(¶meters_with_defaults,
+                           param_and_default.first.GetRenameTo(), ": ",
+                           type_map[param_and_default.first.GetName()], " ",
+                           "= ", param_and_default.second);
+        continue;
+      }
+    }
+
     strings::StrAppend(¶meters, param_and_default.first.GetRenameTo());
     strings::StrAppend(¶meters_with_defaults,
                        param_and_default.first.GetRenameTo(), "=",
@@ -427,6 +440,9 @@ std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
     if (attr.type() == "type") {
       const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
       type_map[attr.name()] = type_var_name;
+    } else if (attr.type() == "bool" || attr.type() == "float" ||
+               attr.type() == "int" || attr.type() == "bytes") {
+      type_map[attr.name()] = attr.type();
     }
   }
 
@@ -507,7 +523,7 @@ void GenEagerPythonOp::GenerateTypeVars() {
     }
   }
 
-  if(added_typevar) strings::StrAppend(&result_, "\n");
+  if (added_typevar) strings::StrAppend(&result_, "\n");
 }
 
 void GenEagerPythonOp::AddReturnTypeAnnotation(std::unordered_map& type_map) {

From 0bef067291d26baf7cde98e7d3d6fbcf54161d58 Mon Sep 17 00:00:00 2001
From: peng 
Date: Fri, 19 Jun 2020 09:24:20 +0800
Subject: [PATCH 0576/1390] fix test error with no GPU config

---
 tensorflow/core/grappler/optimizers/remapper_test.cc | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/remapper_test.cc b/tensorflow/core/grappler/optimizers/remapper_test.cc
index a947210d8a6..56f31cb49b3 100644
--- a/tensorflow/core/grappler/optimizers/remapper_test.cc
+++ b/tensorflow/core/grappler/optimizers/remapper_test.cc
@@ -485,11 +485,6 @@ TEST_F(RemapperTest, FuseConv2DWithBiasAndActivationOnGPU) {
   item.feed = {{"input", input_t}, {"filter", filter_t}, {"bias", bias_t}};
   TF_ASSERT_OK(s.ToGraphDef(&item.graph));
 
-  // Place all nodes on GPU.
-  for (int i = 0; i < item.graph.node_size(); ++i) {
-    item.graph.mutable_node(i)->set_device("/device:GPU:0");
-  }
-
   Remapper optimizer(RewriterConfig::AGGRESSIVE);  // trust placeholders shape
   //Remapper optimizer(RewriterConfig::ON);
   GraphDef output;

From 4e5c26fb53d6e6467ebbb5a65a73af8872d383d3 Mon Sep 17 00:00:00 2001
From: Jonah Kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Thu, 18 Jun 2020 18:40:17 -0700
Subject: [PATCH 0577/1390] Revert commit 3833402 for modular PRs

---
 tensorflow/python/keras/optimizer_v2/optimizer_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
index d8992bbe3e0..c55b332bfc0 100644
--- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
+++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py
@@ -910,7 +910,7 @@ class OptimizerV2(trackable.Trackable):
       return value()
     if tensor_util.is_tensor(value):
       return backend.get_value(value)
-    return float(value)
+    return value
 
   def variables(self):
     """Returns variables of this Optimizer based on the order created."""

From 8e70fe45468fd1c0b845814907bdf7d930e30f46 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Fri, 19 Jun 2020 02:09:23 +0000
Subject: [PATCH 0578/1390] Generate TypeVar name once, Change input & attr
 enumeration

---
 tensorflow/python/framework/python_op_gen.cc | 21 +++++++++-----------
 1 file changed, 9 insertions(+), 12 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 1ef3f9c342b..4ecfccb611f 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -180,10 +180,10 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
 
   void AddRawOpExport(const string& parameters);
 
-  void GenerateTypeVars();
-
   std::unordered_map GetTypeAnnotationMap();
 
+  void GenerateTypeVars(std::unordered_map& type_map);
+
   void AddReturnTypeAnnotation(std::unordered_map& type_map);
 
   void AddAttrForArg(const string& attr, int arg_index) {
@@ -435,8 +435,7 @@ string GenEagerPythonOp::Code() {
 std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
   std::unordered_map type_map;
   // Mapping attrs to TypeVars
-  for (int i = 0; i GenEagerPythonOp::GetTypeAnnotationMap() {
   }
 
   // Mapping input Tensors to their types
-  for (int i = 0; i < op_def_.input_arg_size(); ++i) {
-    const auto& arg = op_def_.input_arg(i);
+  for (const auto& arg : op_def_.input_arg()) {
     // Do not add type annotations to args that accept a sequence of tensors
     if (!arg.number_attr().empty()) continue;
     string type_annotation;
@@ -466,7 +464,7 @@ std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
     type_map[arg.name()] = type_annotation;
   }
 
-  // Mapping output Tensor to its types
+  // Mapping output Tensor to its type
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
     string type_annotation;
@@ -488,10 +486,9 @@ std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
 }
 
 // Generate TypeVars using attrs
-void GenEagerPythonOp::GenerateTypeVars() {
+void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type_map) {
   bool added_typevar = false;
-  for (int i = 0; i allowed_types;
       for (int t : attr.allowed_values().list().type()) {
@@ -517,7 +514,7 @@ void GenEagerPythonOp::GenerateTypeVars() {
         strings::StrAppend(&typevar_dtypes, *it);
       }
 
-      const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
+      const string type_var_name = type_map[attr.name()];
       strings::StrAppend(&result_, type_var_name, " = TypeVar(\"", type_var_name, "\", ", typevar_dtypes,")\n");
       added_typevar = true;
     }
@@ -863,7 +860,7 @@ bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
     const string& parameters, const std::vector& output_sizes,
     const string& eager_not_allowed_error, std::unordered_map& type_map) {
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    GenerateTypeVars();
+    GenerateTypeVars(type_map);
   }
   if (api_def_.visibility() == ApiDef::VISIBLE) {
     strings::StrAppend(&result_, "@_dispatch.add_dispatch_list\n");

From b7edd44ee0f8c264e457c48138474f6e1bf5b18e Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Thu, 18 Jun 2020 19:08:04 -0700
Subject: [PATCH 0579/1390] Enable type annotations for python/ops.

PiperOrigin-RevId: 317229132
Change-Id: I7055e650308c2fc83969385dd25e86fb5b073d75
---
 tensorflow/python/ops/logging_ops.py | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/tensorflow/python/ops/logging_ops.py b/tensorflow/python/ops/logging_ops.py
index 8ca63f55987..02fce277690 100644
--- a/tensorflow/python/ops/logging_ops.py
+++ b/tensorflow/python/ops/logging_ops.py
@@ -54,11 +54,9 @@ except NameError:
 # call relies on certain conditionals for its dependencies.  Use
 # control_flow_ops.Assert.
 
-# Assert and Print are special symbols in python, so we must
-# have an upper-case version of them.
-#
-# For users with Python 3 or Python 2.7
-# with `from __future__ import print_function`, we could also allow lowercase.
+# Assert and Print are special symbols in Python 2, so we must
+# have an upper-case version of them. When support for it is dropped,
+# we can allow lowercase.
 # See https://github.com/tensorflow/tensorflow/issues/18053
 
 
@@ -83,11 +81,6 @@ def Print(input_, data, message=None, first_n=None, summarize=None, name=None):
     with jupyter notebook (printing to the notebook *server's* output, not into
     the notebook).
 
-  Additionally, to use tf.print in python 2.7, users must make sure to import
-  the following:
-
-  `from __future__ import print_function`
-
   Args:
     input_: A tensor passed through this op.
     data: A list of tensors to print out when op is evaluated.
@@ -148,11 +141,6 @@ def print_v2(*inputs, **kwargs):
   Python objects. Printed tensors will recursively show the first and last
   elements of each dimension to summarize.
 
-  @compatibility(python2)
-  In python 2.7, make sure to import the following:
-  `from __future__ import print_function`
-  @end_compatibility
-
   Example:
     Single-input usage:
 

From 13fe5862de7b95fd91aeec8f2d71e9f2e77b699b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 19:11:15 -0700
Subject: [PATCH 0580/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/c830d517b4e4

PiperOrigin-RevId: 317229564
Change-Id: I10163c3e668996252d294018794081394cc0d25c
---
 third_party/mlir/test.BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/mlir/test.BUILD b/third_party/mlir/test.BUILD
index 23287ce28d6..14c2ba7778e 100644
--- a/third_party/mlir/test.BUILD
+++ b/third_party/mlir/test.BUILD
@@ -166,6 +166,7 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:StandardOps",
+        "@llvm-project//mlir:StandardOpsTransforms",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TargetNVVMIR",
         "@llvm-project//mlir:TargetROCDLIR",

From 4a14e778d64853a236941259693aa3c5813c18d8 Mon Sep 17 00:00:00 2001
From: David Majnemer 
Date: Thu, 18 Jun 2020 19:11:40 -0700
Subject: [PATCH 0581/1390] [XLA] Introduce ManifestCheckingTest

PiperOrigin-RevId: 317229603
Change-Id: Ibcc9ea3895d520024f5d80d52330aeb3b970585d
---
 tensorflow/compiler/xla/tests/BUILD           |  23 +++-
 tensorflow/compiler/xla/tests/build_defs.bzl  |   7 +-
 .../xla/tests/client_library_test_base.h      |   3 +-
 tensorflow/compiler/xla/tests/hlo_test_base.h |   3 +-
 .../xla/tests/local_client_test_base.h        |   3 +-
 .../xla/tests/manifest_checking_test.cc       | 129 ++++++++++++++++++
 .../xla/tests/manifest_checking_test.h        |  35 +++++
 tensorflow/compiler/xla/tests/test_macros.cc  |  89 +-----------
 tensorflow/compiler/xla/tests/test_macros.h   | 118 +---------------
 9 files changed, 201 insertions(+), 209 deletions(-)
 create mode 100644 tensorflow/compiler/xla/tests/manifest_checking_test.cc
 create mode 100644 tensorflow/compiler/xla/tests/manifest_checking_test.h

diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index e1863a8a4cf..9b36117602b 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -52,16 +52,26 @@ cc_library(
     name = "test_macros_header",
     testonly = True,
     hdrs = ["test_macros.h"],
-    deps = [
-        "//tensorflow/compiler/xla:types",
-        "//tensorflow/core:test",
-        "@com_google_absl//absl/strings",
-    ],
 )
 
 # Generate a test_macros_${BACKEND} library per backend with the proper copts.
 generate_backend_test_macros()
 
+cc_library(
+    name = "manifest_checking_test",
+    testonly = True,
+    srcs = ["manifest_checking_test.cc"],
+    hdrs = ["manifest_checking_test.h"],
+    deps = [
+        ":test_macros_header",
+        "//tensorflow/core:regexp_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:logging",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
 cc_library(
     name = "test_utils",
     srcs = ["test_utils.cc"],
@@ -136,6 +146,7 @@ cc_library(
     hdrs = ["hlo_test_base.h"],
     deps = [
         ":literal_test_util",
+        ":manifest_checking_test",
         ":test_utils",
         ":verified_hlo_module",
         "//tensorflow/compiler/xla:debug_options_flags",
@@ -193,6 +204,7 @@ cc_library(
     srcs = ["client_library_test_base.cc"],
     hdrs = ["client_library_test_base.h"],
     deps = [
+        ":manifest_checking_test",
         "//tensorflow/compiler/xla:array2d",
         "//tensorflow/compiler/xla:array3d",
         "//tensorflow/compiler/xla:array4d",
@@ -273,6 +285,7 @@ cc_library(
     hdrs = ["local_client_test_base.h"],
     deps = [
         ":client_library_test_base",
+        ":manifest_checking_test",
         ":verified_hlo_module",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:status_macros",
diff --git a/tensorflow/compiler/xla/tests/build_defs.bzl b/tensorflow/compiler/xla/tests/build_defs.bzl
index c0c0751b0de..94d870aa2ef 100644
--- a/tensorflow/compiler/xla/tests/build_defs.bzl
+++ b/tensorflow/compiler/xla/tests/build_defs.bzl
@@ -266,11 +266,6 @@ def generate_backend_test_macros(backends = []):
                 "-DXLA_DISABLED_MANIFEST=\\\"%s\\\"" % manifest,
             ],
             deps = [
-                "@com_google_absl//absl/container:flat_hash_map",
-                "@com_google_absl//absl/strings",
-                "//tensorflow/compiler/xla:types",
-                "//tensorflow/core:lib",
-                "//tensorflow/core:regexp_internal",
-                "//tensorflow/core:test",
+                "//tensorflow/core/platform:logging",
             ],
         )
diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h
index 790497f888e..17bb70bdb42 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.h
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.h
@@ -35,6 +35,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/tests/manifest_checking_test.h"
 #include "tensorflow/compiler/xla/tests/test_utils.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/core/bitmap.h"
@@ -62,7 +63,7 @@ std::vector ExpandUseBfloat16(
 }
 
 // A client library test establishes an in-process XLA client connection.
-class ClientLibraryTestBase : public ::testing::Test {
+class ClientLibraryTestBase : public ManifestCheckingTest {
  protected:
   explicit ClientLibraryTestBase(se::Platform* platform = nullptr);
 
diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h
index 85b1876dd3c..17c2a55ba5b 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.h
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.h
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/shape_layout.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/tests/manifest_checking_test.h"
 #include "tensorflow/compiler/xla/tests/verified_hlo_module.h"
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
@@ -67,7 +68,7 @@ namespace xla {
 //  )
 //
 // For a more detailed example, see "../tests/sample_text_test.cc".
-class HloTestBase : public ::testing::Test {
+class HloTestBase : public ManifestCheckingTest {
  public:
   // Creates a new HLO module for a test. The module created will have
   // TestName() for its name; it will also automatically populate its debug
diff --git a/tensorflow/compiler/xla/tests/local_client_test_base.h b/tensorflow/compiler/xla/tests/local_client_test_base.h
index ea457024618..c1951ad1021 100644
--- a/tensorflow/compiler/xla/tests/local_client_test_base.h
+++ b/tensorflow/compiler/xla/tests/local_client_test_base.h
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/transfer_manager.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/tests/client_library_test_base.h"
+#include "tensorflow/compiler/xla/tests/manifest_checking_test.h"
 #include "tensorflow/compiler/xla/tests/verified_hlo_module.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/platform/mutex.h"
@@ -75,7 +76,7 @@ class TestAllocator : public se::StreamExecutorMemoryAllocator {
 };
 
 // A base class for tests which exercise the LocalClient interface.
-class LocalClientTestBase : public ::testing::Test {
+class LocalClientTestBase : public ManifestCheckingTest {
  protected:
   struct EigenThreadPoolWrapper;
   explicit LocalClientTestBase(se::Platform* platform = nullptr);
diff --git a/tensorflow/compiler/xla/tests/manifest_checking_test.cc b/tensorflow/compiler/xla/tests/manifest_checking_test.cc
new file mode 100644
index 00000000000..8806290472d
--- /dev/null
+++ b/tensorflow/compiler/xla/tests/manifest_checking_test.cc
@@ -0,0 +1,129 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/tests/manifest_checking_test.h"
+
+#include 
+#include 
+#include 
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_split.h"
+#include "tensorflow/compiler/xla/tests/test_macros.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/regexp.h"
+
+namespace xla {
+
+namespace {
+
+// Mapping from test name; i.e. MyTest.MyTestCase to platforms on which it is
+// disabled - a sequence of regexps.
+using ManifestT = absl::flat_hash_map>;
+
+ManifestT ReadManifest() {
+  ManifestT manifest;
+
+  absl::string_view path = absl::NullSafeStringView(kDisabledManifestPath);
+  if (path.empty()) {
+    return manifest;
+  }
+
+  // Note: parens are required to disambiguate vs function decl.
+  std::ifstream file_stream((std::string(path)));
+  std::string contents((std::istreambuf_iterator(file_stream)),
+                       std::istreambuf_iterator());
+
+  std::vector lines = absl::StrSplit(contents, '\n');
+  for (std::string& line : lines) {
+    auto comment = line.find("//");
+    if (comment != std::string::npos) {
+      line = line.substr(0, comment);
+    }
+    if (line.empty()) {
+      continue;
+    }
+    absl::StripTrailingAsciiWhitespace(&line);
+    std::vector pieces = absl::StrSplit(line, ' ');
+    CHECK_GE(pieces.size(), 1);
+    auto& platforms = manifest[pieces[0]];
+    for (size_t i = 1; i < pieces.size(); ++i) {
+      platforms.push_back(pieces[i]);
+    }
+  }
+  return manifest;
+}
+
+}  // namespace
+
+void ManifestCheckingTest::SetUp() {
+  const testing::TestInfo* test_info =
+      testing::UnitTest::GetInstance()->current_test_info();
+  absl::string_view test_case_name = test_info->test_suite_name();
+  absl::string_view test_name = test_info->name();
+  VLOG(1) << "test_case_name: " << test_case_name;
+  VLOG(1) << "test_name: " << test_name;
+
+  // Remove the type suffix from the test case name.
+  if (const char* type_param = test_info->type_param()) {
+    VLOG(1) << "type_param: " << type_param;
+    size_t last_slash = test_case_name.rfind('/');
+    test_case_name = test_case_name.substr(0, last_slash);
+    VLOG(1) << "test_case_name: " << test_case_name;
+  }
+
+  // Remove the test instantiation name if it is present.
+  auto first_slash = test_case_name.find('/');
+  if (first_slash != test_case_name.npos) {
+    test_case_name.remove_prefix(first_slash + 1);
+    VLOG(1) << "test_case_name: " << test_case_name;
+  }
+
+  ManifestT manifest = ReadManifest();
+
+  // If the test name ends with a slash followed by one or more characters,
+  // strip that off.
+  auto last_slash = test_name.rfind('/');
+  if (last_slash != test_name.npos) {
+    test_name = test_name.substr(0, last_slash);
+    VLOG(1) << "test_name: " << test_name;
+  }
+
+  // First try full match: test_case_name.test_name
+  // If that fails, try to find just the test_case_name; this would disable all
+  // tests in the test case.
+  auto it = manifest.find(absl::StrCat(test_case_name, ".", test_name));
+  if (it == manifest.end()) {
+    it = manifest.find(test_case_name);
+    if (it == manifest.end()) {
+      return;
+    }
+  }
+
+  // Expect a full match vs. one of the platform regexps to disable the test.
+  const std::vector& disabled_platforms = it->second;
+  auto platform_string = kTestPlatform;
+  for (const auto& s : disabled_platforms) {
+    if (RE2::FullMatch(/*text=*/platform_string, /*re=*/s)) {
+      GTEST_SKIP();
+      return;
+    }
+  }
+
+  // We didn't hit in the disabled manifest entries, so don't disable it.
+}
+
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/manifest_checking_test.h b/tensorflow/compiler/xla/tests/manifest_checking_test.h
new file mode 100644
index 00000000000..4f44ed76a3e
--- /dev/null
+++ b/tensorflow/compiler/xla/tests/manifest_checking_test.h
@@ -0,0 +1,35 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
+#define TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
+
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+
+// This class allows us to intercept the test name and use an arbitrary
+// heuristic to decide whether the test case should be disabled. We
+// determine whether the test case should be disabled by resolving the (test
+// case name, test name) in a manifest file.
+class ManifestCheckingTest : public ::testing::Test {
+ protected:
+  // This method runs before each test runs.
+  void SetUp() override;
+};
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
diff --git a/tensorflow/compiler/xla/tests/test_macros.cc b/tensorflow/compiler/xla/tests/test_macros.cc
index dc9ac7b684a..9e85af76e89 100644
--- a/tensorflow/compiler/xla/tests/test_macros.cc
+++ b/tensorflow/compiler/xla/tests/test_macros.cc
@@ -15,93 +15,18 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/tests/test_macros.h"
 
-#include 
-#include 
-#include 
-
-#include "absl/container/flat_hash_map.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_split.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/regexp.h"
 
 namespace xla {
-namespace {
 
-// Mapping from test name; i.e. MyTest.MyTestCase to platforms on which it is
-// disabled - a sequence of regexps.
-using ManifestT = absl::flat_hash_map>;
-
-ManifestT ReadManifest() {
-  ManifestT manifest;
-
-  string path = XLA_DISABLED_MANIFEST;
-  if (path.empty()) {
-    return manifest;
-  }
-
-  std::ifstream file_stream(path);
-  // Note: parens are required to disambiguate vs function decl.
-  string contents((std::istreambuf_iterator(file_stream)),
-                  std::istreambuf_iterator());
-
-  std::vector lines = absl::StrSplit(contents, '\n');
-  for (string& line : lines) {
-    auto comment = line.find("//");
-    if (comment != string::npos) {
-      line = line.substr(0, comment);
-    }
-    if (line.empty()) {
-      continue;
-    }
-    absl::StripTrailingAsciiWhitespace(&line);
-    std::vector pieces = absl::StrSplit(line, ' ');
-    CHECK_GE(pieces.size(), 1);
-    auto& platforms = manifest[pieces[0]];
-    for (int64 i = 1; i < pieces.size(); ++i) {
-      platforms.push_back(pieces[i]);
-    }
-  }
-  return manifest;
+static bool InitModule() {
+  kDisabledManifestPath = XLA_DISABLED_MANIFEST;
+  VLOG(1) << "kDisabledManifestPath: " << kDisabledManifestPath;
+  kTestPlatform = XLA_PLATFORM;
+  VLOG(1) << "kTestPlatform: " << kTestPlatform;
+  return false;
 }
 
-}  // namespace
-
-std::string PrependDisabledIfIndicated(absl::string_view test_case_name,
-                                       absl::string_view test_name) {
-  ManifestT manifest = ReadManifest();
-
-  // If the test name ends with a slash followed by one or more digits, strip
-  // that off; this is just a shard number, and matching on this would be
-  // unstable even if someone wanted to do it.
-  static LazyRE2 shard_num_pattern = {R"(/\d+$)"};
-  absl::string_view suffix;
-  if (RE2::PartialMatch(test_name, *shard_num_pattern, &suffix)) {
-    test_name.remove_suffix(suffix.size());
-  }
-
-  // First try full match: test_case_name.test_name
-  // If that fails, try to find just the test_case_name; this would disable all
-  // tests in the test case.
-  auto it = manifest.find(absl::StrCat(test_case_name, ".", test_name));
-  if (it == manifest.end()) {
-    it = manifest.find(test_case_name);
-    if (it == manifest.end()) {
-      return std::string(test_name);
-    }
-  }
-
-  // Expect a full match vs. one of the platform regexps to disable the test.
-  const std::vector& disabled_platforms = it->second;
-  string platform_string = XLA_PLATFORM;
-  for (const auto& s : disabled_platforms) {
-    if (RE2::FullMatch(/*text=*/platform_string, /*re=*/s)) {
-      return absl::StrCat("DISABLED_", test_name);
-    }
-  }
-
-  // We didn't hit in the disabled manifest entries, so don't disable it.
-  return std::string(test_name);
-}
+static bool module_initialized = InitModule();
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h
index 33d2dff9721..f62bccbe850 100644
--- a/tensorflow/compiler/xla/tests/test_macros.h
+++ b/tensorflow/compiler/xla/tests/test_macros.h
@@ -28,12 +28,6 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_
 #define TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_
 
-#include 
-
-#include "absl/strings/string_view.h"
-#include "tensorflow/compiler/xla/types.h"
-#include "tensorflow/core/platform/test.h"
-
 #define DISABLED_ON_CPU(X) X
 #define DISABLED_ON_GPU(X) X
 #define DISABLED_ON_GPU_ROCM(X) X
@@ -79,117 +73,15 @@ limitations under the License.
 
 namespace xla {
 
-// Reads a disabled manifest file to resolve whether test cases should be
-// disabled on a particular platform. For a test that should be disabled,
-// returns DISABLED_ prepended to its name; otherwise returns the test name
-// unmodified.
-std::string PrependDisabledIfIndicated(absl::string_view test_case_name,
-                                       absl::string_view test_name);
+inline const char *kDisabledManifestPath = nullptr;
+inline const char *kTestPlatform = nullptr;
 
 }  // namespace xla
 
-// This is the internal "gtest" class instantiation -- it is identical to the
-// GTEST_TEST_ macro, except that we intercept the test name for potential
-// modification by PrependDisabledIfIndicated. That file can use an arbitrary
-// heuristic to decide whether the test case should be disabled, and we
-// determine whether the test case should be disabled by resolving the (test
-// case name, test name) in a manifest file.
-#define XLA_GTEST_TEST_(test_case_name, test_name, parent_class)             \
-  class GTEST_TEST_CLASS_NAME_(test_case_name, test_name)                    \
-      : public parent_class {                                                \
-   public:                                                                   \
-    GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {}                   \
-                                                                             \
-   private:                                                                  \
-    virtual void TestBody();                                                 \
-    static ::testing::TestInfo* const test_info_ GTEST_ATTRIBUTE_UNUSED_;    \
-    GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name,   \
-                                                           test_name));      \
-  };                                                                         \
-                                                                             \
-  ::testing::TestInfo* const GTEST_TEST_CLASS_NAME_(test_case_name,          \
-                                                    test_name)::test_info_ = \
-      ::testing::RegisterTest(                                               \
-          #test_case_name,                                                   \
-          ::xla::PrependDisabledIfIndicated(#test_case_name, #test_name)     \
-              .c_str(),                                                      \
-          nullptr, nullptr, __FILE__, __LINE__, []() -> parent_class* {      \
-            return new GTEST_TEST_CLASS_NAME_(test_case_name, test_name)();  \
-          });                                                                \
-  void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody()
+#define XLA_TEST_F(test_fixture, test_name) TEST_F(test_fixture, test_name)
 
-// This is identical to the TEST_F macro from "gtest", but it potentially
-// disables the test based on an external manifest file, DISABLED_MANIFEST.
-//
-// Per usual, you can see what tests are available via --gunit_list_tests and
-// choose to run tests that have been disabled via the manifest via
-// --gunit_also_run_disabled_tests.
-#define XLA_TEST_F(test_fixture, test_name) \
-  XLA_GTEST_TEST_(test_fixture, test_name, test_fixture)
+#define XLA_TEST_P(test_case_name, test_name) TEST_P(test_case_name, test_name)
 
-// Likewise, this is identical to the TEST_P macro from "gtest", but
-// potentially disables the test based on the DISABLED_MANIFEST file.
-//
-// We have to wrap this in an outer layer so that any DISABLED_ON_* macros will
-// be properly expanded before the stringification occurs.
-#define XLA_TEST_P_IMPL_(test_case_name, test_name)                            \
-  class GTEST_TEST_CLASS_NAME_(test_case_name, test_name)                      \
-      : public test_case_name {                                                \
-   public:                                                                     \
-    GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {}                     \
-    virtual void TestBody();                                                   \
-                                                                               \
-   private:                                                                    \
-    static int AddToRegistry() {                                               \
-      ::testing::UnitTest::GetInstance()                                       \
-          ->parameterized_test_registry()                                      \
-          .GetTestCasePatternHolder(                           \
-              #test_case_name,                                                 \
-              ::testing::internal::CodeLocation(__FILE__, __LINE__))           \
-          ->AddTestPattern(                                                    \
-              #test_case_name,                                                 \
-              ::xla::PrependDisabledIfIndicated(#test_case_name, #test_name)   \
-                  .c_str(),                                                    \
-              new ::testing::internal::TestMetaFactory());                              \
-      return 0;                                                                \
-    }                                                                          \
-    static int gtest_registering_dummy_ GTEST_ATTRIBUTE_UNUSED_;               \
-    GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name,     \
-                                                           test_name));        \
-  };                                                                           \
-  int GTEST_TEST_CLASS_NAME_(test_case_name,                                   \
-                             test_name)::gtest_registering_dummy_ =            \
-      GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::AddToRegistry();      \
-  void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody()
-
-#define XLA_TEST_P(test_case_name, test_name) \
-  XLA_TEST_P_IMPL_(test_case_name, test_name)
-
-// This is identical to the TEST_F macro from "gtest", but it potentially
-// disables the test based on an external manifest file, DISABLED_MANIFEST.
-#define XLA_TYPED_TEST(CaseName, TestName)                                     \
-  template                                          \
-  class GTEST_TEST_CLASS_NAME_(CaseName, TestName)                             \
-      : public CaseName {                                    \
-   private:                                                                    \
-    typedef CaseName TestFixture;                            \
-    typedef gtest_TypeParam_ TypeParam;                                        \
-    virtual void TestBody();                                                   \
-  };                                                                           \
-  bool gtest_##CaseName##_##TestName##_registered_ GTEST_ATTRIBUTE_UNUSED_ =   \
-      ::testing::internal::TypeParameterizedTest<                              \
-          CaseName,                                                            \
-          ::testing::internal::TemplateSel,  \
-          GTEST_TYPE_PARAMS_(CaseName)>::                                      \
-          Register(                                                            \
-              "", ::testing::internal::CodeLocation(__FILE__, __LINE__),       \
-              #CaseName,                                                       \
-              ::xla::PrependDisabledIfIndicated(#CaseName, #TestName).c_str(), \
-              0);                                                              \
-  template                                          \
-  void GTEST_TEST_CLASS_NAME_(CaseName,                                        \
-                              TestName)::TestBody()
+#define XLA_TYPED_TEST(CaseName, TestName) TYPED_TEST(CaseName, TestName)
 
 #endif  // TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_

From 9c4b749b09b958c436e0681a4276b47fc9316a8a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 19:17:40 -0700
Subject: [PATCH 0582/1390] Internal change

PiperOrigin-RevId: 317230321
Change-Id: I043dae37768f6e9cf946d4db2a8c36123ed2d6d9
---
 tensorflow/core/platform/BUILD               |   7 -
 tensorflow/core/platform/tf32_utils.cc       |  30 --
 tensorflow/core/platform/tf32_utils.h        |  27 --
 tensorflow/python/BUILD                      |  11 -
 tensorflow/python/framework/config.py        |  31 ---
 tensorflow/python/util/tf32.cc               |  22 --
 tensorflow/stream_executor/cuda/BUILD        |   2 -
 tensorflow/stream_executor/cuda/cuda_blas.cc |  98 ++++---
 tensorflow/stream_executor/cuda/cuda_blas.h  |   8 +-
 tensorflow/stream_executor/cuda/cuda_dnn.cc  | 272 ++++++++-----------
 10 files changed, 172 insertions(+), 336 deletions(-)
 delete mode 100644 tensorflow/core/platform/tf32_utils.cc
 delete mode 100644 tensorflow/core/platform/tf32_utils.h
 delete mode 100644 tensorflow/python/util/tf32.cc

diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD
index 33a1e7cfe0a..70bb8a89417 100644
--- a/tensorflow/core/platform/BUILD
+++ b/tensorflow/core/platform/BUILD
@@ -938,13 +938,6 @@ cc_library(
     alwayslink = 1,
 )
 
-cc_library(
-    name = "tf32_utils",
-    srcs = ["tf32_utils.cc"],
-    hdrs = ["tf32_utils.h"],
-    copts = tf_copts(),
-)
-
 tf_cc_tests(
     name = "low_level_library_tests",
     size = "small",
diff --git a/tensorflow/core/platform/tf32_utils.cc b/tensorflow/core/platform/tf32_utils.cc
deleted file mode 100644
index d2f40ea161a..00000000000
--- a/tensorflow/core/platform/tf32_utils.cc
+++ /dev/null
@@ -1,30 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "tensorflow/core/platform/tf32_utils.h"
-
-#include 
-
-namespace tensorflow {
-
-// Whether TensorFloat-32 should be used where supported.
-// TODO(nluehr): Maybe enable by default after TF32 Ampere testing.
-static std::atomic tf32_allowed{false};
-
-void allow_tf32_execution(bool allowed) { tf32_allowed = allowed; }
-
-bool tf32_execution_allowed() { return tf32_allowed; }
-
-}  // namespace tensorflow
diff --git a/tensorflow/core/platform/tf32_utils.h b/tensorflow/core/platform/tf32_utils.h
deleted file mode 100644
index 7a158d00ad3..00000000000
--- a/tensorflow/core/platform/tf32_utils.h
+++ /dev/null
@@ -1,27 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_
-#define TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_
-
-namespace tensorflow {
-
-void allow_tf32_execution(bool allowed);
-
-bool tf32_execution_allowed();
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 5f9e2dfb1ff..de9cf9a24c7 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -788,16 +788,6 @@ tf_python_pybind_extension(
     ],
 )
 
-tf_python_pybind_extension(
-    name = "_pywrap_tf32_execution",
-    srcs = ["util/tf32.cc"],
-    module_name = "_pywrap_tf32_execution",
-    deps = [
-        "//tensorflow/core/platform:tf32_utils",
-        "@pybind11",
-    ],
-)
-
 tf_python_pybind_extension(
     name = "_pywrap_util_port",
     srcs = ["util/port_wrapper.cc"],
@@ -5688,7 +5678,6 @@ py_library(
         "//tensorflow:composite_tensor_whitelist",
     ],
     deps = [
-        ":_pywrap_tf32_execution",
         ":tf_decorator",
         ":tf_export",
         ":tf_stack",
diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py
index 544b6882618..9ff16f2a327 100644
--- a/tensorflow/python/framework/config.py
+++ b/tensorflow/python/framework/config.py
@@ -18,42 +18,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python import _pywrap_tf32_execution
 from tensorflow.python.eager import context
 from tensorflow.python.util import deprecation
 from tensorflow.python.util.tf_export import tf_export
 
 
-# No tf_export until TF is built against CUDA11 which is required for TF32.
-def tensor_float_32_execution_allowed():
-  """Get if TensorFloat-32 operations are enabled on supported hardware.
-
-  Returns:
-    True if TensorFloat-32 execution is enabled and False otherwise.
-  """
-  return _pywrap_tf32_execution.is_allowed()
-
-
-# No tf_export until TF is built against CUDA 11 which is required for TF32.
-def allow_tensor_float_32_execution(allowed):
-  """Allow use of TensorFloat-32 with float32 ops on supported hardware.
-
-  TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture.
-  TensorFloat-32 kernels take float32 inputs and produce float32 outputs.
-  Internally, the inputs are cast to a custom representation with 10-bit
-  mantissa (similar to float16) and 8-bit exponent (similar to float32) and are
-  executed using TensorCores with float32 accumulation. For more information,
-  see https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/.
-
-  TensorFloat-32 execution is disabled by default, but this may change in a
-  future version.
-
-  Args:
-    allowed: whether to allow TensorFloat-32 execution
-  """
-  _pywrap_tf32_execution.allow(allowed)
-
-
 @tf_export('config.threading.get_intra_op_parallelism_threads')
 def get_intra_op_parallelism_threads():
   """Get number of threads used within an individual op for parallelism.
diff --git a/tensorflow/python/util/tf32.cc b/tensorflow/python/util/tf32.cc
deleted file mode 100644
index 7dece6ccdae..00000000000
--- a/tensorflow/python/util/tf32.cc
+++ /dev/null
@@ -1,22 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "pybind11/pybind11.h"
-#include "tensorflow/core/platform/tf32_utils.h"
-
-PYBIND11_MODULE(_pywrap_tf32_execution, m) {
-  m.def("allow", &tensorflow::allow_tf32_execution);
-  m.def("is_allowed", &tensorflow::tf32_execution_allowed);
-}
diff --git a/tensorflow/stream_executor/cuda/BUILD b/tensorflow/stream_executor/cuda/BUILD
index 3a14be9ad50..c3cf9f5db15 100644
--- a/tensorflow/stream_executor/cuda/BUILD
+++ b/tensorflow/stream_executor/cuda/BUILD
@@ -251,7 +251,6 @@ cc_library(
         "@local_config_cuda//cuda:cuda_headers",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-        "//tensorflow/core/platform:tf32_utils",
         "//tensorflow/stream_executor",
         "//tensorflow/stream_executor:event",
         "//tensorflow/stream_executor:host_or_device_scalar",
@@ -357,7 +356,6 @@ cc_library(
         "@local_config_cuda//cuda:cudnn_header",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
-        "//tensorflow/core/platform:tf32_utils",
         "//tensorflow/stream_executor:dnn",
         "//tensorflow/stream_executor:event",
         "//tensorflow/stream_executor:plugin_registry",
diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index fcd0e7b16fb..c9f0fc462c9 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -49,7 +49,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "third_party/eigen3/Eigen/Core"
-#include "tensorflow/core/platform/tf32_utils.h"
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/cuda/cuda_activation.h"
 #include "tensorflow/stream_executor/cuda/cuda_gpu_executor.h"
@@ -102,6 +101,18 @@ static std::string ToString(cublasStatus_t status) {
   }
 }
 
+// Decide whether to enable TENSOR_OP_MATH
+static bool TensorOpMathEnabled() {
+  static bool is_enabled = [] {
+    bool is_disabled;
+    TF_CHECK_OK(
+        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUBLAS_TENSOR_OP_MATH",
+                                       /*default_val=*/false, &is_disabled));
+    return !is_disabled;
+  }();
+  return is_enabled;
+}
+
 // cuBLAS has interfaces that permit pointers to be passed from either the host
 // memory space or the device memory space; however, you must instruct it as to
 // which address space those pointers are in with cublasSetPointerMode.
@@ -226,19 +237,6 @@ bool CUDABlas::Init() {
     return false;
   }
 
-  absl::MutexLock lock(&mu_);
-#if CUDA_VERSION >= 9000
-#if CUBLAS_VER_MAJOR >= 11
-  ret = cublasSetMathMode(blas_, CUBLAS_TF32_TENSOR_OP_MATH);
-#else
-  ret = cublasSetMathMode(blas_, CUBLAS_TENSOR_OP_MATH);
-#endif
-  if (ret != CUBLAS_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set cublas default math mode: " << ToString(ret);
-    return false;
-  }
-#endif
-
   return true;
 }
 
@@ -401,7 +399,7 @@ cudaDataType_t CUDAComputationType(blas::ComputationType ty) {
 template 
 bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                                   bool pointer_mode_host, bool err_on_failure,
-                                  Args... args) {
+                                  bool use_tensor_op_math, Args... args) {
   absl::MutexLock lock(&mu_);
 
   CHECK(blas_ != nullptr);
@@ -415,10 +413,10 @@ bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                                            : CUBLAS_POINTER_MODE_DEVICE)) {
     return false;
   }
-#if CUBLAS_VER_MAJOR >= 11
+#if CUDA_VERSION >= 9000
   ScopedCublasMathMode math_mode{blas_};
-  if (!tensorflow::tf32_execution_allowed()) {
-    if (!math_mode.Init(CUBLAS_DEFAULT_MATH)) {
+  if (use_tensor_op_math) {
+    if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) {
       return false;
     }
   }
@@ -1635,9 +1633,21 @@ bool CUDABlas::DoBlasGemm(
     }
   }
 
+  bool use_tensor_ops = false;
+#if CUDA_VERSION >= 9000
+  int cc_major, cc_minor;
+  stream->parent()->GetDeviceDescription().cuda_compute_capability(&cc_major,
+                                                                   &cc_minor);
+
+  // GPUs < sm_70 don't support tensor ops.
+  if (cc_major >= 7 && TensorOpMathEnabled()) {
+    use_tensor_ops = true;
+  }
+#endif
+
   return DoBlasInternalImpl(
       cublasSgemmEx, stream, true /* = pointer_mode_host */,
-      true /* = err_on_failure= */, CUDABlasTranspose(transa),
+      true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa),
       CUDABlasTranspose(transb), m, n, k, &alpha, GpuMemory(a),
       SE_CUDA_DATA_HALF, lda, GpuMemory(b), SE_CUDA_DATA_HALF, ldb, &beta,
       GpuMemoryMutable(c), SE_CUDA_DATA_HALF, ldc);
@@ -1911,7 +1921,8 @@ static bool TensorOpsAvailable(int cc_major) {
   // strictly correct.  We can't simply enable it, though, as that would change
   // clients' behavior significantly: Using tensor ops on fp32 inputs cause them
   // to be rounded to fp16.
-  if (cc_major >= 7 && std::is_same::value) {
+  if (cc_major >= 7 && TensorOpMathEnabled() &&
+      std::is_same::value) {
     return true;
   }
 #endif
@@ -2259,8 +2270,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal(
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
           &cc_major, &cc_minor) &&
       cc_major >= 5) {
-    bool use_tensor_ops =
-        data_type == CUDA_R_16F || tensorflow::tf32_execution_allowed();
+    bool use_tensor_ops = TensorOpMathEnabled() && data_type == CUDA_R_16F;
     cublasGemmAlgo_t algo =
         (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
     cudaDataType_t compute_type =
@@ -2274,7 +2284,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal(
     bool ok;
     ok = DoBlasInternalImpl(
         AS_LAMBDA(cublasGemmBatchedEx), stream, true /* = pointer_mode_host */,
-        true /* = err_on_failure */, CUDABlasTranspose(transa),
+        true /* = err_on_failure */, use_tensor_ops, CUDABlasTranspose(transa),
         CUDABlasTranspose(transb), m, n, k, &alpha, a_void_ptrs, data_type, lda,
         b_void_ptrs, data_type, ldb, &beta, c_void_ptrs, data_type, ldc,
         batch_count, compute_type, algo);
@@ -2409,25 +2419,33 @@ bool CUDABlas::DoBlasGemmStridedBatched(
     int lda, int64 stride_a, const DeviceMemory &b, int ldb,
     int64 stride_b, float beta, DeviceMemory *c, int ldc,
     int64 stride_c, int batch_count) {
-#if CUDA_VERSION >= 9010
+  bool use_tensor_ops = false;
+#if CUDA_VERSION >= 9000
   int cc_major, cc_minor;
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
-          &cc_major, &cc_minor) &&
-      cc_major >= 5) {
-    cublasGemmAlgo_t algo =
-        (cc_major >= 7 ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
-    bool ok = DoBlasInternalImpl(
-        AS_LAMBDA(cublasGemmStridedBatchedEx), stream,
-        true /* = pointer_mode_host */, true /* = err_on_failure */,
-        CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha,
-        GpuMemory(a), CUDA_R_16F, lda, stride_a, GpuMemory(b), CUDA_R_16F, ldb,
-        stride_b, &beta, GpuMemoryMutable(c), CUDA_R_16F, ldc, stride_c,
-        batch_count, CUDA_R_32F, algo);
-    if (ok) {
-      return true;
+          &cc_major, &cc_minor)) {
+    // GPUs < sm_70 don't support tensor ops.
+    if (cc_major >= 7 && TensorOpMathEnabled()) {
+      use_tensor_ops = true;
     }
-    LOG(ERROR) << "failed BLAS call, see log for details";
-    return false;
+#if CUDA_VERSION >= 9010
+    if (cc_major >= 5) {
+      cublasGemmAlgo_t algo =
+          (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
+      bool ok = DoBlasInternalImpl(
+          AS_LAMBDA(cublasGemmStridedBatchedEx), stream,
+          true /* = pointer_mode_host */, true /* = err_on_failure */,
+          use_tensor_ops, CUDABlasTranspose(transa), CUDABlasTranspose(transb),
+          m, n, k, &alpha, GpuMemory(a), CUDA_R_16F, lda, stride_a,
+          GpuMemory(b), CUDA_R_16F, ldb, stride_b, &beta, GpuMemoryMutable(c),
+          CUDA_R_16F, ldc, stride_c, batch_count, CUDA_R_32F, algo);
+      if (ok) {
+        return true;
+      }
+      LOG(ERROR) << "failed BLAS call, see log for details";
+      return false;
+    }
+#endif
   }
 #endif
   // Either CUDA_VERSION < 9.1 or SM < 5.0. Fall back to a loop.
@@ -2440,7 +2458,7 @@ bool CUDABlas::DoBlasGemmStridedBatched(
         reinterpret_cast<__half *>(GpuMemoryMutable(c) + batch * stride_c);
     bool ok = DoBlasInternalImpl(
         cublasSgemmEx, stream, true /* = pointer_mode_host */,
-        true /* = err_on_failure= */, CUDABlasTranspose(transa),
+        true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa),
         CUDABlasTranspose(transb), m, n, k, &alpha, a_matrix, SE_CUDA_DATA_HALF,
         lda, b_matrix, SE_CUDA_DATA_HALF, ldb, &beta, c_matrix,
         SE_CUDA_DATA_HALF, ldc);
diff --git a/tensorflow/stream_executor/cuda/cuda_blas.h b/tensorflow/stream_executor/cuda/cuda_blas.h
index 556456c83db..817bdb72777 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.h
+++ b/tensorflow/stream_executor/cuda/cuda_blas.h
@@ -83,7 +83,7 @@ class CUDABlas : public blas::BlasSupport {
   template 
   bool DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                           bool pointer_mode_host, bool err_on_failure,
-                          Args... args);
+                          bool use_tensor_op_math, Args... args);
 
   // Convenience functions that call DoBlasInternalImpl with different values
   // for err_on_failure.
@@ -91,7 +91,8 @@ class CUDABlas : public blas::BlasSupport {
   bool DoBlasInternal(FuncT cublas_func, Stream *stream, bool pointer_mode_host,
                       Args... args) {
     return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host,
-                              /*err_on_failure=*/true, args...);
+                              /*err_on_failure=*/true, /*use_tensor_ops=*/false,
+                              args...);
   }
   template 
   bool DoBlasInternalFailureOK(FuncT cublas_func, Stream *stream,
@@ -99,7 +100,8 @@ class CUDABlas : public blas::BlasSupport {
     // Tensor ops are hard-coded off in this path, but can still be enabled with
     // a specific algorithm choice as in DoBlasGemmWithAlgorithmImpl().
     return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host,
-                              /*err_on_failure=*/false, args...);
+                              /*err_on_failure=*/false,
+                              /*use_tensor_ops=*/false, args...);
   }
 
   // A helper function to implement DoBlasGemmBatched interfaces for generic
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 192bae91572..be18c989861 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -22,7 +22,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "third_party/eigen3/Eigen/Core"
 #include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/platform/tf32_utils.h"
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/cuda/cuda_activation.h"
 #include "tensorflow/stream_executor/cuda/cuda_diagnostics.h"
@@ -602,6 +601,31 @@ class CudnnFilterDescriptor {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnFilterDescriptor);
 };
 
+// A helper function to decide whether to enable the TENSOR_OP_MATH math type
+bool TensorOpMathEnabled() {
+  static bool is_enabled = [] {
+    bool is_disabled = false;
+    TF_CHECK_OK(
+        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_TENSOR_OP_MATH",
+                                       /*default_val=*/false, &is_disabled));
+    return !is_disabled;
+  }();
+  return is_enabled;
+}
+
+// A helper function to decide whether to enable the TENSOR_OP_MATH math type
+// for RNNs.
+bool RnnTensorOpMathEnabled() {
+  static bool is_enabled = [] {
+    bool is_disabled = false;
+    TF_CHECK_OK(
+        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH",
+                                       /*default_val=*/false, &is_disabled));
+    return !is_disabled;
+  }();
+  return is_enabled;
+}
+
 // A helper function to decide whether to use
 // CUDNN_BATCHNORM_SPATIAL_PERSISTENT in batchnorm. This mode can be faster in
 // some tasks because an optimized path may be selected for CUDNN_DATA_FLOAT
@@ -706,6 +730,10 @@ class CudnnConvolutionDescriptor {
             : CUDNN_CROSS_CORRELATION,
         data_type));
 
+    // NOTE(benbarsdell): This only applies if tensor op math is enabled
+    //                      and algo selection is set to Default.
+    this->set_use_tensor_op_math(true);
+
 #if CUDNN_MAJOR >= 7
     VLOG(2) << "Requesting grouped convolution: "
             << convolution_descriptor.group_count();
@@ -717,15 +745,13 @@ class CudnnConvolutionDescriptor {
 #endif
   }
 
-  void set_use_tensor_op_math(bool use_tensor_op_math) {
+  void set_use_tensor_op_math(bool use_tensor_op_math) const {
 #if CUDNN_VERSION >= 7000
     cudnnMathType_t math_type =
-#if CUDNN_VERSION >= 8000
-        (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH);
-#else
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
-#endif
-    CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
+    if (TensorOpMathEnabled()) {
+      CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
+    }
 #endif
   }
 
@@ -737,40 +763,6 @@ class CudnnConvolutionDescriptor {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnConvolutionDescriptor);
 };
 
-// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math
-// set
-static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) {
-  cudnnMathType_t math_type;
-  CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type));
-#if CUDNN_VERSION >= 8000
-  return math_type != CUDNN_FMA_MATH;
-#else
-  return math_type == CUDNN_TENSOR_OP_MATH;
-#endif
-}
-
-static bool TensorOpMathAvailable(int cc_major) {
-  return cc_major >= 7 && CUDNN_VERSION >= 7000;
-}
-
-static bool IsTensorMathAllowed(Stream* stream, dnn::DataType input_type) {
-  int cc_major, cc_minor;
-  std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
-  if (!TensorOpMathAvailable(cc_major)) {
-    return false;
-  }
-  if (input_type == dnn::DataType::kFloat) {
-#if CUDNN_VERSION < 8000
-    return false;
-#else
-    if (!tensorflow::tf32_execution_allowed()) {
-      return false;
-    }
-#endif
-  }
-  return true;
-}
-
 // Turns a PoolingDescriptor structure into a cudnn pooling descriptor handle
 // within a scope.
 class CudnnPoolingDescriptor {
@@ -1163,31 +1155,21 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
     // in profile mode, which is run with algorithms returned from
     // GetRnnAlgorithms() (which are non-default and explicitly set whether to
     // use tensor ops). CuDNN 7.2.1 fixed this issue
-    bool allow_tensor_ops =
-        data_type != CUDNN_DATA_FLOAT || tensorflow::tf32_execution_allowed();
-    bool use_tensor_ops;
-    if (algorithm_config.algorithm().has_value()) {
-      use_tensor_ops = algorithm_config.algorithm()->tensor_ops_enabled();
-    } else {
-      use_tensor_ops = CUDNN_VERSION >= 7201 && allow_tensor_ops;
-    }
-
-    if (use_tensor_ops && !allow_tensor_ops) {
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "Algo requests disallowed tensor op evaluation.");
-    }
-
-    cudnnMathType_t math_type;
-    if (use_tensor_ops) {
-      math_type = CUDNN_TENSOR_OP_MATH;
-    } else {
-#if CUDNN_VERSION >= 8000
-      math_type = CUDNN_FMA_MATH;
+    if (RnnTensorOpMathEnabled()) {
+      cudnnMathType_t math_type;
+      if (algorithm_config.algorithm().has_value()) {
+        math_type = algorithm_config.algorithm()->tensor_ops_enabled()
+                        ? CUDNN_TENSOR_OP_MATH
+                        : CUDNN_DEFAULT_MATH;
+      } else {
+#if CUDNN_VERSION >= 7201
+        math_type = CUDNN_TENSOR_OP_MATH;
 #else
-      math_type = CUDNN_DEFAULT_MATH;
-#endif  // CUDNN_VERSION >= 8000
+        math_type = CUDNN_DEFAULT_MATH;
+#endif  // CUDNN_VERSION >= 7201
+      }
+      CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
     }
-    CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
 #endif  // CUDNN_VERSION >= 7000
 
     return CudnnRnnDescriptor(cudnn, std::move(rnn_desc), std::move(rnn_plan),
@@ -2578,11 +2560,10 @@ port::StatusOr> AllocateCudnnConvolutionForwardWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
-    return port::Status(
-        port::error::INTERNAL,
-        "Mismatch between cudnn conv and algorithm descriptors.");
-  }
+  // TODO(csigg): This has side effects on the convolution descriptor. It is
+  // functionally correct because the convolution is run with the algorithm of
+  // the last call to this function, but should be fixed anyway.
+  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2622,11 +2603,10 @@ AllocateCudnnConvolutionBackwardDataWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
-    return port::Status(
-        port::error::INTERNAL,
-        "Mismatch between cudnn conv and algorithm descriptors.");
-  }
+  // TODO(csigg): This has side effects on the convolution descriptor. It is
+  // functionally correct because the convolution is run with the algorithm of
+  // the last call to this function, but should be fixed anyway.
+  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2668,11 +2648,10 @@ AllocateCudnnConvolutionBackwardFilterWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
-    return port::Status(
-        port::error::INTERNAL,
-        "Mismatch between cudnn conv and algorithm descriptors.");
-  }
+  // TODO(csigg): This has side effects on the convolution descriptor. It is
+  // functionally correct because the convolution is run with the algorithm of
+  // the last call to this function, but should be fixed anyway.
+  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2706,42 +2685,18 @@ AllocateCudnnConvolutionBackwardFilterWorkspace(
   return scratch_allocator->AllocateBytes(size_in_bytes);
 }
 
-port::StatusOr UseTensorOps(Stream* stream, dnn::DataType type,
-                                  absl::optional desc) {
-  bool use_tensor_ops;
-  if (desc.has_value()) {
-    use_tensor_ops = desc->tensor_ops_enabled();
-    if (use_tensor_ops && !IsTensorMathAllowed(stream, type)) {
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "Algo requests disallowed tensor op evaluation.");
-    }
-  } else {
-    use_tensor_ops = IsTensorMathAllowed(stream, type);
-  }
-  return use_tensor_ops;
+static bool TensorOpMathAvailable(int cc_major) {
+  return cc_major >= 7 && CUDNN_VERSION >= 7000 && TensorOpMathEnabled();
 }
 
-cudnnDataType_t GetRnnComputeType(dnn::DataType data_type);
-dnn::DataType GetConvAccumulatorType(dnn::DataType data_type);
-
 port::StatusOr GetCudnnConvolutionForwardAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    dnn::DataType element_type,
-    const dnn::ConvolutionDescriptor& convolution_descriptor,
+    const CudnnConvolutionDescriptor& conv,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
-
-  CudnnConvolutionDescriptor conv(
-      convolution_descriptor,
-      ToCudnnDataType(GetConvAccumulatorType(element_type)));
-  bool use_tensor_ops;
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
-
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2754,7 +2709,10 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm(
                         GetCudnnConvolutionForwardAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
+    int cc_major, cc_minor;
+    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
+    algo_desc = dnn::AlgorithmDesc(
+        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
   }
 
   const auto scratch_or = AllocateCudnnConvolutionForwardWorkspace(
@@ -2778,9 +2736,6 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm(
                      "Returned status: ", scratch_or.status().ToString()));
   }
 
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionForwardWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -2791,19 +2746,10 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    dnn::DataType element_type,
-    const dnn::ConvolutionDescriptor& convolution_descriptor,
+    const CudnnConvolutionDescriptor& conv,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
-  CudnnConvolutionDescriptor conv(
-      convolution_descriptor,
-      ToCudnnDataType(GetConvAccumulatorType(element_type)));
-  bool use_tensor_ops;
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
-
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2816,7 +2762,10 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
                         GetCudnnConvolutionBackwardDataAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
+    int cc_major, cc_minor;
+    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
+    algo_desc = dnn::AlgorithmDesc(
+        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
   }
 
   const auto scratch_or = AllocateCudnnConvolutionBackwardDataWorkspace(
@@ -2839,9 +2788,6 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
         "while a secondary algorithm is not provided.");
   }
 
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardDataWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -2852,19 +2798,10 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    dnn::DataType element_type,
-    const dnn::ConvolutionDescriptor& convolution_descriptor,
+    const CudnnConvolutionDescriptor& conv,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
-  CudnnConvolutionDescriptor conv(
-      convolution_descriptor,
-      ToCudnnDataType(GetConvAccumulatorType(element_type)));
-  bool use_tensor_ops;
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
-
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2877,7 +2814,10 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
                         GetCudnnConvolutionBackwardFilterAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
+    int cc_major, cc_minor;
+    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
+    algo_desc = dnn::AlgorithmDesc(
+        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
   }
 
   auto scratch_or = AllocateCudnnConvolutionBackwardFilterWorkspace(
@@ -2900,9 +2840,6 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
         "while a secondary algorithm is not provided.");
   }
 
-  SE_ASSIGN_OR_RETURN(use_tensor_ops,
-                      UseTensorOps(stream, element_type, algo_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardFilterWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -3067,32 +3004,35 @@ port::Status CudnnSupport::DoPrepareForConvolution(
   CudnnTensorDescriptor output_nd(
       output_descriptor,
       ToCudnnDataType(element_type, output_descriptor.layout()));
+  CudnnConvolutionDescriptor conv(
+      convolution_descriptor,
+      ToCudnnDataType(GetConvAccumulatorType(element_type)));
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
 
   switch (kind) {
     case dnn::ConvolutionKind::FORWARD: {
-      SE_ASSIGN_OR_RETURN(*algorithm_desc,
-                          GetCudnnConvolutionForwardAlgorithm(
-                              stream, cudnn, algorithm_config, input_nd,
-                              filter_nd, element_type, convolution_descriptor,
-                              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(
+          *algorithm_desc,
+          GetCudnnConvolutionForwardAlgorithm(
+              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
+              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     case dnn::ConvolutionKind::BACKWARD_DATA: {
-      SE_ASSIGN_OR_RETURN(*algorithm_desc,
-                          GetCudnnConvolutionBackwardDataAlgorithm(
-                              stream, cudnn, algorithm_config, input_nd,
-                              filter_nd, element_type, convolution_descriptor,
-                              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(
+          *algorithm_desc,
+          GetCudnnConvolutionBackwardDataAlgorithm(
+              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
+              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     case dnn::ConvolutionKind::BACKWARD_FILTER: {
-      SE_ASSIGN_OR_RETURN(*algorithm_desc,
-                          GetCudnnConvolutionBackwardFilterAlgorithm(
-                              stream, cudnn, algorithm_config, input_nd,
-                              filter_nd, element_type, convolution_descriptor,
-                              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(
+          *algorithm_desc,
+          GetCudnnConvolutionBackwardFilterAlgorithm(
+              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
+              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     default:
@@ -3121,9 +3061,8 @@ port::Status CudnnSupport::DoConvolve(
   auto accumulator_type = GetConvAccumulatorType(element_type);
   CudnnConvolutionDescriptor conv(convolution_descriptor,
                                   ToCudnnDataType(accumulator_type));
-  SE_ASSIGN_OR_RETURN(bool use_tensor_ops,
-                      UseTensorOps(stream, element_type, algorithm_desc));
-  conv.set_use_tensor_op_math(use_tensor_ops);
+  // Set use_tensor_math param to correct value
+  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
   // Alpha is the scaling factor for input.
@@ -3356,6 +3295,14 @@ port::Status CudnnSupport::DoConvolve(
   return port::Status::OK();
 }
 
+// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math
+// set
+static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) {
+  cudnnMathType_t math_type;
+  CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type));
+  return math_type == CUDNN_TENSOR_OP_MATH;
+}
+
 template 
 port::Status CudnnSupport::DoFusedConvolveImpl(
@@ -3389,6 +3336,8 @@ port::Status CudnnSupport::DoFusedConvolveImpl(
       filter_descriptor,
       GetCudnnDataType(conv_input_descriptor.layout()));
   CudnnTensorDescriptor bias_nd(bias_descriptor, GetCudnnDataType());
+  CudnnConvolutionDescriptor conv(convolution_descriptor,
+                                  ToCudnnDataType(accumulator_type));
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
 
@@ -3398,14 +3347,9 @@ port::Status CudnnSupport::DoFusedConvolveImpl(
   SE_ASSIGN_OR_RETURN(
       dnn::AlgorithmDesc algo_desc,
       GetCudnnConvolutionForwardAlgorithm(
-          stream, cudnn, algorithm_config, conv_input_nd, filter,
-          dnn::ToDataType::value, convolution_descriptor,
+          stream, cudnn, algorithm_config, conv_input_nd, filter, conv,
           output_nd, scratch_allocator, &scratch));
 
-  CudnnConvolutionDescriptor conv(convolution_descriptor,
-                                  ToCudnnDataType(accumulator_type));
-  conv.set_use_tensor_op_math(algo_desc.tensor_ops_enabled());
-
   std::unique_ptr timer;
   if (is_profiling) {
     timer.reset(new GpuTimer(parent_));  // NOLINT
@@ -3536,7 +3480,9 @@ bool CudnnSupport::GetRnnAlgorithms(
   for (auto i : algo_types) {
     out_algorithms->push_back({i, /*use_tensor_ops=*/false});
 #if CUDNN_VERSION >= 7100
-    out_algorithms->push_back({i, /*use_tensor_ops=*/true});
+    if (RnnTensorOpMathEnabled()) {
+      out_algorithms->push_back({i, /*use_tensor_ops=*/true});
+    }
 #endif
   }
   return true;

From 64f7bdd56a394ecae55c5006e483050569b9b136 Mon Sep 17 00:00:00 2001
From: Thai Nguyen 
Date: Thu, 18 Jun 2020 19:32:16 -0700
Subject: [PATCH 0583/1390] Disable tsan on InterpreterFlexTest and
 SelectiveBuiltInterpreterFlexTest

PiperOrigin-RevId: 317231748
Change-Id: I7ab662fd55024c0ed91bd78bfdc8e9206d78b3b6
---
 tensorflow/lite/delegates/flex/BUILD | 1 +
 tensorflow/lite/java/BUILD           | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/lite/delegates/flex/BUILD b/tensorflow/lite/delegates/flex/BUILD
index 42914bf5ab8..99bcf05ab4a 100644
--- a/tensorflow/lite/delegates/flex/BUILD
+++ b/tensorflow/lite/delegates/flex/BUILD
@@ -279,6 +279,7 @@ java_test(
         "no_oss",  # Currently requires --config=monolithic, b/118895218.
         # TODO(b/121204962): Re-enable test after fixing memory leaks.
         "noasan",
+        "notsan",  # TODO(b/158651814) Re-enable after fixing racing condition.
     ],
     test_class = "org.tensorflow.lite.InterpreterFlexTest",
     visibility = ["//visibility:private"],
diff --git a/tensorflow/lite/java/BUILD b/tensorflow/lite/java/BUILD
index 738d66a0eb1..89be932ab4d 100644
--- a/tensorflow/lite/java/BUILD
+++ b/tensorflow/lite/java/BUILD
@@ -304,6 +304,7 @@ java_test(
         "no_oss",  # Currently requires --config=monolithic, b/118895218.
         # TODO(b/121204962): Re-enable test after fixing memory leaks.
         "noasan",
+        "notsan",  # TODO(b/158651814) Re-enable after fixing racing condition.
     ],
     test_class = "org.tensorflow.lite.InterpreterFlexTest",
     visibility = ["//visibility:private"],

From c159f1599548428660c80dada924d69f269384a3 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Thu, 18 Jun 2020 19:35:14 -0700
Subject: [PATCH 0584/1390] Fork the keras related tpu_strategy_test to keras
 integration test.

PiperOrigin-RevId: 317232048
Change-Id: If05867985ff1ff81ac45bb601b701ee68d4d5279
---
 tensorflow/python/distribute/BUILD            |  1 -
 .../python/distribute/tpu_strategy_test.py    | 19 -----
 .../python/keras/integration_test/BUILD       | 13 ++++
 .../integration_test/tpu_strategy_test.py     | 69 +++++++++++++++++++
 4 files changed, 82 insertions(+), 20 deletions(-)
 create mode 100644 tensorflow/python/keras/integration_test/tpu_strategy_test.py

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 7208807a18c..4d77c12f975 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -654,7 +654,6 @@ tpu_py_test(
         "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//tensorflow/python/eager:remote",
         "//tensorflow/python/eager:test",
-        "//tensorflow/python/keras",
     ],
 )
 
diff --git a/tensorflow/python/distribute/tpu_strategy_test.py b/tensorflow/python/distribute/tpu_strategy_test.py
index 6dd7de500e4..400b12112d6 100644
--- a/tensorflow/python/distribute/tpu_strategy_test.py
+++ b/tensorflow/python/distribute/tpu_strategy_test.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python import keras
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.distribute import distribute_lib
 from tensorflow.python.distribute import distribution_strategy_context
@@ -364,24 +363,6 @@ class TPUStrategyTest(test.TestCase):
         expected_result,
         strategy.experimental_local_results(train_step(next(input_iterator))))
 
-  def test_keras_metric_outside_strategy_scope_per_replica(self):
-    strategy = get_tpu_strategy()
-    metric = keras.metrics.Mean("test_metric", dtype=dtypes.float32)
-
-    dataset = dataset_ops.Dataset.range(strategy.num_replicas_in_sync *
-                                        2).batch(2)
-    dataset = strategy.experimental_distribute_dataset(dataset)
-
-    @def_function.function
-    def step_fn(i):
-      metric.update_state(i)
-
-    with self.assertRaisesRegex(ValueError, "Trying to run metric.update_state "
-                                            "in replica context"):
-      with strategy.scope():
-        for i in dataset:
-          strategy.run(step_fn, args=(i,))
-
   # TODO(b/145574622): Remove this test once it is re-enabled in values_test.py.
   def test_all_reduce_on_sync_on_read_variable(self):
     strategy = get_tpu_strategy()
diff --git a/tensorflow/python/keras/integration_test/BUILD b/tensorflow/python/keras/integration_test/BUILD
index 2ef775a190e..b23dcc59b97 100644
--- a/tensorflow/python/keras/integration_test/BUILD
+++ b/tensorflow/python/keras/integration_test/BUILD
@@ -2,6 +2,7 @@
 #   Contains Keras integration tests that verify with other TF high level APIs.
 
 load("//tensorflow:tensorflow.bzl", "cuda_py_test", "tf_py_test")
+load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_test")
 
 package(
     default_visibility = [
@@ -91,3 +92,15 @@ cuda_py_test(
         "//tensorflow/python:extra_py_tests_deps",
     ],
 )
+
+tpu_py_test(
+    name = "tpu_strategy_test",
+    srcs = ["tpu_strategy_test.py"],
+    disable_experimental = True,
+    python_version = "PY3",
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python:extra_py_tests_deps",
+    ],
+)
diff --git a/tensorflow/python/keras/integration_test/tpu_strategy_test.py b/tensorflow/python/keras/integration_test/tpu_strategy_test.py
new file mode 100644
index 00000000000..d24e96ae855
--- /dev/null
+++ b/tensorflow/python/keras/integration_test/tpu_strategy_test.py
@@ -0,0 +1,69 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for TPUStrategy."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import flags
+
+import tensorflow as tf
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
+flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
+flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")
+
+
+def get_tpu_cluster_resolver():
+  resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
+      tpu=FLAGS.tpu,
+      zone=FLAGS.zone,
+      project=FLAGS.project,
+  )
+  return resolver
+
+
+def get_tpu_strategy():
+  resolver = get_tpu_cluster_resolver()
+  tf.config.experimental_connect_to_cluster(resolver)
+  tf.tpu.experimental.initialize_tpu_system(resolver)
+  return tf.distribute.experimental.TPUStrategy(resolver)
+
+
+class TpuStrategyTest(tf.test.TestCase):
+
+  def test_keras_metric_outside_strategy_scope_per_replica(self):
+    strategy = get_tpu_strategy()
+    metric = tf.keras.metrics.Mean("test_metric", dtype=tf.float32)
+
+    dataset = tf.data.Dataset.range(strategy.num_replicas_in_sync * 2).batch(2)
+    dataset = strategy.experimental_distribute_dataset(dataset)
+
+    @tf.function
+    def step_fn(i):
+      metric.update_state(i)
+
+    with self.assertRaisesRegex(ValueError, "Trying to run metric.update_state "
+                                            "in replica context"):
+      with strategy.scope():
+        for i in dataset:
+          strategy.run(step_fn, args=(i,))
+
+
+if __name__ == "__main__":
+  tf.test.main()

From 2399b25e139c729c0f0efd0efe5a009af04ef773 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Fri, 19 Jun 2020 02:49:05 +0000
Subject: [PATCH 0585/1390] Change variable names, Remove comments

---
 tensorflow/python/framework/python_op_gen.cc | 84 ++++++++++----------
 1 file changed, 41 insertions(+), 43 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 4ecfccb611f..062a9aa01e4 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -45,7 +45,7 @@ const int kRightMargin = 78;
 
 constexpr char kEagerFallbackSuffix[] = "_eager_fallback";
 
-std::unordered_map dtypes_map {
+std::unordered_map dtype_type {
       {"_dtypes.float16", "_dtypes.Float16"},
       {"_dtypes.half", "_dtypes.Half"},
       {"_dtypes.float32", "_dtypes.Float32"},
@@ -164,7 +164,7 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
   bool AddEagerFastPathAndGraphCode(const string& parameters,
                                     const std::vector& output_sizes,
                                     const string& eager_not_allowed_error,
-                                    std::unordered_map& type_map);
+                                    std::unordered_map& type_annotations);
   bool AddEagerFallbackCode(const string& parameters,
                             const std::vector& output_sizes,
                             const string& num_outputs_expr,
@@ -182,9 +182,9 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
 
   std::unordered_map GetTypeAnnotationMap();
 
-  void GenerateTypeVars(std::unordered_map& type_map);
+  void GenerateTypeVars(std::unordered_map& type_annotations);
 
-  void AddReturnTypeAnnotation(std::unordered_map& type_map);
+  void AddReturnTypeAnnotation(std::unordered_map& type_annotations);
 
   void AddAttrForArg(const string& attr, int arg_index) {
     gtl::InsertIfNotPresent(&inferred_attrs_, attr,
@@ -348,10 +348,10 @@ string GenEagerPythonOp::Code() {
     param_names_.push_back(param_and_default.first);
   }
 
-  std::unordered_map type_map;
+  std::unordered_map type_annotations;
   // Only populate map for whitelisted ops
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    type_map = GetTypeAnnotationMap();
+    type_annotations = GetTypeAnnotationMap();
   }
 
   string parameters;
@@ -360,9 +360,9 @@ string GenEagerPythonOp::Code() {
     strings::StrAppend(¶meters, param.GetRenameTo());
 
     // Add type annotations to param
-    if (type_map.find(param.GetName()) != type_map.end()) {
-      if (!type_map[param.GetName()].empty()) {
-        strings::StrAppend(¶meters, ": ", type_map[param.GetName()]);
+    if (type_annotations.find(param.GetName()) != type_annotations.end()) {
+      if (!type_annotations[param.GetName()].empty()) {
+        strings::StrAppend(¶meters, ": ", type_annotations[param.GetName()]);
       }
     }
   }
@@ -374,12 +374,12 @@ string GenEagerPythonOp::Code() {
       strings::StrAppend(¶meters_with_defaults, ", ");
 
     // Add type annotations to param_and_default
-    if (type_map.find(param_and_default.first.GetName()) != type_map.end()) {
-      if (!type_map[param_and_default.first.GetName()].empty()) {
-        strings::StrAppend(¶meters, ": ", type_map[param_and_default.first.GetName()]);
+    if (type_annotations.find(param_and_default.first.GetName()) != type_annotations.end()) {
+      if (!type_annotations[param_and_default.first.GetName()].empty()) {
+        strings::StrAppend(¶meters, ": ", type_annotations[param_and_default.first.GetName()]);
         strings::StrAppend(¶meters_with_defaults,
                            param_and_default.first.GetRenameTo(), ": ",
-                           type_map[param_and_default.first.GetName()], " ",
+                           type_annotations[param_and_default.first.GetName()], " ",
                            "= ", param_and_default.second);
         continue;
       }
@@ -420,7 +420,7 @@ string GenEagerPythonOp::Code() {
   string eager_not_allowed_error = GetEagerNotAllowedError();
 
   if (!AddEagerFastPathAndGraphCode(parameters_with_defaults, output_sizes,
-                                    eager_not_allowed_error, type_map)) {
+                                    eager_not_allowed_error, type_annotations)) {
     return result_;
   }
 
@@ -433,60 +433,58 @@ string GenEagerPythonOp::Code() {
 }
 
 std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
-  std::unordered_map type_map;
+  std::unordered_map type_annotations;
   // Mapping attrs to TypeVars
   for (const auto& attr : op_def_.attr()) {
     if (attr.type() == "type") {
       const string type_var_name = "TV_" + op_def_.name() + "_" + attr.name();
-      type_map[attr.name()] = type_var_name;
+      type_annotations[attr.name()] = type_var_name;
     } else if (attr.type() == "bool" || attr.type() == "float" ||
                attr.type() == "int" || attr.type() == "bytes") {
-      type_map[attr.name()] = attr.type();
+      type_annotations[attr.name()] = attr.type();
     }
   }
 
   // Mapping input Tensors to their types
   for (const auto& arg : op_def_.input_arg()) {
-    // Do not add type annotations to args that accept a sequence of tensors
+    // Do not add type annotations to args that accept a sequence of Tensors
     if (!arg.number_attr().empty()) continue;
     string type_annotation;
-    if (type_map.find(arg.type_attr()) != type_map.end()) {
+    if (type_annotations.find(arg.type_attr()) != type_annotations.end()) {
       // Get the correct TypeVar if input maps to an attr
-      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_map[arg.type_attr()], "]");
+      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_annotations[arg.type_attr()], "]");
     } else {
       // Get the dtype of the Tensor
       const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
-      if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
-        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtypes_map[py_dtype], "]");
+      if (dtype_type.find(py_dtype) != dtype_type.end()) {
+        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtype_type[py_dtype], "]");
       }
     }
 
-    type_map[arg.name()] = type_annotation;
+    type_annotations[arg.name()] = type_annotation;
   }
 
   // Mapping output Tensor to its type
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
     string type_annotation;
-    if (type_map.find(arg.type_attr()) != type_map.end()) {
-      // Get the correct TypeVar if input maps to an attr
-      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_map[arg.type_attr()], "]");
+    if (type_annotations.find(arg.type_attr()) != type_annotations.end()) {
+      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_annotations[arg.type_attr()], "]");
     } else {
-      // Get the dtype of the Tensor
       const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
-      if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
-        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtypes_map[py_dtype], "]");
+      if (dtype_type.find(py_dtype) != dtype_type.end()) {
+        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtype_type[py_dtype], "]");
       }
     }
 
-    type_map[arg.name()] = type_annotation;
+    type_annotations[arg.name()] = type_annotation;
   }
 
-  return type_map;
+  return type_annotations;
 }
 
 // Generate TypeVars using attrs
-void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type_map) {
+void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type_annotations) {
   bool added_typevar = false;
   for (const auto& attr : op_def_.attr()) {
     if (attr.type() == "type") {
@@ -494,14 +492,14 @@ void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type
       for (int t : attr.allowed_values().list().type()) {
         DataType dtype = static_cast(t);
         const string py_dtype = python_op_gen_internal::DataTypeToPython(dtype, "_dtypes.");
-        if (dtypes_map.find(py_dtype) != dtypes_map.end()) {
-          allowed_types.emplace_back(dtypes_map[py_dtype]);
+        if (dtype_type.find(py_dtype) != dtype_type.end()) {
+          allowed_types.emplace_back(dtype_type[py_dtype]);
         }
       }
 
       // If all dtypes are allowed, add them all
       if (allowed_types.empty()) {
-        for (std::pair map_dtype : dtypes_map) {
+        for (std::pair map_dtype : dtype_type) {
           allowed_types.emplace_back(map_dtype.second);
         }
       }
@@ -514,7 +512,7 @@ void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type
         strings::StrAppend(&typevar_dtypes, *it);
       }
 
-      const string type_var_name = type_map[attr.name()];
+      const string type_var_name = type_annotations[attr.name()];
       strings::StrAppend(&result_, type_var_name, " = TypeVar(\"", type_var_name, "\", ", typevar_dtypes,")\n");
       added_typevar = true;
     }
@@ -523,14 +521,14 @@ void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type
   if (added_typevar) strings::StrAppend(&result_, "\n");
 }
 
-void GenEagerPythonOp::AddReturnTypeAnnotation(std::unordered_map& type_map) {
+void GenEagerPythonOp::AddReturnTypeAnnotation(std::unordered_map& type_annotations) {
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
     // Add type annotations to param
-    if (type_map.find(arg.name()) != type_map.end()) {
-      if (!type_map[arg.name()].empty()) {
+    if (type_annotations.find(arg.name()) != type_annotations.end()) {
+      if (!type_annotations[arg.name()].empty()) {
         result_.erase(result_.length() - 2);
-        strings::StrAppend(&result_, " -> ", type_map[arg.name()], ":\n");
+        strings::StrAppend(&result_, " -> ", type_annotations[arg.name()], ":\n");
       }
     }
   }
@@ -858,9 +856,9 @@ void GenEagerPythonOp::AddEagerFunctionTeardown(
 
 bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
     const string& parameters, const std::vector& output_sizes,
-    const string& eager_not_allowed_error, std::unordered_map& type_map) {
+    const string& eager_not_allowed_error, std::unordered_map& type_annotations) {
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    GenerateTypeVars(type_map);
+    GenerateTypeVars(type_annotations);
   }
   if (api_def_.visibility() == ApiDef::VISIBLE) {
     strings::StrAppend(&result_, "@_dispatch.add_dispatch_list\n");
@@ -869,7 +867,7 @@ bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
   AddExport();
   AddDefLine(function_name_, parameters);
   if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
-    AddReturnTypeAnnotation(type_map);
+    AddReturnTypeAnnotation(type_annotations);
   }
   AddDocStringDescription();
   AddDocStringArgs();

From 4d54ef31394aefe270826790164edcc6d687bb63 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 20:02:06 -0700
Subject: [PATCH 0586/1390] Enable type annotations for python/ops.

PiperOrigin-RevId: 317234494
Change-Id: I49a24cd1e2127a3c7b0f2eb217cfe023ce5b439f
---
 tensorflow/python/ops/logging_ops.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/logging_ops.py b/tensorflow/python/ops/logging_ops.py
index 02fce277690..8ca63f55987 100644
--- a/tensorflow/python/ops/logging_ops.py
+++ b/tensorflow/python/ops/logging_ops.py
@@ -54,9 +54,11 @@ except NameError:
 # call relies on certain conditionals for its dependencies.  Use
 # control_flow_ops.Assert.
 
-# Assert and Print are special symbols in Python 2, so we must
-# have an upper-case version of them. When support for it is dropped,
-# we can allow lowercase.
+# Assert and Print are special symbols in python, so we must
+# have an upper-case version of them.
+#
+# For users with Python 3 or Python 2.7
+# with `from __future__ import print_function`, we could also allow lowercase.
 # See https://github.com/tensorflow/tensorflow/issues/18053
 
 
@@ -81,6 +83,11 @@ def Print(input_, data, message=None, first_n=None, summarize=None, name=None):
     with jupyter notebook (printing to the notebook *server's* output, not into
     the notebook).
 
+  Additionally, to use tf.print in python 2.7, users must make sure to import
+  the following:
+
+  `from __future__ import print_function`
+
   Args:
     input_: A tensor passed through this op.
     data: A list of tensors to print out when op is evaluated.
@@ -141,6 +148,11 @@ def print_v2(*inputs, **kwargs):
   Python objects. Printed tensors will recursively show the first and last
   elements of each dimension to summarize.
 
+  @compatibility(python2)
+  In python 2.7, make sure to import the following:
+  `from __future__ import print_function`
+  @end_compatibility
+
   Example:
     Single-input usage:
 

From 7e6e549c461118fbefdb11d03adbc80c27109a8a Mon Sep 17 00:00:00 2001
From: Yujing Zhang 
Date: Thu, 18 Jun 2020 20:03:31 -0700
Subject: [PATCH 0587/1390] Support packed variable in DistributedVariable. Add
 an option to enable packed variable in TPUStrategy.

PiperOrigin-RevId: 317234665
Change-Id: I09e806cb8261815cd87a6d98817556dd8f7e8ed7
---
 tensorflow/python/distribute/BUILD            |   6 +-
 .../python/distribute/checkpointing_test.py   |   2 +
 .../custom_training_loop_input_test.py        |   7 +-
 .../python/distribute/distribute_utils.py     |   3 +
 .../distribute/packed_distributed_variable.py |  13 +-
 .../packed_distributed_variable_test.py       |   6 +-
 .../distribute/saved_model_test_base.py       |   1 +
 .../distribute/strategy_combinations.py       |  20 +-
 tensorflow/python/distribute/tpu_strategy.py  |  50 ++-
 .../python/distribute/tpu_strategy_test.py    | 410 +++++++++---------
 tensorflow/python/distribute/tpu_values.py    |  41 +-
 tensorflow/python/distribute/values.py        |  53 ++-
 tensorflow/python/distribute/values_test.py   |  91 ++--
 tensorflow/python/distribute/values_util.py   |  12 +-
 tensorflow/python/tpu/tpu.py                  |   9 +-
 15 files changed, 454 insertions(+), 270 deletions(-)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 4d77c12f975..0062705126f 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -654,6 +654,7 @@ tpu_py_test(
         "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//tensorflow/python/eager:remote",
         "//tensorflow/python/eager:test",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
@@ -787,6 +788,7 @@ py_library(
     name = "tpu_values",
     srcs = ["tpu_values.py"],
     deps = [
+        ":packed_distributed_variable",
         ":values",
         "//tensorflow/python:framework_ops",
         "//tensorflow/python:math_ops",
@@ -1602,7 +1604,7 @@ distribute_py_test(
     srcs = ["saved_model_save_load_test.py"],
     full_precision = True,
     main = "saved_model_save_load_test.py",
-    shard_count = 5,
+    shard_count = 7,
     tags = [
         "multi_and_single_gpu",
         "no_rocm",
@@ -1635,7 +1637,7 @@ distribute_py_test(
     srcs = ["saved_model_mixed_api_test.py"],
     full_precision = True,
     main = "saved_model_mixed_api_test.py",
-    shard_count = 5,
+    shard_count = 7,
     tags = [
         "multi_and_single_gpu",
         "no_rocm",
diff --git a/tensorflow/python/distribute/checkpointing_test.py b/tensorflow/python/distribute/checkpointing_test.py
index ad646905315..edd4c46c371 100644
--- a/tensorflow/python/distribute/checkpointing_test.py
+++ b/tensorflow/python/distribute/checkpointing_test.py
@@ -103,6 +103,7 @@ class TrainingCheckpointTests(test.TestCase, parameterized.TestCase):
               strategy_combinations.mirrored_strategy_with_one_cpu,
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
               strategy_combinations.central_storage_strategy_with_two_gpus,
           ],
           mode=["eager"]))
@@ -138,6 +139,7 @@ class TrainingCheckpointTests(test.TestCase, parameterized.TestCase):
               strategy_combinations.mirrored_strategy_with_one_cpu,
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
               strategy_combinations.central_storage_strategy_with_two_gpus,
           ],
           mode=["eager"]))
diff --git a/tensorflow/python/distribute/custom_training_loop_input_test.py b/tensorflow/python/distribute/custom_training_loop_input_test.py
index 5d1584f5aa7..e4f782810dd 100644
--- a/tensorflow/python/distribute/custom_training_loop_input_test.py
+++ b/tensorflow/python/distribute/custom_training_loop_input_test.py
@@ -197,7 +197,8 @@ class InputIterationTest(test.TestCase, parameterized.TestCase,
       combinations.combine(
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
-              strategy_combinations.tpu_strategy
+              strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["eager"]))
   def testNestedOutput(self, distribution):
@@ -748,6 +749,10 @@ class InputIterationTest(test.TestCase, parameterized.TestCase,
           mode=["eager"]
       ))
   def testMultiDeviceDataCapturedFunction(self, distribution):
+    if getattr(distribution, "_enable_packed_variable_in_eager_mode", False):
+      self.skipTest(
+          "Dataset captured function doesn't support packed tensors yet "
+          "(b/145922293).")
     inputs = constant_op.constant([2., 3.])
     dataset = lambda _: dataset_ops.Dataset.from_tensor_slices(inputs).repeat(5)
     input_iterator = iter(
diff --git a/tensorflow/python/distribute/distribute_utils.py b/tensorflow/python/distribute/distribute_utils.py
index ccf19521718..14b934b4a0f 100644
--- a/tensorflow/python/distribute/distribute_utils.py
+++ b/tensorflow/python/distribute/distribute_utils.py
@@ -148,6 +148,9 @@ def select_replica_mirrored(replica_id, structured):
         raise TypeError(
             "Expected value to be mirrored across replicas: %s in %s." %
             (x, structured))
+      packed_var = getattr(x, "_packed_variable", None)
+      if packed_var is not None:
+        return packed_var
       return x.values[replica_id]
     else:
       return x
diff --git a/tensorflow/python/distribute/packed_distributed_variable.py b/tensorflow/python/distribute/packed_distributed_variable.py
index 62512cb4414..c249b8efc1c 100644
--- a/tensorflow/python/distribute/packed_distributed_variable.py
+++ b/tensorflow/python/distribute/packed_distributed_variable.py
@@ -42,7 +42,7 @@ class PackedDistributedVariable(resource_variable_ops.BaseResourceVariable):
       name: Optional name for the variable. Defaults to `'Variable'` and gets
         uniquified automatically.
     """
-    if not context.executing_eagerly():
+    if not ops.executing_eagerly_outside_functions():
       raise ValueError(
           "PackedDistributedVariable should be created in eager mode.")
     if not distributed_variables:
@@ -84,6 +84,9 @@ class PackedDistributedVariable(resource_variable_ops.BaseResourceVariable):
   def devices(self):
     return self._devices
 
+  def on_device(self, device):
+    return PackedVarAndDevice(self, device)
+
   def get_var_on_device(self, device):
     for i, d in enumerate(self._devices):
       if d == device:
@@ -100,7 +103,10 @@ class PackedDistributedVariable(resource_variable_ops.BaseResourceVariable):
 
   @property
   def handle(self):
-    return self._handle
+    if context.executing_eagerly():
+      return self.get_var_on_current_device().handle
+    else:
+      return self._handle
 
   def _read_variable_op(self):
     if context.executing_eagerly():
@@ -269,7 +275,8 @@ class PackedVarAndDevice(object):
 
   @property
   def handle(self):
-    return self._var.handle
+    with ops.device(self._device):
+      return self._var.handle
 
   @property
   def op(self):
diff --git a/tensorflow/python/distribute/packed_distributed_variable_test.py b/tensorflow/python/distribute/packed_distributed_variable_test.py
index d29d19960a5..ec2e476e4b8 100644
--- a/tensorflow/python/distribute/packed_distributed_variable_test.py
+++ b/tensorflow/python/distribute/packed_distributed_variable_test.py
@@ -46,7 +46,7 @@ class PackedDistributedVariableTest(test.TestCase):
       v1 = resource_variable_ops.ResourceVariable(2.0, name='var1')
 
     packed_var = packed_distributed_variable.PackedDistributedVariable([v0, v1])
-    self.assertTrue(packed_var.handle.is_packed)
+    self.assertFalse(packed_var.handle.is_packed)
     self.assertTrue(packed_var.is_initialized)
 
     with ops.device('/cpu:0'):
@@ -61,6 +61,7 @@ class PackedDistributedVariableTest(test.TestCase):
 
     @def_function.function
     def update_var():
+      self.assertTrue(packed_var.handle.is_packed)
       with ops.device('/cpu:0'):
         packed_var.assign_add(3.0).assign_sub(1.0)
         read0 = packed_var.value()
@@ -85,7 +86,7 @@ class PackedDistributedVariableTest(test.TestCase):
 
     packed_var0 = packed_distributed_variable.PackedVarAndDevice(
         packed_var, device0)
-    self.assertTrue(packed_var0.handle.is_packed)
+    self.assertFalse(packed_var0.handle.is_packed)
     self.assertAllEqual(math_ops.mul(packed_var0, 2.0), 2.0)
 
     packed_var1 = packed_distributed_variable.PackedVarAndDevice(
@@ -94,6 +95,7 @@ class PackedDistributedVariableTest(test.TestCase):
 
     @def_function.function
     def func():
+      self.assertTrue(packed_var.handle.is_packed)
       var0 = packed_distributed_variable.PackedVarAndDevice(packed_var, device0)
       var0.assign_add(3.0)
       var1 = packed_distributed_variable.PackedVarAndDevice(packed_var, device1)
diff --git a/tensorflow/python/distribute/saved_model_test_base.py b/tensorflow/python/distribute/saved_model_test_base.py
index e544e51cddd..70ea582baff 100644
--- a/tensorflow/python/distribute/saved_model_test_base.py
+++ b/tensorflow/python/distribute/saved_model_test_base.py
@@ -58,6 +58,7 @@ strategies = [
     strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
     strategy_combinations.mirrored_strategy_with_two_gpus,
     strategy_combinations.tpu_strategy,
+    strategy_combinations.tpu_strategy_packed_var,
     strategy_combinations.central_storage_strategy_with_two_gpus,
 ]
 
diff --git a/tensorflow/python/distribute/strategy_combinations.py b/tensorflow/python/distribute/strategy_combinations.py
index 350b187f67f..1fa42cb8645 100644
--- a/tensorflow/python/distribute/strategy_combinations.py
+++ b/tensorflow/python/distribute/strategy_combinations.py
@@ -53,7 +53,11 @@ _did_connect_to_cluster = False
 
 
 # pylint: disable=missing-docstring
-def _get_tpu_strategy_creator(steps_per_run, use_single_core=False, **kwargs):
+def _get_tpu_strategy_creator(steps_per_run,
+                              use_single_core=False,
+                              enable_packed_variable=False,
+                              **kwargs):
+
   def _create_tpu_strategy():
     global _did_connect_to_cluster
 
@@ -87,10 +91,13 @@ def _get_tpu_strategy_creator(steps_per_run, use_single_core=False, **kwargs):
 
     # Steps per run is only supported in TF 1.x
     if tf2.enabled():
-      return tpu_lib.TPUStrategy(resolver, device_assignment, **kwargs)
+      strategy = tpu_lib.TPUStrategy(resolver, device_assignment, **kwargs)
     else:
-      return tpu_lib.TPUStrategyV1(resolver, steps_per_run,
-                                   device_assignment, **kwargs)
+      strategy = tpu_lib.TPUStrategyV1(resolver, steps_per_run,
+                                       device_assignment, **kwargs)
+    strategy._enable_packed_variable_in_eager_mode = enable_packed_variable  # pylint: disable=protected-access
+    return strategy
+
   return _create_tpu_strategy
 
 
@@ -117,6 +124,10 @@ one_device_strategy_gpu_on_worker_1 = combinations.NamedDistribution(
     required_gpus=1)
 tpu_strategy = combinations.NamedDistribution(
     "TPU", _get_tpu_strategy_creator(steps_per_run=2), required_tpu=True)
+tpu_strategy_packed_var = combinations.NamedDistribution(
+    "TPUPackedVar",
+    _get_tpu_strategy_creator(steps_per_run=2, enable_packed_variable=True),
+    required_tpu=True)
 tpu_strategy_one_step = combinations.NamedDistribution(
     "TPUOneStep", _get_tpu_strategy_creator(steps_per_run=1), required_tpu=True)
 tpu_strategy_one_core = combinations.NamedDistribution(
@@ -286,6 +297,7 @@ strategies_minus_default_and_tpu = [
 tpu_strategies = [
     tpu_strategy,  # steps_per_run=2
     tpu_strategy_one_step,
+    tpu_strategy_packed_var,
     cloud_tpu_strategy,
 ]
 
diff --git a/tensorflow/python/distribute/tpu_strategy.py b/tensorflow/python/distribute/tpu_strategy.py
index 9493ecce767..7e8f5b97e7e 100644
--- a/tensorflow/python/distribute/tpu_strategy.py
+++ b/tensorflow/python/distribute/tpu_strategy.py
@@ -141,6 +141,10 @@ class TPUStrategy(distribute_lib.Strategy):
         "num_workers").set(self.extended.num_hosts)
     distribute_lib.distribution_strategy_replica_gauge.get_cell(
         "num_replicas_per_worker").set(self.extended.num_replicas_per_host)
+    # Packed variable is used to reduce the overhead of function execution.
+    # For a DistributedVariable, only one variable handle is captured into a
+    # function graph. It's only supported in eager mode.
+    self._enable_packed_variable_in_eager_mode = False
 
   # TODO(cjfj): Modify `_call_for_each_replica` in `TPUExtended` such that this
   # can use the default implementation.
@@ -185,6 +189,10 @@ class TPUStrategyV1(distribute_lib.StrategyV1):
         "num_workers").set(self.extended.num_hosts)
     distribute_lib.distribution_strategy_replica_gauge.get_cell(
         "num_replicas_per_worker").set(self.extended.num_replicas_per_host)
+    # Packed variable is used to reduce the overhead of function execution.
+    # For a DistributedVariable, only one variable handle is captured into a
+    # function graph. It's only supported in eager mode.
+    self._enable_packed_variable_in_eager_mode = False
 
   @property
   def steps_per_run(self):
@@ -671,20 +679,29 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
       return cross_device_ops_lib.reduce_non_distributed_value(
           reduce_op, value, destinations, self._num_replicas_in_sync)
 
+    value_list = value.values
+    # pylint: disable=protected-access
+    if isinstance(
+        value,
+        values.DistributedVariable) and value._packed_variable is not None:
+      value_list = tuple(
+          value._packed_variable.on_device(d)
+          for d in value._packed_variable.devices)
+    # pylint: enable=protected-access
+
     # Currently XLA op by op mode has a limit for the number of inputs for a
     # single op, thus we break one `add_n` op into a group of `add_n` ops to
     # work around the constraint.
     # TODO(cjfj): Detect when it is possible to use `cross_replica_sum`.
     if len(value.values) <= _XLA_OP_BY_OP_INPUTS_LIMIT:
-      output = math_ops.add_n(value.values)
+      output = math_ops.add_n(value_list)
     else:
-      output = array_ops.zeros_like(
-          value.values[0], dtype=value.values[0].dtype)
-      for i in range(0, len(value.values), _XLA_OP_BY_OP_INPUTS_LIMIT):
-        output += math_ops.add_n(value.values[i:i + _XLA_OP_BY_OP_INPUTS_LIMIT])
+      output = array_ops.zeros_like(value_list[0], dtype=value_list[0].dtype)
+      for i in range(0, len(value_list), _XLA_OP_BY_OP_INPUTS_LIMIT):
+        output += math_ops.add_n(value_list[i:i + _XLA_OP_BY_OP_INPUTS_LIMIT])
 
     if reduce_op == reduce_util.ReduceOp.MEAN:
-      output *= (1. / len(value.values))
+      output *= (1. / len(value_list))
 
     devices = cross_device_ops_lib.get_devices_from(destinations)
 
@@ -710,17 +727,28 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
       else:
         return (fn(var, *args, **kwargs),)
 
-    # Otherwise, we revert to MirroredStrategy behavior and update each variable
-    # directly.
+    # Otherwise, we revert to MirroredStrategy behavior and update the variable
+    # on each replica directly.
     updates = []
-    for i, v in enumerate(var.values):
+    values_and_devices = []
+    packed_var = var._packed_variable  # pylint: disable=protected-access
+    if packed_var is not None:
+      for device in packed_var.devices:
+        values_and_devices.append((packed_var, device))
+    else:
+      for value in var.values:
+        values_and_devices.append((value, value.device))
+
+    for i, value_and_device in enumerate(values_and_devices):
+      value = value_and_device[0]
+      device = value_and_device[1]
       name = "update_%d" % i
-      with ops.device(v.device), \
+      with ops.device(device), \
            distribute_lib.UpdateContext(i), \
            ops.name_scope(name):
         # If args and kwargs are not mirrored, the value is returned as is.
         updates.append(
-            fn(v, *distribute_utils.select_replica_mirrored(i, args),
+            fn(value, *distribute_utils.select_replica_mirrored(i, args),
                **distribute_utils.select_replica_mirrored(i, kwargs)))
     return distribute_utils.update_regroup(self, updates, group)
 
diff --git a/tensorflow/python/distribute/tpu_strategy_test.py b/tensorflow/python/distribute/tpu_strategy_test.py
index 400b12112d6..4070336aae8 100644
--- a/tensorflow/python/distribute/tpu_strategy_test.py
+++ b/tensorflow/python/distribute/tpu_strategy_test.py
@@ -18,6 +18,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
+
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.distribute import distribute_lib
 from tensorflow.python.distribute import distribution_strategy_context
@@ -64,14 +66,17 @@ def get_tpu_cluster_resolver():
   return resolver
 
 
-def get_tpu_strategy():
+def get_tpu_strategy(enable_packed_var=False):
   resolver = get_tpu_cluster_resolver()
   remote.connect_to_cluster(resolver)
   tpu_strategy_util.initialize_tpu_system(resolver)
-  return tpu_lib.TPUStrategy(resolver)
+  strategy = tpu_lib.TPUStrategy(resolver)
+  strategy._enable_packed_variable_in_eager_mode = enable_packed_var
+  return strategy
 
 
-class TPUStrategyTest(test.TestCase):
+# TPU tests which don't use TPUStrategy.
+class TPUTest(test.TestCase):
 
   def test_multiple_initialize_system(self):
     resolver = get_tpu_cluster_resolver()
@@ -82,177 +87,6 @@ class TPUStrategyTest(test.TestCase):
       tpu_strategy_util.initialize_tpu_system(resolver)
       self.assertRegex(str(mock_log.call_args), "already been initialized")
 
-  def test_sequential_experimental_runs(self):
-    resolver = get_tpu_cluster_resolver()
-    remote.connect_to_cluster(resolver)
-    topology = tpu_strategy_util.initialize_tpu_system(resolver)
-    # Computation replicated to all cores.
-    device_assignment = device_assignment_lib.DeviceAssignment.build(
-        topology, num_replicas=2)
-    strategy = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment)
-
-    # Computation on the 1st core.
-    device_assignment2 = device_assignment_lib.DeviceAssignment.build(
-        topology, num_replicas=1)
-    strategy2 = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment2)
-
-    def computation(x):
-      return math_ops.square(x)
-
-    @def_function.function
-    def train_step():
-      outputs = strategy.experimental_local_results(
-          strategy.run(computation, args=([2., 2.],)))
-      outputs2 = strategy2.run(
-          computation, args=([outputs[0]],))
-      return outputs2
-
-    self.assertAllEqual([[16., 16.]], train_step())
-
-  def test_device_switch_case(self):
-    strategy = get_tpu_strategy()
-    with strategy.scope():
-      a = variables.Variable(1)
-
-    inference_iteration = variables.Variable(-1)
-
-    def inference_fn(x, i):
-      return a + x + i
-
-    @def_function.function
-    def run_inference(x):
-
-      def do_inference(device, inference_fn, i):
-        with ops.device(device):
-          return inference_fn(x, i)
-
-      branch_fns = {
-          0: (lambda: do_inference("/device:TPU:0", inference_fn, 0)),
-          1: (lambda: do_inference("/device:TPU:1", inference_fn, 1)),
-      }
-      branch_index = inference_iteration.assign_add(1, use_locking=True) % 2
-      return control_flow_ops.switch_case(branch_index, branch_fns)
-
-    self.assertAllEqual(2., run_inference(1))  # Use TPU core 0.
-    self.assertAllEqual(3., run_inference(1))  # Use TPU core 1.
-
-  def test_recover_from_compilation_failures(self):
-    # TODO(b/148150981): Stop skipping this test once recovery works
-    # for non-local TPU.
-    if FLAGS.tpu:
-      self.skipTest("Recovery fails for non-local TPU, see b/148150981")
-
-    # Disable automatic outside compilation.
-    config.set_soft_device_placement(False)
-    strategy = get_tpu_strategy()
-
-    @def_function.function
-    def compilation_failure_run():
-
-      def computation():
-        return random_ops.random_gamma([10], [0.5, 1.5])
-
-      return strategy.run(computation)
-
-    with self.assertRaisesRegexp(errors.InvalidArgumentError,
-                                 "TPU compilation failed"):
-      compilation_failure_run()
-
-    @def_function.function
-    def good_run():
-
-      def computation():
-        return random_ops.random_normal([10])
-
-      return strategy.run(computation)
-
-    good_run()
-
-  def test_dynamic_shape_with_outside_compilation_failure(self):
-    # Enable automatic outside compilation.
-    config.set_soft_device_placement(True)
-    strategy = get_tpu_strategy()
-    dataset = dataset_ops.Dataset.from_tensors(("string", 1.0)).repeat().batch(
-        2, drop_remainder=False)
-    dataset = strategy.experimental_distribute_dataset(dataset)
-    iterator = iter(dataset)
-
-    @def_function.function
-    def train_fn(iterator):
-
-      def step_fn(inputs):
-        _, inputs = inputs
-        return math_ops.reduce_sum(inputs)
-
-      return strategy.experimental_local_results(
-          strategy.run(step_fn, args=(next(iterator),)))
-
-    with self.assertRaisesRegex(errors.InternalError, "Compilation failure"):
-      logging.info(train_fn(iterator))
-
-  def test_computation_on_subset_cores(self):
-    resolver = get_tpu_cluster_resolver()
-    remote.connect_to_cluster(resolver)
-    topology = tpu_strategy_util.initialize_tpu_system(resolver)
-    all_core_strategy = tpu_lib.TPUStrategy(resolver)
-
-    with all_core_strategy.scope():
-      v = variables.Variable(0.0,
-                             aggregation=variables.VariableAggregation.MEAN)
-
-    # Computation on the 1st core.
-    device_assignment = device_assignment_lib.DeviceAssignment.build(
-        topology, num_replicas=1)
-    first_core_strategy = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment)
-
-    # Computation on the 2nd core.
-    device_assignment2 = device_assignment_lib.DeviceAssignment(
-        topology, [[[0, 0, 0, 1]]])
-    second_core_strategy = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment2)
-
-    @def_function.function
-    def train_step():
-
-      def step_fn():
-        return v + 1.0
-
-      all_core_strategy.run(step_fn)
-      r1 = first_core_strategy.run(step_fn)
-      r2 = second_core_strategy.run(step_fn)
-      return r1 + r2
-
-    train_step()
-    self.assertAllEqual(2., train_step())
-
-  def test_worker_devices_on_subset_cores(self):
-    resolver = get_tpu_cluster_resolver()
-    remote.connect_to_cluster(resolver)
-    topology = tpu_strategy_util.initialize_tpu_system(resolver)
-
-    # Strategy for the 1st core.
-    device_assignment = device_assignment_lib.DeviceAssignment.build(
-        topology, num_replicas=1)
-    first_core_strategy = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment)
-
-    # Strategy for the 2nd core.
-    device_assignment2 = device_assignment_lib.DeviceAssignment(
-        topology, [[[0, 0, 0, 1]]])
-    second_core_strategy = tpu_lib.TPUStrategy(
-        resolver, device_assignment=device_assignment2)
-
-    self.assertLen(first_core_strategy.extended.worker_devices, 1)
-    self.assertEndsWith(first_core_strategy.extended.worker_devices[0],
-                        "device:TPU:0")
-
-    self.assertLen(second_core_strategy.extended.worker_devices, 1)
-    self.assertEndsWith(second_core_strategy.extended.worker_devices[0],
-                        "device:TPU:1")
-
   def test_tpu_tf_function_same_device(self):
     with ops.device("/device:TPU:0"):
       a = variables.Variable(1)
@@ -288,8 +122,194 @@ class TPUStrategyTest(test.TestCase):
       result = bar() + 1
       self.assertAllEqual(result, 2)
 
-  def test_control_output_in_while_body_fn(self):
-    strategy = get_tpu_strategy()
+
+@parameterized.named_parameters([("PackedVar", True), ("", False)])
+class TPUStrategyTest(test.TestCase, parameterized.TestCase):
+
+  def test_sequential_experimental_runs(self, enable_packed_var):
+    resolver = get_tpu_cluster_resolver()
+    remote.connect_to_cluster(resolver)
+    topology = tpu_strategy_util.initialize_tpu_system(resolver)
+    # Computation replicated to all cores.
+    device_assignment = device_assignment_lib.DeviceAssignment.build(
+        topology, num_replicas=2)
+    strategy = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment)
+    strategy._enable_packed_variable_in_eager_mode = enable_packed_var
+
+    # Computation on the 1st core.
+    device_assignment2 = device_assignment_lib.DeviceAssignment.build(
+        topology, num_replicas=1)
+    strategy2 = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment2)
+
+    def computation(x):
+      return math_ops.square(x)
+
+    @def_function.function
+    def train_step():
+      outputs = strategy.experimental_local_results(
+          strategy.run(computation, args=([2., 2.],)))
+      outputs2 = strategy2.run(
+          computation, args=([outputs[0]],))
+      return outputs2
+
+    self.assertAllEqual([[16., 16.]], train_step())
+
+  def test_device_switch_case(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
+    with strategy.scope():
+      a = variables.Variable(1)
+
+    inference_iteration = variables.Variable(-1)
+
+    def inference_fn(x, i):
+      return a + x + i
+
+    @def_function.function
+    def run_inference(x):
+
+      def do_inference(device, inference_fn, i):
+        with ops.device(device):
+          return inference_fn(x, i)
+
+      branch_fns = {
+          0: (lambda: do_inference("/device:TPU:0", inference_fn, 0)),
+          1: (lambda: do_inference("/device:TPU:1", inference_fn, 1)),
+      }
+      branch_index = inference_iteration.assign_add(1, use_locking=True) % 2
+      return control_flow_ops.switch_case(branch_index, branch_fns)
+
+    self.assertAllEqual(2., run_inference(1))  # Use TPU core 0.
+    self.assertAllEqual(3., run_inference(1))  # Use TPU core 1.
+
+  def test_recover_from_compilation_failures(self, enable_packed_var):
+    # TODO(b/148150981): Stop skipping this test once recovery works
+    # for non-local TPU.
+    if FLAGS.tpu:
+      self.skipTest("Recovery fails for non-local TPU, see b/148150981")
+
+    # Disable automatic outside compilation.
+    config.set_soft_device_placement(False)
+    strategy = get_tpu_strategy(enable_packed_var)
+
+    @def_function.function
+    def compilation_failure_run():
+
+      def computation():
+        return random_ops.random_gamma([10], [0.5, 1.5])
+
+      return strategy.run(computation)
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "TPU compilation failed"):
+      compilation_failure_run()
+
+    @def_function.function
+    def good_run():
+
+      def computation():
+        return random_ops.random_normal([10])
+
+      return strategy.run(computation)
+
+    good_run()
+
+  def test_dynamic_shape_with_outside_compilation_failure(
+      self, enable_packed_var):
+    # Enable automatic outside compilation.
+    config.set_soft_device_placement(True)
+    strategy = get_tpu_strategy(enable_packed_var)
+    dataset = dataset_ops.Dataset.from_tensors(("string", 1.0)).repeat().batch(
+        2, drop_remainder=False)
+    dataset = strategy.experimental_distribute_dataset(dataset)
+    iterator = iter(dataset)
+
+    @def_function.function
+    def train_fn(iterator):
+
+      def step_fn(inputs):
+        _, inputs = inputs
+        return math_ops.reduce_sum(inputs)
+
+      return strategy.experimental_local_results(
+          strategy.run(step_fn, args=(next(iterator),)))
+
+    with self.assertRaisesRegex(errors.InternalError, "Compilation failure"):
+      logging.info(train_fn(iterator))
+
+  def test_computation_on_subset_cores(self, enable_packed_var):
+    resolver = get_tpu_cluster_resolver()
+    remote.connect_to_cluster(resolver)
+    topology = tpu_strategy_util.initialize_tpu_system(resolver)
+    all_core_strategy = tpu_lib.TPUStrategy(resolver)
+    all_core_strategy._enable_packed_variable_in_eager_mode = enable_packed_var
+
+    with all_core_strategy.scope():
+      v = variables.Variable(0.0,
+                             aggregation=variables.VariableAggregation.MEAN)
+
+    # Computation on the 1st core.
+    device_assignment = device_assignment_lib.DeviceAssignment.build(
+        topology, num_replicas=1)
+    first_core_strategy = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment)
+    first_core_strategy._enable_packed_variable_in_eager_mode = (
+        enable_packed_var)
+
+    # Computation on the 2nd core.
+    device_assignment2 = device_assignment_lib.DeviceAssignment(
+        topology, [[[0, 0, 0, 1]]])
+    second_core_strategy = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment2)
+    second_core_strategy._enable_packed_variable_in_eager_mode = (
+        enable_packed_var)
+
+    @def_function.function
+    def train_step():
+
+      def step_fn():
+        return v + 1.0
+
+      all_core_strategy.run(step_fn)
+      r1 = first_core_strategy.run(step_fn)
+      r2 = second_core_strategy.run(step_fn)
+      return r1 + r2
+
+    train_step()
+    self.assertAllEqual(2., train_step())
+
+  def test_worker_devices_on_subset_cores(self, enable_packed_var):
+    resolver = get_tpu_cluster_resolver()
+    remote.connect_to_cluster(resolver)
+    topology = tpu_strategy_util.initialize_tpu_system(resolver)
+
+    # Strategy for the 1st core.
+    device_assignment = device_assignment_lib.DeviceAssignment.build(
+        topology, num_replicas=1)
+    first_core_strategy = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment)
+    first_core_strategy._enable_packed_variable_in_eager_mode = (
+        enable_packed_var)
+
+    # Strategy for the 2nd core.
+    device_assignment2 = device_assignment_lib.DeviceAssignment(
+        topology, [[[0, 0, 0, 1]]])
+    second_core_strategy = tpu_lib.TPUStrategy(
+        resolver, device_assignment=device_assignment2)
+    second_core_strategy._enable_packed_variable_in_eager_mode = (
+        enable_packed_var)
+
+    self.assertLen(first_core_strategy.extended.worker_devices, 1)
+    self.assertEndsWith(first_core_strategy.extended.worker_devices[0],
+                        "device:TPU:0")
+
+    self.assertLen(second_core_strategy.extended.worker_devices, 1)
+    self.assertEndsWith(second_core_strategy.extended.worker_devices[0],
+                        "device:TPU:1")
+
+  def test_control_output_in_while_body_fn(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
 
     with strategy.scope():
       v = variables.Variable(
@@ -307,8 +327,8 @@ class TPUStrategyTest(test.TestCase):
     train_step()
     self.assertEqual(2.0, v.numpy())
 
-  def test_cluster_in_graph_and_while_body_fn(self):
-    strategy = get_tpu_strategy()
+  def test_cluster_in_graph_and_while_body_fn(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
 
     @def_function.function
     def train_step():
@@ -328,8 +348,8 @@ class TPUStrategyTest(test.TestCase):
     sum_val = train_step().numpy().astype(float)
     self.assertEqual(sum_val, strategy.num_replicas_in_sync * 10)
 
-  def test_two_clusters_with_same_fn(self):
-    strategy = get_tpu_strategy()
+  def test_two_clusters_with_same_fn(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
 
     @def_function.function
     def foo(x):
@@ -342,8 +362,8 @@ class TPUStrategyTest(test.TestCase):
 
     bar(1)
 
-  def test_using_external_variable_inside_tf_function(self):
-    strategy = get_tpu_strategy()
+  def test_using_external_variable_inside_tf_function(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
     dataset = dataset_ops.Dataset.range(
         strategy.num_replicas_in_sync * 2,
         output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
@@ -364,8 +384,8 @@ class TPUStrategyTest(test.TestCase):
         strategy.experimental_local_results(train_step(next(input_iterator))))
 
   # TODO(b/145574622): Remove this test once it is re-enabled in values_test.py.
-  def test_all_reduce_on_sync_on_read_variable(self):
-    strategy = get_tpu_strategy()
+  def test_all_reduce_on_sync_on_read_variable(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
     dataset = dataset_ops.Dataset.range(
         strategy.num_replicas_in_sync, output_type=dtypes.float32).batch(
             strategy.num_replicas_in_sync, drop_remainder=True)
@@ -404,8 +424,8 @@ class TPUStrategyTest(test.TestCase):
     self.assertAllEqual((0.,), w.read_value())
 
   # TODO(b/140633529): Re-enable the test.
-  def disable_test_experimental_run_output_on_device(self):
-    strategy = get_tpu_strategy()
+  def disable_test_experimental_run_output_on_device(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
 
     def computation(x):
       return math_ops.square(x)
@@ -423,8 +443,8 @@ class TPUStrategyTest(test.TestCase):
     self.assertAllEqual("/job:localhost/replica:0/task:0/device:TPU:1",
                         results[1].backing_device)
 
-  def test_composite_input(self):
-    strategy = get_tpu_strategy()
+  def test_composite_input(self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
     if strategy.num_replicas_in_sync != 2:
       self.skipTest("Test assumes two replicas.")
 
@@ -463,8 +483,9 @@ class TPUStrategyTest(test.TestCase):
     self.assertAllEqual(result,
                         [[[0.0, 1.0], [3.0, 8.0]], [[0.0, 1.0], [3.0, 8.0]]])
 
-  def test_composite_input_dynamic_shapes_outside_compilation(self):
-    strategy = get_tpu_strategy()
+  def test_composite_input_dynamic_shapes_outside_compilation(
+      self, enable_packed_var):
+    strategy = get_tpu_strategy(enable_packed_var)
     if strategy.num_replicas_in_sync != 2:
       self.skipTest("Test assumes two replicas.")
 
@@ -506,11 +527,11 @@ class TPUStrategyTest(test.TestCase):
     result = sparse_lookup(dataset)
     self.assertAllEqual(result, [[0.0, 2.0], [1.5, 5.0]])
 
-  def test_per_device_tracing_of_mirrored_variables(self):
+  def test_per_device_tracing_of_mirrored_variables(self, enable_packed_var):
     # Define trace_count as a list to avoid python scoping error
     trace_count = [0]
 
-    strategy = get_tpu_strategy()
+    strategy = get_tpu_strategy(enable_packed_var)
     with strategy.scope():
       variable = variables.Variable(0.0)
 
@@ -527,7 +548,10 @@ class TPUStrategyTest(test.TestCase):
 
     with strategy.scope():
       update_variable.get_concrete_function()
-      self.assertEqual(trace_count[0], len(strategy.extended.worker_devices))
+      self.assertLen(strategy.extended.worker_devices, trace_count[0])
+
+
+class TPUStrategyDataPrefetchTest(test.TestCase):
 
   def test_prefetch_to_device_default(self):
     strategy = get_tpu_strategy()
diff --git a/tensorflow/python/distribute/tpu_values.py b/tensorflow/python/distribute/tpu_values.py
index 40ab058ac7c..33885531966 100644
--- a/tensorflow/python/distribute/tpu_values.py
+++ b/tensorflow/python/distribute/tpu_values.py
@@ -24,6 +24,7 @@ from __future__ import print_function
 
 import contextlib
 
+from tensorflow.python.distribute import packed_distributed_variable as packed
 from tensorflow.python.distribute import values
 from tensorflow.python.eager import context
 from tensorflow.python.eager import tape
@@ -46,15 +47,27 @@ def _maybe_enter_graph(tensor):
       yield
 
 
+@contextlib.contextmanager
+def _maybe_on_device(var):
+  # Add a device scope for packed variables.
+  if isinstance(var, packed.PackedVarAndDevice):
+    with ops.device(var.device):
+      yield
+  else:
+    yield
+
+
 def _make_raw_assign_fn(raw_assign_fn):  # pylint: disable=missing-docstring
 
   def assign_fn(var, value, use_locking=False, name=None, read_value=True):  # pylint: disable=missing-docstring
     del use_locking  # Unused.
 
-    with _maybe_enter_graph(var.handle):
+    handle = var.handle
+    with _maybe_enter_graph(handle), _maybe_on_device(var):
       op = raw_assign_fn(
-          var.handle, ops.convert_to_tensor(value, dtype=var.dtype), name=name)
-
+          handle,
+          ops.convert_to_tensor(value, dtype=var.dtype),
+          name=name)
       with ops.control_dependencies([op]):
         return var._read_variable_op() if read_value else op  # pylint: disable=protected-access
 
@@ -97,23 +110,37 @@ class TPUVariableMixin(object):
 
   @property
   def handle(self):
+    """The handle by which this variable can be accessed."""
     # If we're in a tpu.rewrite(), return the replicated handle.
     tpu_context = enclosing_tpu_context()
     if tpu_context is None or context.executing_eagerly():
       return self._get_on_device_or_primary().handle
     else:
-      return tpu_context.get_replicated_var_handle(self._handle_id,
-                                                   self._values,
-                                                   self._is_mirrored())
+      is_packed = self._packed_var is not None
+      val = self._values
+      if is_packed:
+        val = [self._packed_var]
+
+      return tpu_context.get_replicated_var_handle(self._handle_id, val,
+                                                   self._is_mirrored(),
+                                                   is_packed)
 
   @property
   def device(self):
     return self.handle.device
 
   def _read_variable_op(self):
+    """Reads the value of this variable."""
     if self.trainable:
       tape.variable_accessed(self)
-    return gen_resource_variable_ops.read_variable_op(self.handle, self.dtype)
+
+    handle = self.handle
+    if getattr(handle, "is_packed", False):
+      # Add a device scope for a packed variable handle.
+      with ops.device(self._get_on_device_or_primary().device):
+        return gen_resource_variable_ops.read_variable_op(handle, self.dtype)
+    else:
+      return gen_resource_variable_ops.read_variable_op(handle, self.dtype)
 
   def read_value(self):
     if enclosing_tpu_context() is None:
diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index 60b2ea4fe31..37643e03b18 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -472,6 +472,12 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
     # variable.
     self._var_policy = var_policy
 
+  @property
+  def _devices(self):
+    if self._packed_var is not None:
+      return tuple(d for d in self._packed_var.devices)
+    return tuple(v.device for v in self._values)
+
   def is_initialized(self, name=None):
     """Identifies if all the component variables are initialized.
 
@@ -482,6 +488,8 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       The op that evaluates to True or False depending on if all the
       component variables are initialized.
     """
+    if self._packed_var is not None:
+      return self._packed_var.is_initialized()
     result = self._primary.is_initialized()
     # We iterate through the list of values except the last one to allow us to
     # name the final `logical_and` op the same name that is passed by the user
@@ -552,6 +560,10 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
   def aggregation(self):
     return self._aggregation
 
+  @property
+  def _packed_variable(self):
+    return self._packed_var
+
   @property
   def handle(self):
     replica_id = values_util.get_current_replica_id_as_int()
@@ -559,6 +571,8 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       raise ValueError("`handle` is not available outside the replica context"
                        " or a `tf.distribute.Strategy.update()` call.")
     else:
+      if self._packed_var is not None:
+        return self._packed_var.handle
       return self._values[replica_id].handle
 
   def eval(self, session=None):
@@ -607,6 +621,33 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
   def _in_graph_mode(self):
     return self._primary._in_graph_mode  # pylint: disable=protected-access
 
+  def _get_replica(self, replica_id):
+    """Returns the value on a device with the given replica_id."""
+    if self._packed_var is not None:
+      return self._packed_var.on_device(self._devices[replica_id])
+    return self._values[replica_id]
+
+  def _get(self):
+    """Returns the value for the current device or raises a ValueError."""
+    replica_id = values_util.get_current_replica_id_as_int()
+    if replica_id is None:
+      return self._get_cross_replica()
+    else:
+      return self._get_replica(replica_id)
+
+  def _get_on_device_or_primary(self):
+    """Returns value in same replica or device if possible, else the _primary."""
+    replica_id = values_util.get_current_replica_id_as_int()
+    if replica_id is None:
+      # Try to find a value on the current device.
+      current_device = device_util.canonicalize(device_util.current())
+      for i, value in enumerate(self._values):
+        if device_util.canonicalize(value.device) == current_device:
+          return self._get_replica(i)
+      return self._get_replica(0)
+    else:
+      return self._get_replica(replica_id)
+
   def read_value(self):
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       return array_ops.identity(self._get())
@@ -778,7 +819,8 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       if ds_context.in_cross_replica_context():
         update_replica_id = distribute_lib.get_update_replica_id()
         if update_replica_id is not None:
-          return update_fn(self._values[update_replica_id], value, **kwargs)
+          replica_value = self._get_replica(update_replica_id)
+          return update_fn(replica_value, value, **kwargs)
         return self._update_cross_replica(update_fn, value, **kwargs)
       else:
         values_util.assert_replica_context(self.distribute_strategy)
@@ -802,6 +844,7 @@ class DistributedVariable(DistributedDelegate, variables_lib.Variable,
       obj_map[v] = new_obj
       resource_map[v.handle] = new_obj.handle
     obj_map[self] = new_obj
+    resource_map[self.handle] = new_obj.handle
     resource_map[self] = new_obj.handle
     return obj_map, resource_map
 
@@ -835,6 +878,12 @@ class _MirroredSaveable(saveable_object_util.ResourceVariableSaveable):
   def restore(self, restored_tensors, restored_shapes):
     """Restore the same value into all variables."""
     tensor, = restored_tensors
+    packed_var = self._mirrored_variable._packed_variable  # pylint: disable=protected-access
+    if packed_var is not None:
+      return control_flow_ops.group(
+          tuple(
+              values_util.assign_on_device(d, packed_var, tensor)
+              for d in packed_var.devices))
     return control_flow_ops.group(
         tuple(
             values_util.assign_on_device(v.device, v, tensor)
@@ -1013,7 +1062,7 @@ class SyncOnReadVariable(DistributedVariable):
 
   def _get_cross_replica(self):
     if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
-      return self._primary
+      return self._get_replica(0)
 
     with ds_context.enter_or_assert_strategy(self._distribute_strategy):
       return self._distribute_strategy.reduce(
diff --git a/tensorflow/python/distribute/values_test.py b/tensorflow/python/distribute/values_test.py
index 0cb4d6ddd2a..d0e3eec22a8 100644
--- a/tensorflow/python/distribute/values_test.py
+++ b/tensorflow/python/distribute/values_test.py
@@ -42,7 +42,6 @@ from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import device
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import indexed_slices
 from tensorflow.python.framework import ops
@@ -234,11 +233,11 @@ class DistributedValuesTest(test.TestCase, parameterized.TestCase):
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
               # TODO(b/137795644): support CentralStroageStrategy
               # strategy_combinations.central_storage_strategy_with_two_gpus,
           ],
-          mode=["eager"]
-      ))
+          mode=["eager"]))
   def testMakeDistributedValueDefaultDevicePlacement(self, distribution):
     if not tf2.enabled():
       self.skipTest("Only V2 is supported.")
@@ -259,11 +258,11 @@ class DistributedValuesTest(test.TestCase, parameterized.TestCase):
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
               # TODO(b/137795644): support CentralStroageStrategy
               # strategy_combinations.central_storage_strategy_with_two_gpus,
           ],
-          mode=["eager"]
-      ))
+          mode=["eager"]))
   def testMakeDistributedValueExplicitDevicePlacement(self, distribution):
     if not tf2.enabled():
       self.skipTest("Only V2 is supported.")
@@ -384,6 +383,16 @@ def _make_mirrored():
   return mirrored
 
 
+def mirrored_and_tpu_strategy_combinations():
+  return combinations.combine(
+      distribution=[
+          strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          strategy_combinations.tpu_strategy,
+          strategy_combinations.tpu_strategy_packed_var,
+      ],
+      mode=["graph", "eager"])
+
+
 class RegroupAndSelectDeviceTest(test.TestCase, parameterized.TestCase):
 
   def _is_per_replica(self, result, expected, klass=values.PerReplica):
@@ -563,6 +572,7 @@ class RegroupAndSelectDeviceTest(test.TestCase, parameterized.TestCase):
             strategy_combinations.mirrored_strategy_with_one_cpu,
             strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
             strategy_combinations.tpu_strategy,
+            strategy_combinations.tpu_strategy_packed_var,
             strategy_combinations.central_storage_strategy_with_two_gpus,
         ],
         synchronization=[
@@ -708,29 +718,40 @@ class DistributedVariableTest(test.TestCase, parameterized.TestCase):
         self.evaluate(
             distribution.experimental_local_results(distribution.run(assign)))
 
-  def testPackedVariable(self, distribution, synchronization, aggregation):
+
+@combinations.generate(
+    combinations.combine(
+        distribution=[
+            strategy_combinations.mirrored_strategy_with_one_cpu,
+            strategy_combinations.tpu_strategy,
+        ],
+        mode=["eager"]))
+class PackedDistributedVariableTest(test.TestCase, parameterized.TestCase):
+
+  def testPackedVariable(self, distribution):
     with distribution.scope():
-      v0 = variables_lib.Variable(
-          0., synchronization=synchronization, aggregation=aggregation)
-      if not isinstance(v0, values.DistributedVariable):
-        self.skipTest("This test doesn't apply to non DistributedVariables")
-
-    self.assertEqual(v0._packed_var, None)
-
-    device_type = device.DeviceSpec.from_string(v0._devices[0]).device_type
-    for d in v0._devices:
-      if device.DeviceSpec.from_string(d).device_type != device_type:
-        self.skipTest("Packing variables on devices of different types "
-                      "is not supported yet.")
+      v0 = variables_lib.Variable(0.)
+    self.assertIsNone(v0._packed_var)
 
     distribution._enable_packed_variable_in_eager_mode = True
     with distribution.scope():
-      v1 = variables_lib.Variable(
-          0., synchronization=synchronization, aggregation=aggregation)
-    if ops.executing_eagerly_outside_functions():
+      v1 = variables_lib.Variable(0)
       self.assertIsInstance(v1._packed_var, packed.PackedDistributedVariable)
-    else:
-      self.assertEqual(v1._packed_var, None)
+
+    devices = v1._devices
+    for i in range(1, len(devices)):
+      with distribute_lib.ReplicaContext(distribution, i):
+        v1.assign(i)
+    val = v1._get()
+    self.assertIsInstance(val, packed.PackedVarAndDevice)
+    self.assertEqual(val.device, devices[0])
+    self.assertEqual(self.evaluate(val.read_value()), 0)
+    for i in range(0, len(devices)):
+      with distribute_lib.ReplicaContext(distribution, i):
+        val = v1._get()
+        self.assertIsInstance(val, packed.PackedVarAndDevice)
+        self.assertEqual(val.device, devices[i])
+        self.assertEqual(self.evaluate(val.read_value()), i)
 
 
 class MirroredVariableTest(test.TestCase, parameterized.TestCase):
@@ -920,6 +941,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["eager"]))
   def testAssignValueInReplicaContextWithoutAggregation(self, distribution):
@@ -943,6 +965,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
               strategy_combinations.mirrored_strategy_with_one_cpu,
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["graph", "eager"]))
   def testValueInReplicaContext(self, distribution):
@@ -968,6 +991,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
               strategy_combinations.mirrored_strategy_with_one_cpu,
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["graph", "eager"]))
   def testAssignOutOfScope(self, distribution):
@@ -1041,6 +1065,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["eager"]))
   def testInitializedToSameValueInsideEagerRun(self, distribution):
@@ -1066,6 +1091,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
               strategy_combinations.mirrored_strategy_with_one_cpu,
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["graph", "eager"]))
   def testAggregationOnlyFirstReplica(self, distribution):
@@ -1093,6 +1119,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
           distribution=[
               strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["eager"]))
   def testInitScope(self, distribution):
@@ -1143,13 +1170,7 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
         distribution.experimental_local_results(distribution.run(add)))
     self.assertAllEqual([2, 2], per_replica_results)
 
-  @combinations.generate(
-      combinations.combine(
-          distribution=[
-              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
-              strategy_combinations.tpu_strategy,
-          ],
-          mode=["graph", "eager"]))
+  @combinations.generate(mirrored_and_tpu_strategy_combinations())
   def testAssignAdd(self, distribution):
     with distribution.scope():
       v = variable_scope.variable(
@@ -1456,15 +1477,6 @@ class SyncOnReadVariablePropertiesTest(test.TestCase):
     self.assertEqual(2., self.evaluate(add1(replica_local)))
 
 
-def mirrored_and_tpu_strategy_combinations():
-  return combinations.combine(
-      distribution=[
-          strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
-          strategy_combinations.tpu_strategy,
-      ],
-      mode=["graph", "eager"])
-
-
 # TODO(b/144432582): Add variable aggregation type to combinations to simplify
 # tests.
 def strategy_and_run_tf_function_combinations():
@@ -1478,6 +1490,7 @@ def strategy_and_run_tf_function_combinations():
       experimental_run_tf_function=[True, False]) + combinations.combine(
           distribution=[
               strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
           ],
           mode=["graph", "eager"],
           experimental_run_tf_function=[True])
diff --git a/tensorflow/python/distribute/values_util.py b/tensorflow/python/distribute/values_util.py
index ddb0d2d0401..5909bdd229e 100644
--- a/tensorflow/python/distribute/values_util.py
+++ b/tensorflow/python/distribute/values_util.py
@@ -61,8 +61,14 @@ def on_write_assign_sub(var, value, use_locking=False, name=None,
 
 
 def assign_on_each_device(var, assign_func, value, read_value):
-  update = control_flow_ops.group(
-      tuple(assign_func(v.device, v, value) for v in var._values))  # pylint: disable=protected-access
+  """Update the variable on each replica with the given assign_func and value."""
+  if var._packed_variable is not None:  # pylint: disable=protected-access
+    update = control_flow_ops.group(
+        tuple(
+            assign_func(d, var._packed_variable, value) for d in var._devices))  # pylint: disable=protected-access
+  else:
+    update = control_flow_ops.group(
+        tuple(assign_func(v.device, v, value) for v in var._values))  # pylint: disable=protected-access
   if not read_value:
     return update
   with ops.control_dependencies([update] if update else []):
@@ -104,7 +110,7 @@ def on_read_assign_cross_replica(var, value, read_value=True):
       # TODO(anjs): Should this be over all the replicas in sync since we
       # call `reduce` on the variable during read?
       if var.aggregation == vs.VariableAggregation.SUM:
-        tensor = math_ops.cast(tensor / len(var._values), var.dtype)  # pylint: disable=protected-access
+        tensor = math_ops.cast(tensor / len(var._devices), var.dtype)  # pylint: disable=protected-access
       return assign_on_each_device(var, assign_on_device, tensor,
                                    read_value)
 
diff --git a/tensorflow/python/tpu/tpu.py b/tensorflow/python/tpu/tpu.py
index ce3aaa8a058..6f5f0bc26c2 100644
--- a/tensorflow/python/tpu/tpu.py
+++ b/tensorflow/python/tpu/tpu.py
@@ -298,7 +298,8 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
     self._pivot = pivot
     self._replicated_vars = {}
 
-  def get_replicated_var_handle(self, name, vars_, is_mirrored=False):
+  def get_replicated_var_handle(self, name, vars_, is_mirrored=False,
+                                is_packed=False):
     """Returns a variable handle for replicated TPU variable 'var'.
 
     This is a method used by an experimental replicated variable implementation
@@ -309,6 +310,7 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
       vars_: The replicated TPU variables.
       is_mirrored: Whether the variables are mirrored, which guarantees the
         values in each replica are always the same.
+      is_packed: Whether the replicated variables are packed into one variable.
 
     Returns:
       The handle of the TPU replicated input node.
@@ -320,7 +322,7 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
     if handle is not None:
       return handle
 
-    if device_assignment is not None:
+    if device_assignment is not None and not is_packed:
       # Find a variable copy for each replica in the device assignment.
       # Note that the order of devices for replicas for the variable and the
       # device assignment might not match.
@@ -356,7 +358,8 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext):
       graph._set_control_flow_context(self.outer_context)
       handle = tpu_ops.tpu_replicated_input([v.handle for v in replicated_vars],
                                             name=name + "/handle",
-                                            is_mirrored_variable=is_mirrored)
+                                            is_mirrored_variable=is_mirrored,
+                                            is_packed=is_packed)
       graph._set_control_flow_context(saved_context)
       # pylint: enable=protected-access
     self._replicated_vars[name] = handle

From 62082d40720b56974436dd17625c247a5fce2a6b Mon Sep 17 00:00:00 2001
From: YoungSeok Yoon 
Date: Thu, 18 Jun 2020 20:16:19 -0700
Subject: [PATCH 0588/1390] Add build flags for objc libraries

PiperOrigin-RevId: 317235962
Change-Id: I976ccd1ce3db49be3acac44f60b2dc44ed25d767
---
 tensorflow/lite/experimental/objc/BUILD.apple | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/experimental/objc/BUILD.apple b/tensorflow/lite/experimental/objc/BUILD.apple
index ff7e8fa58e9..09d4547813a 100644
--- a/tensorflow/lite/experimental/objc/BUILD.apple
+++ b/tensorflow/lite/experimental/objc/BUILD.apple
@@ -97,7 +97,7 @@ objc_library(
         "//tensorflow/lite:testdata/add.bin",
         "//tensorflow/lite:testdata/add_quantized.bin",
     ],
-    tags = TFL_DEFAULT_TAGS,
+    tags = TFL_DEFAULT_TAGS + ["builder_default_ios_x86_64"],
     deps = [
         ":TensorFlowLite",
     ],
@@ -135,7 +135,10 @@ objc_library(
         "apis",
     ],
     module_name = "TestApp",
-    tags = TFL_DEFAULT_TAGS + ["manual"],
+    tags = TFL_DEFAULT_TAGS + [
+        "manual",
+        "builder_default_ios_x86_64",
+    ],
     deps = [
         ":TensorFlowLite",
     ],

From 85ad8031f60536361de71dd689c9d88848fefed6 Mon Sep 17 00:00:00 2001
From: Gaurav Jain 
Date: Thu, 18 Jun 2020 20:27:03 -0700
Subject: [PATCH 0589/1390] Expand dtype support for Neg

PiperOrigin-RevId: 317237033
Change-Id: I59c5e45d469f7bf704976b66bc122aaac3982b5e
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    |  4 +--
 tensorflow/core/kernels/BUILD                 |  3 ++-
 .../core/kernels/cwise_op_gpu_neg.cu.cc       |  4 +--
 .../{cwise_op_neg.cc => cwise_op_neg_1.cc}    |  6 ++---
 tensorflow/core/kernels/cwise_op_neg_2.cc     | 26 +++++++++++++++++++
 tensorflow/core/ops/math_ops.cc               | 12 ++++-----
 .../kernel_tests/cwise_ops_unary_test.py      |  6 +++++
 7 files changed, 46 insertions(+), 15 deletions(-)
 rename tensorflow/core/kernels/{cwise_op_neg.cc => cwise_op_neg_1.cc} (87%)
 create mode 100644 tensorflow/core/kernels/cwise_op_neg_2.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index dcd083fc398..3b1f3eec699 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -6059,11 +6059,11 @@ I.e., \\(y = -x\\).
   }];
 
   let arguments = (ins
-    TensorOf<[BF16, F16, F32, F64, I32, I64, TF_Complex128, TF_Complex64]>:$x
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64]>:$x
   );
 
   let results = (outs
-    TensorOf<[BF16, F16, F32, F64, I32, I64, TF_Complex128, TF_Complex64]>:$y
+    TensorOf<[BF16, F16, F32, F64, I16, I32, I64, I8, TF_Complex128, TF_Complex64]>:$y
   );
 
   TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index ffe2a035591..279dff92c58 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -6802,7 +6802,8 @@ filegroup(
         "cwise_op_minimum.cc",
         "cwise_op_mul_1.cc",
         "cwise_op_mul_2.cc",
-        "cwise_op_neg.cc",
+        "cwise_op_neg_1.cc",
+        "cwise_op_neg_2.cc",
         "cwise_op_pow.cc",
         "cwise_op_real.cc",
         "cwise_op_reciprocal.cc",
diff --git a/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc b/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc
index ea1ca623560..4f7bb9b2075 100644
--- a/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc
+++ b/tensorflow/core/kernels/cwise_op_gpu_neg.cu.cc
@@ -19,8 +19,8 @@ limitations under the License.
 
 namespace tensorflow {
 namespace functor {
-DEFINE_UNARY7(neg, Eigen::half, float, double, int32, int64, complex64,
-              complex128);
+DEFINE_UNARY4(neg, int8, int16, int32, int64);
+DEFINE_UNARY6(neg, Eigen::half, float, double, bfloat16, complex64, complex128);
 }  // namespace functor
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/kernels/cwise_op_neg.cc b/tensorflow/core/kernels/cwise_op_neg_1.cc
similarity index 87%
rename from tensorflow/core/kernels/cwise_op_neg.cc
rename to tensorflow/core/kernels/cwise_op_neg_1.cc
index f52cf6c8b91..18a7c61be90 100644
--- a/tensorflow/core/kernels/cwise_op_neg.cc
+++ b/tensorflow/core/kernels/cwise_op_neg_1.cc
@@ -16,8 +16,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/cwise_ops_common.h"
 
 namespace tensorflow {
-REGISTER8(UnaryOp, CPU, "Neg", functor::neg, float, Eigen::half, double, int32,
-          complex64, int64, complex128, bfloat16);
+REGISTER4(UnaryOp, CPU, "Neg", functor::neg, int8, int16, int32, int64);
 
 #ifdef TENSORFLOW_USE_SYCL
 REGISTER3(UnaryOp, SYCL, "Neg", functor::neg, float, double, int64);
@@ -30,8 +29,7 @@ REGISTER_KERNEL_BUILDER(Name("Neg")
 #endif  // TENSORFLOW_USE_SYCL
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-REGISTER6(UnaryOp, GPU, "Neg", functor::neg, float, Eigen::half, double, int64,
-          complex64, complex128);
+REGISTER3(UnaryOp, GPU, "Neg", functor::neg, int8, int16, int64);
 
 // A special GPU kernel for int32.
 // TODO(b/25387198): Also enable int32 in device memory. This kernel
diff --git a/tensorflow/core/kernels/cwise_op_neg_2.cc b/tensorflow/core/kernels/cwise_op_neg_2.cc
new file mode 100644
index 00000000000..5ea78ad665c
--- /dev/null
+++ b/tensorflow/core/kernels/cwise_op_neg_2.cc
@@ -0,0 +1,26 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/kernels/cwise_ops_common.h"
+
+namespace tensorflow {
+REGISTER6(UnaryOp, CPU, "Neg", functor::neg, Eigen::half, float, double,
+          bfloat16, complex64, complex128);
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+REGISTER6(UnaryOp, GPU, "Neg", functor::neg, Eigen::half, float, double,
+          bfloat16, complex64, complex128);
+#endif
+}  // namespace tensorflow
diff --git a/tensorflow/core/ops/math_ops.cc b/tensorflow/core/ops/math_ops.cc
index b81bb9d3afc..2a70f420260 100644
--- a/tensorflow/core/ops/math_ops.cc
+++ b/tensorflow/core/ops/math_ops.cc
@@ -201,12 +201,12 @@ REGISTER_OP("ComplexAbs")
     .SetShapeFn(shape_inference::UnchangedShape);
 
 // Declares cwise unary operations signature: 't -> 't
-#define UNARY()                                                          \
-  Input("x: T")                                                          \
-      .Output("y: T")                                                    \
-      .Attr(                                                             \
-          "T: {bfloat16, half, float, double, int32, int64, complex64, " \
-          "complex128}")                                                 \
+#define UNARY()                                                            \
+  Input("x: T")                                                            \
+      .Output("y: T")                                                      \
+      .Attr(                                                               \
+          "T: {bfloat16, half, float, double, int8, int16, int32, int64, " \
+          "complex64, complex128}")                                        \
       .SetShapeFn(shape_inference::UnchangedShape)
 
 #define UNARY_REAL()                              \
diff --git a/tensorflow/python/kernel_tests/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/cwise_ops_unary_test.py
index f4beaabc29a..df848a653d4 100644
--- a/tensorflow/python/kernel_tests/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/cwise_ops_unary_test.py
@@ -389,16 +389,22 @@ class UnaryOpTest(test.TestCase):
                   2).reshape(1, 3, 2).astype(dtypes_lib.bfloat16.as_numpy_dtype)
     self._compareCpu(x, np.abs, math_ops.abs)
     self._compareCpu(x, np.abs, _ABS)
+    self._compareBoth(x, np.negative, math_ops.negative)
+    self._compareBoth(x, np.negative, _NEG)
 
   def testInt8Basic(self):
     x = np.arange(-6, 6, 2).reshape(1, 3, 2).astype(np.int8)
     self._compareCpu(x, np.abs, math_ops.abs)
     self._compareCpu(x, np.abs, _ABS)
+    self._compareBoth(x, np.negative, math_ops.negative)
+    self._compareBoth(x, np.negative, _NEG)
 
   def testInt16Basic(self):
     x = np.arange(-6, 6, 2).reshape(1, 3, 2).astype(np.int16)
     self._compareCpu(x, np.abs, math_ops.abs)
     self._compareCpu(x, np.abs, _ABS)
+    self._compareBoth(x, np.negative, math_ops.negative)
+    self._compareBoth(x, np.negative, _NEG)
 
   def testInt32Basic(self):
     x = np.arange(-6, 6, 2).reshape(1, 3, 2).astype(np.int32)

From 2a05589bd4f5e3042d1baf539e564d7ab9bd6287 Mon Sep 17 00:00:00 2001
From: Taehee Jeong 
Date: Thu, 18 Jun 2020 20:48:49 -0700
Subject: [PATCH 0590/1390] Add inference instruction for iOS

PiperOrigin-RevId: 317239235
Change-Id: I55bd7e43bc286f34024ccfc27db61d28304a651d
---
 tensorflow/lite/g3doc/guide/inference.md | 425 +++++++++++++++--------
 1 file changed, 287 insertions(+), 138 deletions(-)

diff --git a/tensorflow/lite/g3doc/guide/inference.md b/tensorflow/lite/g3doc/guide/inference.md
index 5f3fba98cff..6e47d6d5190 100644
--- a/tensorflow/lite/g3doc/guide/inference.md
+++ b/tensorflow/lite/g3doc/guide/inference.md
@@ -7,9 +7,9 @@ inference with a TensorFlow Lite model, you must run it through an
 The interpreter uses a static graph ordering and a custom (less-dynamic) memory
 allocator to ensure minimal load, initialization, and execution latency.
 
-This page describes how to access to the TensorFlow Lite interpreter and
-perform an inference using C++, Java, and Python, plus links to other resources
-for each [supported platform](#supported-platforms).
+This page describes how to access to the TensorFlow Lite interpreter and perform
+an inference using C++, Java, and Python, plus links to other resources for each
+[supported platform](#supported-platforms).
 
 [TOC]
 
@@ -17,31 +17,31 @@ for each [supported platform](#supported-platforms).
 
 TensorFlow Lite inference typically follows the following steps:
 
-1. **Loading a model**
+1.  **Loading a model**
 
-   You must load the `.tflite` model into memory, which contains the model's
-   execution graph.
+    You must load the `.tflite` model into memory, which contains the model's
+    execution graph.
 
-1. **Transforming data**
+1.  **Transforming data**
 
-   Raw input data for the model generally does not match the input data format
-   expected by the model. For example, you might need to resize an image or
-   change the image format to be compatible with the model.
+    Raw input data for the model generally does not match the input data format
+    expected by the model. For example, you might need to resize an image or
+    change the image format to be compatible with the model.
 
-1. **Running inference**
+1.  **Running inference**
 
-   This step involves using the TensorFlow Lite API to execute the model. It
-   involves a few steps such as building the interpreter, and allocating
-   tensors, as described in the following sections.
+    This step involves using the TensorFlow Lite API to execute the model. It
+    involves a few steps such as building the interpreter, and allocating
+    tensors, as described in the following sections.
 
-1. **Interpreting output**
+1.  **Interpreting output**
 
-   When you receive results from the model inference, you must interpret the
-   tensors in a meaningful way that's useful in your application.
+    When you receive results from the model inference, you must interpret the
+    tensors in a meaningful way that's useful in your application.
 
-   For example, a model might return only a list of probabilities. It's up to
-   you to map the probabilities to relevant categories and present it to your
-   end-user.
+    For example, a model might return only a list of probabilities. It's up to
+    you to map the probabilities to relevant categories and present it to your
+    end-user.
 
 ## Supported platforms
 
@@ -54,8 +54,8 @@ should be no surprise that the APIs try to avoid unnecessary copies at the
 expense of convenience. Similarly, consistency with TensorFlow APIs was not an
 explicit goal and some variance between languages is to be expected.
 
-Across all libraries, the TensorFlow Lite API enables you to load models,
-feed inputs, and retrieve inference outputs.
+Across all libraries, the TensorFlow Lite API enables you to load models, feed
+inputs, and retrieve inference outputs.
 
 ### Android
 
@@ -64,8 +64,8 @@ APIs. The Java APIs provide convenience and can be used directly within your
 Android Activity classes. The C++ APIs offer more flexibility and speed, but may
 require writing JNI wrappers to move data between Java and C++ layers.
 
-See below for details about using C++ and Java, or
-follow the [Android quickstart](android.md) for a tutorial and example code.
+See below for details about using C++ and Java, or follow the
+[Android quickstart](android.md) for a tutorial and example code.
 
 #### TensorFlow Lite Android wrapper code generator
 
@@ -86,103 +86,36 @@ On iOS, TensorFlow Lite is available with native iOS libraries written in
 [Swift](https://www.tensorflow.org/code/tensorflow/lite/experimental/swift)
 and
 [Objective-C](https://www.tensorflow.org/code/tensorflow/lite/experimental/objc).
+You can also use
+[C API](https://www.tensorflow.org/code/tensorflow/lite/c/c_api.h)
+directly in Objective-C codes.
 
-This page doesn't include a discussion for about these languages, so you should
-refer to the [iOS quickstart](ios.md) for a tutorial and example code.
+See below for details about using Swift, Objective-C and C API, or follow the
+[iOS quickstart](ios.md) for a tutorial and example code.
 
 ### Linux
 
 On Linux platforms (including [Raspberry Pi](build_rpi.md)), you can run
-inferences using TensorFlow Lite APIs available in C++ and Python, as shown
-in the following sections.
+inferences using TensorFlow Lite APIs available in C++ and Python, as shown in
+the following sections.
 
+## Running a model
 
-## Load and run a model in C++
+Running a TensorFlow Lite model involves a few simple steps:
 
-Running a TensorFlow Lite model with C++ involves a few simple steps:
-
-  1. Load the model into memory as a `FlatBufferModel`.
-  2. Build an `Interpreter` based on an existing `FlatBufferModel`.
-  3. Set input tensor values. (Optionally resize input tensors if the
-     predefined sizes are not desired.)
-  4. Invoke inference.
-  5. Read output tensor values.
-
-The [`FlatBufferModel`](
-https://www.tensorflow.org/lite/api_docs/cc/class/tflite/flat-buffer-model.html)
-class encapsulates a TensorFlow Lite model and you can
-build it in a couple of different ways, depending on where the model is stored:
-
-```c++
-class FlatBufferModel {
-  // Build a model based on a file. Return a nullptr in case of failure.
-  static std::unique_ptr BuildFromFile(
-      const char* filename,
-      ErrorReporter* error_reporter);
-
-  // Build a model based on a pre-loaded flatbuffer. The caller retains
-  // ownership of the buffer and should keep it alive until the returned object
-  // is destroyed. Return a nullptr in case of failure.
-  static std::unique_ptr BuildFromBuffer(
-      const char* buffer,
-      size_t buffer_size,
-      ErrorReporter* error_reporter);
-};
-```
-
-Note: If TensorFlow Lite detects the presence of the [Android NNAPI](
-https://developer.android.com/ndk/guides/neuralnetworks), it will
-automatically try to use shared memory to store the `FlatBufferModel`.
-
-Now that you have the model as a `FlatBufferModel` object, you can execute it
-with an [`Interpreter`](
-https://www.tensorflow.org/lite/api_docs/cc/class/tflite/interpreter.html).
-A single `FlatBufferModel` can be used
-simultaneously by more than one `Interpreter`.
-
-Caution: The `FlatBufferModel` object must remain valid until
-all instances of `Interpreter` using it have been destroyed.
-
-The important parts of the `Interpreter` API are shown in the
-code snippet below. It should be noted that:
-
-  * Tensors are represented by integers, in order to avoid string comparisons
-    (and any fixed dependency on string libraries).
-  * An interpreter must not be accessed from concurrent threads.
-  * Memory allocation for input and output tensors must be triggered
-    by calling `AllocateTensors()` right after resizing tensors.
-
-The simplest usage of TensorFlow Lite with C++ looks like this:
-
-```c++
-// Load the model
-std::unique_ptr model =
-    tflite::FlatBufferModel::BuildFromFile(filename);
-
-// Build the interpreter
-tflite::ops::builtin::BuiltinOpResolver resolver;
-std::unique_ptr interpreter;
-tflite::InterpreterBuilder(*model, resolver)(&interpreter);
-
-// Resize input tensors, if desired.
-interpreter->AllocateTensors();
-
-float* input = interpreter->typed_input_tensor(0);
-// Fill `input`.
-
-interpreter->Invoke();
-
-float* output = interpreter->typed_output_tensor(0);
-```
-
-For more example code, see [`minimal.cc`](
-https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/minimal/minimal.cc)
-and [`label_image.cc`](
-https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/label_image/label_image.cc).
+1.  Load the model into memory.
+2.  Build an `Interpreter` based on an existing model.
+3.  Set input tensor values. (Optionally resize input tensors if the predefined
+    sizes are not desired.)
+4.  Invoke inference.
+5.  Read output tensor values.
 
+Following sections describe how these steps can be done in each language.
 
 ## Load and run a model in Java
 
+*Platform: Android*
+
 The Java API for running an inference with TensorFlow Lite is primarily designed
 for use with Android, so it's available as an Android library dependency:
 `org.tensorflow:tensorflow-lite`.
@@ -203,12 +136,12 @@ public Interpreter(@NotNull MappedByteBuffer mappedByteBuffer);
 ```
 
 In both cases, you must provide a valid TensorFlow Lite model or the API throws
-`IllegalArgumentException`. If you use `MappedByteBuffer` to
-initialize an `Interpreter`, it must remain unchanged for the whole lifetime
-of the `Interpreter`.
+`IllegalArgumentException`. If you use `MappedByteBuffer` to initialize an
+`Interpreter`, it must remain unchanged for the whole lifetime of the
+`Interpreter`.
 
-To then run an inference with the model, simply call `Interpreter.run()`.
-For example:
+To then run an inference with the model, simply call `Interpreter.run()`. For
+example:
 
 ```java
 try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
@@ -228,9 +161,9 @@ In this case, each entry in `inputs` corresponds to an input tensor and
 output data.
 
 In both cases, the tensor indices should correspond to the values you gave to
-the [TensorFlow Lite Converter](../convert/) when you created the model.
-Be aware that the order of tensors in `input` must match the
-order given to the TensorFlow Lite Converter.
+the [TensorFlow Lite Converter](../convert/) when you created the model. Be
+aware that the order of tensors in `input` must match the order given to the
+TensorFlow Lite Converter.
 
 The `Interpreter` class also provides convenient functions for you to get the
 index of any model input or output using an operation name:
@@ -250,8 +183,8 @@ resources must be released after use by:
 interpreter.close();
 ```
 
-For an example project with Java, see the [Android image classification sample](
-https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/android).
+For an example project with Java, see the
+[Android image classification sample](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/android).
 
 ### Supported data types (in Java)
 
@@ -295,13 +228,231 @@ have dynamic outputs, where the shape of output tensors can vary depending on
 the input. There's no straightforward way of handling this with the existing
 Java inference API, but planned extensions will make this possible.
 
+## Load and run a model in Swift
+
+*Platform: iOS*
+
+The
+[Swift API](https://www.tensorflow.org/code/tensorflow/lite/experimental/swift)
+is available in `TensorFlowLiteSwift` Pod from Cocoapods.
+
+First, you need to import `TensorFlowLite` module.
+
+```swift
+import TensorFlowLite
+```
+
+```swift
+// Getting model path
+guard
+  let modelPath = Bundle.main.path(forResource: "model", ofType: "tflite")
+else {
+  // Error handling...
+}
+
+do {
+  // Initialize an interpreter with the model.
+  let interpreter = try Interpreter(modelPath: modelPath)
+
+  // Allocate memory for the model's input `Tensor`s.
+  try interpreter.allocateTensors()
+
+  let inputData: Data  // Should be initialized
+
+  // input data preparation...
+
+  // Copy the input data to the input `Tensor`.
+  try self.interpreter.copy(inputData, toInputAt: 0)
+
+  // Run inference by invoking the `Interpreter`.
+  try self.interpreter.invoke()
+
+  // Get the output `Tensor`
+  let outputTensor = try self.interpreter.output(at: 0)
+
+  // Copy output to `Data` to process the inference results.
+  let outputSize = outputTensor.shape.dimensions.reduce(1, {x, y in x * y})
+  let outputData =
+        UnsafeMutableBufferPointer.allocate(capacity: outputSize)
+  outputTensor.data.copyBytes(to: outputData)
+
+  if (error != nil) { /* Error handling... */ }
+} catch error {
+  // Error handling...
+}
+```
+
+## Load and run a model in Objective-C
+
+*Platform: iOS*
+
+The
+[Objective-C API](https://www.tensorflow.org/code/tensorflow/lite/experimental/objc)
+is available in `TensorFlowLiteObjC` Pod from Cocoapods.
+
+First, you need to import `TensorFlowLite` module.
+
+```objc
+@import TensorFlowLite;
+```
+
+```objc
+NSString *modelPath = [[NSBundle mainBundle] pathForResource:@"model"
+                                                      ofType:@"tflite"];
+NSError *error;
+
+// Initialize an interpreter with the model.
+TFLInterpreter *interpreter = [[TFLInterpreter alloc] initWithModelPath:modelPath
+                                                                  error:&error];
+if (error != nil) { /* Error handling... */ }
+
+// Allocate memory for the model's input `TFLTensor`s.
+[interpreter allocateTensorsWithError:&error];
+if (error != nil) { /* Error handling... */ }
+
+NSMutableData *inputData;  // Should be initialized
+// input data preparation...
+
+// Copy the input data to the input `TFLTensor`.
+[interpreter copyData:inputData toInputTensorAtIndex:0 error:&error];
+if (error != nil) { /* Error handling... */ }
+
+// Run inference by invoking the `TFLInterpreter`.
+[interpreter invokeWithError:&error];
+if (error != nil) { /* Error handling... */ }
+
+// Get the output `TFLTensor`
+TFLTensor *outputTensor = [interpreter outputTensorAtIndex:0 error:&error];
+if (error != nil) { /* Error handling... */ }
+
+// Copy output to `NSData` to process the inference results.
+NSData *outputData = [outputTensor dataWithError:&error];
+if (error != nil) { /* Error handling... */ }
+```
+
+### Using C API in Objective-C code
+
+Currently Objective-C API does not support delegates. In order to use delegates
+with Objective-C code, you need to directly call underlying
+[C API](https://www.tensorflow.org/code/tensorflow/lite/c/c_api.h).
+
+```c
+#include "tensorflow/lite/c/c_api.h"
+```
+
+```c
+TfLiteModel* model = TfLiteModelCreateFromFile([modelPath UTF8String]);
+TfLiteInterpreterOptions* options = TfLiteInterpreterOptionsCreate();
+
+// Create the interpreter.
+TfLiteInterpreter* interpreter = TfLiteInterpreterCreate(model, options);
+
+// Allocate tensors and populate the input tensor data.
+TfLiteInterpreterAllocateTensors(interpreter);
+TfLiteTensor* input_tensor =
+    TfLiteInterpreterGetInputTensor(interpreter, 0);
+TfLiteTensorCopyFromBuffer(input_tensor, input.data(),
+                           input.size() * sizeof(float));
+
+// Execute inference.
+TfLiteInterpreterInvoke(interpreter);
+
+// Extract the output tensor data.
+const TfLiteTensor* output_tensor =
+//      TfLiteInterpreterGetOutputTensor(interpreter, 0);
+TfLiteTensorCopyToBuffer(output_tensor, output.data(),
+                         output.size() * sizeof(float));
+
+// Dispose of the model and interpreter objects.
+TfLiteInterpreterDelete(interpreter);
+TfLiteInterpreterOptionsDelete(options);
+TfLiteModelDelete(model);
+```
+
+## Load and run a model in C++
+
+*Platforms: Android and Linux*
+
+In C++, the model is stored in
+[`FlatBufferModel`](https://www.tensorflow.org/lite/api_docs/cc/class/tflite/flat-buffer-model.html)
+class. It encapsulates a TensorFlow Lite model and you can build it in a couple
+of different ways, depending on where the model is stored:
+
+```c++
+class FlatBufferModel {
+  // Build a model based on a file. Return a nullptr in case of failure.
+  static std::unique_ptr BuildFromFile(
+      const char* filename,
+      ErrorReporter* error_reporter);
+
+  // Build a model based on a pre-loaded flatbuffer. The caller retains
+  // ownership of the buffer and should keep it alive until the returned object
+  // is destroyed. Return a nullptr in case of failure.
+  static std::unique_ptr BuildFromBuffer(
+      const char* buffer,
+      size_t buffer_size,
+      ErrorReporter* error_reporter);
+};
+```
+
+Note: If TensorFlow Lite detects the presence of the
+[Android NNAPI](https://developer.android.com/ndk/guides/neuralnetworks), it
+will automatically try to use shared memory to store the `FlatBufferModel`.
+
+Now that you have the model as a `FlatBufferModel` object, you can execute it
+with an
+[`Interpreter`](https://www.tensorflow.org/lite/api_docs/cc/class/tflite/interpreter.html).
+A single `FlatBufferModel` can be used simultaneously by more than one
+`Interpreter`.
+
+Caution: The `FlatBufferModel` object must remain valid until all instances of
+`Interpreter` using it have been destroyed.
+
+The important parts of the `Interpreter` API are shown in the code snippet
+below. It should be noted that:
+
+*   Tensors are represented by integers, in order to avoid string comparisons
+    (and any fixed dependency on string libraries).
+*   An interpreter must not be accessed from concurrent threads.
+*   Memory allocation for input and output tensors must be triggered by calling
+    `AllocateTensors()` right after resizing tensors.
+
+The simplest usage of TensorFlow Lite with C++ looks like this:
+
+```c++
+// Load the model
+std::unique_ptr model =
+    tflite::FlatBufferModel::BuildFromFile(filename);
+
+// Build the interpreter
+tflite::ops::builtin::BuiltinOpResolver resolver;
+std::unique_ptr interpreter;
+tflite::InterpreterBuilder(*model, resolver)(&interpreter);
+
+// Resize input tensors, if desired.
+interpreter->AllocateTensors();
+
+float* input = interpreter->typed_input_tensor(0);
+// Fill `input`.
+
+interpreter->Invoke();
+
+float* output = interpreter->typed_output_tensor(0);
+```
+
+For more example code, see
+[`minimal.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/minimal/minimal.cc)
+and
+[`label_image.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/label_image/label_image.cc).
 
 ## Load and run a model in Python
 
-The Python API for running an inference is provided in the `tf.lite`
-module. From which, you mostly need only [`tf.lite.Interpreter`](
-https://www.tensorflow.org/api_docs/python/tf/lite/Interpreter) to load
-a model and run an inference.
+*Platform: Linux*
+
+The Python API for running an inference is provided in the `tf.lite` module.
+From which, you mostly need only
+[`tf.lite.Interpreter`](https://www.tensorflow.org/api_docs/python/tf/lite/Interpreter)
+to load a model and run an inference.
 
 The following example shows how to use the Python interpreter to load a
 `.tflite` file and run inference with random input data:
@@ -358,13 +509,12 @@ interpreter.allocate_tensors()
 # Continue to get tensors and so forth, as shown above...
 ```
 
-For more Python sample code, see [`label_image.py`](
-https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/python/label_image.py).
+For more Python sample code, see
+[`label_image.py`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/python/label_image.py).
 
 Tip: Run `help(tf.lite.Interpreter)` in the Python terminal to get detailed
 documentation about the interpreter.
 
-
 ## Write a custom operator
 
 All TensorFlow Lite operators (both custom and builtin) are defined using a
@@ -379,10 +529,10 @@ typedef struct {
 } TfLiteRegistration;
 ```
 
-Refer to `context.h` for details on `TfLiteContext` and `TfLiteNode`. The
-former provides error reporting facilities and access to global objects,
-including all the tensors. The latter allows implementations to access their
-inputs and outputs.
+Refer to `context.h` for details on `TfLiteContext` and `TfLiteNode`. The former
+provides error reporting facilities and access to global objects, including all
+the tensors. The latter allows implementations to access their inputs and
+outputs.
 
 When the interpreter loads a model, it calls `init()` once for each node in the
 graph. A given `init()` will be called more than once if the op is used multiple
@@ -403,9 +553,9 @@ implementations can access their state using `node->user_data`.
 Finally, each time inference runs, the interpreter traverses the graph calling
 `invoke()`, and here too the state is available as `node->user_data`.
 
-Custom ops can be implemented in exactly the same way as builtin ops, by
-defined those four functions and a global registration function that usually
-looks like this:
+Custom ops can be implemented in exactly the same way as builtin ops, by defined
+those four functions and a global registration function that usually looks like
+this:
 
 ```c++
 namespace tflite {
@@ -461,8 +611,7 @@ You can optionally register custom ops (before you pass the resolver to the
 resolver.AddOp("MY_CUSTOM_OP", Register_MY_CUSTOM_OP());
 ```
 
-If the set of builtin ops is deemed to be too large, a new `OpResolver` could
-be code-generated  based on a given subset of ops, possibly only the ones
-contained in a given model. This is the equivalent of TensorFlow's selective
-registration (and a simple version of it is available in the `tools`
-directory).
+If the set of builtin ops is deemed to be too large, a new `OpResolver` could be
+code-generated based on a given subset of ops, possibly only the ones contained
+in a given model. This is the equivalent of TensorFlow's selective registration
+(and a simple version of it is available in the `tools` directory).

From 397494a2313aa51fe0b87b4e51d3a2349e4f8ecc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 20:51:42 -0700
Subject: [PATCH 0591/1390] Hoisting unconditional converts from conditional
 branch computations.

PiperOrigin-RevId: 317239618
Change-Id: If3b16ff4f2bbcf38ee1ca51f5e8b187c58ab8e91
---
 .../xla/service/conditional_code_motion.cc    | 285 ++++++++++++++++--
 .../xla/service/conditional_code_motion.h     |  15 +-
 .../service/conditional_code_motion_test.cc   | 101 ++++++-
 3 files changed, 369 insertions(+), 32 deletions(-)

diff --git a/tensorflow/compiler/xla/service/conditional_code_motion.cc b/tensorflow/compiler/xla/service/conditional_code_motion.cc
index eecdcc851e9..6db4c3eb6d4 100644
--- a/tensorflow/compiler/xla/service/conditional_code_motion.cc
+++ b/tensorflow/compiler/xla/service/conditional_code_motion.cc
@@ -106,7 +106,6 @@ class BranchVisitor {
         boundaries_.emplace_back(operand, i, inst);
         continue;
       }
-
       worklist_.push_back(operand);
       visited_.insert(operand);
     }
@@ -197,6 +196,7 @@ bool WorthHoisting(HloInstruction* instruction) {
     case HloOpcode::kMultiply:
     case HloOpcode::kDivide:
     case HloOpcode::kTuple:
+    case HloOpcode::kSqrt:
     case HloOpcode::kGetTupleElement:
       return true;
     default:
@@ -206,10 +206,11 @@ bool WorthHoisting(HloInstruction* instruction) {
 
 // Compare if the instructions to be visited at each branches are identical.
 bool InstructionWithinBranchIdentical(
-    const std::vector& instructions, bool is_layout_senstive) {
+    const std::vector& instructions,
+    bool is_layout_sensitive) {
   // Identical includes the shape of each operands are equal.
   auto eq_operand = [&](const HloInstruction* a, const HloInstruction* b) {
-    bool eq_operands = is_layout_senstive
+    bool eq_operands = is_layout_sensitive
                            ? ShapeUtil::Equal(a->shape(), b->shape())
                            : ShapeUtil::Compatible(a->shape(), b->shape());
     return eq_operands;
@@ -233,7 +234,7 @@ bool InstructionWithinBranchIdentical(
           auto old_channel_id = instruction->channel_id();
           instruction->set_channel_id(instructions[0]->channel_id());
           bool eq_instructions = instructions[0]->Identical(
-              *instruction, eq_operand, eq_computations, is_layout_senstive);
+              *instruction, eq_operand, eq_computations, is_layout_sensitive);
           instruction->set_channel_id(old_channel_id);
           return eq_instructions;
         });
@@ -243,7 +244,7 @@ bool InstructionWithinBranchIdentical(
                      [&](HloInstruction* instruction) {
                        return instructions[0]->Identical(
                            *instruction, eq_operand, eq_computations,
-                           is_layout_senstive);
+                           is_layout_sensitive);
                      });
 }
 
@@ -354,12 +355,228 @@ Status RemoveInstructionFromComputation(
   return Status::OK();
 }
 
+// Identify converts to be hoisted/rematerialized out of the branch
+// computations.
+absl::flat_hash_set FindSpecialConverts(HloInstruction* old_root,
+                                               int branch_count,
+                                               HloInstruction* conditional,
+                                               bool is_layout_sensitive) {
+  absl::flat_hash_set kspecial_convert;
+  for (int64 operand_num = 0; operand_num < old_root->operand_count();
+       ++operand_num) {
+    if (old_root->operand(operand_num)->opcode() != HloOpcode::kConvert) {
+      continue;
+    }
+    bool replica = true;
+    HloInstruction* kspecial_convert_candidate =
+        old_root->mutable_operand(operand_num);
+    // Check whether an identical candidate appears in other branches
+    for (int others = 1; others < branch_count; ++others) {
+      HloInstruction* others_root =
+          conditional->branch_computation(others)->root_instruction();
+      bool eq_shape =
+          is_layout_sensitive
+              ? ShapeUtil::Equal(others_root->operand(operand_num)->shape(),
+                                 kspecial_convert_candidate->shape())
+              : ShapeUtil::Compatible(
+                    others_root->operand(operand_num)->shape(),
+                    kspecial_convert_candidate->shape());
+      if ((others_root->operand(operand_num)->opcode() ==
+           HloOpcode::kConvert) &&
+          eq_shape) {
+        // Nothing to be done.
+      } else {
+        replica = false;
+        break;
+      }
+    }
+    if (replica) {
+      kspecial_convert.insert(operand_num);
+    }
+  }
+  return kspecial_convert;
+}
+
+// Restructuring the conditional instruction as follows:
+// i.e., %result = conditional() becomes
+// x = conditional()
+// y.{0..n} = gte(x, {0..n})
+// z = tuple(y.0, y.1, ...y.n)
+// Doing so ensures that we can accommodate the possible shape-change of the
+// conditional when the instructions are hoisted.
+Status RestructureConditionalInstruction(HloComputation* computation,
+                                         HloInstruction* conditional) {
+  HloInstruction* old_root = computation->root_instruction();
+  std::vector new_operands;
+  int cur_index = 0;
+  for (; cur_index < ShapeUtil::TupleElementCount(conditional->shape());
+       ++cur_index) {
+    new_operands.push_back(
+        computation->AddInstruction(HloInstruction::CreateGetTupleElement(
+            ShapeUtil::GetTupleElementShape(conditional->shape(), cur_index),
+            conditional, cur_index)));
+  }
+  HloInstruction* new_tuple =
+      computation->AddInstruction(HloInstruction::CreateTuple(new_operands));
+  if (old_root == conditional) {
+    computation->set_root_instruction(new_tuple);
+  } else {
+    std::vector new_tuple_users;
+    for (auto conditional_user : conditional->users()) {
+      auto is_new_gte = absl::c_find_if(
+          new_operands,
+          [&](HloInstruction* instr) { return instr == conditional_user; });
+      if (is_new_gte == new_operands.end()) {
+        new_tuple_users.push_back(conditional_user);
+      }
+    }
+    for (auto new_tuple_user : new_tuple_users) {
+      TF_RETURN_IF_ERROR(
+          conditional->ReplaceUseWith(new_tuple_user, new_tuple));
+    }
+  }
+  VLOG(2) << "computation after root restructure:\n" << computation->ToString();
+  return Status::OK();
+}
+
+StatusOr ConvertSpecialMove(HloInstruction* conditional,
+                                  bool is_layout_sensitive) {
+  int branch_count = conditional->branch_count();
+  if (branch_count <= 0) {
+    return false;
+  }
+
+  HloInstruction* old_root =
+      conditional->branch_computation(0)->root_instruction();
+  if (old_root->opcode() != HloOpcode::kTuple) {
+    return false;
+  } else {
+    VLOG(2) << "BEFORE :" << conditional->parent()->parent()->ToString();
+    // Identify the gte using `index'.
+    auto find_gte = [](const HloInstruction* conditional_result,
+                       int64 index) -> HloInstruction* {
+      for (HloInstruction* instr : conditional_result->users()) {
+        if (instr->opcode() != HloOpcode::kGetTupleElement) {
+          return nullptr;
+        }
+        if (instr->tuple_index() == index) {
+          return instr;
+        }
+      }
+      return nullptr;
+    };
+
+    // Captures tuple indices refering to converts to be rematerialized/hoisted.
+    absl::flat_hash_set kspecial_convert = FindSpecialConverts(
+        old_root, branch_count, conditional, is_layout_sensitive);
+
+    // Exit if we cannot find any converts to be hoisted.
+    if (kspecial_convert.empty()) {
+      return false;
+    }
+
+    TF_RETURN_IF_ERROR(
+        RestructureConditionalInstruction(conditional->parent(), conditional));
+
+    for (int branch = 0; branch < branch_count; branch++) {
+      old_root = conditional->branch_computation(branch)->root_instruction();
+      absl::flat_hash_map map_inst_to_tuple_index;
+      std::vector new_operands(old_root->operand_count());
+      std::unordered_set to_hoist_set;
+
+      for (int64 operand_num = 0; operand_num < old_root->operand_count();
+           ++operand_num) {
+        map_inst_to_tuple_index[old_root->mutable_operand(operand_num)] =
+            operand_num;
+      }
+      for (int64 operand_num = 0; operand_num < old_root->operand_count();
+           ++operand_num) {
+        HloInstruction* hoist = old_root->mutable_operand(operand_num);
+        if (!kspecial_convert.contains(operand_num)) {
+          new_operands[operand_num] = old_root->mutable_operand(operand_num);
+          continue;
+        }
+
+        to_hoist_set.insert(hoist);
+        int64 new_tuple_count = old_root->operand_count();
+
+        // Replace the hoisted instr in the tuple with the operand/operands.
+        // We will replace at least one of the operands of the hoist at the
+        // tuple place; the rest will be added at the end.
+        bool inplace = true;
+        CHECK(!hoist->operands().empty());
+        for (HloInstruction* prod : hoist->operands()) {
+          if (inplace) {
+            map_inst_to_tuple_index[prod] = map_inst_to_tuple_index[hoist];
+            new_operands[map_inst_to_tuple_index[hoist]] = prod;
+            inplace = false;
+          } else {
+            map_inst_to_tuple_index[prod] = new_tuple_count++;
+            new_operands.push_back(prod);
+          }
+        }
+      }
+
+      // Create the new root instruction.
+      HloComputation* cur_branch = conditional->branch_computation(branch);
+      HloInstruction* new_branch_root =
+          cur_branch->AddInstruction(HloInstruction::CreateTuple(new_operands));
+      // The shape can vary since the operands to convert are now
+      // being returned through the branches' root.
+      cur_branch->set_root_instruction(new_branch_root, true /*new shape*/);
+      TF_CHECK_OK(cur_branch->RemoveInstruction(old_root));
+
+      // Only one of the branches needs to change the conditional->parent().
+      if (branch != 0) {
+        continue;
+      }
+      HloComputation* conditional_parent = conditional->parent();
+      HloInstruction* newconditional =
+          conditional_parent->AddInstruction(HloInstruction::CreateConditional(
+              cur_branch->root_instruction()->shape(),
+              conditional->mutable_operand(0),
+              absl::MakeSpan(conditional->branch_computations()),
+              absl::MakeSpan(conditional->operands()).subspan(1)));
+      // Ensure that all the users of conditional refer to the new one.
+      TF_RETURN_IF_ERROR(
+          conditional->ReplaceAllUsesWithDifferentShape(newconditional));
+      TF_CHECK_OK(conditional_parent->RemoveInstruction(conditional));
+      conditional = newconditional;
+      // Add the hoisted instructions in the parent.
+      for (HloInstruction* hoist : to_hoist_set) {
+        VLOG(2) << "Hoisting instruction:" << hoist->ToString();
+        int64 hoist_index = map_inst_to_tuple_index[hoist];
+        // Find out the gte that captured the hoisted instr result.
+        HloInstruction* gte_hoist = find_gte(conditional, hoist_index);
+        CHECK(gte_hoist != nullptr);
+        std::vector new_operands;
+        for (HloInstruction* op : hoist->operands()) {
+          HloInstruction* gte = conditional_parent->AddInstruction(
+              HloInstruction::CreateGetTupleElement(
+                  op->shape(), conditional, map_inst_to_tuple_index[op]));
+          new_operands.push_back(gte);
+        }
+        HloInstruction* hoisted = conditional_parent->AddInstruction(
+            hoist->CloneWithNewOperands(hoist->shape(), new_operands));
+        VLOG(2) << "Hoisted instruction in parent:" << hoisted->ToString();
+        TF_RETURN_IF_ERROR(gte_hoist->ReplaceAllUsesWith(hoisted));
+        TF_CHECK_OK(conditional_parent->RemoveInstruction(gte_hoist));
+      }
+      // No need to explicitly delete a hoisted instruction since if its dead
+      // then the subsequent DCE will remove it.
+    }
+  }
+  VLOG(2) << "AFTER :" << conditional->parent()->parent()->ToString();
+  return true;
+}
+
 // Hoist identical ops out of the conditional. The definition of identical
 // are the shape of the operands are identical and their properties are
 // identical. Will start from the root instruction of each branch and get
 // the identical ops to hoist.
 StatusOr MergeIdenticalElements(HloInstruction* conditional,
                                       bool is_layout_sensitive) {
+  VLOG(1) << " visiting conditional:" << conditional->ToString();
   int branch_count = conditional->branch_count();
   if (branch_count <= 0) {
     return false;
@@ -399,7 +616,7 @@ StatusOr MergeIdenticalElements(HloInstruction* conditional,
     }
   }
 
-  if (visitors[0].HoistInstructionSize() <= 1) {
+  if (visitors[0].HoistInstructionSize() < 1) {
     return false;
   }
 
@@ -442,7 +659,6 @@ StatusOr MergeIdenticalElements(HloInstruction* conditional,
         RemoveInstructionFromComputation(visitors[i].instructions_to_hoist(),
                                          conditional->branch_computation(i)));
   }
-
   return true;
 }
 
@@ -451,26 +667,55 @@ StatusOr MergeIdenticalElements(HloInstruction* conditional,
 StatusOr ConditionalCodeMotion::Run(HloModule* module) {
   bool changed = false;
 
-  // Gather all the conditional ops in our module. We do this ahead of time so
-  // we don't have to worry about mutating the lists of computations or
-  // instructions as we iterate.
-  std::vector conditional_ops;
-  for (auto* comp : module->MakeComputationPostOrder()) {
-    for (auto* instr : comp->MakeInstructionPostOrder()) {
-      if (instr->opcode() == HloOpcode::kConditional) {
-        conditional_ops.push_back(instr);
+  if (pursue_full_conditional_code_motion_) {
+    std::vector conditional_ops;
+    for (auto* comp : module->MakeComputationPostOrder()) {
+      for (auto* instr : comp->MakeInstructionPostOrder()) {
+        if (instr->opcode() == HloOpcode::kConditional) {
+          conditional_ops.push_back(instr);
+        }
       }
     }
+
+    for (HloInstruction* conditional_op : conditional_ops) {
+      TF_ASSIGN_OR_RETURN(
+          bool result,
+          MergeIdenticalElements(conditional_op, is_layout_sensitive_));
+      changed |= result;
+    }
+
+    if (changed) {
+      HloPassPipeline subpipeline("after_conditional_code_motion");
+      subpipeline.AddPass();
+      subpipeline.AddPass();
+      subpipeline.AddPass();
+      TF_ASSIGN_OR_RETURN(bool cleanup_changed, subpipeline.Run(module));
+      changed |= cleanup_changed;
+    }
   }
 
-  for (HloInstruction* conditional_op : conditional_ops) {
-    TF_ASSIGN_OR_RETURN(bool result, MergeIdenticalElements(
-                                         conditional_op, is_layout_sensitive_));
-    changed |= result;
+  // handling convert rematerialization/hoisting
+  {
+    std::vector conditional_ops;
+    for (auto* comp : module->MakeComputationPostOrder()) {
+      for (auto* instr : comp->MakeInstructionPostOrder()) {
+        if (instr->opcode() == HloOpcode::kConditional) {
+          conditional_ops.push_back(instr);
+        }
+      }
+    }
+    for (HloInstruction* conditional_op : conditional_ops) {
+      TF_ASSIGN_OR_RETURN(
+          bool convert_result,
+          ConvertSpecialMove(conditional_op, is_layout_sensitive_));
+      changed |= convert_result;
+    }
   }
 
   if (changed) {
-    HloPassPipeline subpipeline("after_conditional_code_motion");
+    HloPassPipeline subpipeline(
+        "after_conditional_code_motion_after_convert_hoisting");
+    subpipeline.AddPass();
     subpipeline.AddPass();
     subpipeline.AddPass();
     TF_ASSIGN_OR_RETURN(bool cleanup_changed, subpipeline.Run(module));
diff --git a/tensorflow/compiler/xla/service/conditional_code_motion.h b/tensorflow/compiler/xla/service/conditional_code_motion.h
index 1197a8b3620..95f02833e15 100644
--- a/tensorflow/compiler/xla/service/conditional_code_motion.h
+++ b/tensorflow/compiler/xla/service/conditional_code_motion.h
@@ -23,7 +23,11 @@ limitations under the License.
 
 namespace xla {
 
-// HLO pass that moves identical ops out of conditional.
+// ConditionalCodeMotion specializes in hoisting/rematerializing
+// unconditional converts in the default mode.
+// When pursue_full_conditional_code_motion_ is set to true, the
+// full HLO pass moves identical ops out of a conditional in addition to moving
+// converts.
 // - The definition of identical are the shape of the operands are identical
 // and their properties are identical.
 // - Currently, only some types of instructions is supported.
@@ -35,13 +39,18 @@ class ConditionalCodeMotion : public HloModulePass {
  public:
   // If is_layout_sensitive is true, then the hoist process preserves layout
   // during identical comparison. Otherwise, layout is ignored.
-  explicit ConditionalCodeMotion(bool is_layout_sensitive = true)
-      : is_layout_sensitive_(is_layout_sensitive) {}
+  explicit ConditionalCodeMotion(
+      bool is_layout_sensitive = true,
+      bool pursue_full_conditional_code_motion = false)
+      : is_layout_sensitive_(is_layout_sensitive),
+        pursue_full_conditional_code_motion_(
+            pursue_full_conditional_code_motion) {}
   absl::string_view name() const override { return "conditional-code-motion"; }
   StatusOr Run(HloModule* module) override;
 
  private:
   const bool is_layout_sensitive_;
+  const bool pursue_full_conditional_code_motion_;
 };
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/conditional_code_motion_test.cc b/tensorflow/compiler/xla/service/conditional_code_motion_test.cc
index 4a52303a42a..38b2b515fa0 100644
--- a/tensorflow/compiler/xla/service/conditional_code_motion_test.cc
+++ b/tensorflow/compiler/xla/service/conditional_code_motion_test.cc
@@ -38,7 +38,86 @@ namespace {
 using ConditionalCodeMotionTest = HloTestBase;
 namespace op = xla::testing::opcode_matchers;
 
-TEST_F(ConditionalCodeMotionTest, DoNotMoveConvertOut) {
+TEST_F(ConditionalCodeMotionTest, MoveSubsetTupleOut) {
+  absl::string_view hlo_string =
+      R"(
+HloModule RemoveDotOpOut
+
+on_true {
+  %arg_tuple.1 = (f32[93184,4]{1,0}) parameter(0)
+  %get-tuple-element.1 = f32[93184,4]{1,0} get-tuple-element(%arg_tuple.1), index=0
+  %reshape.8493 = f32[2,512,364]{2,1,0} reshape(f32[93184,4]{1,0} %get-tuple-element.1)
+  %convert.2894 = bf16[2,512,364]{2,1,0} convert(f32[2,512,364]{2,1,0} %reshape.8493)
+  ROOT %tuple.1 = ( bf16[2,512,364]{2,1,0}, f32[2,512,364]{2,1,0}) tuple(%convert.2894, %reshape.8493)
+}
+
+on_false {
+  %arg_tuple.2 = (f32[93184,4]{1,0}) parameter(0)
+  %get-tuple-element.3 = f32[93184,4]{1,0} get-tuple-element(%arg_tuple.2), index=0
+  %reshape.9717 = f32[2,512,364]{2,1,0} reshape(f32[93184,4]{1,0} %get-tuple-element.3)
+  %add = f32[2,512,364]{2,1,0} add(f32[2,512,364]{2,1,0} %reshape.9717, f32[2,512,364]{2,1,0} %reshape.9717)
+  %convert.3604 = bf16[2,512,364]{2,1,0} convert(f32[2,512,364]{2,1,0} %reshape.9717), metadata={op_type="Cast" op_name="gradients/Cast_125_grad/Cast"}
+  ROOT %tuple.2 = (bf16[2,512,364]{2,1,0}, f32[2,512,364]{2,1,0}) tuple(%convert.3604, %add)
+}
+
+ENTRY main {
+  pred.1 = pred[] parameter(0)
+  arg_tuple.11 = (f32[93184,4]{1,0}) parameter(1)
+  arg_tuple.22 = (f32[93184,4]{1,0}) parameter(2)
+  conditional = (bf16[2,512,364]{2,1,0}, f32[2,512,364]{2,1,0}) conditional(pred.1, arg_tuple.11, arg_tuple.22), true_computation=on_true, false_computation=on_false
+  get-first-index = bf16[2,512,364]{2,1,0} get-tuple-element(conditional), index=0
+  get-first-index.2 = f32[2,512,364]{2,1,0} get-tuple-element(conditional), index=1
+  ROOT result = (bf16[2,512,364]{2,1,0}, f32[2,512,364]{2,1,0}) tuple(get-first-index, get-first-index.2)
+}
+)";
+  auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
+  ConditionalCodeMotion pass(true, true);
+  ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Tuple(op::Convert(), op::GetTupleElement())));
+}
+
+TEST_F(ConditionalCodeMotionTest, MoveConvertOutConditionalRoot) {
+  absl::string_view hlo_string =
+      R"(
+HloModule RemoveDotOpOut
+
+on_true {
+  %arg_tuple.1 = (f32[93184,4]{1,0}) parameter(0)
+  %get-tuple-element.1 = f32[93184,4]{1,0} get-tuple-element(%arg_tuple.1), index=0
+  %reshape.8493 = f32[2,512,364]{2,1,0} reshape(f32[93184,4]{1,0} %get-tuple-element.1)
+  %add.8493 = f32[2,512,364]{2,1,0} add(f32[2,512,364]{2,1,0} %reshape.8493, f32[2,512,364]{2,1,0} %reshape.8493)
+  %convert.2894 = bf16[2,512,364]{2,1,0} convert(f32[2,512,364]{2,1,0} %add.8493)
+  ROOT %tuple.1 = ( bf16[2,512,364]{2,1,0}) tuple(%convert.2894)
+}
+
+on_false {
+  %arg_tuple.2 = (f32[93184,4]{1,0}) parameter(0)
+  %get-tuple-element.3 = f32[93184,4]{1,0} get-tuple-element(%arg_tuple.2), index=0
+  %reshape.9717 = f32[2,512,364]{2,1,0} reshape(f32[93184,4]{1,0} %get-tuple-element.3)
+  %add.8493 = f32[2,512,364]{2,1,0} add(f32[2,512,364]{2,1,0} %reshape.9717, f32[2,512,364]{2,1,0} %reshape.9717)
+  %sub.8493 = f32[2,512,364]{2,1,0} subtract(f32[2,512,364]{2,1,0} %add.8493, f32[2,512,364]{2,1,0} %reshape.9717)
+  %convert.3604 = bf16[2,512,364]{2,1,0} convert(f32[2,512,364]{2,1,0} %reshape.9717), metadata={op_type="Cast" op_name="gradients/Cast_125_grad/Cast"}
+  ROOT %tuple.2 = (bf16[2,512,364]{2,1,0}) tuple(%convert.3604)
+}
+
+ENTRY main {
+  pred.1 = pred[] parameter(0)
+  arg_tuple.11 = (f32[93184,4]{1,0}) parameter(1)
+  arg_tuple.22 = (f32[93184,4]{1,0}) parameter(2)
+  ROOT conditional = (bf16[2,512,364]{2,1,0}) conditional(pred.1, arg_tuple.11, arg_tuple.22), true_computation=on_true, false_computation=on_false
+}
+)";
+  auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
+  ConditionalCodeMotion pass(true, true);
+  ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Tuple(op::Convert())));
+}
+
+TEST_F(ConditionalCodeMotionTest, MoveConvertOut) {
   absl::string_view hlo_string =
       R"(
 HloModule RemoveDotOpOut
@@ -65,12 +144,16 @@ ENTRY main {
   arg_tuple.22 = (f32[93184,4]{1,0}) parameter(2)
   conditional = (bf16[2,512,364]{2,1,0}) conditional(pred.1, arg_tuple.11, arg_tuple.22), true_computation=on_true, false_computation=on_false
   get-first-index = bf16[2,512,364]{2,1,0} get-tuple-element(conditional), index=0
-  ROOT result = (bf16[2,512,364]{2,1,0}) tuple(get-first-index)
+  add.1 = bf16[2,512,364]{2,1,0} add(bf16[2,512,364]{2,1,0} get-first-index, bf16[2,512,364]{2,1,0} get-first-index)
+  ROOT result = (bf16[2,512,364]{2,1,0}) tuple(add.1)
 }
 )";
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
-  ASSERT_FALSE(pass.Run(&*module).ValueOrDie());
+  ConditionalCodeMotion pass(true, true);
+  ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Tuple(op::Add(op::Convert(), op::Convert()))));
 }
 
 TEST_F(ConditionalCodeMotionTest, UserShareOperandCannotBeMoved) {
@@ -123,7 +206,7 @@ ENTRY main {
 }
 )";
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
+  ConditionalCodeMotion pass(true, true);
   ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
 
   const HloInstruction* conditional =
@@ -181,7 +264,7 @@ ENTRY main {
 }
 )";
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
+  ConditionalCodeMotion pass(true, true);
   ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
   const HloInstruction* conditional =
       FindInstruction(module.get(), "conditional");
@@ -245,7 +328,7 @@ ENTRY main {
 }
 )";
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
+  ConditionalCodeMotion pass(true, true);
   ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
 
   const HloInstruction* conditional =
@@ -317,7 +400,7 @@ ENTRY main {
 )";
 
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
+  ConditionalCodeMotion pass(true, true);
   ASSERT_FALSE(pass.Run(&*module).ValueOrDie());
 }
 
@@ -390,7 +473,7 @@ ENTRY main {
 }
 )";
   auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie();
-  ConditionalCodeMotion pass;
+  ConditionalCodeMotion pass(true, true);
   ASSERT_TRUE(pass.Run(&*module).ValueOrDie());
   const HloInstruction* conditional =
       FindInstruction(module.get(), "conditional");

From 0b0eef4031fa2674a2c5d32aa7570a82c3def6a8 Mon Sep 17 00:00:00 2001
From: Chao Mei 
Date: Thu, 18 Jun 2020 20:59:28 -0700
Subject: [PATCH 0592/1390] Move enabling xnnpack delegate to AllocateTensors
 to allow other delegates to be applied first.

PiperOrigin-RevId: 317240424
Change-Id: I89b616f891f65f7cff6beedbf5c2a372f7456592
---
 tensorflow/lite/interpreter.cc         | 15 +++++++++++++--
 tensorflow/lite/interpreter.h          | 15 +++++++++++----
 tensorflow/lite/interpreter_builder.cc | 17 ++++++-----------
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/tensorflow/lite/interpreter.cc b/tensorflow/lite/interpreter.cc
index cae2ca7dde0..b49aa5031bf 100644
--- a/tensorflow/lite/interpreter.cc
+++ b/tensorflow/lite/interpreter.cc
@@ -86,8 +86,9 @@ TfLiteQuantization GetQuantizationFromLegacy(
 }  // namespace
 
 Interpreter::Interpreter(ErrorReporter* error_reporter)
-    : error_reporter_(error_reporter ? error_reporter
-                                     : DefaultErrorReporter()) {
+    : error_reporter_(error_reporter ? error_reporter : DefaultErrorReporter()),
+      lazy_delegate_provider_(
+          TfLiteDelegatePtr(nullptr, [](TfLiteDelegate*) {})) {
   // TODO(b/128420794): Include the TFLite runtime version in the log.
   // Prod logging is useful for mobile platforms where scraping console logs is
   // critical for debugging.
@@ -175,6 +176,16 @@ TfLiteStatus Interpreter::SetVariables(std::vector variables) {
 }
 
 TfLiteStatus Interpreter::AllocateTensors() {
+  // Apply the default delegate that TFLite will enable at this point to allow
+  // other user-level delegates to be applied first.
+  if (lazy_delegate_provider_) {
+    // The execution will fall back to default implementation if the XNNPACK
+    // delegate fails to be applied. Therefore, we ignore the return status
+    // here and let it fall through the rest of the code.
+    ModifyGraphWithDelegate(std::move(lazy_delegate_provider_));
+    lazy_delegate_provider_.reset();
+  }
+
   return primary_subgraph().AllocateTensors();
 }
 
diff --git a/tensorflow/lite/interpreter.h b/tensorflow/lite/interpreter.h
index 59cab6add6d..41377c4ce1f 100644
--- a/tensorflow/lite/interpreter.h
+++ b/tensorflow/lite/interpreter.h
@@ -347,10 +347,12 @@ class Interpreter {
   /// WARNING: Experimental interface, subject to change
   TfLiteStatus ReleaseNonPersistentMemory();
 
-  /// Update allocations for all tensors. This will redim dependent tensors
-  /// using the input tensor dimensionality as given. This is relatively
-  /// expensive. If you know that your sizes are not changing, you need not call
-  /// this. Returns status of success or failure.
+  // Update allocations for all tensors. This will redim dependent tensors
+  // using the input tensor dimensionality as given. This is relatively
+  // expensive. This *must be* called after the interpreter has been created
+  // and before running inference (and accessing tensor buffers), and *must be*
+  // called again if (and only if) an input tensor is resized. Returns status of
+  // success or failure.
   TfLiteStatus AllocateTensors();
 
   /// Invoke the interpreter (run the whole graph in dependency order).
@@ -594,6 +596,11 @@ class Interpreter {
 
   // A map of resources. Owned by interpreter and shared by multiple subgraphs.
   resource::ResourceMap resources_;
+
+  // Indicating a delegate that the TFLite interpreter will apply by default.
+  // A nullptr value means there's no delegate to be applied by default or the
+  // delegate has been applied and doesn't need to be applied again.
+  TfLiteDelegatePtr lazy_delegate_provider_;
 };
 
 }  // namespace impl
diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index d73b298e595..4b491d41881 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -545,17 +545,7 @@ TfLiteStatus InterpreterBuilder::ParseTensors(
 
 TfLiteStatus InterpreterBuilder::ApplyDelegates(Interpreter* interpreter,
                                                 int num_threads) {
-  // First, apply XNNPACK delegate if applicable.
-  if (num_fp32_tensors_ > 0) {
-    // The execution will fall back to default implementation if the XNNPACK
-    // delegate fails to be applied. Therefore, we ignore the return status
-    // here and let it fall through the rest of the code.
-    if (auto xnnpack_delegate = MaybeCreateXNNPACKDelegate(num_threads)) {
-      interpreter->ModifyGraphWithDelegate(std::move(xnnpack_delegate));
-    }
-  }
-
-  // Secondly, apply Flex delegate if applicable.
+  // Apply Flex delegate if applicable.
   if (has_flex_op_) {
     if (auto flex_delegate = AcquireFlexDelegate()) {
       return interpreter->ModifyGraphWithDelegate(std::move(flex_delegate));
@@ -672,6 +662,11 @@ TfLiteStatus InterpreterBuilder::operator()(
     modified_subgraph->SetVariables(std::move(variables));
   }
 
+  if (num_fp32_tensors_ > 0) {
+    (*interpreter)->lazy_delegate_provider_ =
+        MaybeCreateXNNPACKDelegate(num_threads);
+  }
+
   if (ApplyDelegates(interpreter->get(), num_threads) != kTfLiteOk)
     return cleanup_and_error();
 

From b5bb616121f1805c5ff5391daf00c86b6bcad1ae Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 21:06:35 -0700
Subject: [PATCH 0593/1390] *** Reason for rollback ***

 CL316577397 prevents H2D prefetch overlapping for TF2 MLPerf Resnet model on V100x8
More details in b/159372996

  With the rollback, the prefetching overlapping is back, and the training speed also recovered.
*** Original change description ***

Add DT_BOOL support to GPU variable ops

This is a follow-on to PR #38848 & PR #39172 and resolves remaining ask
in Issue #35994. The original PR tried to add many variable ops on the
GPU including DT_BOOL. However, this caused testCondModifyBoolPred to
fail and thus the DT_BOOL type was removed. The reason for the test
failure is once DT_BOOL variables are supported on the GPU, we need to
ensure the switch ops are also updated to not have ho...

***

PiperOrigin-RevId: 317241338
Change-Id: Id7b7d79622e0537ccb677f081b487014ac4d2395
---
 tensorflow/core/kernels/control_flow_ops.cc            | 10 +++++-----
 tensorflow/core/kernels/variable_ops.cc                |  3 ++-
 .../debug/lib/debug_graph_reconstruction_test.py       |  6 +++---
 tensorflow/python/ops/control_flow_ops_test.py         |  6 +++---
 4 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/kernels/control_flow_ops.cc b/tensorflow/core/kernels/control_flow_ops.cc
index 435de3c5954..c8e83b6f672 100644
--- a/tensorflow/core/kernels/control_flow_ops.cc
+++ b/tensorflow/core/kernels/control_flow_ops.cc
@@ -111,17 +111,15 @@ REGISTER_GPU_SWITCH(uint64);
 TF_CALL_variant(REGISTER_GPU_SWITCH);
 TF_CALL_uint32(REGISTER_GPU_SWITCH);
 TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
-TF_CALL_bool(REGISTER_GPU_SWITCH);
-TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
 
 #undef REGISTER_CPU_SWITCH
 #undef REGISTER_CPU_REF_SWITCH
 #undef REGISTER_GPU_SWITCH
 #undef REGISTER_GPU_REF_SWITCH
 
-// Special GPU kernels for int32, string & resource handles. Requiring all
-// inputs and outputs to be in host memory.
-// TODO(b/25387198): Also enable int32 in device memory.
+// Special GPU kernels for int32 and string.
+// TODO(b/25387198): Also enable int32 in device memory. This kernel
+// registration requires all int32 inputs and outputs to be in host memory.
 #define REGISTER_GPU_HOST_KERNEL(type)                    \
   REGISTER_KERNEL_BUILDER(Name("Switch")                  \
                               .Device(DEVICE_GPU)         \
@@ -151,6 +149,8 @@ TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
 
 REGISTER_GPU_HOST_KERNEL(int32);
 REGISTER_GPU_HOST_REF_KERNEL(int32);
+REGISTER_GPU_HOST_KERNEL(bool);
+REGISTER_GPU_HOST_REF_KERNEL(bool);
 REGISTER_GPU_HOST_KERNEL(tstring);
 REGISTER_GPU_HOST_REF_KERNEL(tstring);
 REGISTER_GPU_HOST_KERNEL(ResourceHandle);
diff --git a/tensorflow/core/kernels/variable_ops.cc b/tensorflow/core/kernels/variable_ops.cc
index ccd33e8c75a..6f5e0b94eca 100644
--- a/tensorflow/core/kernels/variable_ops.cc
+++ b/tensorflow/core/kernels/variable_ops.cc
@@ -252,7 +252,8 @@ TF_CALL_GPU_NUMBER_TYPES_NO_HALF(REGISTER_SYCL_KERNEL);
 
 TF_CALL_int64(REGISTER_GPU_KERNELS);
 TF_CALL_uint32(REGISTER_GPU_KERNELS);
-TF_CALL_GPU_ALL_TYPES(REGISTER_GPU_KERNELS);
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);
+TF_CALL_COMPLEX_TYPES(REGISTER_GPU_KERNELS);
 #undef REGISTER_GPU_KERNELS
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
diff --git a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
index b3baa6e7bc2..fb722efab4e 100644
--- a/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
+++ b/tensorflow/python/debug/lib/debug_graph_reconstruction_test.py
@@ -73,9 +73,9 @@ class ReconstructNonDebugGraphTest(test_util.TensorFlowTestCase):
           for attr_key in new_node.attr:
             if attr_key == "parallel_iterations":
               new_node.attr[attr_key].i = 1
-        elif new_node.op == "Switch" or new_node.op == "Identity":
-          # We don't check the inputs to Switch or Identity ops as their inputs
-          # may be Send/Recv nodes.
+        elif new_node.op == "Switch":
+          # We don't check the inputs to Switch ops as their inputs may be
+          # Send/Recv nodes.
           del new_node.input[:]
 
     return output_graph_def
diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py
index 3ca9bda82f2..9254695d988 100644
--- a/tensorflow/python/ops/control_flow_ops_test.py
+++ b/tensorflow/python/ops/control_flow_ops_test.py
@@ -396,10 +396,10 @@ class CondTest(test_util.TensorFlowTestCase):
         fn2=lambda: math_ops.add(y, 23))
     self.assertEquals(self.evaluate(z), 24)
 
-  @test_util.run_v1_only("Exercises Ref variables")
+  @test_util.run_deprecated_v1
   def testCondModifyBoolPred(self):
-    # We want to use the GPU here because we want to ensure that we can update
-    # a boolean ref variable on the GPU.
+    # This test in particular used to fail only when running in GPU, hence
+    # use_gpu=True.
     with test_util.use_gpu():
       bool_var = variable_scope.get_variable(
           "bool_var", dtype=dtypes.bool, initializer=True)

From cfbdd27fe3f2b904609e1551490a01640ae4fcac Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 21:17:54 -0700
Subject: [PATCH 0594/1390] Update ops-related pbtxt files.

PiperOrigin-RevId: 317242544
Change-Id: I36000fdb2d595b5006ea111105ece5ca6f537732
---
 .../core/ops/compat/ops_history_v2/Acos.pbtxt | 29 +++++++++++++++++++
 .../core/ops/compat/ops_history_v2/Asin.pbtxt | 29 +++++++++++++++++++
 .../core/ops/compat/ops_history_v2/Atan.pbtxt | 29 +++++++++++++++++++
 .../core/ops/compat/ops_history_v2/Inv.pbtxt  | 29 +++++++++++++++++++
 .../core/ops/compat/ops_history_v2/Neg.pbtxt  | 29 +++++++++++++++++++
 .../compat/ops_history_v2/Reciprocal.pbtxt    | 29 +++++++++++++++++++
 .../ops/compat/ops_history_v2/Round.pbtxt     | 29 +++++++++++++++++++
 .../ops/compat/ops_history_v2/Square.pbtxt    | 29 +++++++++++++++++++
 .../core/ops/compat/ops_history_v2/Tan.pbtxt  | 29 +++++++++++++++++++
 tensorflow/core/ops/ops.pbtxt                 | 18 ++++++++++++
 10 files changed, 279 insertions(+)

diff --git a/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
index 3ed45186f6e..417dbfc7e7d 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Acos.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Acos"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
index 7df768f7c66..c799ff99169 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Asin.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Asin"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
index 86f0628ab53..4a80c7a751e 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Atan.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Atan"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
index ca208664617..0c191790030 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Inv.pbtxt
@@ -168,3 +168,32 @@ op {
     }
   }
 }
+op {
+  name: "Inv"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
index 77bb4a5872d..864d0257fe4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Neg.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Neg"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
index 5ea1abe4c9c..7e03554871a 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Reciprocal.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Reciprocal"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
index 4f59b21afd5..c5685dc6143 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Round.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Round"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
index 4d07faf4fd0..6af75b3ddc1 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Square.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Square"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt b/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
index 7dc7f84fd38..80e0b1e22c4 100644
--- a/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history_v2/Tan.pbtxt
@@ -78,3 +78,32 @@ op {
     }
   }
 }
+op {
+  name: "Tan"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  output_arg {
+    name: "y"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_BFLOAT16
+        type: DT_HALF
+        type: DT_FLOAT
+        type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
+        type: DT_INT32
+        type: DT_INT64
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 1f1cf7444fb..dbd91c91b65 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -216,6 +216,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -2333,6 +2335,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -2646,6 +2650,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -19442,6 +19448,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -25498,6 +25506,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -35191,6 +35201,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -40686,6 +40698,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -48071,6 +48085,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64
@@ -50832,6 +50848,8 @@ op {
         type: DT_HALF
         type: DT_FLOAT
         type: DT_DOUBLE
+        type: DT_INT8
+        type: DT_INT16
         type: DT_INT32
         type: DT_INT64
         type: DT_COMPLEX64

From b7caba2c42285a8e1cb875bec1664d5b0e6c65e9 Mon Sep 17 00:00:00 2001
From: Ashwin Murthy 
Date: Thu, 18 Jun 2020 22:10:41 -0700
Subject: [PATCH 0595/1390] Update RNN conversion tflite g3doc

This uses the content from the blog post/dogfood announcement email

PiperOrigin-RevId: 317248288
Change-Id: I210c64bd54c70aa5b68742d59d6d36fa154e856c
---
 tensorflow/lite/g3doc/convert/rnn.md | 240 +++++++++++++++++++--------
 1 file changed, 167 insertions(+), 73 deletions(-)

diff --git a/tensorflow/lite/g3doc/convert/rnn.md b/tensorflow/lite/g3doc/convert/rnn.md
index 52bc287c151..734992c0904 100644
--- a/tensorflow/lite/g3doc/convert/rnn.md
+++ b/tensorflow/lite/g3doc/convert/rnn.md
@@ -1,99 +1,193 @@
-# Convert RNN models
+# TensorFlow RNN conversion to TensorFlow Lite
 
-The TensorFlow Lite interpreter currently implements a subset of TensorFlow
-operations, meaning some model architectures cannot immediately be converted due
-to missing operations.
+## Overview
 
-Some RNN-based architectures are affected by this. The following document
-outlines the current state of play and provides strategies for converting RNN
-models.
+TensorFlow Lite supports converting TensorFlow RNN models to TensorFlow Lite’s
+fused LSTM operators. Fused operators exist to maximize the performance of their
+underlying kernel implementations, as well as provide a higher level interface
+to define complex transformations like quantizatization.
 
-## Currently supported
+Since there are many variants of RNN APIs in TensorFlow, our approach has been
+two fold:
 
-Currently, RNN models using
-[`tf.compat.v1.nn.static_rnn`](https://www.tensorflow.org/api_docs/python/tf/nn/static_rnn)
-can be converted successfully as long as no `sequence_length` is specified.
+1.  Provide **native support for standard TensorFlow RNN APIs** like Keras LSTM.
+    This is the recommended option.
+1.  Provide an **interface** **into the conversion infrastructure for**
+    **user-defined** **RNN implementations** to plug in and get converted to
+    TensorFlow Lite. We provide a couple of out of box examples of such
+    conversion using lingvo’s
+    [LSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L123)
+    and
+    [LayerNormalizedLSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/utils/lstm_utils.cc#L519)
+    RNN interfaces.
 
-The following `tf.compat.v1.nn.rnn_cell` operations work with
-`tf.compat.v1.nn.static_rnn`:
+## Converter API
 
-*   [tf.compat.v1.nn.rnn_cell.LSTMCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/LSTMCell)
-*   [tf.compat.v1.nn.rnn_cell.RNNCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/RNNCell)
-*   [tf.compat.v1.nn.rnn_cell.GRUCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/GRUCell)
-*   [tf.compat.v1.nn.rnn_cell.BasicLSTMCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/BasicLSTMCell)
-*   [tf.compat.v1.nn.rnn_cell.BasicRNNCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/BasicRNNCell)
+Currently this feature is available through the
+[tf-nightly](https://pypi.org/project/tf-nightly/) pip or from head. This will
+be available in the TensorFlow 2.3 release.
 
-In addition, TensorFlow Lite provides some experimental drop-in replacements for
-RNN operations that enable dynamic RNN architectures with TensorFlow Lite.
+This conversion functionality is available when converting to TensorFlow Lite
+via a SavedModel or from the Keras model directly. See example usages.
 
-Drop-in replacements are available for the following:
+### From saved model
 
-*   [tf.compat.v1.nn.dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn)
-*   [tf.compat.v1.nn.bidirectional_dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/bidirectional_dynamic_rnn)
-*   [tf.compat.v1.nn.rnn_cell.RNNCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/RNNCell)
-*   [tf.compat.v1.nn.rnn_cell.LSTMCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/LSTMCell)
+```
+# build a saved model. Here concrete_function is the exported function
+# corresponding to the TensorFlow model containing one or more
+# Keras LSTM layers.
+saved_model, saved_model_dir = build_saved_model_lstm(...)
+saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)
 
-## Not currently supported
+# Convert the model.
+converter = TFLiteConverter.from_saved_model(saved_model_dir)
+tflite_model = converter.convert()
+```
 
-TensorFlow Lite does not currently support
-[Control Flow](https://www.tensorflow.org/api_docs/cc/group/control-flow-ops)
-operations. This means that, unless one of the conversion strategies discussed
-in the next section are employed, models built with the following TensorFlow
-functions will not convert successfully:
+### From Keras model
 
-*   [tf.compat.v1.nn.static_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/static_rnn)
-    where a `sequence_length` is specified
-*   [tf.compat.v1.nn.dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn)
-*   [tf.compat.v1.nn.bidirectional_dynamic_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/bidirectional_dynamic_rnn)
+```
+# build a Keras model
+keras_model = build_keras_lstm(...)
 
-Note: TensorFlow Lite plans to implement all required Control Flow operations by
-the end of 2019. At this point, all RNN architectures will convert successfully.
+# Convert the model.
+converter = TFLiteConverter.from_keras_model(keras_model)
+tflite_model = converter.convert()
 
-## Conversion strategies
+```
 
-To convert an RNN model that uses the functions specified above, you will have
-to modify its architecture and retrain it. The following strategies can be used.
+## Example
 
-### 1. Refactoring
+Keras LSTM to TensorFlow Lite
+[Colab](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/experimental_new_converter/Keras_LSTM_fusion_Codelab.ipynb)
+illustrates the end to end usage with the TensorFlow Lite interpreter.
 
-The simplest approach, if possible, is to refactor the model architecture to use
-[tf.compat.v1.nn.static_rnn](https://www.tensorflow.org/api_docs/python/tf/nn/static_rnn)
-without `sequence_length`.
+## TensorFlow RNNs APIs supported
 
-### 2. Drop-in replacements that use op hints and fused ops
+### Keras LSTM conversion (recommended)
 
-TensorFlow Lite provides the some experimental drop-in replacements for RNN
-operations that enable dynamic RNN architectures with TensorFlow Lite. Using
-[OpHints](https://www.tensorflow.org/lite/guide/ops_custom#converting_tensorflow_models_to_convert_graphs),
-they run normally during training, but are substituted with special fused ops
-when run by the Lite interpreter.
+We support out-of-the-box conversion of Keras LSTM to TensorFlow Lite. For
+details on how this works please refer to the
+[Keras LSTM interface](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/examples/experimental_new_converter/Keras_LSTM_fusion_Codelab.ipynb)
+and to the conversion logic
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/utils/lstm_utils.cc#L627).
 
-The following drop-in replacements are available:
+Also important is to highlight the TensorFlow Lite’s LSTM contract with respect
+to the Keras operation definition:
 
-*   [tf.compat.v1.lite.experimental.nn.dynamic_rnn](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/rnn.py#L41)
-    *   replacement for tf.nn.dynamic_rnn
-*   [tf.compat.v1.lite.experimental.nn.bidirectional_dynamic_rnn](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/rnn.py#L279)
-    *   replacement for tf.nn.bidirectional_dynamic_rnn
-*   [tf.compat.v1.lite.experimental.nn.TfLiteRNNCell](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/rnn_cell.py#L39)
-    *   replacement for tf.nn.rnn_cell.RNNCell
-*   [tf.compat.v1.lite.experimental.nn.TfLiteLSTMCell](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/rnn_cell.py#L159)
-    *   replacement for tf.nn.rnn_cell.LSTMCell
+1.  The dimension 0 of the input tensor is the batch size.
+1.  The dimension 0 of the recurrent\_weight tensor is the number of outputs.
+1.  The **weight** and **recurrent\_kernel** tensors are transposed.
+1.  The transposed weight, transposed recurrent\_kernel and bias tensors are
+    split into 4 equal sized tensors along the dimension 0. These correspond to
+    **input gate, forget gate, cell, and output gate**.
 
-Note: These replacements must be used together. For example, if you are using
-`tf.compat.v1.lite.experimental.nn.dynamic_rnn`, you must combine it with
-`tf.compat.v1.lite.experimental.nn.TfLiteRNNCell` instead of using
-`tf.compat.v1.nn.rnn_cell.RNNCell`.
+#### Keras LSTM Variants
 
-Instead of
-[tf.compat.v1.nn.rnn_cell.MultiRNNCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/MultiRNNCell),
-you should use
-[tf.compat.v1.keras.layers.StackedRNNCells](https://www.tensorflow.org/api_docs/python/tf/keras/layers/StackedRNNCells).
+##### Time major
 
-For a tutorial on using these replacements, see
-[TensorFlow Lite LSTM ops API](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/g3doc/README.md).
+Users may choose time-major or no time-major. Keras LSTM adds a time-major
+attribute in the function def attributes. For Unidirectional sequence LSTM, we
+can simply map to unidirecional\_sequence\_lstm's
+[time major attribute](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/ir/tfl_ops.td#L3508).
 
-For a Colab demonstrating these classes, refer to
-[TensorFlowLite_LSTM_Keras_Tutorial](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/experimental/examples/lstm/TensorFlowLite_LSTM_Keras_Tutorial.ipynb).
+##### BiDirectional LSTM
 
-Note: There is no replacement available for
-[tf.compat.v1.nn.rnn_cell.GRUCell](https://www.tensorflow.org/api_docs/python/tf/nn/rnn_cell/GRUCell).
+Bidirectional LSTM can be implemented with two Keras LSTM layers, one for
+forward and one for backward, see examples
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/wrappers.py#L381).
+Once we see the go\_backward attribute, we recognize it as backward LSTM, then
+we group forward & backward LSTM together. **This is future work.** Currently,
+this creates two UnidirectionalSequenceLSTM operators in the TensorFlow Lite
+model.
+
+### User-defined LSTM conversion examples
+
+TensorFlow Lite also provides a way to convert user defined LSTM
+implementations. Here we use Lingvo’s LSTM as an example of how that can be
+implemented. For details please refer to the
+[lingvo.LSTMCellSimple interface](https://github.com/tensorflow/lingvo/blob/master/lingvo/core/rnn_cell.py#L230)
+and the conversion logic
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L123).
+We also provide an example for another of Lingvo’s LSTM definitions in
+[lingvo.LayerNormalizedLSTMCellSimple interface](https://github.com/tensorflow/lingvo/blob/master/lingvo/core/rnn_cell.py#L1179)
+and its convertion logic
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L130).
+
+## “Bring your own TensorFlow RNN” to TensorFlow Lite
+
+If a user's RNN interface is different from the standard supported ones, there
+are a couple of options:
+
+**Option 1:** Write adapter code in TensorFlow python to adapt the RNN interface
+to the Keras RNN interface. This means a tf.function with
+[tf\_implements annotation](https://github.com/tensorflow/community/pull/113) on
+the generated RNN interface’s function that is identical to the one generated by
+the Keras LSTM layer. After this, the same conversion API used for Keras LSTM
+will work.
+
+**Option 2:** If the above is not possible (e.g. the Keras LSTM is missing some
+functionality that is currently exposed by TensorFlow Lite’s fused LSTM op like
+layer normalization), then extend the TensorFlow Lite converter by writing
+custom conversion code and plug it into the prepare-composite-functions
+MLIR-pass
+[here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L108).
+The function’s interface should be treated like an API contract and should
+contain the arguments needed to convert to fused TensorFlow Lite LSTM
+operators - i.e. input, bias, weights, projection, layer normalization, etc. It
+is preferable for the tensors passed as arguments to this function to have known
+rank (i.e. RankedTensorType in MLIR). This makes it much easier to write
+conversion code that can assume these tensors as RankedTensorType and helps
+transform them to ranked tensors corresponding to the fused TensorFlow Lite
+operator’s operands.
+
+A complete example of such conversion flow is Lingvo’s LSTMCellSimple to
+TensorFlow Lite conversion.
+
+The LSTMCellSimple in Lingvo is defined
+[here](https://github.com/tensorflow/lingvo/blob/master/lingvo/core/rnn_cell.py#L230).
+Models trained with this LSTM cell can be converted to TensorFlow Lite as
+follows:
+
+1.  Wrap all uses of LSTMCellSimple in a tf.function with a tf\_implements
+    annotation that is labelled as such (e.g. lingvo.LSTMCellSimple would be a
+    good annotation name here). Make sure the tf.function that is generated
+    matches the interface of the function expected in the conversion code. This
+    is a contract between the model author adding the annotation and the
+    conversion code.
+1.  Extend the prepare-composite-functions pass to plug in a custom composite op
+    to TensorFlow Lite fused LSTM op conversion. See
+    [LSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc#L123)
+    conversion code.
+
+    The conversion contract:
+
+1.  **Weight** and **projection** tensors are transposed.
+
+1.  The **{input, recurrent}** to **{cell, input gate, forget gate, output
+    gate}** are extracted by slicing the transposed weight tensor.
+
+1.  The **{bias}** to **{cell, input gate, forget gate, output gate}** are
+    extracted by slicing the bias tensor.
+
+1.  The **projection** is extracted by slicing the transposed projection tensor.
+
+1.  Similar conversion is written for
+    [LayerNormalizedLSTMCellSimple](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/utils/lstm_utils.cc#L519).
+
+1.  The rest of the TensorFlow Lite conversion infrastructure, including all the
+    [MLIR passes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/tf_tfl_passes.cc#L58)
+    defined as well as the final export to TensorFlow Lite flatbuffer can be
+    reused.
+
+## Known issues/limitations
+
+1.  Currently there is support only for converting stateless Keras LSTM (default
+    behavior in Keras). Stateful Keras LSTM conversion is future work.
+1.  It is still possible to model a stateful Keras LSTM layer using the
+    underlying stateless Keras LSTM layer and managing the state explicitly in
+    the user program. Such a TensorFlow program can still be converted to
+    TensorFlow Lite using the feature being described here.
+1.  Bidirectional LSTM is currently modelled as two UnidirectionalSequenceLSTM
+    operators in TensorFlow Lite. This will be replaced with a single
+    BidirectionalSequenceLSTM op.

From 158d4be42d7aea11a395d2f79483ac93289e1bb8 Mon Sep 17 00:00:00 2001
From: Xinyi Wang 
Date: Thu, 18 Jun 2020 22:15:44 -0700
Subject: [PATCH 0596/1390] Add get_next_as_optional method for a distributed
 iterator

The function is called on a distributed iterator and returns an `Optional` that contains the next value, the PerReplica input, from Distributed iterator or no value if this `iterator` has reached the end of the sequence.

PiperOrigin-RevId: 317248910
Change-Id: Ide217da1aff1d62f8d0d8f43423be2d859d933d3
---
 .../custom_training_loop_input_test.py        | 49 ++++++++++-
 .../python/distribute/distribute_lib.py       |  4 +
 tensorflow/python/distribute/input_lib.py     | 60 +++++++++++++
 .../python/distribute/input_lib_test.py       | 86 +++++++++++++------
 ...low.distribute.-distributed-iterator.pbtxt |  4 +
 5 files changed, 177 insertions(+), 26 deletions(-)

diff --git a/tensorflow/python/distribute/custom_training_loop_input_test.py b/tensorflow/python/distribute/custom_training_loop_input_test.py
index e4f782810dd..5660b5839ce 100644
--- a/tensorflow/python/distribute/custom_training_loop_input_test.py
+++ b/tensorflow/python/distribute/custom_training_loop_input_test.py
@@ -30,6 +30,7 @@ from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.eager import def_function
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
@@ -136,8 +137,52 @@ class InputIterationTest(test.TestCase, parameterized.TestCase,
 
   @combinations.generate(
       combinations.combine(
-          distribution=strategy_combinations.tpu_strategies,
-          mode=["eager"]))
+          distribution=strategy_combinations.all_strategies, mode=["eager"]))
+  def testGetNextAsOptional(self, distribution):
+    data = [5., 6., 7., 8.]
+    dataset = get_dataset_from_tensor_slices(data).batch(2)
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    iterator = iter(dist_dataset)
+
+    def train_step(data):
+      return math_ops.square(data)
+
+    @def_function.function
+    def run(iterator):
+      return distribution.experimental_local_results(
+          distribution.run(
+              train_step, args=(iterator.get_next_as_optional().get_value(),)))
+
+    self.assert_equal_flattened([[25., 36.]], [run(iterator)])
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=strategy_combinations.all_strategies, mode=["eager"]))
+  def testGetNextAsOptionalExampleUsage(self, distribution):
+    global_batch_size = 2
+    steps_per_loop = 6
+    dataset = dataset_ops.Dataset.range(
+        8, output_type=dtypes.int32).batch(global_batch_size)
+    distributed_iterator = iter(
+        distribution.experimental_distribute_dataset(dataset))
+
+    @def_function.function
+    def train_fn(distributed_iterator):
+
+      def step_fn(x):
+        return x
+
+      for _ in math_ops.range(steps_per_loop):
+        optional_data = distributed_iterator.get_next_as_optional()
+        if not optional_data.has_value():
+          break
+        distribution.run(step_fn, args=(optional_data.get_value(),))
+
+    train_fn(distributed_iterator)
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=strategy_combinations.tpu_strategies, mode=["eager"]))
   def testFullEagerTPU(self, distribution):
     dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
 
diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index b6a89463426..ec0b911ebe0 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -200,6 +200,7 @@ import six
 from tensorflow.python.autograph.core import ag_ctx as autograph_ctx
 from tensorflow.python.autograph.impl import api as autograph
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
 from tensorflow.python.distribute import collective_util
 from tensorflow.python.distribute import device_util
 from tensorflow.python.distribute import distribution_strategy_context
@@ -2879,6 +2880,9 @@ class _DefaultDistributionExtended(StrategyExtendedV1):
     def get_next(self):
       return self._iterator.get_next()
 
+    def get_next_as_optional(self):
+      return iterator_ops.get_next_as_optional(self._iterator)
+
     @deprecated(None, "Use the iterator's `initializer` property instead.")
     def initialize(self):
       """Initialize underlying iterators.
diff --git a/tensorflow/python/distribute/input_lib.py b/tensorflow/python/distribute/input_lib.py
index ff468af7f87..e4a362a92c6 100644
--- a/tensorflow/python/distribute/input_lib.py
+++ b/tensorflow/python/distribute/input_lib.py
@@ -29,6 +29,7 @@ from tensorflow.python.data.experimental.ops import batching
 from tensorflow.python.data.experimental.ops import distribute
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.data.ops import multi_device_iterator_ops
+from tensorflow.python.data.ops import optional_ops
 from tensorflow.python.distribute import device_util
 from tensorflow.python.distribute import distribute_utils
 from tensorflow.python.distribute import distribution_strategy_context
@@ -235,6 +236,40 @@ class DistributedIteratorInterface(collections.Iterator,
     raise NotImplementedError(
         "DistributedIterator.element_spec() must be implemented in descendants")
 
+  def get_next_as_optional(self):
+    """Returns a `tf.experimental.Optional` that contains the next value for all replicas.
+
+    If the `tf.distribute.DistributedIterator` has reached the end of the
+    sequence, the returned `tf.experimental.Optional` will have no value.
+
+    Example usage:
+
+    >>> strategy = tf.distribute.MirroredStrategy()
+    >>> global_batch_size = 2
+    >>> steps_per_loop = 2
+    >>> dataset = tf.data.Dataset.range(10).batch(global_batch_size)
+    >>> distributed_iterator = iter(
+    ...     strategy.experimental_distribute_dataset(dataset))
+    >>> def step_fn(x):
+    ...   return x
+    >>> @tf.function
+    ... def train_fn(distributed_iterator):
+    ...   for _ in tf.range(steps_per_loop):
+    ...     optional_data = distributed_iterator.get_next_as_optional()
+    ...     if not optional_data.has_value():
+    ...       break
+    ...     tf.print(strategy.run(step_fn, args=(optional_data.get_value(),)))
+    >>> train_fn(distributed_iterator)
+    ... # ([0 1],)
+    ... # ([2 3],)
+
+    Returns:
+      An `tf.experimental.Optional` object representing the next value from the
+      `tf.distribute.DistributedIterator` (if it has one) or no value.
+    """
+    raise NotImplementedError(
+        "get_next_as_optional() not implemented in descendants")
+
 
 @tf_export("distribute.DistributedDataset", v1=[])
 class DistributedDatasetInterface(collections.Iterable,
@@ -622,6 +657,31 @@ class DistributedIteratorBase(DistributedIteratorInterface):
   def __iter__(self):
     return self
 
+  def get_next_as_optional(self):
+    global_has_value, replicas = _get_next_as_optional(self, self._strategy)
+
+    def return_none():
+      return optional_ops.Optional.empty(self._element_spec)
+
+    def return_value(replicas):
+      """Wraps the inputs for replicas in an `tf.experimental.Optional`."""
+      results = []
+      for i, worker in enumerate(self._input_workers.worker_devices):
+        with ops.device(worker):
+          devices = self._input_workers.compute_devices_for_worker(i)
+          for j, device in enumerate(devices):
+            with ops.device(device):
+              result = replicas[i][j]
+              results.append(result)
+      replicas = results
+
+      return optional_ops.Optional.from_value(
+          distribute_utils.regroup(replicas))
+
+    return control_flow_ops.cond(global_has_value,
+                                 lambda: return_value(replicas),
+                                 lambda: return_none())  # pylint: disable=unnecessary-lambda
+
   def get_next(self, name=None):
     """Returns the next input from the iterator for all replicas."""
     if not self._enable_get_next_as_optional:
diff --git a/tensorflow/python/distribute/input_lib_test.py b/tensorflow/python/distribute/input_lib_test.py
index ff4436c4c8c..7f02d0121d0 100644
--- a/tensorflow/python/distribute/input_lib_test.py
+++ b/tensorflow/python/distribute/input_lib_test.py
@@ -185,38 +185,76 @@ class DistributedIteratorTestBase(test.TestCase):
       if not ops.executing_eagerly_outside_functions():
         evaluate(control_flow_ops.group(iterator.initializer))
 
-      for expected_value in expected_values:
-        next_element = iterator.get_next()
-        computed_value = evaluate(
-            [distribute_utils.select_replica(r, next_element)
-             for r in range(len(devices))])
-        self.assertEqual(len(expected_value), len(computed_value))
-        for i in range(len(expected_value)):
-          self.assertAllEqual(expected_value[i], computed_value[i])
+      def test_get_next(iterator):
+        for expected_value in expected_values:
+          next_element = iterator.get_next()
+          computed_value = evaluate([
+              distribute_utils.select_replica(r, next_element)
+              for r in range(len(devices))
+          ])
 
-      with self.assertRaises(errors.OutOfRangeError):
-        next_element = iterator.get_next()
-        evaluate(
-            [distribute_utils.select_replica(r, next_element)
-             for r in range(len(devices))])
+          self.assertEqual(len(expected_value), len(computed_value))
+          for i in range(len(expected_value)):
+            self.assertAllEqual(expected_value[i], computed_value[i])
 
-      # After re-initializing the iterator, should be able to iterate again.
-      if not ops.executing_eagerly_outside_functions():
-        evaluate(control_flow_ops.group(iterator.initializer))
+        with self.assertRaises(errors.OutOfRangeError):
+          next_element = iterator.get_next()
+          evaluate([
+              distribute_utils.select_replica(r, next_element)
+              for r in range(len(devices))
+          ])
+
+        # After re-initializing the iterator, should be able to iterate again.
+        if not ops.executing_eagerly_outside_functions():
+          evaluate(control_flow_ops.group(iterator.initializer))
+        else:
+          if api_type == "wrap_into_iterator":
+            self.skipTest("unsupported test combination")
+          else:
+            iterator = iter(dataset)
+
+        for expected_value in expected_values:
+          next_element = iterator.get_next()
+          computed_value = evaluate([
+              distribute_utils.select_replica(r, next_element)
+              for r in range(len(devices))
+          ])
+          self.assertEqual(len(expected_value), len(computed_value))
+          for i in range(len(expected_value)):
+            self.assertAllEqual(expected_value[i], computed_value[i])
+
+      def test_get_next_as_optional(iterator):
+        for expected_value in expected_values:
+          next_element = iterator.get_next_as_optional()
+          computed_value = evaluate([
+              distribute_utils.select_replica(r, next_element.get_value())
+              for r in range(len(devices))
+          ])
+
+          self.assertEqual(len(expected_value), len(computed_value))
+          for i in range(len(expected_value)):
+            self.assertAllEqual(expected_value[i], computed_value[i])
+
+        next_element = iterator.get_next_as_optional()
+        self.assertFalse(self.evaluate(next_element.has_value()))
+        with self.assertRaises(errors.InvalidArgumentError):
+          evaluate([
+              distribute_utils.select_replica(r, next_element.get_value())
+              for r in range(len(devices))
+          ])
+
+      test_get_next(iterator)
+
+      # re-initializing the iterator
+      if not tf2.enabled():
+        self.skipTest("Not testing get_next_as_optional in TF1")
       else:
         if api_type == "wrap_into_iterator":
           self.skipTest("unsupported test combination")
         else:
           iterator = iter(dataset)
 
-      for expected_value in expected_values:
-        next_element = iterator.get_next()
-        computed_value = evaluate(
-            [distribute_utils.select_replica(r, next_element)
-             for r in range(len(devices))])
-        self.assertEqual(len(expected_value), len(computed_value))
-        for i in range(len(expected_value)):
-          self.assertAllEqual(expected_value[i], computed_value[i])
+      test_get_next_as_optional(iterator)
 
     if iteration_type == "for_loop" and context.executing_eagerly():
       actual_values = []
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
index f712d9058b9..47899cc4188 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
@@ -13,4 +13,8 @@ tf_class {
     name: "get_next"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
   }
+  member_method {
+    name: "get_next_as_optional"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
 }

From 539e9cb3a22793aad5d2df885e016f43b81a6a9f Mon Sep 17 00:00:00 2001
From: Meghna Natraj 
Date: Thu, 18 Jun 2020 22:33:44 -0700
Subject: [PATCH 0597/1390] Update quantization docs to use
 TFLiteConverter.from_saved_model() API instead of .from_keras_model() API

PiperOrigin-RevId: 317251205
Change-Id: Ia8166decfa76327e3fd44871b194ffcae0f049f8
---
 .../lite/g3doc/performance/post_training_quantization.md  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/g3doc/performance/post_training_quantization.md b/tensorflow/lite/g3doc/performance/post_training_quantization.md
index ac584dd4c1c..dcf251e6d3d 100644
--- a/tensorflow/lite/g3doc/performance/post_training_quantization.md
+++ b/tensorflow/lite/g3doc/performance/post_training_quantization.md
@@ -34,7 +34,7 @@ weights from floating point to integer, which has 8-bits of precision:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 tflite_quant_model = converter.convert()
 

@@ -68,7 +68,7 @@ the following steps:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 def representative_dataset_gen():
   for _ in range(num_calibration_steps):
@@ -96,7 +96,7 @@ the following steps:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 def representative_dataset_gen():
   for _ in range(num_calibration_steps):
@@ -120,7 +120,7 @@ quantization of weights, use the following steps:
 
 
 import tensorflow as tf
-converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
+converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
 converter.optimizations = [tf.lite.Optimize.DEFAULT]
 converter.target_spec.supported_types = [tf.float16]
 tflite_quant_model = converter.convert()

From 8e654afea4adba36b94b0f7a3d33a23e788612e0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Thu, 18 Jun 2020 23:35:25 -0700
Subject: [PATCH 0598/1390] tf.numpy: Improve ndarray.__getitem__ to match
 numpy semantics.

PiperOrigin-RevId: 317256717
Change-Id: Ie89b81689f96242e3e9b01568e13937b80aaffc7
---
 .../python/ops/numpy_ops/np_array_ops.py      | 261 ++++++++++++++++--
 tensorflow/python/ops/numpy_ops/np_arrays.py  | 137 ---------
 2 files changed, 245 insertions(+), 153 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 906e53c556d..47236d45561 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -20,12 +20,15 @@ from __future__ import division
 from __future__ import print_function
 
 import math
+import numbers
+from typing import Sequence
 import numpy as np
 import six
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import clip_ops
 from tensorflow.python.ops import control_flow_ops
@@ -164,9 +167,11 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  #
 @np_utils.np_doc_only(np.array)
 def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
   """Since Tensors are immutable, a copy is made only if val is placed on a
+
   different device than the current one. Even if `copy` is False, a new Tensor
   may need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
-  is an ndarray or a Tensor."""  # pylint:disable=g-docstring-missing-newline
+  is an ndarray or a Tensor.
+  """  # pylint:disable=g-docstring-missing-newline
   if dtype:
     dtype = np_utils.result_type(dtype)
   if isinstance(val, np_arrays.ndarray):
@@ -215,6 +220,8 @@ def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-out
   result_t = np_utils.cond(
       np_utils.greater(ndmin, ndims), true_fn, lambda: result_t)
   return np_arrays.tensor_to_ndarray(result_t)
+
+
 # pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
 
 
@@ -1446,14 +1453,13 @@ def take_along_axis(arr, indices, axis):  # pylint: disable=missing-docstring
   # broadcast.
   arr_shape_original = array_ops.shape(arr)
   indices_shape_original = array_ops.shape(indices)
-  arr_shape = array_ops.tensor_scatter_update(
-      arr_shape_original, [[axis]], [1])
-  indices_shape = array_ops.tensor_scatter_update(
-      indices_shape_original, [[axis]], [1])
-  broadcasted_shape = array_ops.broadcast_dynamic_shape(
-      arr_shape, indices_shape)
-  arr_shape = array_ops.tensor_scatter_update(
-      broadcasted_shape, [[axis]], [arr_shape_original[axis]])
+  arr_shape = array_ops.tensor_scatter_update(arr_shape_original, [[axis]], [1])
+  indices_shape = array_ops.tensor_scatter_update(indices_shape_original,
+                                                  [[axis]], [1])
+  broadcasted_shape = array_ops.broadcast_dynamic_shape(arr_shape,
+                                                        indices_shape)
+  arr_shape = array_ops.tensor_scatter_update(broadcasted_shape, [[axis]],
+                                              [arr_shape_original[axis]])
   indices_shape = array_ops.tensor_scatter_update(
       broadcasted_shape, [[axis]], [indices_shape_original[axis]])
   arr = array_ops.broadcast_to(arr, arr_shape)
@@ -1468,10 +1474,10 @@ def take_along_axis(arr, indices, axis):  # pylint: disable=missing-docstring
   swapaxes_ = lambda t: swapaxes(np_utils.tensor_to_ndarray(t), axis, -1).data
 
   dont_move_axis_to_end = math_ops.equal(axis, rank - 1)
-  arr = np_utils.cond(
-      dont_move_axis_to_end, lambda: arr, lambda: swapaxes_(arr))
-  indices = np_utils.cond(
-      dont_move_axis_to_end, lambda: indices, lambda: swapaxes_(indices))
+  arr = np_utils.cond(dont_move_axis_to_end, lambda: arr,
+                      lambda: swapaxes_(arr))
+  indices = np_utils.cond(dont_move_axis_to_end, lambda: indices,
+                          lambda: swapaxes_(indices))
 
   arr_shape = array_ops.shape(arr)
   arr = array_ops.reshape(arr, [-1, arr_shape[-1]])
@@ -1481,8 +1487,231 @@ def take_along_axis(arr, indices, axis):  # pylint: disable=missing-docstring
 
   result = array_ops.gather(arr, indices, batch_dims=1)
   result = array_ops.reshape(result, indices_shape)
-  result = np_utils.cond(
-      dont_move_axis_to_end, lambda: result, lambda: swapaxes_(result))
+  result = np_utils.cond(dont_move_axis_to_end, lambda: result,
+                         lambda: swapaxes_(result))
   result.set_shape(possible_result_shape)
 
-  return  np_utils.tensor_to_ndarray(result)
+  return np_utils.tensor_to_ndarray(result)
+
+
+_SLICE_ERORR = (
+    'only integers, slices (`:`), ellipsis (`...`), '
+    'numpy.newaxis (`None`) and integer or boolean arrays are valid indices')
+
+
+def _as_index(idx, need_scalar=True):
+  """Helper function to parse idx as an index.
+
+  Args:
+    idx: index
+    need_scalar: If idx needs to be a scalar value.
+
+  Returns:
+    A pair, (indx, bool). First one is the parsed index and can be a tensor,
+    or scalar integer / Dimension. Second one is True if rank is known to be 0.
+
+  Raises:
+    IndexError: For incorrect indices.
+  """
+  if isinstance(idx, (numbers.Integral, tensor_shape.Dimension)):
+    return idx, True
+  data = asarray(idx).data
+  if data.dtype == dtypes.bool:
+    if data.shape.ndims != 1:
+      # TODO(agarwal): handle higher rank boolean masks.
+      raise NotImplementedError('Need rank 1 for bool index %s' % idx)
+    data = array_ops.where_v2(data)
+    data = array_ops.reshape(data, [-1])
+  if need_scalar and data.shape.rank not in (None, 0):
+    raise IndexError(_SLICE_ERORR + ', got {!r}'.format(idx))
+  np_dtype = data.dtype.as_numpy_dtype
+  if not np.issubdtype(np_dtype, np.integer):
+    raise IndexError(_SLICE_ERORR + ', got {!r}'.format(idx))
+  if data.dtype not in (dtypes.int64, dtypes.int32):
+    # TF slicing can only handle int32/int64. So we need to cast.
+    promoted_dtype = np.promote_types(np.int32, np_dtype)
+    if promoted_dtype == np.int32:
+      data = math_ops.cast(data, dtypes.int32)
+    elif promoted_dtype == np.int64:
+      data = math_ops.cast(data, dtypes.int64)
+    else:
+      raise IndexError(_SLICE_ERORR + ', got {!r}'.format(idx))
+  return data, data.shape.rank == 0
+
+
+def _slice_helper(tensor, slice_spec):
+  """Helper function for __getitem__."""
+  begin, end, strides = [], [], []
+  new_axis_mask, shrink_axis_mask = 0, 0
+  begin_mask, end_mask = 0, 0
+  ellipsis_mask = 0
+  advanced_indices = []
+  shrink_indices = []
+  for index, s in enumerate(slice_spec):
+    if isinstance(s, slice):
+      if s.start is not None:
+        begin.append(_as_index(s.start)[0])
+      else:
+        begin.append(0)
+        begin_mask |= (1 << index)
+      if s.stop is not None:
+        end.append(_as_index(s.stop)[0])
+      else:
+        end.append(0)
+        end_mask |= (1 << index)
+      if s.step is not None:
+        strides.append(_as_index(s.step)[0])
+      else:
+        strides.append(1)
+    elif s is Ellipsis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      ellipsis_mask |= (1 << index)
+    elif s is array_ops.newaxis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      new_axis_mask |= (1 << index)
+    else:
+      s, is_scalar = _as_index(s, False)
+      if is_scalar:
+        begin.append(s)
+        end.append(s + 1)
+        strides.append(1)
+        shrink_axis_mask |= (1 << index)
+        shrink_indices.append(index)
+      else:
+        begin.append(0)
+        end.append(0)
+        strides.append(1)
+        begin_mask |= (1 << index)
+        end_mask |= (1 << index)
+        advanced_indices.append((index, s, ellipsis_mask != 0))
+
+  # stack possibly involves no tensors, so we must use op_scope correct graph.
+  with ops.name_scope(
+      None,
+      'strided_slice', [tensor] + begin + end + strides,
+      skip_on_eager=False) as name:
+    if begin:
+      packed_begin, packed_end, packed_strides = (array_ops.stack(begin),
+                                                  array_ops.stack(end),
+                                                  array_ops.stack(strides))
+      if (packed_begin.dtype == dtypes.int64 or
+          packed_end.dtype == dtypes.int64 or
+          packed_strides.dtype == dtypes.int64):
+        if packed_begin.dtype != dtypes.int64:
+          packed_begin = math_ops.cast(packed_begin, dtypes.int64)
+        if packed_end.dtype != dtypes.int64:
+          packed_end = math_ops.cast(packed_end, dtypes.int64)
+        if packed_strides.dtype != dtypes.int64:
+          packed_strides = math_ops.cast(packed_strides, dtypes.int64)
+    else:
+      var_empty = constant_op.constant([], dtype=dtypes.int32)
+      packed_begin = packed_end = packed_strides = var_empty
+    # TODO(agarwal): set_shape on tensor to set rank.
+    tensor = array_ops.strided_slice(
+        tensor,
+        packed_begin,
+        packed_end,
+        packed_strides,
+        begin_mask=begin_mask,
+        end_mask=end_mask,
+        shrink_axis_mask=shrink_axis_mask,
+        new_axis_mask=new_axis_mask,
+        ellipsis_mask=ellipsis_mask,
+        name=name)
+    if not advanced_indices:
+      return tensor
+    advanced_indices_map = {}
+    for index, data, had_ellipsis in advanced_indices:
+      if had_ellipsis:
+        num_shrink = len([x for x in shrink_indices if x > index])
+        dim = index - len(slice_spec) + num_shrink
+      else:
+        num_shrink = len([x for x in shrink_indices if x < index])
+        dim = index - num_shrink
+      advanced_indices_map[dim] = data
+    dims = sorted(advanced_indices_map.keys())
+    dims_contiguous = True
+    if len(dims) > 1:
+      if dims[0] < 0 and dims[-1] >= 0:  # not all same sign
+        dims_contiguous = False
+      else:
+        for i in range(len(dims) - 1):
+          if dims[i] + 1 != dims[i + 1]:
+            dims_contiguous = False
+            break
+    indices = [advanced_indices_map[x] for x in dims]
+    indices = [x.data for x in _promote_dtype(*indices)]
+    indices = np_utils.tf_broadcast(*indices)
+    stacked_indices = array_ops.stack(indices, axis=-1)
+    if not dims_contiguous:
+      tensor = moveaxis(tensor, dims, range(len(dims))).data
+      tensor_shape_prefix = array_ops.shape(
+          tensor, out_type=stacked_indices.dtype)[:len(dims)]
+      stacked_indices = array_ops.where_v2(
+          stacked_indices < 0, stacked_indices + tensor_shape_prefix,
+          stacked_indices)
+      return array_ops.gather_nd(tensor, stacked_indices)
+    # Note that gather_nd does not support gathering from inside the array.
+    # To avoid shuffling data back and forth, we transform the indices and
+    # do a gather instead.
+    rank = np_utils._maybe_static(array_ops.rank(tensor))  # pylint: disable=protected-access
+    dims = [(x + rank if x < 0 else x) for x in dims]
+    shape_tensor = array_ops.shape(tensor, out_type=stacked_indices.dtype)
+    dim_sizes = array_ops.gather(shape_tensor, dims)
+    if len(dims) == 1:
+      stacked_indices = indices[0]
+    stacked_indices = array_ops.where_v2(stacked_indices < 0,
+                                         stacked_indices + dim_sizes,
+                                         stacked_indices)
+    axis = dims[0]
+    if len(dims) > 1:
+      index_scaling = math_ops.cumprod(
+          dim_sizes, reverse=True, exclusive=True)
+      stacked_indices = math_ops.tensordot(
+          stacked_indices, index_scaling, axes=1)
+      flat_shape = array_ops.concat(
+          [shape_tensor[:axis], [-1], shape_tensor[axis + len(dims):]],
+          axis=0)
+      tensor = array_ops.reshape(tensor, flat_shape)
+
+    return array_ops.gather(tensor, stacked_indices, axis=axis)
+
+
+def _as_spec_tuple(slice_spec):
+  """Convert slice_spec to tuple."""
+  if isinstance(slice_spec,
+                Sequence) and not isinstance(slice_spec, np.ndarray):
+    is_index = True
+    for s in slice_spec:
+      if s is None or s is Ellipsis or isinstance(s, (Sequence, slice)):
+        is_index = False
+        break
+      elif isinstance(s, (np_arrays.ndarray, np.ndarray)) and s.ndim != 0:
+        is_index = False
+        break
+    if not is_index:
+      return tuple(slice_spec)
+  return (slice_spec,)
+
+
+def _getitem(self, slice_spec):
+  """Implementation of ndarray.__getitem__."""
+  if (isinstance(slice_spec, bool) or (isinstance(slice_spec, ops.Tensor) and
+                                       slice_spec.dtype == dtypes.bool) or
+      (isinstance(slice_spec, (np.ndarray, np_arrays.ndarray)) and
+       slice_spec.dtype == np.bool)):
+    return np_utils.tensor_to_ndarray(
+        array_ops.boolean_mask(tensor=self.data, mask=slice_spec))
+
+  if not isinstance(slice_spec, tuple):
+    slice_spec = _as_spec_tuple(slice_spec)
+
+  result_t = _slice_helper(self.data, slice_spec)
+  return np_utils.tensor_to_ndarray(result_t)
+
+
+setattr(np_arrays.ndarray, '__getitem__', _getitem)
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index 8bec8a469a2..88bf4e7499a 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -20,138 +20,17 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import numbers
 import numpy as np
 import six
 
 from tensorflow.python.framework import composite_tensor
-from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import type_spec
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.numpy_ops import np_dtypes
-from tensorflow.python.util import nest
-
-
-_SLICE_TYPE_ERROR = (
-    'Only integers, slices (`:`), ellipsis (`...`), '
-    'tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid '
-    'indices')
-
-_SUPPORTED_SLICE_DTYPES = (dtypes.int32, dtypes.int32_ref, dtypes.int64,
-                           dtypes.int64_ref)
-
-
-def _check_index(idx):
-  """Check if a given value is a valid index into a tensor."""
-  if isinstance(idx, (numbers.Integral, tensor_shape.Dimension)):
-    return
-
-  # Optimistic check. Assumptions:
-  # * any object with a dtype is supported
-  # * any object with a dtype has a sizeable shape attribute.
-  dtype = getattr(idx, 'dtype', None)
-  if (dtype is None or dtypes.as_dtype(dtype) not in _SUPPORTED_SLICE_DTYPES or
-      idx.shape and len(idx.shape) == 1):
-    # TODO(slebedev): IndexError seems more appropriate here, but it
-    # will break `_slice_helper` contract.
-    raise TypeError(_SLICE_TYPE_ERROR + ', got {!r}'.format(idx))
-
-
-def _is_undefined_dimension(d):
-  return isinstance(d, tensor_shape.Dimension) and d.value is None
-
-
-def _slice_helper(tensor, slice_spec, var=None):
-  """Copied from array_ops._slice_helper, will be merged back later."""
-  if isinstance(slice_spec, bool) or \
-  (isinstance(slice_spec, ops.Tensor) and slice_spec.dtype == dtypes.bool) or \
-  (isinstance(slice_spec, np.ndarray) and slice_spec.dtype == bool):
-    return array_ops.boolean_mask(tensor=tensor, mask=slice_spec)
-
-  if not isinstance(slice_spec, (list, tuple)):
-    slice_spec = [slice_spec]
-
-  begin, end, strides = [], [], []
-  index = 0
-
-  new_axis_mask, shrink_axis_mask = 0, 0
-  begin_mask, end_mask = 0, 0
-  ellipsis_mask = 0
-  for s in slice_spec:
-    if isinstance(s, slice):
-      if s.start is not None and not _is_undefined_dimension(s.start):
-        _check_index(s.start)
-        begin.append(s.start)
-      else:
-        begin.append(0)
-        begin_mask |= (1 << index)
-      if s.stop is not None and not _is_undefined_dimension(s.stop):
-        _check_index(s.stop)
-        end.append(s.stop)
-      else:
-        end.append(0)
-        end_mask |= (1 << index)
-      if s.step is not None and not _is_undefined_dimension(s.step):
-        _check_index(s.step)
-        strides.append(s.step)
-      else:
-        strides.append(1)
-    elif s is Ellipsis:
-      begin.append(0)
-      end.append(0)
-      strides.append(1)
-      ellipsis_mask |= (1 << index)
-    elif s is array_ops.newaxis:
-      begin.append(0)
-      end.append(0)
-      strides.append(1)
-      new_axis_mask |= (1 << index)
-    else:
-      _check_index(s)
-      begin.append(s)
-      end.append(s + 1)
-      strides.append(1)
-      shrink_axis_mask |= (1 << index)
-    index += 1
-
-  # stack possibly involves no tensors, so we must use op_scope correct graph.
-  with ops.name_scope(
-      None,
-      'strided_slice', [tensor] + begin + end + strides,
-      skip_on_eager=False) as name:
-    if begin:
-      packed_begin, packed_end, packed_strides = (array_ops.stack(begin),
-                                                  array_ops.stack(end),
-                                                  array_ops.stack(strides))
-      if (packed_begin.dtype == dtypes.int64 or
-          packed_end.dtype == dtypes.int64 or
-          packed_strides.dtype == dtypes.int64):
-        if packed_begin.dtype != dtypes.int64:
-          packed_begin = math_ops.cast(packed_begin, dtypes.int64)
-        if packed_end.dtype != dtypes.int64:
-          packed_end = math_ops.cast(packed_end, dtypes.int64)
-        if packed_strides.dtype != dtypes.int64:
-          packed_strides = math_ops.cast(packed_strides, dtypes.int64)
-    else:
-      var_empty = constant_op.constant([], dtype=dtypes.int32)
-      packed_begin = packed_end = packed_strides = var_empty
-    return array_ops.strided_slice(
-        tensor,
-        packed_begin,
-        packed_end,
-        packed_strides,
-        begin_mask=begin_mask,
-        end_mask=end_mask,
-        shrink_axis_mask=shrink_axis_mask,
-        new_axis_mask=new_axis_mask,
-        ellipsis_mask=ellipsis_mask,
-        var=var,
-        name=name)
 
 
 def convert_to_tensor(value, dtype=None, dtype_hint=None):
@@ -361,22 +240,6 @@ class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
   def __bool__(self):
     return self.__nonzero__()
 
-  def __getitem__(self, slice_spec):
-    # TODO(srbs): Need to support better indexing.
-    def _gettensor(x):
-      if isinstance(x, ndarray):
-        x = x.data
-      if isinstance(x, ops.Tensor) and x.dtype not in (
-          dtypes.int32, dtypes.int64):
-        # Currently _slice_helper will only work with int32/int64 tensors, but
-        # type inference by numpy can create {u,}int{8,16}, so just cast.
-        x = math_ops.cast(x, dtypes.int32)
-      return x
-    slice_spec = nest.map_structure(_gettensor, slice_spec)
-
-    result_t = _slice_helper(self.data, slice_spec)
-    return tensor_to_ndarray(result_t)
-
   def __iter__(self):
     if not isinstance(self.data, ops.EagerTensor):
       raise TypeError('Iteration over symbolic tensor is not allowed')

From e972c5572634efd188696038e9241b75cdcd69bc Mon Sep 17 00:00:00 2001
From: Gaurav Jain 
Date: Fri, 19 Jun 2020 00:07:20 -0700
Subject: [PATCH 0599/1390] Add uint32 & uint64 to TF_CALL_INTEGRAL_TYPES

Both uint32 & uint64 had been omitted from TF_CALL_INTEGRAL_TYPES due to
suggested concerns of size bloat. In reality it seems that the size
increase is only around 2MB. Further, this fixes #39649 since we are no
longer inadvertently using the XLA_CPU device to perform tf.reduce_mean.

PiperOrigin-RevId: 317259372
Change-Id: Iacf75eaedce198fbef4bd9fd59b6fefa584cbf34
---
 tensorflow/core/framework/register_types.h    | 21 +++++---------
 tensorflow/core/framework/types.cc            |  5 ----
 tensorflow/core/kernels/BUILD                 |  2 ++
 tensorflow/core/kernels/concat_lib_cpu.cc     |  2 --
 tensorflow/core/kernels/concat_op.cc          |  2 --
 tensorflow/core/kernels/constant_op.cc        |  1 -
 tensorflow/core/kernels/control_flow_ops.cc   |  5 ----
 .../core/kernels/data/dataset_test_base.cc    |  2 --
 tensorflow/core/kernels/dense_update_ops.cc   |  1 -
 .../core/kernels/dynamic_partition_op.cc      |  2 --
 tensorflow/core/kernels/fill_functor.cc       |  5 +++-
 tensorflow/core/kernels/gather_op.cc          |  2 --
 tensorflow/core/kernels/identity_op.cc        |  1 -
 tensorflow/core/kernels/ragged_gather_op.cc   |  2 --
 .../kernels/ragged_tensor_from_variant_op.cc  |  2 --
 .../kernels/ragged_tensor_to_tensor_op.cc     |  2 --
 .../kernels/ragged_tensor_to_variant_op.cc    |  2 --
 .../core/kernels/resource_variable_ops.cc     |  1 -
 tensorflow/core/kernels/split_lib_cpu.cc      |  1 -
 tensorflow/core/kernels/split_op.cc           |  1 -
 tensorflow/core/kernels/strided_slice_op.cc   |  2 --
 .../core/kernels/strided_slice_op_impl.h      |  2 --
 tensorflow/core/kernels/topk_op.cc            |  2 --
 .../core/kernels/topk_op_gpu_uint32.cu.cc     | 28 +++++++++++++++++++
 .../core/kernels/topk_op_gpu_uint64.cu.cc     | 28 +++++++++++++++++++
 tensorflow/core/util/batch_util.cc            |  8 ------
 .../core/util/saved_tensor_slice_util.h       |  2 ++
 27 files changed, 71 insertions(+), 63 deletions(-)
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc

diff --git a/tensorflow/core/framework/register_types.h b/tensorflow/core/framework/register_types.h
index bc3e5e1743b..0cf6536e8c2 100644
--- a/tensorflow/core/framework/register_types.h
+++ b/tensorflow/core/framework/register_types.h
@@ -153,16 +153,9 @@ limitations under the License.
 #endif  // defined(IS_MOBILE_PLATFORM)  - end of TF_CALL_type defines
 
 // Defines for sets of types.
-
-// TODO(b/111604096): Add uint32 and uint64 to TF_CALL_INTEGRAL_TYPES.
-//
-// The uint32 and uint64 types were introduced in 10/2017 to be used via XLA and
-// thus were not included in TF_CALL_INTEGRAL_TYPES. Including them in
-// TF_CALL_INTEGRAL_TYPES should only happen after evaluating the effect on the
-// TF binary size and performance.
-#define TF_CALL_INTEGRAL_TYPES(m)                                      \
-  TF_CALL_int64(m) TF_CALL_int32(m) TF_CALL_uint16(m) TF_CALL_int16(m) \
-      TF_CALL_uint8(m) TF_CALL_int8(m)
+#define TF_CALL_INTEGRAL_TYPES(m)                                       \
+  TF_CALL_uint64(m) TF_CALL_int64(m) TF_CALL_uint32(m) TF_CALL_int32(m) \
+      TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)
 
 #define TF_CALL_FLOAT_TYPES(m) \
   TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)
@@ -174,10 +167,10 @@ limitations under the License.
 #define TF_CALL_REAL_NUMBER_TYPES_NO_BFLOAT16(m) \
   TF_CALL_INTEGRAL_TYPES(m) TF_CALL_half(m) TF_CALL_float(m) TF_CALL_double(m)
 
-#define TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m)                              \
-  TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)   \
-      TF_CALL_int64(m) TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) \
-          TF_CALL_int8(m)
+#define TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m)                                \
+  TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)     \
+      TF_CALL_uint64(m) TF_CALL_int64(m) TF_CALL_uint32(m) TF_CALL_uint16(m) \
+          TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)
 
 #define TF_CALL_COMPLEX_TYPES(m) TF_CALL_complex64(m) TF_CALL_complex128(m)
 
diff --git a/tensorflow/core/framework/types.cc b/tensorflow/core/framework/types.cc
index 97eaec98ffe..d6455e012d0 100644
--- a/tensorflow/core/framework/types.cc
+++ b/tensorflow/core/framework/types.cc
@@ -238,11 +238,6 @@ int DataTypeSize(DataType dt) {
     TF_CALL_qint16(CASE);
     TF_CALL_quint16(CASE);
 
-    // uint32 and uint64 aren't included in TF_CALL_POD_TYPES because we
-    // don't want to define kernels for them at this stage to avoid binary
-    // bloat.
-    TF_CALL_uint32(CASE);
-    TF_CALL_uint64(CASE);
     default:
       return 0;
   }
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 279dff92c58..97f974c6af4 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -4900,7 +4900,9 @@ tf_kernel_library(
         "topk_op_gpu_double.cu.cc",
         "topk_op_gpu_float.cu.cc",
         "topk_op_gpu_half.cu.cc",
+        "topk_op_gpu_uint64.cu.cc",
         "topk_op_gpu_int64.cu.cc",
+        "topk_op_gpu_uint32.cu.cc",
         "topk_op_gpu_int32.cu.cc",
         "topk_op_gpu_int16.cu.cc",
         "topk_op_gpu_uint16.cu.cc",
diff --git a/tensorflow/core/kernels/concat_lib_cpu.cc b/tensorflow/core/kernels/concat_lib_cpu.cc
index da73d3d2c56..1dec589d3ff 100644
--- a/tensorflow/core/kernels/concat_lib_cpu.cc
+++ b/tensorflow/core/kernels/concat_lib_cpu.cc
@@ -116,8 +116,6 @@ REGISTER(qint8)
 REGISTER(quint16)
 REGISTER(qint16)
 REGISTER(qint32)
-REGISTER(uint32)
-REGISTER(uint64)
 
 #if defined(IS_MOBILE_PLATFORM) && !defined(SUPPORT_SELECTIVE_REGISTRATION) && \
     !defined(__ANDROID_TYPES_FULL__)
diff --git a/tensorflow/core/kernels/concat_op.cc b/tensorflow/core/kernels/concat_op.cc
index be3e9a67c5f..d3f3a04f33b 100644
--- a/tensorflow/core/kernels/concat_op.cc
+++ b/tensorflow/core/kernels/concat_op.cc
@@ -208,8 +208,6 @@ REGISTER_CONCAT(qint8);
 REGISTER_CONCAT(quint16);
 REGISTER_CONCAT(qint16);
 REGISTER_CONCAT(qint32);
-REGISTER_CONCAT(uint32);
-REGISTER_CONCAT(uint64);
 
 #undef REGISTER_CONCAT
 
diff --git a/tensorflow/core/kernels/constant_op.cc b/tensorflow/core/kernels/constant_op.cc
index 4bcbc076446..dc178d17d49 100644
--- a/tensorflow/core/kernels/constant_op.cc
+++ b/tensorflow/core/kernels/constant_op.cc
@@ -211,7 +211,6 @@ TF_CALL_ALL_TYPES(REGISTER_CPU_KERNEL);
 // the conversion from uint8 to quint8.
 REGISTER_KERNEL(CPU, quint8);
 REGISTER_KERNEL(CPU, quint16);
-REGISTER_KERNEL(CPU, uint32);
 #undef REGISTER_CPU_KERNEL
 
 #ifdef TENSORFLOW_USE_SYCL
diff --git a/tensorflow/core/kernels/control_flow_ops.cc b/tensorflow/core/kernels/control_flow_ops.cc
index c8e83b6f672..accb2c59540 100644
--- a/tensorflow/core/kernels/control_flow_ops.cc
+++ b/tensorflow/core/kernels/control_flow_ops.cc
@@ -101,16 +101,12 @@ TF_CALL_ALL_TYPES(REGISTER_CPU_SWITCH);
 TF_CALL_ALL_TYPES(REGISTER_CPU_REF_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_REF_SWITCH);
-REGISTER_CPU_SWITCH(uint64);
 
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_SWITCH);
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_REF_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_REF_SWITCH);
-REGISTER_GPU_SWITCH(uint64);
 TF_CALL_variant(REGISTER_GPU_SWITCH);
-TF_CALL_uint32(REGISTER_GPU_SWITCH);
-TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
 
 #undef REGISTER_CPU_SWITCH
 #undef REGISTER_CPU_REF_SWITCH
@@ -311,7 +307,6 @@ TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_REF_KERNEL);
 REGISTER_GPU_KERNEL(bool);
 REGISTER_GPU_REF_KERNEL(bool);
-REGISTER_GPU_KERNEL(uint64);
 TF_CALL_variant(REGISTER_GPU_KERNEL);
 
 #undef REGISTER_GPU_KERNEL
diff --git a/tensorflow/core/kernels/data/dataset_test_base.cc b/tensorflow/core/kernels/data/dataset_test_base.cc
index b91ab9b733c..e41e35be1e9 100644
--- a/tensorflow/core/kernels/data/dataset_test_base.cc
+++ b/tensorflow/core/kernels/data/dataset_test_base.cc
@@ -220,8 +220,6 @@ Status DatasetOpsTestBase::ExpectEqual(const Tensor& a, const Tensor& b) {
     break;
     TF_CALL_NUMBER_TYPES(CASE);
     TF_CALL_tstring(CASE);
-    TF_CALL_uint32(CASE);
-    TF_CALL_uint64(CASE);
     // TODO(feihugis): figure out how to support variant tensors.
 #undef CASE
     default:
diff --git a/tensorflow/core/kernels/dense_update_ops.cc b/tensorflow/core/kernels/dense_update_ops.cc
index 55e4cd7606a..71235fca143 100644
--- a/tensorflow/core/kernels/dense_update_ops.cc
+++ b/tensorflow/core/kernels/dense_update_ops.cc
@@ -98,7 +98,6 @@ typedef Eigen::SyclDevice SYCLDevice;
 
 TF_CALL_ALL_TYPES(REGISTER_KERNELS);
 // uint32 not included in ALL_TYPES
-TF_CALL_uint32(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 // quint16 not included in QUANTIZIED_TYPES
 TF_CALL_quint16(REGISTER_KERNELS);
diff --git a/tensorflow/core/kernels/dynamic_partition_op.cc b/tensorflow/core/kernels/dynamic_partition_op.cc
index 90ed71dccce..95af19c4c48 100644
--- a/tensorflow/core/kernels/dynamic_partition_op.cc
+++ b/tensorflow/core/kernels/dynamic_partition_op.cc
@@ -164,8 +164,6 @@ class DynamicPartitionOp : public DynamicPartitionOp_Shared {
       DynamicPartitionOp)
 
 TF_CALL_ALL_TYPES(REGISTER_DYNAMIC_PARTITION);
-// For partitioning fingerprints.
-TF_CALL_uint64(REGISTER_DYNAMIC_PARTITION);
 #undef REGISTER_DYNAMIC_PARTITION
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/fill_functor.cc b/tensorflow/core/kernels/fill_functor.cc
index 10dd3df1915..174a4e45a79 100644
--- a/tensorflow/core/kernels/fill_functor.cc
+++ b/tensorflow/core/kernels/fill_functor.cc
@@ -45,6 +45,8 @@ DEFINE_SETZERO_CPU(Eigen::half);
 DEFINE_SETZERO_CPU(bfloat16);
 DEFINE_SETZERO_CPU(float);
 DEFINE_SETZERO_CPU(double);
+DEFINE_SETZERO_CPU(uint32);
+DEFINE_SETZERO_CPU(uint64);
 DEFINE_SETZERO_CPU(uint8);
 DEFINE_SETZERO_CPU(int8);
 DEFINE_SETZERO_CPU(uint16);
@@ -96,6 +98,8 @@ DEFINE_SETONE_CPU(Eigen::half);
 DEFINE_SETONE_CPU(bfloat16);
 DEFINE_SETONE_CPU(float);
 DEFINE_SETONE_CPU(double);
+DEFINE_SETONE_CPU(uint32);
+DEFINE_SETONE_CPU(uint64);
 DEFINE_SETONE_CPU(uint8);
 DEFINE_SETONE_CPU(int8);
 DEFINE_SETONE_CPU(uint16);
@@ -137,7 +141,6 @@ struct FillFunctor {
 TF_CALL_ALL_TYPES(DEFINE_FILL_CPU);
 DEFINE_FILL_CPU(quint8);
 DEFINE_FILL_CPU(quint16);
-DEFINE_FILL_CPU(uint32);
 #undef DEFINE_FILL_CPU
 
 #ifdef TENSORFLOW_USE_SYCL
diff --git a/tensorflow/core/kernels/gather_op.cc b/tensorflow/core/kernels/gather_op.cc
index 6d493a5f2ea..948567e019a 100644
--- a/tensorflow/core/kernels/gather_op.cc
+++ b/tensorflow/core/kernels/gather_op.cc
@@ -211,8 +211,6 @@ TF_CALL_ALL_TYPES(REGISTER_GATHER_CPU);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GATHER_CPU);
 TF_CALL_quint16(REGISTER_GATHER_CPU);
 TF_CALL_qint16(REGISTER_GATHER_CPU);
-TF_CALL_uint32(REGISTER_GATHER_CPU);
-TF_CALL_uint64(REGISTER_GATHER_CPU);
 
 #undef REGISTER_GATHER_CPU
 
diff --git a/tensorflow/core/kernels/identity_op.cc b/tensorflow/core/kernels/identity_op.cc
index fd94df9a768..daa8a1ddb25 100644
--- a/tensorflow/core/kernels/identity_op.cc
+++ b/tensorflow/core/kernels/identity_op.cc
@@ -122,7 +122,6 @@ REGISTER_SYCL_HOST_KERNEL(bool);
 
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL);
 REGISTER_GPU_KERNEL(Variant);
-TF_CALL_uint32(REGISTER_GPU_KERNEL);
 
 #undef REGISTER_GPU_KERNEL
 
diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc
index 88c0d1ebd69..3bf82cba050 100644
--- a/tensorflow/core/kernels/ragged_gather_op.cc
+++ b/tensorflow/core/kernels/ragged_gather_op.cc
@@ -296,8 +296,6 @@ TF_CALL_tstring(REGISTER_CPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_KERNEL);
 TF_CALL_quint16(REGISTER_CPU_KERNEL);
 TF_CALL_qint16(REGISTER_CPU_KERNEL);
-TF_CALL_uint32(REGISTER_CPU_KERNEL);
-TF_CALL_uint64(REGISTER_CPU_KERNEL);
 #undef REGISTER_CPU_KERNEL
 #undef REGISTER_CPU_KERNEL_WITH_INDEX_TYPE
 
diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
index f83bcb38c6c..ad0712e6fd0 100644
--- a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
@@ -308,8 +308,6 @@ TF_CALL_tstring(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 TF_CALL_quint16(REGISTER_KERNELS);
 TF_CALL_qint16(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
-TF_CALL_uint64(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 #undef REGISTER_KERNELS_WITH_SPLIT_TYPE
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
index d729c43f25a..9ae5d7ffbdc 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
@@ -561,8 +561,6 @@ TF_CALL_string(REGISTER_CPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_KERNEL);
 TF_CALL_quint16(REGISTER_CPU_KERNEL);
 TF_CALL_qint16(REGISTER_CPU_KERNEL);
-TF_CALL_uint32(REGISTER_CPU_KERNEL);
-TF_CALL_uint64(REGISTER_CPU_KERNEL);
 
 #undef REGISTER_CPU_KERNEL
 
diff --git a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
index 7a5ae1c6240..64c372b005e 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
@@ -213,8 +213,6 @@ TF_CALL_tstring(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 TF_CALL_quint16(REGISTER_KERNELS);
 TF_CALL_qint16(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
-TF_CALL_uint64(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 #undef REGISTER_KERNELS_WITH_SPLIT_TYPE
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 0fc1d53749f..79a64cb9219 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -512,7 +512,6 @@ class AssignVariableOp : public OpKernel {
 
 TF_CALL_ALL_TYPES(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/tensorflow/core/kernels/split_lib_cpu.cc b/tensorflow/core/kernels/split_lib_cpu.cc
index 0cb0a94d498..a3060e4e90d 100644
--- a/tensorflow/core/kernels/split_lib_cpu.cc
+++ b/tensorflow/core/kernels/split_lib_cpu.cc
@@ -43,7 +43,6 @@ void Split::operator()(
 
 TF_CALL_ALL_TYPES(DEFINE_CPU_KERNELS)
 DEFINE_CPU_KERNELS(quint8)
-DEFINE_CPU_KERNELS(uint64)
 
 #ifdef TENSORFLOW_USE_SYCL
 template 
diff --git a/tensorflow/core/kernels/split_op.cc b/tensorflow/core/kernels/split_op.cc
index f09740c6198..08575f01f67 100644
--- a/tensorflow/core/kernels/split_op.cc
+++ b/tensorflow/core/kernels/split_op.cc
@@ -404,7 +404,6 @@ class SplitOpSYCL : public SplitOpBase {
 
 TF_CALL_ALL_TYPES(REGISTER_SPLIT);
 REGISTER_SPLIT(quint8);
-REGISTER_SPLIT(uint64);
 
 #undef REGISTER_SPLIT
 
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index ccc1984bb98..b4099213303 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -440,8 +440,6 @@ class StridedSliceAssignOp : public OpKernel {
                           StridedSliceAssignOp)
 
 TF_CALL_ALL_TYPES(REGISTER_STRIDED_SLICE);
-TF_CALL_uint32(REGISTER_STRIDED_SLICE);
-TF_CALL_uint64(REGISTER_STRIDED_SLICE);
 
 #undef REGISTER_STRIDED_SLICE
 
diff --git a/tensorflow/core/kernels/strided_slice_op_impl.h b/tensorflow/core/kernels/strided_slice_op_impl.h
index 1ae959b7b3f..5ce1d773e33 100644
--- a/tensorflow/core/kernels/strided_slice_op_impl.h
+++ b/tensorflow/core/kernels/strided_slice_op_impl.h
@@ -287,8 +287,6 @@ TF_CALL_GPU_ALL_TYPES(DECLARE_FOR_N_GPU);
 #endif  // END GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
 TF_CALL_ALL_TYPES(DECLARE_FOR_N_CPU);
-TF_CALL_uint32(DECLARE_FOR_N_CPU);
-TF_CALL_uint64(DECLARE_FOR_N_CPU);
 
 #ifdef TENSORFLOW_USE_SYCL
 #define PREVENT_FOR_N_SYCL(T) \
diff --git a/tensorflow/core/kernels/topk_op.cc b/tensorflow/core/kernels/topk_op.cc
index c555b42f005..50325b7bcfe 100644
--- a/tensorflow/core/kernels/topk_op.cc
+++ b/tensorflow/core/kernels/topk_op.cc
@@ -258,7 +258,6 @@ namespace functor {
 
 TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
 TF_CALL_INTEGRAL_TYPES(DECLARE_GPU_SPEC);
-TF_CALL_uint32(DECLARE_GPU_SPEC);
 
 #undef DECLARE_GPU_SPEC
 
@@ -276,7 +275,6 @@ TF_CALL_uint32(DECLARE_GPU_SPEC);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER_KERNELS);
 TF_CALL_INTEGRAL_TYPES(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS)
 #undef REGISTER_KERNELS
 
 #endif  // end GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc b/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
new file mode 100644
index 00000000000..16e2e0e9420
--- /dev/null
+++ b/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include "tensorflow/core/kernels/topk_op.h"
+#include "tensorflow/core/kernels/topk_op_gpu.h"
+
+namespace tensorflow {
+using Eigen::GpuDevice;
+
+template struct functor::TopKFunctor;
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc b/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc
new file mode 100644
index 00000000000..895247a63a2
--- /dev/null
+++ b/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include "tensorflow/core/kernels/topk_op.h"
+#include "tensorflow/core/kernels/topk_op_gpu.h"
+
+namespace tensorflow {
+using Eigen::GpuDevice;
+
+template struct functor::TopKFunctor;
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/core/util/batch_util.cc b/tensorflow/core/util/batch_util.cc
index b88c365ced0..e03188b04da 100644
--- a/tensorflow/core/util/batch_util.cc
+++ b/tensorflow/core/util/batch_util.cc
@@ -182,8 +182,6 @@ Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index) {
   switch (element.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopyElementToSlice Unhandled data type: ",
@@ -207,8 +205,6 @@ Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index) {
   switch (parent.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopySliceToElement Unhandled data type: ",
@@ -280,8 +276,6 @@ Status CopyContiguousSlices(const Tensor& src, int64 src_offset,
   switch (src.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopyContiguousSlices unhandled data type: ",
@@ -308,8 +302,6 @@ Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element, int64 index) {
   switch (parent->dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented(
diff --git a/tensorflow/core/util/saved_tensor_slice_util.h b/tensorflow/core/util/saved_tensor_slice_util.h
index 09b9235b711..1f9768f5163 100644
--- a/tensorflow/core/util/saved_tensor_slice_util.h
+++ b/tensorflow/core/util/saved_tensor_slice_util.h
@@ -116,7 +116,9 @@ TENSOR_PROTO_EXTRACT_TYPE(double, double, double);
 TENSOR_PROTO_EXTRACT_TYPE_COMPLEX(complex64, scomplex, float);
 TENSOR_PROTO_EXTRACT_TYPE_COMPLEX(complex128, dcomplex, double);
 TENSOR_PROTO_EXTRACT_TYPE(int32, int, int32);
+TENSOR_PROTO_EXTRACT_TYPE(uint32, uint32, uint32);
 TENSOR_PROTO_EXTRACT_TYPE(int64, int64, protobuf_int64);
+TENSOR_PROTO_EXTRACT_TYPE(uint64, uint64, protobuf_uint64);
 TENSOR_PROTO_EXTRACT_TYPE(uint16, int, int32);
 TENSOR_PROTO_EXTRACT_TYPE(uint8, int, int32);
 TENSOR_PROTO_EXTRACT_TYPE(int8, int, int32);

From 9f20b156bc7862fb621756fd5d6744255b1f3735 Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Fri, 19 Jun 2020 01:13:01 -0700
Subject: [PATCH 0600/1390] [XLA:GPU] [NFC] Clarify the precondition for the
 fast reduction emitter

PiperOrigin-RevId: 317266013
Change-Id: I384acac279f0db53f195d5b43318c38c87a1739c
---
 tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index b97aa3651c6..01bcf456f75 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -226,6 +226,11 @@ bool IsReductionFromOrToContiguousDimensions(const HloInstruction& reduce) {
       dims_to_keep.push_back(dim);
     }
   }
+
+  // We support fast codegen for three cases:
+  // 1) Row reduction: (K, R)
+  // 2) Column reduction: (K, R, K)
+  // 3) "Batched" row reduction: (R, K, R)
   if (!LayoutUtil::AreDimensionsConsecutive(input->shape().layout(),
                                             dims_to_keep) &&
       !LayoutUtil::AreDimensionsConsecutive(input->shape().layout(),

From 051d1b70f5f0636316e2651630f0ade554f192c0 Mon Sep 17 00:00:00 2001
From: Stefano Galarraga 
Date: Fri, 19 Jun 2020 01:21:24 -0700
Subject: [PATCH 0601/1390] Fix NNAPI delegation error on models with MAX/MIN
 operations with scalar quantized operators

PiperOrigin-RevId: 317266736
Change-Id: Ieed8a77685d4ca0d51389b5976addf0de167cfcf
---
 tensorflow/lite/delegates/nnapi/nnapi_delegate.cc |  2 ++
 tensorflow/lite/kernels/maximum_minimum_test.cc   | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
index a3a3f9fda4d..1c35ee370c2 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
@@ -160,6 +160,8 @@ bool IsScalarInputSupported(int builtin_code) {
     case kTfLiteBuiltinLess:
     case kTfLiteBuiltinLessEqual:
     case kTfLiteBuiltinPow:
+    case kTfLiteBuiltinMaximum:
+    case kTfLiteBuiltinMinimum:
       return true;
     default:
       return false;
diff --git a/tensorflow/lite/kernels/maximum_minimum_test.cc b/tensorflow/lite/kernels/maximum_minimum_test.cc
index 2c036e369bd..803fe91c460 100644
--- a/tensorflow/lite/kernels/maximum_minimum_test.cc
+++ b/tensorflow/lite/kernels/maximum_minimum_test.cc
@@ -190,6 +190,17 @@ TEST(MaximumOpTest, Int32WithBroadcastTest_ScalarY) {
                      data1, data2, {1, 0, -1, -2, 2, 2}, /*is_constant=*/true);
 }
 
+TEST(MaximumOpTest, Int8WithBroadcastTest_ScalarY) {
+  std::initializer_list data1 = {1, 0, -1, -2, 3, 11};
+  std::initializer_list data2 = {2};
+  TestModel(BuiltinOperator_MAXIMUM, {TensorType_INT8, {3, 1, 2}},
+                    {TensorType_INT8, {}}, {TensorType_INT8, {3, 1, 2}}, data1,
+                    data2, {2, 2, 2, 2, 3, 11}, /*is_constant=*/true);
+  TestModel(BuiltinOperator_MINIMUM, {TensorType_INT8, {3, 1, 2}},
+                    {TensorType_INT8, {}}, {TensorType_INT8, {3, 1, 2}}, data1,
+                    data2, {1, 0, -1, -2, 2, 2}, /*is_constant=*/true);
+}
+
 TEST(MaxMinOpTest, Int8Test8D) {
   std::initializer_list data1 = {1, 0, 2, 11, 2, 23};
   std::initializer_list data2 = {0, 0, 1, 12, 123, 1};

From e51b17f4582183a216d3a47450117c5e8cdd387d Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Fri, 19 Jun 2020 01:49:14 -0700
Subject: [PATCH 0602/1390] Add a small test to cover the mlir generated Tanh
 GPU kernel.

This test is a first step towards being able to ensure that we don't
accidentally break the kernel generation.

PiperOrigin-RevId: 317269120
Change-Id: Iad6bdd7ab7e9fb819a478c947ba6294a191f1099
---
 tensorflow/core/kernels/BUILD                 | 19 +++++
 .../mlir_generated_op_gpu_tanh_test.cc        | 85 +++++++++++++++++++
 2 files changed, 104 insertions(+)
 create mode 100644 tensorflow/core/kernels/mlir_generated_op_gpu_tanh_test.cc

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 97f974c6af4..0b7a092033b 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -4168,6 +4168,25 @@ tf_kernel_library(
     ]),
 )
 
+tf_cuda_cc_test(
+    name = "mlir_generated_op_gpu_tanh_test",
+    size = "small",
+    srcs = if_mlir_generated_gpu_kernels_enabled(["mlir_generated_op_gpu_tanh_test.cc"]),
+    tags = tf_cuda_tests_tags() + ["no_rocm"],
+    deps = [
+        ":cwise_op",
+        ":ops_testutil",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "//tensorflow/core/common_runtime:device",
+        "//tensorflow/core/common_runtime:device_factory",
+    ],
+)
+
 tf_kernel_library(
     name = "nextafter_op",
     prefix = "nextafter_op",
diff --git a/tensorflow/core/kernels/mlir_generated_op_gpu_tanh_test.cc b/tensorflow/core/kernels/mlir_generated_op_gpu_tanh_test.cc
new file mode 100644
index 00000000000..39c1d709b1e
--- /dev/null
+++ b/tensorflow/core/kernels/mlir_generated_op_gpu_tanh_test.cc
@@ -0,0 +1,85 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include 
+#include 
+#include 
+
+#include "tensorflow/core/common_runtime/device.h"
+#include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+class MlirGeneratedOpGpuTanhTest : public OpsTestBase {
+ protected:
+  void SetUp() override {
+    std::unique_ptr device_gpu(
+        tensorflow::DeviceFactory::NewDevice("GPU", {},
+                                             "/job:a/replica:0/task:0"));
+    SetDevice(tensorflow::DEVICE_GPU, std::move(device_gpu));
+  }
+  template 
+  void RunTanhOp(std::initializer_list input) {
+    TensorShape shape({2, 7});
+    TF_ASSERT_OK(NodeDefBuilder("tanh_op", "Tanh")
+                     .Input(FakeInput(DataTypeToEnum::v()))
+                     .Attr("T", DataTypeToEnum::v())
+                     .Finalize(node_def()));
+
+    TF_ASSERT_OK(InitOp());
+    AddInputFromArray(shape, input);
+    TF_ASSERT_OK(RunOpKernel());
+
+    Tensor expected_tensor(allocator(), DataTypeToEnum::value, shape);
+    std::vector expected;
+    expected.reserve(input.size());
+    for (const T& inp : input) {
+      expected.push_back(static_cast(std::tanh(static_cast(inp))));
+    }
+    test::FillValues(&expected_tensor, expected);
+    test::ExpectClose(expected_tensor, *GetOutput(0));
+  }
+};
+
+TEST_F(MlirGeneratedOpGpuTanhTest, TanhFloat) {
+  RunTanhOp({-18.0f, -9.0f, -1e-6f, -0.0f, 0.0f, 1e-6, 0.1f, 0.2f, 0.3f,
+                    0.5f, 0.7f, 0.9f, 9.0f, 18.0f});
+}
+
+TEST_F(MlirGeneratedOpGpuTanhTest, TanhDouble) {
+  RunTanhOp({-18.0, -9.0, -1e-6, -0.0, 0.0, 1e-6, 0.1, 0.2, 0.3, 0.5,
+                     0.7, 0.9, 9.0, 18.0});
+}
+
+TEST_F(MlirGeneratedOpGpuTanhTest, TanhHalf) {
+  RunTanhOp(
+      {static_cast(-18.0), static_cast(-9.0),
+       static_cast(-1e-6), static_cast(-0.0),
+       static_cast(0.0), static_cast(1e-6),
+       static_cast(0.1), static_cast(0.2),
+       static_cast(0.3), static_cast(0.5),
+       static_cast(0.7), static_cast(0.9),
+       static_cast(9.0), static_cast(18.0)});
+}
+
+}  // namespace
+}  // end namespace tensorflow

From 772433a2a2120d0aefc6c3628c6254d5a1aaf19d Mon Sep 17 00:00:00 2001
From: YoungSeok Yoon 
Date: Fri, 19 Jun 2020 01:58:58 -0700
Subject: [PATCH 0603/1390] Add flag for using optimized TFLite CPU kernels on
 iOS

This adds new experimental flags to the interpreter options of TFLite Obj-C and
Swift APIs, which can be used for opting in to a set of highly optimized
floating point kernels provided via the XNNPACK delegate. The flags can be used
as follows.

Obj-C:

    TFLInterpreterOptions *options = [[TFLInterpreterOptions alloc] init];
    options.useXNNPACK = YES;
    NSError *error;
    TFLInterpreter *interpreter =
        [[TFLInterpreter alloc] initWithModelPath:@"model/path"
                                          options:options
                                            error:&error];

Swift:

    var options = InterpreterOptions()
    options.isXNNPackEnabled = true
    var interpreter = try Interpreter(modelPath: "model/path", options: options)

PiperOrigin-RevId: 317270012
Change-Id: I82aae43c3de13ab08af3c70513e2a458e807b0f1
---
 tensorflow/lite/delegates/xnnpack/BUILD       |  4 ++
 tensorflow/lite/experimental/ios/BUILD.apple  | 18 +++++
 tensorflow/lite/experimental/objc/BUILD.apple |  1 +
 .../objc/TensorFlowLiteObjC-nightly.podspec   |  1 +
 .../objc/TensorFlowLiteObjC.podspec           |  1 +
 .../objc/TensorFlowLiteObjC.podspec.template  |  1 +
 .../objc/apis/TFLInterpreterOptions.h         | 21 ++++++
 .../objc/sources/TFLInterpreter.mm            | 15 +++++
 .../objc/tests/TFLInterpreterOptionsTests.m   |  9 +++
 .../swift/Sources/Interpreter.swift           | 67 ++++++++++++++++---
 .../swift/Tests/InterpreterTests.swift        | 62 +++++++++++------
 11 files changed, 171 insertions(+), 29 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index 97e6aea2a6b..eaf7d8f6f03 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -14,6 +14,10 @@ EMSCRIPTEN_LINKOPTS = [
     "-s TOTAL_MEMORY=134217728",
 ]
 
+exports_files([
+    "xnnpack_delegate.h",
+])
+
 cc_library(
     name = "xnnpack_delegate",
     srcs = ["xnnpack_delegate.cc"],
diff --git a/tensorflow/lite/experimental/ios/BUILD.apple b/tensorflow/lite/experimental/ios/BUILD.apple
index 1a85b604f9b..7a40ca7b8e7 100644
--- a/tensorflow/lite/experimental/ios/BUILD.apple
+++ b/tensorflow/lite/experimental/ios/BUILD.apple
@@ -18,10 +18,26 @@ sh_binary(
     ],
 )
 
+# When the static framework is built with bazel, the all header files are moved
+# to the "Headers" directory with no header path prefixes. This auxiliary rule
+# is used for stripping the path prefix to the "common.h" file included by the
+# "xnnpack_delegate.h" header.
+genrule(
+    name = "strip_xnnpack_include_hdr",
+    srcs = ["//tensorflow/lite/delegates/xnnpack:xnnpack_delegate.h"],
+    outs = ["xnnpack_delegate.h"],
+    cmd = """
+    sed 's|#include ".*common.h"|#include "common.h"|'\
+    "$(location //tensorflow/lite/delegates/xnnpack:xnnpack_delegate.h)"\
+    > "$@"
+    """,
+)
+
 # bazel build -c opt --config=ios_fat //tensorflow/lite/experimental/ios:TensorFlowLiteC_framework
 tflite_ios_static_framework(
     name = "TensorFlowLiteC_framework",
     hdrs = [
+        ":xnnpack_delegate.h",
         "//tensorflow/lite/c:c_api.h",
         "//tensorflow/lite/c:common.h",
     ],
@@ -105,6 +121,7 @@ cc_library(
     hdrs = [
         "//tensorflow/lite/c:c_api.h",
         "//tensorflow/lite/c:common.h",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate.h",
     ],
     tags = [
         "nobuilder",
@@ -112,6 +129,7 @@ cc_library(
     ],
     deps = [
         "//tensorflow/lite/c:c_api",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
     ],
 )
 
diff --git a/tensorflow/lite/experimental/objc/BUILD.apple b/tensorflow/lite/experimental/objc/BUILD.apple
index 09d4547813a..d26d90c46a1 100644
--- a/tensorflow/lite/experimental/objc/BUILD.apple
+++ b/tensorflow/lite/experimental/objc/BUILD.apple
@@ -64,6 +64,7 @@ objc_library(
     visibility = ios_visibility_whitelist(),
     deps = [
         "//tensorflow/lite/c:c_api",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
     ],
     alwayslink = 1,
 )
diff --git a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC-nightly.podspec b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC-nightly.podspec
index e039fb57114..eed0f087f44 100644
--- a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC-nightly.podspec
+++ b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC-nightly.podspec
@@ -26,6 +26,7 @@ Pod::Spec.new do |s|
     objc_dir + '{apis,sources}/*.{h,m,mm}',
     tfl_dir + 'c/c_api.h',
     tfl_dir + 'c/common.h',
+    tfl_dir + 'delegates/xnnpack/xnnpack_delegate.h',
   ]
   s.module_map = objc_dir + 'apis/framework.modulemap'
   s.dependency 'TensorFlowLiteC', "~> #{s.version}"
diff --git a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec
index c673cfad759..5817619a58f 100644
--- a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec
+++ b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec
@@ -26,6 +26,7 @@ Pod::Spec.new do |s|
     objc_dir + '{apis,sources}/*.{h,m,mm}',
     tfl_dir + 'c/c_api.h',
     tfl_dir + 'c/common.h',
+    tfl_dir + 'delegates/xnnpack/xnnpack_delegate.h',
   ]
   s.module_map = objc_dir + 'apis/framework.modulemap'
   s.dependency 'TensorFlowLiteC', "#{s.version}"
diff --git a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec.template b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec.template
index fc9e10e4a2c..4ab5753e016 100644
--- a/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec.template
+++ b/tensorflow/lite/experimental/objc/TensorFlowLiteObjC.podspec.template
@@ -26,6 +26,7 @@ Pod::Spec.new do |s|
     objc_dir + '{apis,sources}/*.{h,m,mm}',
     tfl_dir + 'c/c_api.h',
     tfl_dir + 'c/common.h',
+    tfl_dir + 'delegates/xnnpack/xnnpack_delegate.h',
   ]
   s.module_map = objc_dir + 'apis/framework.modulemap'
   s.dependency 'TensorFlowLiteC', '~> 0.0.1-nightly'
diff --git a/tensorflow/lite/experimental/objc/apis/TFLInterpreterOptions.h b/tensorflow/lite/experimental/objc/apis/TFLInterpreterOptions.h
index 6461fbf0178..d7dbb2bd970 100644
--- a/tensorflow/lite/experimental/objc/apis/TFLInterpreterOptions.h
+++ b/tensorflow/lite/experimental/objc/apis/TFLInterpreterOptions.h
@@ -25,6 +25,27 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property(nonatomic) NSUInteger numberOfThreads;
 
+/**
+ * Experimental: Enable an optimized set of floating point CPU kernels (provided by XNNPACK).
+ *
+ * Enabling this flag will enable use of a new, highly optimized set of CPU kernels provided via the
+ * XNNPACK delegate. Currently, this is restricted to a subset of floating point operations.
+ * Eventually, we plan to enable this by default, as it can provide significant performance benefits
+ * for many classes of floating point models. See
+ * https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md
+ * for more details.
+ *
+ * Things to keep in mind when enabling this flag:
+ *
+ *     * Startup time and resize time may increase.
+ *     * Baseline memory consumption may increase.
+ *     * Compatibility with other delegates (e.g., GPU) has not been fully validated.
+ *     * Quantized models will not see any benefit.
+ *
+ * WARNING: This is an experimental interface that is subject to change.
+ */
+@property(nonatomic) BOOL useXNNPACK;
+
 /**
  * Initializes a new instance of `TFLInterpreterOptions`.
  *
diff --git a/tensorflow/lite/experimental/objc/sources/TFLInterpreter.mm b/tensorflow/lite/experimental/objc/sources/TFLInterpreter.mm
index 94031ee5428..34dd119885d 100644
--- a/tensorflow/lite/experimental/objc/sources/TFLInterpreter.mm
+++ b/tensorflow/lite/experimental/objc/sources/TFLInterpreter.mm
@@ -23,6 +23,7 @@
 #import "tensorflow/lite/experimental/objc/apis/TFLTensor.h"
 
 #include "tensorflow/lite/c/c_api.h"
+#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -45,6 +46,9 @@ static void TFLInterpreterErrorReporter(void *user_data, const char *format, va_
 /** TfLiteInterpreter backed by C API. */
 @property(nonatomic, nullable) TfLiteInterpreter *interpreter;
 
+/** TfLiteDelegate backed by C API. */
+@property(nonatomic, nullable) TfLiteDelegate *xnnpack_delegate;
+
 @end
 
 @implementation TFLInterpreter
@@ -53,6 +57,7 @@ static void TFLInterpreterErrorReporter(void *user_data, const char *format, va_
 
 - (void)dealloc {
   TfLiteInterpreterDelete(_interpreter);
+  TfLiteXNNPackDelegateDelete(_xnnpack_delegate);
 }
 
 #pragma mark - Public
@@ -104,6 +109,16 @@ static void TFLInterpreterErrorReporter(void *user_data, const char *format, va_
       }
       TfLiteInterpreterOptionsSetErrorReporter(cOptions, TFLInterpreterErrorReporter, nullptr);
 
+      if (options.useXNNPACK) {
+        TfLiteXNNPackDelegateOptions xnnpack_options = TfLiteXNNPackDelegateOptionsDefault();
+        if (options.numberOfThreads > 0) {
+          xnnpack_options.num_threads = (int32_t)options.numberOfThreads;
+        }
+
+        _xnnpack_delegate = TfLiteXNNPackDelegateCreate(&xnnpack_options);
+        TfLiteInterpreterOptionsAddDelegate(cOptions, _xnnpack_delegate);
+      }
+
       _interpreter = TfLiteInterpreterCreate(model, cOptions);
       if (_interpreter == nullptr) {
         [TFLErrorUtil saveInterpreterErrorWithCode:TFLInterpreterErrorCodeFailedToCreateInterpreter
diff --git a/tensorflow/lite/experimental/objc/tests/TFLInterpreterOptionsTests.m b/tensorflow/lite/experimental/objc/tests/TFLInterpreterOptionsTests.m
index 00b800d6af9..286cba98b49 100644
--- a/tensorflow/lite/experimental/objc/tests/TFLInterpreterOptionsTests.m
+++ b/tensorflow/lite/experimental/objc/tests/TFLInterpreterOptionsTests.m
@@ -32,6 +32,7 @@ NS_ASSUME_NONNULL_BEGIN
   TFLInterpreterOptions *options = [[TFLInterpreterOptions alloc] init];
   XCTAssertNotNil(options);
   XCTAssertEqual(options.numberOfThreads, 0);
+  XCTAssertFalse(options.useXNNPACK);
 }
 
 - (void)testSetNumberOfThread {
@@ -44,6 +45,14 @@ NS_ASSUME_NONNULL_BEGIN
   XCTAssertEqual(options.numberOfThreads, 3);
 }
 
+- (void)testUseXNNPACK {
+  TFLInterpreterOptions *options = [[TFLInterpreterOptions alloc] init];
+  options.useXNNPACK = YES;
+  XCTAssertTrue(options.useXNNPACK);
+  options.useXNNPACK = NO;
+  XCTAssertFalse(options.useXNNPACK);
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/tensorflow/lite/experimental/swift/Sources/Interpreter.swift b/tensorflow/lite/experimental/swift/Sources/Interpreter.swift
index b83c36c4e1d..3567822208d 100644
--- a/tensorflow/lite/experimental/swift/Sources/Interpreter.swift
+++ b/tensorflow/lite/experimental/swift/Sources/Interpreter.swift
@@ -39,6 +39,9 @@ public final class Interpreter {
   /// The underlying `TfLiteInterpreter` C pointer.
   private var cInterpreter: CInterpreter?
 
+  /// The underlying `TfLiteDelegate` C pointer for XNNPACK delegate.
+  private var cXNNPackDelegate: Delegate.CDelegate?
+
   /// Creates a new instance with the given values.
   ///
   /// - Parameters:
@@ -78,6 +81,14 @@ public final class Interpreter {
       )
     }
     delegates?.forEach { TfLiteInterpreterOptionsAddDelegate(cInterpreterOptions, $0.cDelegate) }
+
+    // Configure the XNNPack delegate after the other delegates explicitly added by the user.
+    options.map {
+      if $0.isXNNPackEnabled {
+        configureXNNPack(options: $0, cInterpreterOptions: cInterpreterOptions)
+      }
+    }
+
     guard let cInterpreter = TfLiteInterpreterCreate(model.cModel, cInterpreterOptions) else {
       throw InterpreterError.failedToCreateInterpreter
     }
@@ -86,6 +97,7 @@ public final class Interpreter {
 
   deinit {
     TfLiteInterpreterDelete(cInterpreter)
+    TfLiteXNNPackDelegateDelete(cXNNPackDelegate)
   }
 
   /// Invokes the interpreter to perform inference from the loaded graph.
@@ -201,12 +213,13 @@ public final class Interpreter {
     guard case 0...maxIndex = index else {
       throw InterpreterError.invalidTensorIndex(index: index, maxIndex: maxIndex)
     }
-    guard TfLiteInterpreterResizeInputTensor(
-      cInterpreter,
-      Int32(index),
-      shape.int32Dimensions,
-      Int32(shape.rank)
-    ) == kTfLiteOk
+    guard
+      TfLiteInterpreterResizeInputTensor(
+        cInterpreter,
+        Int32(index),
+        shape.int32Dimensions,
+        Int32(shape.rank)
+      ) == kTfLiteOk
     else {
       throw InterpreterError.failedToResizeInputTensor(index: index)
     }
@@ -236,11 +249,11 @@ public final class Interpreter {
     }
 
     #if swift(>=5.0)
-    let status = data.withUnsafeBytes {
-      TfLiteTensorCopyFromBuffer(cTensor, $0.baseAddress, data.count)
-    }
+      let status = data.withUnsafeBytes {
+        TfLiteTensorCopyFromBuffer(cTensor, $0.baseAddress, data.count)
+      }
     #else
-    let status = data.withUnsafeBytes { TfLiteTensorCopyFromBuffer(cTensor, $0, data.count) }
+      let status = data.withUnsafeBytes { TfLiteTensorCopyFromBuffer(cTensor, $0, data.count) }
     #endif  // swift(>=5.0)
     guard status == kTfLiteOk else { throw InterpreterError.failedToCopyDataToInputTensor }
     return try input(at: index)
@@ -256,6 +269,18 @@ public final class Interpreter {
       throw InterpreterError.failedToAllocateTensors
     }
   }
+
+  // MARK: - Private
+
+  private func configureXNNPack(options: Options, cInterpreterOptions: OpaquePointer) {
+    var cXNNPackOptions = TfLiteXNNPackDelegateOptionsDefault()
+    if let threadCount = options.threadCount, threadCount > 0 {
+      cXNNPackOptions.num_threads = Int32(threadCount)
+    }
+
+    cXNNPackDelegate = TfLiteXNNPackDelegateCreate(&cXNNPackOptions)
+    TfLiteInterpreterOptionsAddDelegate(cInterpreterOptions, cXNNPackDelegate)
+  }
 }
 
 extension Interpreter {
@@ -265,6 +290,28 @@ extension Interpreter {
     /// indicating that the `Interpreter` will decide the number of threads to use.
     public var threadCount: Int? = nil
 
+    /// Indicates whether an optimized set of floating point CPU kernels, provided by XNNPACK, is
+    /// enabled.
+    ///
+    /// - Experiment:
+    /// Enabling this flag will enable use of a new, highly optimized set of CPU kernels provided
+    /// via the XNNPACK delegate. Currently, this is restricted to a subset of floating point
+    /// operations. Eventually, we plan to enable this by default, as it can provide significant
+    /// performance benefits for many classes of floating point models. See
+    /// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md
+    /// for more details.
+    ///
+    /// - Important:
+    /// Things to keep in mind when enabling this flag:
+    ///
+    ///     * Startup time and resize time may increase.
+    ///     * Baseline memory consumption may increase.
+    ///     * Compatibility with other delegates (e.g., GPU) has not been fully validated.
+    ///     * Quantized models will not see any benefit.
+    ///
+    /// - Warning: This is an experimental interface that is subject to change.
+    public var isXNNPackEnabled: Bool = false
+
     /// Creates a new instance with the default values.
     public init() {}
   }
diff --git a/tensorflow/lite/experimental/swift/Tests/InterpreterTests.swift b/tensorflow/lite/experimental/swift/Tests/InterpreterTests.swift
index 8d0140279af..67d8120df4d 100644
--- a/tensorflow/lite/experimental/swift/Tests/InterpreterTests.swift
+++ b/tensorflow/lite/experimental/swift/Tests/InterpreterTests.swift
@@ -142,10 +142,12 @@ class InterpreterTests: XCTestCase {
   }
 
   func testResizeInputTensorAtIndexToShape_ThrowsInvalidIndex() {
-    XCTAssertThrowsError(try interpreter.resizeInput(
-      at: AddModel.invalidIndex,
-      to: [2, 2, 3]
-    )) { error in
+    XCTAssertThrowsError(
+      try interpreter.resizeInput(
+        at: AddModel.invalidIndex,
+        to: [2, 2, 3]
+      )
+    ) { error in
       let maxIndex = AddModel.inputTensorCount - 1
       self.assertEqualErrors(
         actual: error,
@@ -162,10 +164,12 @@ class InterpreterTests: XCTestCase {
   }
 
   func testCopyDataToInputTensorAtIndex_ThrowsInvalidIndex() {
-    XCTAssertThrowsError(try interpreter.copy(
-      AddModel.inputData,
-      toInputAt: AddModel.invalidIndex
-    )) { error in
+    XCTAssertThrowsError(
+      try interpreter.copy(
+        AddModel.inputData,
+        toInputAt: AddModel.invalidIndex
+      )
+    ) { error in
       let maxIndex = AddModel.inputTensorCount - 1
       self.assertEqualErrors(
         actual: error,
@@ -178,10 +182,12 @@ class InterpreterTests: XCTestCase {
     try interpreter.resizeInput(at: AddModel.validIndex, to: AddModel.shape)
     try interpreter.allocateTensors()
     let invalidData = Data(count: AddModel.dataCount - 1)
-    XCTAssertThrowsError(try interpreter.copy(
-      invalidData,
-      toInputAt: AddModel.validIndex
-    )) { error in
+    XCTAssertThrowsError(
+      try interpreter.copy(
+        invalidData,
+        toInputAt: AddModel.validIndex
+      )
+    ) { error in
       self.assertEqualErrors(
         actual: error,
         expected: .invalidTensorDataCount(provided: invalidData.count, required: AddModel.dataCount)
@@ -223,12 +229,20 @@ class InterpreterOptionsTests: XCTestCase {
   func testInitWithDefaultValues() {
     let options = Interpreter.Options()
     XCTAssertNil(options.threadCount)
+    XCTAssertFalse(options.isXNNPackEnabled)
   }
 
   func testInitWithCustomValues() {
     var options = Interpreter.Options()
+
     options.threadCount = 2
     XCTAssertEqual(options.threadCount, 2)
+
+    options.isXNNPackEnabled = false
+    XCTAssertFalse(options.isXNNPackEnabled)
+
+    options.isXNNPackEnabled = true
+    XCTAssertTrue(options.isXNNPackEnabled)
   }
 
   func testEquatable() {
@@ -242,6 +256,15 @@ class InterpreterOptionsTests: XCTestCase {
 
     options2.threadCount = 3
     XCTAssertNotEqual(options1, options2)
+
+    options2.threadCount = 2
+    XCTAssertEqual(options1, options2)
+
+    options2.isXNNPackEnabled = true
+    XCTAssertNotEqual(options1, options2)
+
+    options1.isXNNPackEnabled = true
+    XCTAssertEqual(options1, options2)
   }
 }
 
@@ -326,14 +349,15 @@ extension Array {
   init?(unsafeData: Data) {
     guard unsafeData.count % MemoryLayout.stride == 0 else { return nil }
     #if swift(>=5.0)
-    self = unsafeData.withUnsafeBytes { .init($0.bindMemory(to: Element.self)) }
+      self = unsafeData.withUnsafeBytes { .init($0.bindMemory(to: Element.self)) }
     #else
-    self = unsafeData.withUnsafeBytes {
-      .init(UnsafeBufferPointer(
-        start: $0,
-        count: unsafeData.count / MemoryLayout.stride
-      ))
-    }
+      self = unsafeData.withUnsafeBytes {
+        .init(
+          UnsafeBufferPointer(
+            start: $0,
+            count: unsafeData.count / MemoryLayout.stride
+          ))
+      }
     #endif  // swift(>=5.0)
   }
 }

From 4b0a6f818fa8e3c38fd0cf68d9a647f82cf6c93a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 02:01:50 -0700
Subject: [PATCH 0604/1390] Update GraphDef version to 437.

PiperOrigin-RevId: 317270285
Change-Id: Ib8d1e6dbb565c01d2bdf0304a03be1c1eebbde41
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 546d86e58fa..9a79fc1eddf 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 436  // Updated: 2020/6/18
+#define TF_GRAPH_DEF_VERSION 437  // Updated: 2020/6/19
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From d41d28120e6aac1efbad27523a78cd254434dc4e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 02:02:01 -0700
Subject: [PATCH 0605/1390] compat: Update forward compatibility horizon to
 2020-06-19

PiperOrigin-RevId: 317270310
Change-Id: Idc8188172496af9f2494c580cdab27558b16e4a8
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 32545ac8463..22988d26cfc 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 18)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 19)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 9d1ec55aed0a4d9baf7302974fefe08546bfad25 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 02:04:36 -0700
Subject: [PATCH 0606/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/7f0d7f326316

PiperOrigin-RevId: 317270655
Change-Id: Ic80ab697da45212c8d58bcda989e5ee0a330b565
---
 .../mlir/xla/transforms/hlo_legalize_to_lhlo.cc     | 13 +++++++------
 .../compiler/mlir/xla/transforms/legalize_tf.cc     |  4 ++--
 .../mlir/xla/transforms/lhlo_legalize_to_gpu.cc     |  2 +-
 .../xla/transforms/lhlo_legalize_to_llvm_pass.cc    |  2 +-
 .../transforms/lhlo_legalize_to_parallel_loops.cc   |  2 +-
 .../mlir/xla/transforms/xla_legalize_to_linalg.cc   |  4 ++--
 .../xla/service/mlir_gpu/kernel_lowering.cc         |  2 +-
 tensorflow/workspace.bzl                            |  4 ++--
 8 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
index 1cfe0c12e20..a11b08e0ea6 100644
--- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
@@ -389,10 +389,13 @@ struct HloLegalizeToLhlo
     target.addLegalOp();
     target.addLegalOp();
     target.addIllegalDialect();
+
+    BufferAssignmentTypeConverter converter;
     target.addDynamicallyLegalOp([&](FuncOp op) {
       auto inputs = op.getType().getInputs();
-      return std::all_of(inputs.begin(), inputs.end(),
-                         [](Type input) { return input.isa(); });
+      return llvm::all_of(inputs,
+                          [](Type input) { return input.isa(); }) &&
+             converter.isLegal(&op.getBody());
     });
     target.addDynamicallyLegalOp([&](mlir::ReturnOp returnOp) {
       return std::all_of(returnOp.operand_type_begin(),
@@ -401,8 +404,7 @@ struct HloLegalizeToLhlo
     });
 
     auto module = getOperation();
-    BufferAssignmentTypeConverter converter;
-    module.walk([&](FuncOp func) {
+    module.walk([&](FuncOp func) -> WalkResult {
       BufferAssignmentPlacer bufferAssignment(func);
       OwningRewritePatternList patterns;
       populateHLOToLHLOConversionPattern(func.getContext(), &bufferAssignment,
@@ -418,8 +420,7 @@ struct HloLegalizeToLhlo
             /*allowMemrefFunctionResults=*/false>(&context, &bufferAssignment,
                                                   &converter, &patterns);
       }
-      return WalkResult(
-          applyPartialConversion(func, target, patterns, &converter));
+      return applyPartialConversion(func, target, patterns);
     });
   }
 
diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc b/tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc
index b7cad554043..1788cd1b270 100644
--- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc
@@ -5238,8 +5238,8 @@ LogicalResult legalizeTF(Operation *op, bool allow_partial_conversion,
     // Fully qualify ReturnOp here as xla_hlo dialect also defines a ReturnOp.
     target.addLegalOp();
     DenseSet nonlegalized_ops;
-    LogicalResult result = applyPartialConversion(
-        op, target, patterns, /*converter=*/nullptr, &nonlegalized_ops);
+    LogicalResult result =
+        applyPartialConversion(op, target, patterns, &nonlegalized_ops);
     // In order to enforce that the conversion result is fully converted,
     // fail if there are any nonlegalized ops in the set.
     if (failed(result) || !nonlegalized_ops.empty()) {
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_gpu.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_gpu.cc
index f0eb3cc1a0f..c23b8b49268 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_gpu.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_gpu.cc
@@ -177,7 +177,7 @@ struct LhloLegalizeToGpu : public PassWrapper {
     target.addIllegalOp();
     auto func = getFunction();
     patterns.insert(func.getContext());
-    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+    if (failed(applyPartialConversion(func, target, patterns))) {
       signalPassFailure();
     }
   }
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
index 9b809049290..63265c4a7e7 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
@@ -43,7 +43,7 @@ class TestLhloToLLVMPass
     target.addLegalOp();
     target.addIllegalDialect();
 
-    if (failed(applyFullConversion(m, target, patterns, &converter))) {
+    if (failed(applyFullConversion(m, target, patterns))) {
       signalPassFailure();
     }
   }
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc
index b3112d49103..65962c5b7a5 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_parallel_loops.cc
@@ -711,7 +711,7 @@ struct LhloLegalizeToParallelLoops
     target.addIllegalOp();
 
-    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+    if (failed(applyPartialConversion(func, target, patterns))) {
       signalPassFailure();
     }
   }
diff --git a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
index ad78a01100b..8a2f8ce7d04 100644
--- a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
@@ -867,7 +867,7 @@ struct LhloLegalizeToLinalg
 
     auto func = getFunction();
     populateLHLOToLinalgConversionPattern(func.getContext(), &patterns);
-    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+    if (failed(applyPartialConversion(func, target, patterns))) {
       signalPassFailure();
     }
   }
@@ -882,7 +882,7 @@ struct HloLegalizeToLinalg
 
     auto func = getFunction();
     xla_hlo::populateHLOToLinalgConversionPattern(func.getContext(), &patterns);
-    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
+    if (failed(applyPartialConversion(func, target, patterns))) {
       signalPassFailure();
     }
   }
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index 9d5b52df010..ecd1308be4b 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -552,7 +552,7 @@ class LowerToNVVMPass
     // TODO(csigg): Remove once we support replacing non-root ops.
     target.addLegalOp<::mlir::gpu::GPUModuleOp, ::mlir::gpu::ModuleEndOp,
                       ::mlir::gpu::YieldOp>();
-    if (failed(mlir::applyFullConversion(m, target, patterns, &converter))) {
+    if (failed(mlir::applyFullConversion(m, target, patterns))) {
       signalPassFailure();
     }
   }
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 52c573628ac..27eca0ee54f 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "92d8ad02e92fed3884169ba5d98056fe4fa5660d"
-    LLVM_SHA256 = "a4995ace7ddaef0c49293dc65771f58ef1fea96ebe1f39aa0a2d6d75d07f6cc7"
+    LLVM_COMMIT = "7f0d7f32631648acf48bc23047635ab5e2058a1a"
+    LLVM_SHA256 = "2f1dbae231b3b8f9c67d6a4f578c8ce29f3aa2831313b34c40ff2edb4014476a"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),

From c662daf4891a1e6efe64797615c3bd2bebedc5f5 Mon Sep 17 00:00:00 2001
From: Smit Hinsu 
Date: Fri, 19 Jun 2020 02:22:43 -0700
Subject: [PATCH 0607/1390] Override CustomCall in MlirHloBuilder

Also, enable mlir bridge for image ops compilers test. ResizeBilinear op
lowering usese CustomCall in case of TPU lowerings.

PiperOrigin-RevId: 317272443
Change-Id: I134c828cdc76552a0cbfdeb7c65532aa986314e2
---
 .../compiler/mlir/xla/ir/mlir_hlo_builder.cc  | 16 ++++++++++++
 .../compiler/mlir/xla/ir/mlir_hlo_builder.h   |  6 +++++
 .../xla/transforms/legalize_tf_with_tf2xla.cc |  8 ++++++
 tensorflow/compiler/tests/BUILD               |  1 +
 tensorflow/compiler/xla/client/xla_builder.cc | 26 ++++++++++++++-----
 tensorflow/compiler/xla/client/xla_builder.h  |  8 ++++++
 6 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc b/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc
index 21b1ac5f0ea..3c11d8e590d 100644
--- a/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc
+++ b/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc
@@ -132,6 +132,22 @@ StatusOr MlirHloBuilder::FftInternal(
   return MakeXlaOp(op);
 }
 
+StatusOr MlirHloBuilder::CustomCallInternal(
+    const string& call_target_name, absl::Span operands,
+    const Shape& shape, const string& opaque,
+    absl::optional> operand_shapes_with_layout) {
+  if (operand_shapes_with_layout.has_value())
+    return Unimplemented(
+        "CustomCall doesn't support operands shapes with layout");
+  TF_ASSIGN_OR_RETURN(mlir::Type ty, ConvertShapeToType(
+                                         shape, builder_));
+  auto op = builder_.create(
+      loc_, ty, GetValues(operands), builder_.getStringAttr(call_target_name),
+      /*has_side_effect=*/builder_.getBoolAttr(false),
+      builder_.getStringAttr(opaque));
+  return MakeXlaOp(op);
+}
+
 StatusOr MlirHloBuilder::ReduceInternal(
     const Shape& shape, absl::Span all_operands,
     const XlaComputation& computation,
diff --git a/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.h b/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.h
index 4b28c32db99..4d7d93af7a7 100644
--- a/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.h
+++ b/tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.h
@@ -124,6 +124,12 @@ class MlirHloBuilder : public XlaBuilder {
                               FftType fft_type,
                               absl::Span fft_length) override;
 
+  StatusOr CustomCallInternal(const string& call_target_name,
+                                     absl::Span operands,
+                                     const Shape& shape, const string& opaque,
+                                     absl::optional>
+                                         operand_shapes_with_layout) override;
+
   StatusOr ReduceInternal(
       const Shape& shape, absl::Span all_operands,
       const XlaComputation& computation,
diff --git a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
index ef79c8868bb..8f96f4d1305 100644
--- a/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc
@@ -88,6 +88,9 @@ static bool IsOpWhitelisted(Operation* op) {
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
+    TypeID::get(),
+    TypeID::get(),
+    TypeID::get(),
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
@@ -127,6 +130,7 @@ static bool IsOpWhitelisted(Operation* op) {
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
+    TypeID::get(),
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
@@ -157,10 +161,14 @@ static bool IsOpWhitelisted(Operation* op) {
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
+    TypeID::get(),
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
+    TypeID::get(),
+    TypeID::get(),
+    TypeID::get(),
     TypeID::get(),
     TypeID::get(),
     TypeID::get(),
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index b574622efce..034ec82de10 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -770,6 +770,7 @@ tf_xla_py_test(
     size = "small",
     timeout = "long",
     srcs = ["image_ops_test.py"],
+    enable_mlir_bridge = True,
     python_version = "PY3",
     shard_count = 10,
     tags = [
diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index c7b6a7f9491..03ae23ea18b 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -1564,16 +1564,12 @@ XlaOp XlaBuilder::CustomCall(
     const Shape& shape, const string& opaque,
     absl::optional> operand_shapes_with_layout) {
   return ReportErrorOrReturn([&]() -> StatusOr {
-    HloInstructionProto instr;
     if (absl::StartsWith(call_target_name, "$")) {
       return InvalidArgument(
           "Invalid custom_call_target \"%s\": Call targets that start with '$' "
           "are reserved for internal use.",
           call_target_name);
     }
-    *instr.mutable_shape() = shape.ToProto();
-    instr.set_custom_call_target(call_target_name);
-    instr.set_backend_config(opaque);
     if (operand_shapes_with_layout.has_value()) {
       if (!LayoutUtil::HasLayout(shape)) {
         return InvalidArgument(
@@ -1586,7 +1582,6 @@ XlaOp XlaBuilder::CustomCall(
             "with constrained layout; given %d shapes, expected %d",
             operand_shapes_with_layout->size(), operands.size());
       }
-      instr.set_constrain_layout(true);
       int64 operand_num = 0;
       for (const Shape& operand_shape : *operand_shapes_with_layout) {
         if (!LayoutUtil::HasLayout(operand_shape)) {
@@ -1595,14 +1590,31 @@ XlaOp XlaBuilder::CustomCall(
               "constrained layout.",
               operand_num);
         }
-        *instr.add_operand_shapes_with_layout() = operand_shape.ToProto();
         ++operand_num;
       }
     }
-    return AddInstruction(std::move(instr), HloOpcode::kCustomCall, operands);
+    return CustomCallInternal(call_target_name, operands, shape, opaque,
+                              operand_shapes_with_layout);
   });
 }
 
+StatusOr XlaBuilder::CustomCallInternal(
+    const string& call_target_name, absl::Span operands,
+    const Shape& shape, const string& opaque,
+    absl::optional> operand_shapes_with_layout) {
+  HloInstructionProto instr;
+  *instr.mutable_shape() = shape.ToProto();
+  instr.set_custom_call_target(call_target_name);
+  instr.set_backend_config(opaque);
+  if (operand_shapes_with_layout.has_value()) {
+    instr.set_constrain_layout(true);
+    for (const Shape& operand_shape : *operand_shapes_with_layout) {
+      *instr.add_operand_shapes_with_layout() = operand_shape.ToProto();
+    }
+  }
+  return AddInstruction(std::move(instr), HloOpcode::kCustomCall, operands);
+}
+
 XlaOp XlaBuilder::CustomCall(
     const string& call_target_name, absl::Span operands,
     const XlaComputation& computation, const Shape& shape, const string& opaque,
diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h
index b8af180b83e..3fc26747468 100644
--- a/tensorflow/compiler/xla/client/xla_builder.h
+++ b/tensorflow/compiler/xla/client/xla_builder.h
@@ -527,6 +527,14 @@ class XlaBuilder {
       const Shape& shape_with_layout, const string& opaque,
       absl::optional> operand_shapes_with_layout);
 
+  // Internal version of CustomCall without computation that doesn't do op
+  // specific error handling and expects arguments to be legal. CustomCall
+  // method above calls this method after error handling.
+  virtual StatusOr CustomCallInternal(
+      const string& call_target_name, absl::Span operands,
+      const Shape& shape_with_layout, const string& opaque,
+      absl::optional> operand_shapes_with_layout);
+
   XlaOp CustomCall(
       const string& call_target_name, absl::Span operands,
       const XlaComputation& computation, const Shape& shape_with_layout,

From 8ad5bc80e71921c3c2530d93d3856ba59e524c60 Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Fri, 19 Jun 2020 11:50:24 +0200
Subject: [PATCH 0608/1390] Remove unnecessary assert

---
 .../mixed_precision/experimental/autocast_variable_test.py      | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 2fa7c103258..9036109af96 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -372,7 +372,6 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
         # Variable should be increased, despite it appearing to be the same
         # float16 value.
         self.evaluate(x.assign(1. + small_tensor))
-        self.assertEqual(1. + small_val, self.evaluate(x._variable))
         self.assertEqual(1., self.evaluate(x.value()))
       self.assertEqual(1. + small_val, self.evaluate(x))
 
@@ -380,7 +379,6 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
       with ops.get_default_graph()._enable_auto_casting_variables(
           dtypes.float16):
         self.evaluate(x.assign_add(small_tensor))
-        self.assertEqual(1. + small_val, self.evaluate(x._variable))
         self.assertEqual(1., self.evaluate(x.value()))
       self.assertEqual(1. + small_val, self.evaluate(x))
 

From dc8d42922b9ff89e717f130515c968186ec4504c Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Fri, 19 Jun 2020 12:33:01 +0200
Subject: [PATCH 0609/1390] Remove unnecessary control_dependencies

---
 .../mixed_precision/experimental/autocast_variable.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
index ca6420f0c0b..b60100c7b48 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
@@ -190,25 +190,20 @@ class AutoCastVariable(variables.Variable, core.Tensor):
 
   def _apply_assign_update(
       self, update_fn, value, use_locking=None, name=None, read_value=True):
-    if not read_value:
-      return update_fn(value, use_locking, name, read_value)
-
     if context.executing_eagerly() or ops.inside_function():
       assign_op = update_fn(value, use_locking, name, False)
-      with ops.control_dependencies([assign_op]):
-        return self
+      return self if read_value else assign_op
 
     # Fallback to wrapping the returned variable in graph mode if possible
     assign_var = update_fn(value, use_locking, name, read_value)
-    if resource_variable_ops.is_resource_variable(assign_var):
+    if read_value and resource_variable_ops.is_resource_variable(assign_var):
       return create_autocast_variable(assign_var)
     return assign_var
 
   def _apply_update(self, update_fn, *args, **kwargs):
     update_var = update_fn(*args, **kwargs)
     if context.executing_eagerly() or ops.inside_function():
-      with ops.control_dependencies([update_var]):
-        return self
+      return self
 
     # Fallback to wrapping the returned variable in graph mode if possible
     if resource_variable_ops.is_resource_variable(update_var):

From b58e6000457f26c7a53a9b945642fbe2baddbf20 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev 
Date: Fri, 19 Jun 2020 04:25:24 -0700
Subject: [PATCH 0610/1390] [XLA][MLIR] Enable xla_hlo.ReshapeOp ->
 xla_lhlo.ReshapeOp conversion.

PiperOrigin-RevId: 317284676
Change-Id: Ia845183efcfabe77f6eb66d8c56dcbfc82653982
---
 tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc | 1 +
 tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h    | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
index a11b08e0ea6..446f2aae833 100644
--- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
@@ -464,6 +464,7 @@ void populateHLOToLHLOConversionPattern(
       HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
+      HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
diff --git a/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h b/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
index 4b9397795a1..8d5f27474a5 100644
--- a/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
+++ b/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
@@ -61,6 +61,7 @@ MAP_HLO_TO_LHLO(MulOp);
 MAP_HLO_TO_LHLO(NegOp);
 MAP_HLO_TO_LHLO(RealOp);
 MAP_HLO_TO_LHLO(ReduceOp);
+MAP_HLO_TO_LHLO(ReshapeOp);
 MAP_HLO_TO_LHLO(RemOp);
 MAP_HLO_TO_LHLO(RsqrtOp);
 MAP_HLO_TO_LHLO(SelectOp);

From 7b1a726ec3ab2d831167d32b9b820bcf14aba6f7 Mon Sep 17 00:00:00 2001
From: Devi Sandeep Endluri 
Date: Fri, 19 Jun 2020 06:23:39 -0500
Subject: [PATCH 0611/1390] Ensure there are test samples for imdb dataset,
 when maxlen is low

With the current imdb.load_data(), the following results are seen
for different values of maxlen.

	load_data                (len(x_train), len(x_test))
------------------------------------------------------------
imdb.load_data(maxlen=50)    -->    (1035, 0)
imdb.load_data(maxlen=100)   -->    (5736, 0)
imdb.load_data(maxlen=200)   -->    (25000, 3913)
imdb.load_data()             -->    (25000, 25000)

Analysis: We can observe that when maxlen is low, the number
of test samples can be 0. This is because the train and test data is
concatenated, then the samples with length > maxlen are removed, and
the first 25,000 are considered as training data.

Fix: This can be fixed when data can be filtered first to remove the
ones with length > maxlen, and then concatenate to process further.
The following are the results after the fix.

     fixed load_data              (len(x_train), len(x_test))
------------------------------------------------------------
imdb.load_data(maxlen=50)    -->    (477, 558)
imdb.load_data(maxlen=100)   -->    (2773, 2963)
imdb.load_data(maxlen=200)   -->    (14244, 14669)
imdb.load_data()             -->    (25000, 25000)
---
 tensorflow/python/keras/datasets/imdb.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tensorflow/python/keras/datasets/imdb.py b/tensorflow/python/keras/datasets/imdb.py
index 37403228edf..e359d691a5d 100644
--- a/tensorflow/python/keras/datasets/imdb.py
+++ b/tensorflow/python/keras/datasets/imdb.py
@@ -124,20 +124,24 @@ def load_data(path='imdb.npz',
   x_test = x_test[indices]
   labels_test = labels_test[indices]
 
-  xs = np.concatenate([x_train, x_test])
-  labels = np.concatenate([labels_train, labels_test])
-
   if start_char is not None:
-    xs = [[start_char] + [w + index_from for w in x] for x in xs]
+    x_train = [[start_char] + [w + index_from for w in x] for x in x_train]
+    x_test = [[start_char] + [w + index_from for w in x] for x in x_test]
   elif index_from:
-    xs = [[w + index_from for w in x] for x in xs]
+    x_train = [[w + index_from for w in x] for x in x_train]
+    x_test = [[w + index_from for w in x] for x in x_test]
 
   if maxlen:
-    xs, labels = _remove_long_seq(maxlen, xs, labels)
-    if not xs:
+    x_train, labels_train = _remove_long_seq(maxlen, x_train, labels_train)
+    x_test, labels_test = _remove_long_seq(maxlen, x_test, labels_test)
+    if not x_train or not x_test:
       raise ValueError('After filtering for sequences shorter than maxlen=' +
                        str(maxlen) + ', no sequence was kept. '
                        'Increase maxlen.')
+
+  xs = np.concatenate([x_train, x_test])
+  labels = np.concatenate([labels_train, labels_test])
+
   if not num_words:
     num_words = max(max(x) for x in xs)
 

From 6ddd920c44540284ab28cc7cdf2faef18b85ff11 Mon Sep 17 00:00:00 2001
From: Xinan Jiang 
Date: Fri, 19 Jun 2020 20:51:30 +0800
Subject: [PATCH 0612/1390] [MLIR][XLA] Modify for FileCheck commands

---
 tensorflow/compiler/xla/service/mlir_gpu/tests/gather.hlo | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/tests/gather.hlo b/tensorflow/compiler/xla/service/mlir_gpu/tests/gather.hlo
index 99faa319bf6..8a24daae0f4 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/tests/gather.hlo
+++ b/tensorflow/compiler/xla/service/mlir_gpu/tests/gather.hlo
@@ -11,5 +11,5 @@ ENTRY %Gather (x: f32[100,10], y: s64[4,6]) -> f32[4,6,10] {
 
 // CHECK: func @gather(%[[ARG0:.*]]: [[TYPE0:.*]], %[[ARG1:.*]]: [[TYPE1:.*]], %[[RESULT:.*]]: [[RTYPE:.*]]) {
 // CHECK:   "xla_lhlo.gather"(%[[ARG0]], %[[ARG1]], %[[RESULT]])
-// CHECK:   {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, slice_sizes = dense<[1, 10]> : tensor<2xi64>, start_index_map = dense<0> : tensor<1xi64>} : ([[TYPE0]], [[TYPE1]], [[RTYPE]]) -> ()
+// CHECK-SAME:   {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, slice_sizes = dense<[1, 10]> : tensor<2xi64>, start_index_map = dense<0> : tensor<1xi64>} : ([[TYPE0]], [[TYPE1]], [[RTYPE]]) -> ()
 // CHECK: }

From 42579858f9cda701c7c69d4a1f89035f0a68b258 Mon Sep 17 00:00:00 2001
From: Xinan Jiang 
Date: Fri, 19 Jun 2020 20:56:08 +0800
Subject: [PATCH 0613/1390] [MLIR][XLA] Add GatherOp to HLO to LHLO converters

---
 tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc | 1 +
 tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h    | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
index 1cfe0c12e20..524965cfacd 100644
--- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc
@@ -453,6 +453,7 @@ void populateHLOToLHLOConversionPattern(
       HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
+      HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
       HloToLhloOpConverter,
diff --git a/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h b/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
index 4b9397795a1..5e3d1cb9302 100644
--- a/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
+++ b/tensorflow/compiler/mlir/xla/transforms/map_hlo_to_lhlo_op.h
@@ -52,6 +52,7 @@ MAP_HLO_TO_LHLO(CosOp);
 MAP_HLO_TO_LHLO(DivOp);
 MAP_HLO_TO_LHLO(DotOp);
 MAP_HLO_TO_LHLO(ExpOp);
+MAP_HLO_TO_LHLO(GatherOp);
 MAP_HLO_TO_LHLO(ImagOp);
 MAP_HLO_TO_LHLO(IotaOp);
 MAP_HLO_TO_LHLO(LogOp);

From 0c7e61d6608e29324357fb5f79a9b925281521a0 Mon Sep 17 00:00:00 2001
From: Hanhan Wang 
Date: Fri, 19 Jun 2020 07:09:13 -0700
Subject: [PATCH 0614/1390] Remove the canonicalize pattern for folding a pad
 op into the following conv op.

Basically rolledback for cl/305641881, the pattern could hurt performance
because the operation can't be fully tiled in Linalg transformation. In this
context, not everyone wants this pattern, so remove it from canonicalize
patterns.

PiperOrigin-RevId: 317302072
Change-Id: I19aa64e14eecccfd738ad3f775f3670974bc68f9
---
 tensorflow/compiler/mlir/xla/ir/hlo_ops.cc    | 56 ----------------
 tensorflow/compiler/mlir/xla/ir/hlo_ops.td    |  2 -
 .../compiler/mlir/xla/tests/canonicalize.mlir | 65 -------------------
 .../mlir/xla/transforms/canonicalize.td       | 51 ---------------
 4 files changed, 174 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
index d7950919883..e0fa1da93b8 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.cc
@@ -106,53 +106,6 @@ DenseIntElementsAttr BuildSliceLimits(DenseIntElementsAttr start_indices,
   return GetI64ElementsAttr(slice_limits, builder);
 }
 
-// Returns the padding value of the given position. If padding_attr is a
-// nullptr, returns 0.
-static int64_t GetPaddingValue(DenseIntElementsAttr padding_attr,
-                               ArrayRef index) {
-  if (!padding_attr) return 0;
-  return padding_attr.getValue(index);
-}
-
-static bool IsOnlyPaddingSpatialDims(Value lhs,
-                                     ConvDimensionNumbers dimension_numbers,
-                                     DenseIntElementsAttr edge_padding_low,
-                                     DenseIntElementsAttr edge_padding_high) {
-  const int64_t batch_dim = dimension_numbers.input_batch_dimension().getInt();
-  const int64_t feature_dim =
-      dimension_numbers.input_feature_dimension().getInt();
-  if (edge_padding_low.getValue(batch_dim) ||
-      edge_padding_high.getValue(batch_dim))
-    return false;
-  if (edge_padding_low.getValue(feature_dim) ||
-      edge_padding_high.getValue(feature_dim))
-    return false;
-  return true;
-}
-
-DenseIntElementsAttr BuildConvPaddingAttrs(
-    DenseIntElementsAttr edge_padding_low,
-    DenseIntElementsAttr edge_padding_high, DenseIntElementsAttr padding_attr,
-    ConvDimensionNumbers dimension_numbers, Builder* builder) {
-  SmallVector padding_low, padding_high;
-  for (const auto& dim : dimension_numbers.input_spatial_dimensions()) {
-    unsigned i = dim.getZExtValue();
-    padding_low.push_back(edge_padding_low.getValue(i));
-    padding_high.push_back(edge_padding_high.getValue(i));
-  }
-
-  int rank = padding_low.size();
-  SmallVector padding;
-  for (unsigned i = 0, e = rank; i < e; ++i) {
-    padding.push_back(GetPaddingValue(padding_attr, {i, 0}) + padding_low[i]);
-    padding.push_back(GetPaddingValue(padding_attr, {i, 1}) + padding_high[i]);
-  }
-  // padding_attr.getType() doesn't work because it is an optional attribute,
-  // which can be a nullptr.
-  auto type = RankedTensorType::get({rank, 2}, builder->getIntegerType(64));
-  return DenseIntElementsAttr::get(type, padding);
-}
-
 #include "tensorflow/compiler/mlir/xla/transforms/generated_canonicalize.inc"
 }  // namespace
 
@@ -2153,14 +2106,5 @@ LogicalResult deriveShapeFromFirstOperand(
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// ConvOp
-//===----------------------------------------------------------------------===//
-
-void ConvOp::getCanonicalizationPatterns(OwningRewritePatternList& results,
-                                         MLIRContext* context) {
-  results.insert(context);
-}
-
 }  // namespace xla_hlo
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
index b1745c73fbf..f92d1c5b85c 100644
--- a/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
+++ b/tensorflow/compiler/mlir/xla/ir/hlo_ops.td
@@ -929,8 +929,6 @@ def HLO_ConvOp : HLO_Op<"convolution", [NoSideEffect]>, BASE_HLO_ConvOp {
   );
 
   let results = (outs HLO_Tensor);
-
-  let hasCanonicalizer = 1;
 }
 
 def HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CopyOp {
diff --git a/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir b/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir
index ef0f8c4d200..1954c3344df 100644
--- a/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir
+++ b/tensorflow/compiler/mlir/xla/tests/canonicalize.mlir
@@ -415,71 +415,6 @@ func @fold_copy(%arg : tensor<1x4xf32>) -> tensor<1x4xf32> {
   return %0 : tensor<1x4xf32>
 }
 
-// CHECK-LABEL: func @fold_pad_into_conv_f32
-func @fold_pad_into_conv_f32(%arg0 : tensor<1x32x32x3xf32>,
-                         %arg1 : tensor<7x7x3x64xf32>)
-    -> tensor<1x16x16x64xf32> {
-  //  CHECK-NOT: xla_hlo.pad
-  //      CHECK: xla_hlo.convolution
-  // CHECK-SAME: padding = dense<3> : tensor<2x2xi64>
-  %0 = xla_hlo.constant dense<0.000000e+00> : tensor
-  %1 = "xla_hlo.pad"(%arg0, %0) {
-    edge_padding_high = dense<[0, 3, 3, 0]> : tensor<4xi64>,
-    edge_padding_low = dense<[0, 3, 3, 0]> : tensor<4xi64>,
-    interior_padding = dense<0> : tensor<4xi64>
-  } : (tensor<1x32x32x3xf32>, tensor) -> tensor<1x38x38x3xf32>
-  %2 = "xla_hlo.convolution"(%1, %arg1) {
-    batch_group_count = 1 : i64,
-    dimension_numbers = {
-      input_batch_dimension = 0 : i64,
-      input_feature_dimension = 3 : i64,
-      input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-      kernel_input_feature_dimension = 2 : i64,
-      kernel_output_feature_dimension = 3 : i64,
-      kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-      output_batch_dimension = 0 : i64,
-      output_feature_dimension = 3 : i64,
-      output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-    },
-    feature_group_count = 1 : i64,
-    padding = dense<0> : tensor<2x2xi64>,
-    window_strides = dense<2> : tensor<2xi64>
-  } : (tensor<1x38x38x3xf32>, tensor<7x7x3x64xf32>) -> tensor<1x16x16x64xf32>
-  return %2 : tensor<1x16x16x64xf32>
-}
-
-// CHECK-LABEL: func @fold_pad_into_conv_i32
-func @fold_pad_into_conv_i32(%arg0 : tensor<1x32x32x3xi32>,
-                         %arg1 : tensor<7x7x3x64xi32>)
-    -> tensor<1x16x16x64xi32> {
-  //  CHECK-NOT: xla_hlo.pad
-  //      CHECK: xla_hlo.convolution
-  // CHECK-SAME: padding = dense<3> : tensor<2x2xi64>
-  %0 = xla_hlo.constant dense<0> : tensor
-  %1 = "xla_hlo.pad"(%arg0, %0) {
-    edge_padding_high = dense<[0, 3, 3, 0]> : tensor<4xi64>,
-    edge_padding_low = dense<[0, 3, 3, 0]> : tensor<4xi64>,
-    interior_padding = dense<0> : tensor<4xi64>
-  } : (tensor<1x32x32x3xi32>, tensor) -> tensor<1x38x38x3xi32>
-  %2 = "xla_hlo.convolution"(%1, %arg1) {
-    batch_group_count = 1 : i64,
-    dimension_numbers = {
-      input_batch_dimension = 0 : i64,
-      input_feature_dimension = 3 : i64,
-      input_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>,
-      kernel_input_feature_dimension = 2 : i64,
-      kernel_output_feature_dimension = 3 : i64,
-      kernel_spatial_dimensions = dense<[0, 1]> : tensor<2xi64>,
-      output_batch_dimension = 0 : i64,
-      output_feature_dimension = 3 : i64,
-      output_spatial_dimensions = dense<[1, 2]> : tensor<2xi64>
-    },
-    feature_group_count = 1 : i64,
-    window_strides = dense<2> : tensor<2xi64>
-  } : (tensor<1x38x38x3xi32>, tensor<7x7x3x64xi32>) -> tensor<1x16x16x64xi32>
-  return %2 : tensor<1x16x16x64xi32>
-}
-
 // CHECK-LABEL: func @dynamic_reshape_not_actually_dynamic
 func @dynamic_reshape_not_actually_dynamic(%arg0: tensor<4xf32>, %shape: tensor<2xindex>) -> tensor<4x1xf32> {
   // CHECK: xla_hlo.reshape
diff --git a/tensorflow/compiler/mlir/xla/transforms/canonicalize.td b/tensorflow/compiler/mlir/xla/transforms/canonicalize.td
index b788cb80380..c319551d92a 100644
--- a/tensorflow/compiler/mlir/xla/transforms/canonicalize.td
+++ b/tensorflow/compiler/mlir/xla/transforms/canonicalize.td
@@ -28,54 +28,3 @@ def UnaryEinsumToEinsum : Pat<
   (HLO_UnaryEinsumOp $operand, $equation),
   (HLO_EinsumOp (HLO_ConstOp (GetScalarOfType<1> $operand)),
                 $operand, (UnaryToBinaryEinsumEq $equation))>;
-
-//===----------------------------------------------------------------------===//
-// Conv op patterns.
-//===----------------------------------------------------------------------===//
-
-def IsZero : Attr() &&"
-  "$_self.cast().isSplat() &&"
-  "$_self.cast().getSplatValue()"
-  ".getValue().isZero()) ||"
-  "($_self.isa() &&"
-  "$_self.cast().isSplat() &&"
-  "$_self.cast().getSplatValue()"
-  ".getInt() == 0)">>;
-
-def IsOnlyPaddingSpatialDims
-  : Constraint>;
-
-def BuildConvPaddingAttrs : NativeCodeCall<
-  "BuildConvPaddingAttrs($0, $1, $2, $3, &$_builder)">;
-
-def FoldPadIntoConv : Pat<
-  (HLO_ConvOp
-    (HLO_PadOp $lhs,
-      (HLO_ConstOp IsZero:$padding_value),
-      $edge_padding_low,
-      $edge_padding_high,
-      IsZero:$interior_padding),
-    $rhs,
-    $window_strides,
-    $padding,
-    $lhs_dilation,
-    $rhs_dilation,
-    $dimension_numbers,
-    $feature_group_count,
-    $batch_group_count,
-    $precision_config),
-  (HLO_ConvOp
-    $lhs,
-    $rhs,
-    $window_strides,
-    (BuildConvPaddingAttrs $edge_padding_low, $edge_padding_high, $padding,
-      $dimension_numbers),
-    $lhs_dilation,
-    $rhs_dilation,
-    $dimension_numbers,
-    $feature_group_count,
-    $batch_group_count,
-    $precision_config),
-    [(IsOnlyPaddingSpatialDims $lhs, $dimension_numbers, $edge_padding_low,
-      $edge_padding_high)]>;

From 9e7d5ef6f25e436fffae03597838294d872404f0 Mon Sep 17 00:00:00 2001
From: "T.J. Alumbaugh" 
Date: Fri, 19 Jun 2020 07:29:02 -0700
Subject: [PATCH 0615/1390] Full int8 quantization BatchMatMul

PiperOrigin-RevId: 317304259
Change-Id: Icf96d9d129db30b965e36f5c8befd27762b173b2
---
 tensorflow/compiler/mlir/lite/ir/tfl_ops.td   |   6 +-
 tensorflow/lite/kernels/batch_matmul.cc       |  92 ++++++++++-
 tensorflow/lite/kernels/batch_matmul_test.cc  | 156 ++++++++++++++++--
 .../kernels/internal/optimized/batch_matmul.h | 106 ++++++++++++
 .../kernels/internal/reference/batch_matmul.h |  93 +++++++++++
 tensorflow/lite/kernels/register.cc           |   4 +-
 .../lite/tools/optimize/operator_property.cc  |   6 +
 .../lite/tools/versioning/op_version.cc       |   2 +
 .../lite/tools/versioning/runtime_version.cc  |   1 +
 9 files changed, 439 insertions(+), 27 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
index 509c13ae161..33281cc58fb 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
@@ -953,14 +953,14 @@ in the batch dimensions and broadcasting.
   }];
 
   let arguments = (ins
-    TFL_TensorOf<[F32]>:$x,
-    TFL_TensorOf<[F32]>:$y,
+    TFL_TensorOf<[F32, QI8]>:$x,
+    TFL_TensorOf<[F32, QI8]>:$y,
     DefaultValuedAttr:$adj_x,
     DefaultValuedAttr:$adj_y
   );
 
    let results = (outs
-    TFL_TensorOf<[F32]>:$output
+    TFL_TensorOf<[F32, QI8]>:$output
   );
 
   let hasOptions = 1;
diff --git a/tensorflow/lite/kernels/batch_matmul.cc b/tensorflow/lite/kernels/batch_matmul.cc
index 8bc23c9c94a..a414a226504 100644
--- a/tensorflow/lite/kernels/batch_matmul.cc
+++ b/tensorflow/lite/kernels/batch_matmul.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include 
 #include 
+#include 
 
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
@@ -52,6 +53,14 @@ enum KernelType {
 };
 
 struct OpData {
+  // The scaling factor from input to output (aka the 'real multiplier') can
+  // be represented as a fixed point multiplier plus a left shift.
+  int32_t output_multiplier;
+  int output_shift;
+  // The range of the fused activation layer. For example for kNone and
+  // uint8_t these would be 0 and 255.
+  int32_t output_activation_min;
+  int32_t output_activation_max;
   // The index of the temporary tensors where we store transposed LHS/RHS.
   int scratch_tensor_index;
   bool rhs_transposed;
@@ -274,6 +283,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   OpContext op_context(context, node);
   TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context));
+  OpData* op_data = reinterpret_cast(node->user_data);
 
   bool adj_x = op_context.params->adj_x;
   bool adj_y = op_context.params->adj_y;
@@ -282,7 +292,24 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* rhs_data = GetInput(context, node, kInputRHSTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TF_LITE_ENSURE_TYPES_EQ(context, lhs_data->type, kTfLiteFloat32);
+  // Note that quantized inference requires that all tensors have their
+  // parameters set. This is usually done during quantized training.
+  if (lhs_data->type == kTfLiteInt8) {
+    double real_multiplier = 0.0;
+    TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler(
+        context, lhs_data, rhs_data, output, &real_multiplier));
+    int exponent;
+    QuantizeMultiplier(real_multiplier, &op_data->output_multiplier, &exponent);
+    op_data->output_shift = exponent;
+    // BatchMatMul has no fused activation functions. Therefore, set
+    // output activation min and max to min and max of int8_t type,
+    // respecitvely.
+    op_data->output_activation_min = std::numeric_limits::min();
+    op_data->output_activation_max = std::numeric_limits::max();
+  }
+
+  TF_LITE_ENSURE(context, lhs_data->type == kTfLiteFloat32 ||
+                              lhs_data->type == kTfLiteInt8);
   TF_LITE_ENSURE(context, rhs_data->type == kTfLiteFloat32 ||
                               rhs_data->type == kTfLiteInt8);
   // Support dimensions between 2 and 4, inclusive.
@@ -433,6 +460,41 @@ TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, OpData* data,
   return kTfLiteOk;
 }
 
+template 
+TfLiteStatus EvalInt8(TfLiteContext* context, const OpData* data,
+                      const RuntimeShape& lhs_shape, const TfLiteTensor* lhs,
+                      const RuntimeShape& rhs_shape, const TfLiteTensor* rhs,
+                      const RuntimeShape& output_shape, TfLiteTensor* output) {
+  // Reuse params struct from FullyConnected Op.
+  FullyConnectedParams op_params;
+  int32_t input_offset = -lhs->params.zero_point;
+  int32_t filter_offset = -rhs->params.zero_point;
+  int32_t output_offset = output->params.zero_point;
+  op_params.input_offset = input_offset;
+  op_params.weights_offset = filter_offset;
+  op_params.output_offset = output_offset;
+  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_shift = data->output_shift;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+  op_params.lhs_cacheable = IsConstantTensor(lhs);
+  op_params.rhs_cacheable = IsConstantTensor(rhs);
+
+  if (kernel_type == kReference) {
+    reference_ops::BatchMatMul(op_params, rhs_shape, GetTensorData(rhs),
+                               lhs_shape, GetTensorData(lhs),
+                               GetTensorShape(output),
+                               GetTensorData(output));
+  } else {
+    optimized_ops::BatchMatMul(op_params, rhs_shape, GetTensorData(rhs),
+                               lhs_shape, GetTensorData(lhs),
+                               GetTensorShape(output),
+                               GetTensorData(output),
+                               CpuBackendContext::GetFromContext(context));
+  }
+  return kTfLiteOk;
+}
+
 template 
 TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
                            OpData* data, const RuntimeShape& lhs_shape,
@@ -448,25 +510,39 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
     return EvalHybrid(
         context, node, data, lhs_shape, lhs, rhs_shape, rhs, input_quantized,
         scaling_factors, accum_scratch, row_sums, input_offsets, output);
+  } else if (lhs->type == kTfLiteInt8) {
+    return EvalInt8(context, data, lhs_shape, lhs, rhs_shape, rhs,
+                                 GetTensorShape(output), output);
   } else {
-    TF_LITE_KERNEL_LOG(context,
-                       "Currently only hybrid quantization is supported.\n");
+    TF_LITE_KERNEL_LOG(
+        context, "Currently only hybrid and int8 quantization is supported.\n");
     return kTfLiteError;
   }
   return kTfLiteOk;
 }
 
-TfLiteTensor* GetRhs(TfLiteContext* context, TfLiteNode* node,
-                     const TfLiteTensor* rhs) {
+TfLiteTensor* GetTempRhs(TfLiteContext* context, TfLiteNode* node,
+                         const TfLiteTensor* rhs) {
   TfLiteTensor* transposed_rhs = GetTemporary(context, node, 1);
   if (rhs->type == kTfLiteInt8) {
-    // Get the quantization params from the weights tensors.
+    // Get the quantization params from the RHS tensor.
     transposed_rhs->params.scale = rhs->params.scale;
     transposed_rhs->params.zero_point = rhs->params.zero_point;
   }
   return transposed_rhs;
 }
 
+TfLiteTensor* GetTempLhs(TfLiteContext* context, TfLiteNode* node,
+                         const TfLiteTensor* lhs) {
+  TfLiteTensor* transposed_lhs = GetTemporary(context, node, 0);
+  if (lhs->type == kTfLiteInt8) {
+    // Get the quantization params from the LHS tensor.
+    transposed_lhs->params.scale = lhs->params.scale;
+    transposed_lhs->params.zero_point = lhs->params.zero_point;
+  }
+  return transposed_lhs;
+}
+
 // Perform a batch matrix multiply on
 // LHS <..., A, B>  X  RHS<..., B, C>
 // where the leading dimensions of LHS and RHS obey broadcasting rules
@@ -491,8 +567,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   bool adj_y = op_context.params->adj_y;
   bool adj_x = op_context.params->adj_x;
 
-  const TfLiteTensor* rhs_tensor = adj_y ? rhs : GetRhs(context, node, rhs);
-  const TfLiteTensor* lhs_tensor = adj_x ? GetTemporary(context, node, 0) : lhs;
+  const TfLiteTensor* rhs_tensor = adj_y ? rhs : GetTempRhs(context, node, rhs);
+  const TfLiteTensor* lhs_tensor = adj_x ? GetTempLhs(context, node, lhs) : lhs;
   if (!adj_y) {
     // TODO(b/154760341) Constant tensors should already be transposed, but
     // we transpose once if necessary for now.
diff --git a/tensorflow/lite/kernels/batch_matmul_test.cc b/tensorflow/lite/kernels/batch_matmul_test.cc
index 5e52479f49b..98df8ebe3db 100644
--- a/tensorflow/lite/kernels/batch_matmul_test.cc
+++ b/tensorflow/lite/kernels/batch_matmul_test.cc
@@ -24,8 +24,19 @@ limitations under the License.
 #include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
+
+namespace ops {
+namespace builtin {
+
+TfLiteRegistration* Register_BATCH_MATMUL_REF();
+TfLiteRegistration* Register_BATCH_MATMUL_GENERIC_OPTIMIZED();
+
+}  // namespace builtin
+}  // namespace ops
+
 namespace {
 
+using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
 
 template 
@@ -53,7 +64,20 @@ class BatchMatMulOpModel : public SingleOpModel {
   int output_id_;
 };
 
-TEST(BatchMatMulOpModelTest, Float32Test_Simple) {
+const auto kKernelMap = new std::map({
+    {"Reference", ops::builtin::Register_BATCH_MATMUL_REF()},
+    {"GenericOptimized",
+     ops::builtin::Register_BATCH_MATMUL_GENERIC_OPTIMIZED()},
+});
+
+class BatchMatMulOpTest : public SingleOpTest {
+ protected:
+  const std::map& GetKernelMap() override {
+    return *kKernelMap;
+  }
+};
+
+TEST_P(BatchMatMulOpTest, Float32Test_Simple) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {1, 2, 3}},
                                   {TensorType_FLOAT32, {1, 3, 4}});
   model.PopulateTensor(model.lhs(), {1, 2, 3, 4, 5, 6});
@@ -65,7 +89,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Simple) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_SimpleRHSAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_SimpleRHSAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {1, 2, 3}},
                                   {TensorType_FLOAT32, {1, 4, 3}}, false, true);
   model.PopulateTensor(model.lhs(), {1, 2, 3, 4, 5, 6});
@@ -77,7 +101,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_SimpleRHSAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_SimpleLHSAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_SimpleLHSAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {1, 3, 2}},
                                   {TensorType_FLOAT32, {1, 3, 4}}, true, false);
   model.PopulateTensor(model.lhs(), {1, 4, 2, 5, 3, 6});
@@ -89,7 +113,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_SimpleLHSAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_BatchSizeTwo) {
+TEST_P(BatchMatMulOpTest, Float32Test_BatchSizeTwo) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 2, 3}},
                                   {TensorType_FLOAT32, {2, 3, 4}});
   model.PopulateTensor(model.lhs(),
@@ -105,7 +129,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_BatchSizeTwo) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_Broadcast) {
+TEST_P(BatchMatMulOpTest, Float32Test_Broadcast) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 2, 3}},
                                   {TensorType_FLOAT32, {3, 4}});
   model.PopulateTensor(model.lhs(),
@@ -121,7 +145,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Broadcast) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_BroadcastLHSAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_BroadcastLHSAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 3, 2}},
                                   {TensorType_FLOAT32, {3, 4}}, true, false);
   model.PopulateTensor(model.lhs(),
@@ -137,7 +161,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_BroadcastLHSAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 2, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2) {
+TEST_P(BatchMatMulOpTest, Float32Test_Broadcast2) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 1, 3, 2}},
                                   {TensorType_FLOAT32, {3, 2, 4}});
   model.PopulateTensor(model.lhs(),
@@ -161,7 +185,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 3, 3, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2LHSAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_Broadcast2LHSAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 1, 2, 3}},
                                   {TensorType_FLOAT32, {3, 2, 4}}, true, false);
   model.PopulateTensor(model.lhs(),
@@ -185,7 +209,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2LHSAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 3, 3, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2RHSAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_Broadcast2RHSAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 1, 3, 2}},
                                   {TensorType_FLOAT32, {3, 4, 2}}, false, true);
   model.PopulateTensor(model.lhs(),
@@ -208,7 +232,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2RHSAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 3, 3, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2BothAdjoint) {
+TEST_P(BatchMatMulOpTest, Float32Test_Broadcast2BothAdjoint) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {2, 1, 2, 3}},
                                   {TensorType_FLOAT32, {3, 4, 2}}, true, true);
   model.PopulateTensor(model.lhs(),
@@ -231,7 +255,7 @@ TEST(BatchMatMulOpModelTest, Float32Test_Broadcast2BothAdjoint) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 3, 3, 4}));
 }
 
-TEST(BatchMatMulOpModelTest, Float32Test_BroadcastFromRHS) {
+TEST_P(BatchMatMulOpTest, Float32Test_BroadcastFromRHS) {
   BatchMatMulOpModel model({TensorType_FLOAT32, {4, 5}},
                                   {TensorType_FLOAT32, {3, 1, 5, 2}});
   model.PopulateTensor(
@@ -251,6 +275,10 @@ TEST(BatchMatMulOpModelTest, Float32Test_BroadcastFromRHS) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({3, 1, 4, 2}));
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    BatchMatMulOpTest, BatchMatMulOpTest,
+    ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
+
 // In the hybrid model the weights are quantized int8. But the input
 // and output are expected to be in float precision.
 class HybridAsymmetricBatchMatMulOpModel : public SingleOpModel {
@@ -304,7 +332,14 @@ class HybridAsymmetricBatchMatMulOpModel : public SingleOpModel {
   int input_size_;
 };
 
-TEST(HybridAsymmetricBatchMatMulOpTest, SimpleTestQuantizedInt8) {
+class HybridAsymmetricBatchMatMulOpTest : public SingleOpTest {
+ protected:
+  const std::map& GetKernelMap() override {
+    return *kKernelMap;
+  }
+};
+
+TEST_P(HybridAsymmetricBatchMatMulOpTest, SimpleTestQuantizedInt8) {
   HybridAsymmetricBatchMatMulOpModel m(
       /*units=*/3, /*batches=*/2,
       /*lhs=*/{TensorType_FLOAT32, {2, 10}},
@@ -335,7 +370,7 @@ TEST(HybridAsymmetricBatchMatMulOpTest, SimpleTestQuantizedInt8) {
   EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3}));
 }
 
-TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastWeights) {
+TEST_P(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastWeights) {
   HybridAsymmetricBatchMatMulOpModel m(
       /*units=*/3, /*batches=*/2,
       /*lhs=*/{TensorType_FLOAT32, {2, 2, 10}},
@@ -366,7 +401,7 @@ TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastWeights) {
   EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 3}));
 }
 
-TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastBigWeights) {
+TEST_P(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastBigWeights) {
   HybridAsymmetricBatchMatMulOpModel m(
       /*units=*/9, /*batches=*/2,
       /*lhs=*/{TensorType_FLOAT32, {2, 2, 10}},
@@ -401,7 +436,7 @@ TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastBigWeights) {
   EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 9}));
 }
 
-TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastInputs) {
+TEST_P(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastInputs) {
   HybridAsymmetricBatchMatMulOpModel m(
       /*units=*/3, /*batches=*/2,
       /*lhs=*/{TensorType_FLOAT32, {2, 10}},
@@ -431,5 +466,96 @@ TEST(HybridAsymmetricBatchMatMulOpTest, QuantizedInt8BroadcastInputs) {
   EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 3}));
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    HybridAsymmetricBatchMatMulOpTest, HybridAsymmetricBatchMatMulOpTest,
+    ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
+
+class QuantizedBatchMatMulOpModel : public SingleOpModel {
+ public:
+  QuantizedBatchMatMulOpModel(int units, int batches, const TensorData& lhs,
+                              const TensorData& output = {TensorType_INT8},
+                              bool adj_x = false, bool adj_y = false)
+      : units_(units), batches_(batches) {
+    int total_input_size = 1;
+    for (size_t i = 0; i < lhs.shape.size(); ++i) {
+      total_input_size *= lhs.shape[i];
+    }
+    input_size_ = total_input_size / batches_;
+
+    lhs_id_ = AddInput(lhs);
+    rhs_id_ = AddInput({lhs.type, {input_size_, units_}, lhs.min, lhs.max});
+
+    output_id_ = AddOutput(output);
+
+    SetBuiltinOp(BuiltinOperator_BATCH_MATMUL,
+                 BuiltinOptions_BatchMatMulOptions,
+                 CreateBatchMatMulOptions(builder_, adj_x, adj_y).Union());
+    BuildInterpreter({GetShape(lhs_id_), GetShape(rhs_id_)});
+  }
+
+  template 
+  void SetWeights(const std::vector& data) {
+    QuantizeAndPopulate(rhs_id_, data);
+  }
+
+  template 
+  void SetInput(const std::vector& data) {
+    QuantizeAndPopulate(lhs_id_, data);
+  }
+
+  template 
+  std::vector GetOutput() {
+    return ExtractVector(output_id_);
+  }
+
+  template 
+  std::vector GetDequantizedOutput() {
+    return Dequantize(ExtractVector(output_id_), GetScale(output_id_),
+                         GetZeroPoint(output_id_));
+  }
+
+ protected:
+  int lhs_id_;
+  int rhs_id_;
+  int output_id_;
+  int units_;
+  int batches_;
+  int input_size_;
+};
+
+class QuantizedBatchMatMulOpTest : public SingleOpTest {
+ protected:
+  const std::map& GetKernelMap() override {
+    return *kKernelMap;
+  }
+};
+
+TEST_P(QuantizedBatchMatMulOpTest, SimpleTestQuantizedInt8) {
+  QuantizedBatchMatMulOpModel m(
+      /*units=*/3, /*batches*/ 2,
+      /*lhs=*/{TensorType_INT8, {2, 10}, -63.5, 64},
+      /*output=*/{TensorType_INT8, {}, -127, 128});
+
+  m.SetWeights({
+      1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5,  5,  5,
+      6, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10,
+  });
+
+  m.SetInput({
+      1, 2, 3, 4, 5, 6, 7, 8,  -9, -10,  // b = 0
+      1, 2, 3, 4, 5, 6, 7, -8, 9,  -10,  // b = 1
+  });
+
+  m.Invoke();
+
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(ArrayFloatNear({23, 23, 23, 57, 57, 57})));
+  EXPECT_THAT(m.GetOutput(), ElementsAre(22, 22, 22, 56, 56, 56));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    QuantizedBatchMatMulOpTest, QuantizedBatchMatMulOpTest,
+    ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
+
 }  // namespace
 }  // namespace tflite
diff --git a/tensorflow/lite/kernels/internal/optimized/batch_matmul.h b/tensorflow/lite/kernels/internal/optimized/batch_matmul.h
index 24b5012304f..5e622154d60 100644
--- a/tensorflow/lite/kernels/internal/optimized/batch_matmul.h
+++ b/tensorflow/lite/kernels/internal/optimized/batch_matmul.h
@@ -272,6 +272,112 @@ inline void BatchMatMul(const RuntimeShape& lhs_shape, const int8_t* lhs_data,
   }
 }
 
+inline void BatchMatMul(const FullyConnectedParams& params,
+                        const RuntimeShape& lhs_shape, const int8_t* lhs_data,
+                        const RuntimeShape& rhs_shape, const int8_t* rhs_data,
+                        const RuntimeShape& output_shape, int8_t* output_data,
+                        CpuBackendContext* context) {
+  using ::tflite::cpu_backend_gemm::Gemm;
+  using ::tflite::cpu_backend_gemm::GemmParams;
+  using ::tflite::cpu_backend_gemm::MatrixParams;
+
+  const RuntimeShape extended_lhs_shape =
+      RuntimeShape::ExtendedShape(5, lhs_shape);
+  const RuntimeShape extended_rhs_shape =
+      RuntimeShape::ExtendedShape(5, rhs_shape);
+
+  // Determine which dimension is the broadcast dimension.
+  auto broadcast_dim = [](int lhs_dim, int rhs_dim) {
+    if (lhs_dim == rhs_dim) return lhs_dim;
+    if (lhs_dim == 1) return rhs_dim;
+    TFLITE_DCHECK_EQ(rhs_dim, 1);
+    return lhs_dim;
+  };
+
+  // Compute the "extent" for iterating on this dimension.
+  // If we are broadcasting, then don't advance (i.e return 0).
+  auto extent = [](const RuntimeShape& shape, int x) {
+    if (shape.Dims(x) == 1) {
+      return 0;
+    }
+    int prod = 1;
+    for (int i = x + 1; i < shape.DimensionsCount(); ++i) {
+      prod *= shape.Dims(i);
+    }
+    return prod;
+  };
+
+  const int batch_dim0 =
+      broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0));
+  const int batch_dim1 =
+      broadcast_dim(extended_lhs_shape.Dims(1), extended_rhs_shape.Dims(1));
+  const int batch_dim2 =
+      broadcast_dim(extended_lhs_shape.Dims(2), extended_rhs_shape.Dims(2));
+
+  const int lhs_ext0 = extent(extended_lhs_shape, 0);
+  const int lhs_ext1 = extent(extended_lhs_shape, 1);
+  const int lhs_ext2 = extent(extended_lhs_shape, 2);
+  const int rhs_ext0 = extent(extended_rhs_shape, 0);
+  const int rhs_ext1 = extent(extended_rhs_shape, 1);
+  const int rhs_ext2 = extent(extended_rhs_shape, 2);
+
+  // Set params for each matrix multiply.
+  const int lhs_rows = extended_lhs_shape.Dims(3);
+  const int rhs_cols = extended_rhs_shape.Dims(4);
+  const int accum_depth = extended_lhs_shape.Dims(4);
+
+  const int32 input_offset = params.input_offset;
+  const int32 filter_offset = params.weights_offset;
+  const int32 output_offset = params.output_offset;
+  const int32 output_multiplier = params.output_multiplier;
+  const int output_shift = params.output_shift;
+  const int32 output_activation_min = params.quantized_activation_min;
+  const int32 output_activation_max = params.quantized_activation_max;
+  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
+
+  MatrixParams lhs_params;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.rows = lhs_rows;
+  lhs_params.cols = accum_depth;
+  lhs_params.zero_point = -filter_offset;
+
+  MatrixParams rhs_params;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.rows = accum_depth;
+  rhs_params.cols = rhs_cols;
+  rhs_params.zero_point = -input_offset;
+
+  MatrixParams dst_params;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.rows = lhs_rows;
+  dst_params.cols = rhs_cols;
+  dst_params.zero_point = output_offset;
+
+  for (int b0 = 0; b0 < batch_dim0; ++b0) {
+    const int8_t* lhs_ptr0 = lhs_data + (b0 * lhs_ext0);
+    const int8_t* rhs_ptr0 = rhs_data + (b0 * rhs_ext0);
+    for (int b1 = 0; b1 < batch_dim1; ++b1) {
+      const int8_t* lhs_ptr1 = lhs_ptr0 + b1 * lhs_ext1;
+      const int8_t* rhs_ptr1 = rhs_ptr0 + b1 * rhs_ext1;
+      for (int b2 = 0; b2 < batch_dim2; ++b2) {
+        const int8_t* lhs_ptr2 = lhs_ptr1 + b2 * lhs_ext2;
+        const int8_t* rhs_ptr2 = rhs_ptr1 + b2 * rhs_ext2;
+        int8_t* out_ptr = output_data + ((b0 * batch_dim1 * batch_dim2) +
+                                         b1 * batch_dim2 + b2) *
+                                            lhs_rows * rhs_cols;
+
+        GemmParams gemm_params;
+        gemm_params.clamp_min = output_activation_min;
+        gemm_params.clamp_max = output_activation_max;
+        gemm_params.multiplier_fixedpoint = output_multiplier;
+        gemm_params.multiplier_exponent = output_shift;
+        cpu_backend_gemm::Gemm(lhs_params, lhs_ptr2, rhs_params, rhs_ptr2,
+                               dst_params, out_ptr, gemm_params, context);
+      }
+    }
+  }
+}
+
 }  // namespace optimized_ops
 }  // namespace tflite
 
diff --git a/tensorflow/lite/kernels/internal/reference/batch_matmul.h b/tensorflow/lite/kernels/internal/reference/batch_matmul.h
index 1394bd9da64..05caefaca5d 100644
--- a/tensorflow/lite/kernels/internal/reference/batch_matmul.h
+++ b/tensorflow/lite/kernels/internal/reference/batch_matmul.h
@@ -217,6 +217,99 @@ inline void BatchMatMul(const RuntimeShape& lhs_shape, const int8_t* lhs_data,
   }
 }
 
+inline void BatchMatMul(const FullyConnectedParams& params,
+                        const RuntimeShape& lhs_shape, const int8_t* lhs_data,
+                        const RuntimeShape& rhs_shape, const int8_t* rhs_data,
+                        const RuntimeShape& output_shape, int8_t* output_data) {
+  const RuntimeShape extended_lhs_shape =
+      RuntimeShape::ExtendedShape(5, lhs_shape);
+  const RuntimeShape extended_rhs_shape =
+      RuntimeShape::ExtendedShape(5, rhs_shape);
+
+  // Determine which dimension is the broadcast dimension.
+  auto broadcast_dim = [](int lhs_dim, int rhs_dim) {
+    if (lhs_dim == rhs_dim) return lhs_dim;
+    if (lhs_dim == 1) return rhs_dim;
+    TFLITE_DCHECK_EQ(rhs_dim, 1);
+    return lhs_dim;
+  };
+
+  // Compute the "extent" for iterating on this dimension.
+  // If we are broadcasting, then don't advance (i.e return 0).
+  auto extent = [](const RuntimeShape& shape, int x) {
+    if (shape.Dims(x) == 1) {
+      return 0;
+    }
+    int prod = 1;
+    for (int i = x + 1; i < shape.DimensionsCount(); ++i) {
+      prod *= shape.Dims(i);
+    }
+    return prod;
+  };
+
+  const int batch_dim0 =
+      broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0));
+  const int batch_dim1 =
+      broadcast_dim(extended_lhs_shape.Dims(1), extended_rhs_shape.Dims(1));
+  const int batch_dim2 =
+      broadcast_dim(extended_lhs_shape.Dims(2), extended_rhs_shape.Dims(2));
+
+  const int lhs_ext0 = extent(extended_lhs_shape, 0);
+  const int lhs_ext1 = extent(extended_lhs_shape, 1);
+  const int lhs_ext2 = extent(extended_lhs_shape, 2);
+  const int rhs_ext0 = extent(extended_rhs_shape, 0);
+  const int rhs_ext1 = extent(extended_rhs_shape, 1);
+  const int rhs_ext2 = extent(extended_rhs_shape, 2);
+
+  // Set params for each matrix multiply.
+  const int lhs_rows = extended_lhs_shape.Dims(3);
+  const int rhs_cols = extended_rhs_shape.Dims(4);
+  const int accum_depth = extended_lhs_shape.Dims(4);
+
+  const int32 input_offset = params.input_offset;
+  const int32 filter_offset = params.weights_offset;
+  const int32 output_offset = params.output_offset;
+  const int32 output_multiplier = params.output_multiplier;
+  const int output_shift = params.output_shift;
+  const int32 output_activation_min = params.quantized_activation_min;
+  const int32 output_activation_max = params.quantized_activation_max;
+  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
+
+  for (int b0 = 0; b0 < batch_dim0; ++b0) {
+    const int8_t* lhs_ptr0 = lhs_data + (b0 * lhs_ext0);
+    const int8_t* rhs_ptr0 = rhs_data + (b0 * rhs_ext0);
+    for (int b1 = 0; b1 < batch_dim1; ++b1) {
+      const int8_t* lhs_ptr1 = lhs_ptr0 + b1 * lhs_ext1;
+      const int8_t* rhs_ptr1 = rhs_ptr0 + b1 * rhs_ext1;
+      for (int b2 = 0; b2 < batch_dim2; ++b2) {
+        const int8_t* lhs_ptr2 = lhs_ptr1 + b2 * lhs_ext2;
+        const int8_t* rhs_ptr2 = rhs_ptr1 + b2 * rhs_ext2;
+        int8_t* out_ptr = output_data + ((b0 * batch_dim1 * batch_dim2) +
+                                         b1 * batch_dim2 + b2) *
+                                            lhs_rows * rhs_cols;
+
+        for (int j = 0; j < rhs_cols; ++j) {
+          for (int i = 0; i < lhs_rows; ++i) {
+            int32_t total = 0;
+            for (int k = 0; k < accum_depth; ++k) {
+              int32 lhs_val = lhs_ptr2[accum_depth * i + k];
+              int32 rhs_val = rhs_ptr2[accum_depth * j + k];
+              total += (lhs_val + filter_offset) * (rhs_val + input_offset);
+            }
+            total = MultiplyByQuantizedMultiplier(total, output_multiplier,
+                                                  output_shift);
+            total += output_offset;
+            total = std::max(total, output_activation_min);
+            total = std::min(total, output_activation_max);
+            const int idx = lhs_rows * j + i;
+            out_ptr[idx] = static_cast(total);
+          }
+        }
+      }
+    }
+  }
+}
+
 }  // namespace reference_ops
 }  // namespace tflite
 
diff --git a/tensorflow/lite/kernels/register.cc b/tensorflow/lite/kernels/register.cc
index 90688a2aa1f..c3a4aaad16d 100644
--- a/tensorflow/lite/kernels/register.cc
+++ b/tensorflow/lite/kernels/register.cc
@@ -289,7 +289,9 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_SCATTER_ND, Register_SCATTER_ND());
   AddBuiltin(BuiltinOperator_DENSIFY, Register_DENSIFY());
   AddBuiltin(BuiltinOperator_SEGMENT_SUM, Register_SEGMENT_SUM());
-  AddBuiltin(BuiltinOperator_BATCH_MATMUL, Register_BATCH_MATMUL());
+  AddBuiltin(BuiltinOperator_BATCH_MATMUL, Register_BATCH_MATMUL(),
+             /* min_version = */ 1,
+             /* max_version = */ 2);
   AddCustom("NumericVerify", tflite::ops::custom::Register_NUMERIC_VERIFY());
   // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that
   // custom ops aren't always included by default.
diff --git a/tensorflow/lite/tools/optimize/operator_property.cc b/tensorflow/lite/tools/optimize/operator_property.cc
index 8a0cbca29e2..f2cb98ef31a 100644
--- a/tensorflow/lite/tools/optimize/operator_property.cc
+++ b/tensorflow/lite/tools/optimize/operator_property.cc
@@ -88,6 +88,12 @@ OperatorProperty GetOperatorProperty(const ModelT* model, int subgraph_index,
       property.restrict_same_input_output_scale = true;
       property.version = 2;
       break;
+    case BuiltinOperator_BATCH_MATMUL: {
+      property.inputs = {{0, {}}, {1, {}}};
+      property.outputs = {{0, {}}};
+      property.version = 2;
+      break;
+    }
     case BuiltinOperator_BATCH_TO_SPACE_ND:
     case BuiltinOperator_SPACE_TO_BATCH_ND:
     case BuiltinOperator_SPACE_TO_DEPTH:
diff --git a/tensorflow/lite/tools/versioning/op_version.cc b/tensorflow/lite/tools/versioning/op_version.cc
index 118e2d420f8..a97b9da47f1 100644
--- a/tensorflow/lite/tools/versioning/op_version.cc
+++ b/tensorflow/lite/tools/versioning/op_version.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/strings/str_split.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/schema/schema_generated.h"
 
 namespace tflite {
 namespace {
@@ -518,6 +519,7 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
     case BuiltinOperator_LESS:
     case BuiltinOperator_LESS_EQUAL:
     case BuiltinOperator_SELECT:
+    case BuiltinOperator_BATCH_MATMUL:
       if (op_sig.input_types.at(0) == TensorType_INT8) {
         return 2;
       }
diff --git a/tensorflow/lite/tools/versioning/runtime_version.cc b/tensorflow/lite/tools/versioning/runtime_version.cc
index 92a7001606f..36976354685 100644
--- a/tensorflow/lite/tools/versioning/runtime_version.cc
+++ b/tensorflow/lite/tools/versioning/runtime_version.cc
@@ -58,6 +58,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_AVERAGE_POOL_2D, 2}, "1.14.0"},
               {{BuiltinOperator_AVERAGE_POOL_2D, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_BATCH_MATMUL, 1}, kPendingReleaseVersion},
+              {{BuiltinOperator_BATCH_MATMUL, 2}, kPendingReleaseVersion},
               {{BuiltinOperator_CONV_2D, 1}, "1.5.0"},
               {{BuiltinOperator_CONV_2D, 2}, "1.14.0"},
               {{BuiltinOperator_CONV_2D, 3}, "1.14.0"},

From 07c54454eec55c1279c243a3c148eeee81b41ed5 Mon Sep 17 00:00:00 2001
From: Tamara Norman 
Date: Fri, 19 Jun 2020 08:48:51 -0700
Subject: [PATCH 0616/1390] Add an option such that the cached host_value can
 be discarded

PiperOrigin-RevId: 317315157
Change-Id: I9d7145390a526003069321c7e04794e139a53c09
---
 tensorflow/compiler/xla/pjrt/pjrt_client.cc | 6 +++++-
 tensorflow/compiler/xla/pjrt/pjrt_client.h  | 8 ++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.cc b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
index 46f592100c9..b4f0363e69a 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.cc
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.cc
@@ -1077,13 +1077,17 @@ Status PjRtBuffer::CopyToHostAsync() {
   return Status::OK();
 }
 
-StatusOr> PjRtBuffer::ToLiteral() {
+StatusOr> PjRtBuffer::ToLiteral(
+    const bool discard_cached_copy) {
   tensorflow::profiler::TraceMe traceme("PjRtBuffer::ToLiteral");
   TF_RETURN_IF_ERROR(CopyToHostAsync());
   std::shared_ptr host_value;
   {
     absl::MutexLock lock(&mu_);
     host_value = host_value_;
+    if (discard_cached_copy) {
+      host_value_ = nullptr;
+    }
   }
   if (host_value == nullptr) {
     return InvalidArgument("ToLiteral called on invalid buffer");
diff --git a/tensorflow/compiler/xla/pjrt/pjrt_client.h b/tensorflow/compiler/xla/pjrt/pjrt_client.h
index 754eb19bec6..8f74e6244d6 100644
--- a/tensorflow/compiler/xla/pjrt/pjrt_client.h
+++ b/tensorflow/compiler/xla/pjrt/pjrt_client.h
@@ -478,8 +478,12 @@ class PjRtBuffer {
 
   // Returns the buffer's value as an XLA Literal. If the value has previously
   // been prefetched to the host, then returns the prefetched version, otherwise
-  // copies the buffer to the host. Blocks until the value is ready.
-  StatusOr> ToLiteral();
+  // copies the buffer to the host. Blocks until the value is ready. If
+  // `discard_cached_copy` is true then buffer will no longer keep hold of a
+  // cached copy of the literal (i.e. The reference to the host value will be
+  // removed.)
+  StatusOr> ToLiteral(
+      bool discard_cached_copy = false);
 
   // Initiates a copy of the buffer to the host. Does not block waiting for
   // the transfer to complete. The value can be retrieved by a later call to

From 16cb89bd7b40fc816aa7440f62b443ca480bcf6d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 09:15:04 -0700
Subject: [PATCH 0617/1390] Qualify uses of std::string

PiperOrigin-RevId: 317319501
Change-Id: Ib75a31ad89fa1a6bda81450f2ab5ba07d7338ada
---
 tensorflow/lite/toco/tflite/export.cc         | 58 +++++++++----------
 tensorflow/lite/toco/tflite/export.h          | 12 ++--
 tensorflow/lite/toco/tflite/export_test.cc    | 56 ++++++++++--------
 tensorflow/lite/toco/tflite/import.cc         | 18 +++---
 tensorflow/lite/toco/tflite/import.h          |  6 +-
 tensorflow/lite/toco/tflite/import_test.cc    |  6 +-
 tensorflow/lite/toco/tflite/op_version.cc     |  8 +--
 .../lite/toco/tflite/op_version_test.cc       | 32 +++++-----
 tensorflow/lite/toco/tflite/operator.cc       | 37 ++++++------
 tensorflow/lite/toco/tflite/operator.h        | 12 ++--
 tensorflow/lite/toco/tflite/operator_test.cc  |  8 +--
 tensorflow/lite/toco/tflite/types.cc          |  6 +-
 12 files changed, 133 insertions(+), 126 deletions(-)

diff --git a/tensorflow/lite/toco/tflite/export.cc b/tensorflow/lite/toco/tflite/export.cc
index d72a902001d..d109ab875b5 100644
--- a/tensorflow/lite/toco/tflite/export.cc
+++ b/tensorflow/lite/toco/tflite/export.cc
@@ -52,7 +52,7 @@ using ::tflite::Tensor;
 namespace {
 
 // Check if a TensorFlow Op is a control flow op by its name.
-bool IsControlFlowOp(const string& tensorflow_op) {
+bool IsControlFlowOp(const std::string& tensorflow_op) {
   // Technically this is equivalent to `::tensorflow::Node::IsControlFlow()`.
   // It requires to construct a `::tensorflow::Graph` to use that helper
   // function, so we simply hardcode the list of control flow ops here.
@@ -68,7 +68,7 @@ bool IsControlFlowOp(const string& tensorflow_op) {
 }
 
 // Check if a TensorFlow Op is unsupported by the Flex runtime.
-bool IsUnsupportedFlexOp(const string& tensorflow_op) {
+bool IsUnsupportedFlexOp(const std::string& tensorflow_op) {
   if (IsControlFlowOp(tensorflow_op)) {
     return true;
   }
@@ -82,14 +82,14 @@ bool IsUnsupportedFlexOp(const string& tensorflow_op) {
 }
 
 // Map from operator name to TF Lite enum value, for all builtins.
-const std::map& GetBuiltinOpsMap() {
-  static std::map* builtin_ops = nullptr;
+const std::map& GetBuiltinOpsMap() {
+  static std::map* builtin_ops = nullptr;
   if (builtin_ops == nullptr) {
-    builtin_ops = new std::map();
+    builtin_ops = new std::map();
 
     for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
       BuiltinOperator op = static_cast(i);
-      string name = EnumNameBuiltinOperator(op);
+      std::string name = EnumNameBuiltinOperator(op);
       if (op != BuiltinOperator_CUSTOM && !name.empty()) {
         (*builtin_ops)[name] = op;
       }
@@ -99,10 +99,10 @@ const std::map& GetBuiltinOpsMap() {
 }
 
 void WriteModelToString(const flatbuffers::FlatBufferBuilder& builder,
-                        string* file_contents) {
+                        std::string* file_contents) {
   const uint8_t* buffer = builder.GetBufferPointer();
   int size = builder.GetSize();
-  *file_contents = string(reinterpret_cast(buffer), size);
+  *file_contents = std::string(reinterpret_cast(buffer), size);
 }
 
 }  // Anonymous namespace.
@@ -115,7 +115,7 @@ OperatorKey::OperatorKey(
     bool enable_select_tf_ops) {
   // Get the op name (by Toco definition).
   const ::toco::Operator& op = *op_signature.op;
-  string name = HelpfulOperatorTypeName(op);
+  std::string name = HelpfulOperatorTypeName(op);
 
   bool is_builtin = false;
   const auto& builtin_ops = GetBuiltinOpsMap();
@@ -146,7 +146,7 @@ OperatorKey::OperatorKey(
       is_flex_op_ = true;
       flex_tensorflow_op_ = tensorflow_op;
       custom_code_ =
-          string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
+          std::string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
     } else {
       custom_code_ = tensorflow_op;
     }
@@ -158,7 +158,7 @@ OperatorKey::OperatorKey(
     is_flex_op_ = true;
     flex_tensorflow_op_ = name;
     custom_code_ =
-        string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
+        std::string(::tflite::kFlexCustomCodePrefix) + flex_tensorflow_op_;
   } else {
     // If Flex is disabled or the original TensorFlow NodeDef isn't available,
     // we produce a custom op. This gives developers a chance to implement
@@ -175,7 +175,7 @@ OperatorKey::OperatorKey(
 
 void LoadTensorsMap(const Model& model, TensorsMap* tensors_map) {
   // First find a list of unique array names.
-  std::set names;
+  std::set names;
   for (const auto& array_pair : model.GetArrayMap()) {
     names.insert(array_pair.first);
   }
@@ -218,7 +218,7 @@ Offset>> ExportTensors(
   std::map> ordered_tensors;
 
   for (const auto& array_pair : model.GetArrayMap()) {
-    const string& tensor_name = array_pair.first;
+    const std::string& tensor_name = array_pair.first;
     const toco::Array& array = *array_pair.second;
 
     int buffer_index = buffers_to_write->size();
@@ -283,7 +283,7 @@ Offset> ExportOutputTensors(
     const Model& model, const details::TensorsMap& tensors_map,
     FlatBufferBuilder* builder) {
   std::vector outputs;
-  for (const string& output : model.flags.output_arrays()) {
+  for (const std::string& output : model.flags.output_arrays()) {
     outputs.push_back(tensors_map.at(output));
   }
   return builder->CreateVector(outputs);
@@ -295,10 +295,10 @@ Offset>> ExportOperatorCodes(
     const details::OperatorsMap& operators_map, FlatBufferBuilder* builder,
     const ExportParams& params) {
   // Map from operator name to TF Lite enum value, for all builtins.
-  std::map builtin_ops;
+  std::map builtin_ops;
   for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) {
     BuiltinOperator op = static_cast(i);
-    string name = EnumNameBuiltinOperator(op);
+    std::string name = EnumNameBuiltinOperator(op);
     if (op != BuiltinOperator_CUSTOM && !name.empty()) {
       builtin_ops[name] = op;
     }
@@ -349,13 +349,13 @@ Offset>> ExportOperators(
   std::vector> op_vector;
   for (const auto& op : model.operators) {
     std::vector inputs;
-    for (const string& input : op->inputs) {
+    for (const std::string& input : op->inputs) {
       // -1 is the ID for optional tensor in TFLite output
       int id = model.IsOptionalArray(input) ? -1 : tensors_map.at(input);
       inputs.push_back(id);
     }
     std::vector outputs;
-    for (const string& output : op->outputs) {
+    for (const std::string& output : op->outputs) {
       outputs.push_back(tensors_map.at(output));
     }
     const toco::OperatorSignature op_signature = {op.get(), &model};
@@ -428,15 +428,15 @@ Offset>> ExportBuffers(
   return builder->CreateVector(buffer_vector);
 }
 
-tensorflow::Status Export(const Model& model, string* output_file_contents,
+tensorflow::Status Export(const Model& model, std::string* output_file_contents,
                           const ExportParams& params) {
   const auto ops_by_type = BuildOperatorByTypeMap(params.enable_select_tf_ops);
   return Export(model, output_file_contents, params, ops_by_type);
 }
 
-void ParseControlFlowErrors(std::set* custom_ops,
-                            std::vector* error_msgs) {
-  std::set unsupported_control_flow_ops;
+void ParseControlFlowErrors(std::set* custom_ops,
+                            std::vector* error_msgs) {
+  std::set unsupported_control_flow_ops;
   // Check if unsupported ops contains control flow ops. It's impossible
   // to implement these ops as custom ops at the moment.
   for (const auto& op : *custom_ops) {
@@ -471,10 +471,10 @@ void ExportModelVersionBuffer(
 }
 
 tensorflow::Status Export(
-    const Model& model, string* output_file_contents,
+    const Model& model, std::string* output_file_contents,
     const ExportParams& params,
     const std::map>& ops_by_type) {
-  for (const string& input_array : model.GetInvalidInputArrays()) {
+  for (const std::string& input_array : model.GetInvalidInputArrays()) {
     if (model.HasArray(input_array)) {
       return tensorflow::errors::InvalidArgument(
           absl::StrCat("Placeholder ", input_array,
@@ -509,11 +509,11 @@ tensorflow::Status Export(
   }
 
   // The set of used builtin ops.
-  std::set builtin_ops;
+  std::set builtin_ops;
   // The set of custom ops (not including Flex ops).
-  std::set custom_ops;
+  std::set custom_ops;
   // The set of Flex ops which are not supported.
-  std::set unsupported_flex_ops;
+  std::set unsupported_flex_ops;
 
   for (const auto& it : operators_map) {
     const details::OperatorKey& key = it.first;
@@ -540,7 +540,7 @@ tensorflow::Status Export(
                "40-tflite-op-request.md\n and pasting the following:\n\n";
       };
 
-      std::vector error_msgs;
+      std::vector error_msgs;
       ParseControlFlowErrors(&custom_ops, &error_msgs);
 
       // Remove ExpandDims and ReorderAxes from unimplemented list unless they
@@ -549,7 +549,7 @@ tensorflow::Status Export(
       // transformation is unable to run because the output shape is not
       // defined. This causes unnecessary confusion during model conversion
       // time.
-      std::set custom_ops_final;
+      std::set custom_ops_final;
       for (const auto& op_type : custom_ops) {
         if (op_type != "ReorderAxes" && op_type != "ExpandDims") {
           custom_ops_final.insert(op_type);
diff --git a/tensorflow/lite/toco/tflite/export.h b/tensorflow/lite/toco/tflite/export.h
index 3af77ffcf43..64f7c7b128f 100644
--- a/tensorflow/lite/toco/tflite/export.h
+++ b/tensorflow/lite/toco/tflite/export.h
@@ -35,19 +35,19 @@ struct ExportParams {
 
 // Transform the given tf.mini model into a TF Lite flatbuffer and deposit the
 // result in the given string.
-tensorflow::Status Export(const Model& model, string* output_file_contents,
+tensorflow::Status Export(const Model& model, std::string* output_file_contents,
                           const ExportParams& params);
 
 // Export API with custom TFLite operator mapping.
 tensorflow::Status Export(
-    const Model& model, string* output_file_contents,
+    const Model& model, std::string* output_file_contents,
     const ExportParams& params,
     const std::map>& ops_by_type);
 
 // This is for backward-compatibility.
 // TODO(ycling): Remove the deprecated entry functions.
 inline void Export(const Model& model, bool allow_custom_ops,
-                   bool quantize_weights, string* output_file_contents) {
+                   bool quantize_weights, std::string* output_file_contents) {
   ExportParams params;
   params.allow_custom_ops = allow_custom_ops;
   params.quantize_weights =
@@ -60,7 +60,7 @@ inline void Export(const Model& model, bool allow_custom_ops,
 // TODO(ycling): Remove the deprecated entry functions.
 inline void Export(
     const Model& model, bool allow_custom_ops, bool quantize_weights,
-    string* output_file_contents,
+    std::string* output_file_contents,
     const std::map>& ops_by_type) {
   ExportParams params;
   params.allow_custom_ops = allow_custom_ops;
@@ -72,7 +72,7 @@ inline void Export(
 
 // This is for backward-compatibility.
 // TODO(ycling): Remove the deprecated entry functions.
-inline void Export(const Model& model, string* output_file_contents) {
+inline void Export(const Model& model, std::string* output_file_contents) {
   ExportParams params;
   params.allow_custom_ops = true;
   auto status = Export(model, output_file_contents, params);
@@ -82,7 +82,7 @@ inline void Export(const Model& model, string* output_file_contents) {
 namespace details {
 
 // A map from tensor name to its final position in the TF Lite buffer.
-using TensorsMap = std::unordered_map;
+using TensorsMap = std::unordered_map;
 
 // A key to identify an operator.
 // Only when `type` is `kUnsupported`, `custom_code` is filled to
diff --git a/tensorflow/lite/toco/tflite/export_test.cc b/tensorflow/lite/toco/tflite/export_test.cc
index 19b77543c66..ed347a28d51 100644
--- a/tensorflow/lite/toco/tflite/export_test.cc
+++ b/tensorflow/lite/toco/tflite/export_test.cc
@@ -34,13 +34,13 @@ using ::testing::HasSubstr;
 class ExportTest : public ::testing::Test {
  protected:
   void ResetOperators() { input_model_.operators.clear(); }
-  void AddTensorsByName(std::initializer_list names) {
-    for (const string& name : names) {
+  void AddTensorsByName(std::initializer_list names) {
+    for (const std::string& name : names) {
       input_model_.GetOrCreateArray(name);
     }
   }
-  void AddOperatorsByName(std::initializer_list names) {
-    for (const string& name : names) {
+  void AddOperatorsByName(std::initializer_list names) {
+    for (const std::string& name : names) {
       if (name == "Conv") {
         auto* op = new ConvOperator;
         op->padding.type = PaddingType::kSame;
@@ -153,14 +153,15 @@ class ExportTest : public ::testing::Test {
   }
 
   tensorflow::Status ExportAndReturnStatus(const ExportParams& params) {
-    string result;
+    std::string result;
     return Export(input_model_, &result, params);
   }
 
-  std::vector ExportAndSummarizeOperators(const ExportParams& params) {
-    std::vector names;
+  std::vector ExportAndSummarizeOperators(
+      const ExportParams& params) {
+    std::vector names;
 
-    string result;
+    std::string result;
     auto status = Export(input_model_, &result, params);
     if (!status.ok()) {
       LOG(INFO) << status.error_message();
@@ -171,10 +172,12 @@ class ExportTest : public ::testing::Test {
 
     for (const ::tflite::OperatorCode* opcode : *model->operator_codes()) {
       if (opcode->builtin_code() != ::tflite::BuiltinOperator_CUSTOM) {
-        names.push_back(string("builtin:") + ::tflite::EnumNameBuiltinOperator(
-                                                 opcode->builtin_code()));
+        names.push_back(
+            std::string("builtin:") +
+            ::tflite::EnumNameBuiltinOperator(opcode->builtin_code()));
       } else {
-        names.push_back(string("custom:") + opcode->custom_code()->c_str());
+        names.push_back(std::string("custom:") +
+                        opcode->custom_code()->c_str());
       }
     }
 
@@ -185,7 +188,7 @@ class ExportTest : public ::testing::Test {
       const ExportParams& params) {
     std::vector indices;
 
-    string result;
+    std::string result;
     if (!Export(input_model_, &result, params).ok()) return indices;
     auto* model = ::tflite::GetModel(result.data());
 
@@ -257,7 +260,7 @@ TEST_F(ExportTest, ExportMinRuntime) {
   params.enable_select_tf_ops = false;
   params.quantize_weights = QuantizedBufferType::NONE;
 
-  string output;
+  std::string output;
   auto status = Export(input_model_, &output, params);
   auto* model = ::tflite::GetModel(output.data());
   EXPECT_EQ(model->metadata()->size(), 1);
@@ -265,7 +268,8 @@ TEST_F(ExportTest, ExportMinRuntime) {
   auto buf = model->metadata()->Get(0)->buffer();
   auto* buffer = (*model->buffers())[buf];
   auto* array = buffer->data();
-  string version(reinterpret_cast(array->data()), array->size());
+  std::string version(reinterpret_cast(array->data()),
+                      array->size());
   EXPECT_EQ(version, "1.6.0");
 }
 
@@ -275,7 +279,7 @@ TEST_F(ExportTest, ExportEmptyMinRuntime) {
   ExportParams params;
   params.allow_custom_ops = true;
 
-  string output;
+  std::string output;
   auto status = Export(input_model_, &output, params);
   auto* model = ::tflite::GetModel(output.data());
   EXPECT_EQ(model->metadata()->size(), 1);
@@ -283,7 +287,8 @@ TEST_F(ExportTest, ExportEmptyMinRuntime) {
   auto buf = model->metadata()->Get(0)->buffer();
   auto* buffer = (*model->buffers())[buf];
   auto* array = buffer->data();
-  string version(reinterpret_cast(array->data()), array->size());
+  std::string version(reinterpret_cast(array->data()),
+                      array->size());
   EXPECT_EQ(version, "");
 }
 
@@ -296,7 +301,7 @@ TEST_F(ExportTest, UnsupportedControlFlowErrors) {
   // The model contains control flow ops which are not convertible, so we should
   // check the returned error message.
 
-  string output;
+  std::string output;
   const auto ops_by_type = BuildOperatorByTypeMap();
   auto status = Export(input_model_, &output, params, ops_by_type);
   EXPECT_EQ(status.error_message(),
@@ -318,7 +323,7 @@ TEST_F(ExportTest, UnsupportedOpsAndNeedEnableFlex) {
   params.allow_custom_ops = false;
   params.enable_select_tf_ops = false;
 
-  string output;
+  std::string output;
   const auto ops_by_type = BuildOperatorByTypeMap();
   auto status = Export(input_model_, &output, params, ops_by_type);
   EXPECT_EQ(
@@ -348,7 +353,7 @@ TEST_F(ExportTest, UnsupportedOpsNeedCustomImplementation) {
   params.allow_custom_ops = false;
   params.enable_select_tf_ops = true;
 
-  string output;
+  std::string output;
   const auto ops_by_type = BuildOperatorByTypeMap();
   auto status = Export(input_model_, &output, params, ops_by_type);
   EXPECT_EQ(
@@ -378,7 +383,7 @@ TEST_F(ExportTest, UnsupportedControlFlowAndCustomOpsErrors) {
   // The model contains control flow ops which are not convertible, so we should
   // check the returned error message.
 
-  string output;
+  std::string output;
   const auto ops_by_type = BuildOperatorByTypeMap();
   auto status = Export(input_model_, &output, params, ops_by_type);
   EXPECT_EQ(
@@ -407,11 +412,11 @@ TEST_F(ExportTest, UnsupportedControlFlowAndCustomOpsErrors) {
 TEST_F(ExportTest, QuantizeWeights) {
   // Sanity check for quantize_weights parameter.
   BuildQuantizableTestModel();
-  string unquantized_result;
+  std::string unquantized_result;
   Export(input_model_, true, /*quantize_weights*/ false, &unquantized_result);
 
   BuildQuantizableTestModel();
-  string quantized_result;
+  std::string quantized_result;
   Export(input_model_, true, /*quantize_weights*/ true, &quantized_result);
 
   // The quantized models should be smaller.
@@ -443,12 +448,13 @@ class OpSetsTest : public ExportTest {
     }
   }
 
-  std::vector ImportExport(std::initializer_list op_names) {
+  std::vector ImportExport(
+      std::initializer_list op_names) {
     ResetOperators();
     if (!import_all_ops_as_unsupported_) {
       AddOperatorsByName(op_names);
     } else {
-      for (const string& name : op_names) {
+      for (const std::string& name : op_names) {
         auto* op = new TensorFlowUnsupportedOperator;
         op->tensorflow_op = name;
         input_model_.operators.emplace_back(op);
@@ -644,7 +650,7 @@ TEST_F(VersionedOpExportTest, Export) {
   AddConvOp(false);
   AddConvOp(true);
 
-  string result;
+  std::string result;
   const auto ops_by_type = BuildFakeOperatorByTypeMap();
   Export(input_model_, true, false, &result, ops_by_type);
 
diff --git a/tensorflow/lite/toco/tflite/import.cc b/tensorflow/lite/toco/tflite/import.cc
index 0f3dd48652e..136aa4ffaa8 100644
--- a/tensorflow/lite/toco/tflite/import.cc
+++ b/tensorflow/lite/toco/tflite/import.cc
@@ -99,7 +99,7 @@ void ImportTensors(const ::tflite::Model& input_model, Model* model) {
 
 void ImportOperators(
     const ::tflite::Model& input_model,
-    const std::map>& ops_by_name,
+    const std::map>& ops_by_name,
     const details::TensorsTable& tensors_table,
     const details::OperatorsTable& operators_table, Model* model) {
   // TODO(aselle): add support for multiple subgraphs.
@@ -112,12 +112,12 @@ void ImportOperators(
       LOG(FATAL) << "Index " << index << " must be between zero and "
                  << operators_table.size();
     }
-    string opname = operators_table.at(index);
+    std::string opname = operators_table.at(index);
 
     // Find and use the appropriate operator deserialization factory.
     std::unique_ptr new_op = nullptr;
     if (ops_by_name.count(opname) == 0) {
-      string effective_opname = "TENSORFLOW_UNSUPPORTED";
+      std::string effective_opname = "TENSORFLOW_UNSUPPORTED";
       if (ops_by_name.count(effective_opname) == 0) {
         LOG(FATAL) << "Internal logic error: TENSORFLOW_UNSUPPORTED not found.";
       }
@@ -147,10 +147,10 @@ void ImportOperators(
       auto input_index = inputs->Get(i);
       // input_index == -1 indicates optional tensor.
       if (input_index != -1) {
-        const string& input_name = tensors_table.at(input_index);
+        const std::string& input_name = tensors_table.at(input_index);
         op->inputs.push_back(input_name);
       } else {
-        const string& tensor_name =
+        const std::string& tensor_name =
             toco::AvailableArrayName(*model, "OptionalTensor");
         model->CreateOptionalArray(tensor_name);
         op->inputs.push_back(tensor_name);
@@ -159,7 +159,7 @@ void ImportOperators(
     auto outputs = input_op->outputs();
     for (int i = 0; i < outputs->Length(); i++) {
       auto output_index = outputs->Get(i);
-      const string& output_name = tensors_table.at(output_index);
+      const std::string& output_name = tensors_table.at(output_index);
       op->outputs.push_back(output_name);
     }
   }
@@ -173,7 +173,7 @@ void ImportIOTensors(const ModelFlags& model_flags,
     auto inputs = (*input_model.subgraphs())[0]->inputs();
     if (inputs) {
       for (int input : *inputs) {
-        const string& input_name = tensors_table.at(input);
+        const std::string& input_name = tensors_table.at(input);
         model->flags.add_input_arrays()->set_name(input_name);
       }
     }
@@ -184,7 +184,7 @@ void ImportIOTensors(const ModelFlags& model_flags,
     auto outputs = (*input_model.subgraphs())[0]->outputs();
     if (outputs) {
       for (int output : *outputs) {
-        const string& output_name = tensors_table.at(output);
+        const std::string& output_name = tensors_table.at(output);
         model->flags.add_output_arrays(output_name);
       }
     }
@@ -199,7 +199,7 @@ bool Verify(const void* buf, size_t len) {
 }  // namespace
 
 std::unique_ptr Import(const ModelFlags& model_flags,
-                              const string& input_file_contents) {
+                              const std::string& input_file_contents) {
   ::tflite::AlwaysTrueResolver r;
   if (!::tflite::Verify(input_file_contents.data(), input_file_contents.size(),
                         r, ::tflite::DefaultErrorReporter())) {
diff --git a/tensorflow/lite/toco/tflite/import.h b/tensorflow/lite/toco/tflite/import.h
index f5de3b53b5b..bac55aae8b9 100644
--- a/tensorflow/lite/toco/tflite/import.h
+++ b/tensorflow/lite/toco/tflite/import.h
@@ -24,17 +24,17 @@ namespace tflite {
 
 // Parse the given string as TF Lite flatbuffer and return a new tf.mini model.
 std::unique_ptr Import(const ModelFlags &model_flags,
-                              const string &input_file_contents);
+                              const std::string &input_file_contents);
 
 namespace details {
 
 // The names of all tensors found in a TF Lite model.
-using TensorsTable = std::vector;
+using TensorsTable = std::vector;
 
 // The names of all operators found in TF Lite model. If the operator is
 // builtin, the string representation of the corresponding enum value is used
 // as name.
-using OperatorsTable = std::vector;
+using OperatorsTable = std::vector;
 
 void LoadTensorsTable(const ::tflite::Model &input_model,
                       TensorsTable *tensors_table);
diff --git a/tensorflow/lite/toco/tflite/import_test.cc b/tensorflow/lite/toco/tflite/import_test.cc
index b00c4124d83..6163ebab45b 100644
--- a/tensorflow/lite/toco/tflite/import_test.cc
+++ b/tensorflow/lite/toco/tflite/import_test.cc
@@ -134,9 +134,9 @@ class ImportTest : public ::testing::Test {
 
     input_model_ = ::tflite::GetModel(builder_.GetBufferPointer());
   }
-  string InputModelAsString() {
-    return string(reinterpret_cast(builder_.GetBufferPointer()),
-                  builder_.GetSize());
+  std::string InputModelAsString() {
+    return std::string(reinterpret_cast(builder_.GetBufferPointer()),
+                       builder_.GetSize());
   }
   flatbuffers::FlatBufferBuilder builder_;
   const ::tflite::Model* input_model_ = nullptr;
diff --git a/tensorflow/lite/toco/tflite/op_version.cc b/tensorflow/lite/toco/tflite/op_version.cc
index cf127a9f459..efa53c69cae 100644
--- a/tensorflow/lite/toco/tflite/op_version.cc
+++ b/tensorflow/lite/toco/tflite/op_version.cc
@@ -29,7 +29,7 @@ namespace tflite {
 
 // Deprecated and please register new ops/versions in
 // tflite/tools/versioning/op_version.cc".
-string GetMinimumRuntimeVersionForModel(const Model& model) {
+std::string GetMinimumRuntimeVersionForModel(const Model& model) {
   // Use this as the placeholder string if a particular op is not yet included
   // in any Tensorflow's RC/Final release source package. Once that op is
   // included in the release, please update this with the real version string.
@@ -37,8 +37,8 @@ string GetMinimumRuntimeVersionForModel(const Model& model) {
   // A map from the version key of an op to its minimum runtime version.
   // For example, {{kAveragePool, 1}, "1.5.0"},  means the 1st version of
   // AveragePool requires a minimum TF Lite runtime version '1.5.0`.
-  static const std::map, string>* op_version_map =
-      new std::map, string>({
+  static const std::map, std::string>*
+      op_version_map = new std::map, std::string>({
           {{OperatorType::kAveragePool, 1}, "1.5.0"},
           {{OperatorType::kAveragePool, 2}, "1.14.0"},
           {{OperatorType::kAveragePool, 3}, kPendingReleaseOpVersion},
@@ -253,7 +253,7 @@ string GetMinimumRuntimeVersionForModel(const Model& model) {
       tflite::BuildOperatorByTypeMap(false /*enable_select_tf_ops=*/);
   OperatorSignature op_signature;
   op_signature.model = &model;
-  string model_min_version;
+  std::string model_min_version;
   for (const auto& op : model.operators) {
     if (op_types_map.find(op->type) == op_types_map.end()) continue;
     op_signature.op = op.get();
diff --git a/tensorflow/lite/toco/tflite/op_version_test.cc b/tensorflow/lite/toco/tflite/op_version_test.cc
index 14b086471b7..8466fc35ad7 100644
--- a/tensorflow/lite/toco/tflite/op_version_test.cc
+++ b/tensorflow/lite/toco/tflite/op_version_test.cc
@@ -27,9 +27,9 @@ TEST(OpVersionTest, MinimumVersionForSameOpVersions) {
   Model model;
   // Float convolutional kernel is introduced since '1.5.0'.
   std::unique_ptr conv(new ConvOperator());
-  const string conv_input = "conv_input";
-  const string conv_filter = "conv_filter";
-  const string conv_output = "conv_output";
+  const std::string conv_input = "conv_input";
+  const std::string conv_filter = "conv_filter";
+  const std::string conv_output = "conv_output";
   conv->inputs.push_back(conv_input);
   conv->inputs.push_back(conv_filter);
   conv->outputs.push_back(conv_output);
@@ -44,8 +44,8 @@ TEST(OpVersionTest, MinimumVersionForSameOpVersions) {
 
   // Float softmax kernel is introduced since '1.5.0'.
   std::unique_ptr softmax(new SoftmaxOperator());
-  const string softmax_input = "softmax_input";
-  const string softmax_output = "softmax_output";
+  const std::string softmax_input = "softmax_input";
+  const std::string softmax_output = "softmax_output";
   softmax->inputs.push_back(softmax_input);
   softmax->outputs.push_back(softmax_output);
   array_map[softmax_input] = std::unique_ptr(new Array);
@@ -60,9 +60,9 @@ TEST(OpVersionTest, MinimumVersionForMultipleOpVersions) {
   Model model;
   // Dilated DepthWiseConvolution is introduced since '1.12.0'.
   std::unique_ptr conv(new DepthwiseConvOperator());
-  const string conv_input = "conv_input";
-  const string conv_filter = "conv_filter";
-  const string conv_output = "conv_output";
+  const std::string conv_input = "conv_input";
+  const std::string conv_filter = "conv_filter";
+  const std::string conv_output = "conv_output";
   conv->inputs.push_back(conv_input);
   conv->inputs.push_back(conv_filter);
   conv->outputs.push_back(conv_output);
@@ -77,10 +77,10 @@ TEST(OpVersionTest, MinimumVersionForMultipleOpVersions) {
   // FullyConnected op with kShuffled4x16Int8 weight format is introduced from
   // '1.10.0'.
   std::unique_ptr fc(new FullyConnectedOperator());
-  const string fc_input = "fc_input";
-  const string fc_weights = "fc_weights";
-  const string fc_bias = "fc_bias";
-  const string fc_output = "fc_output";
+  const std::string fc_input = "fc_input";
+  const std::string fc_weights = "fc_weights";
+  const std::string fc_bias = "fc_bias";
+  const std::string fc_output = "fc_output";
   fc->inputs.push_back(fc_input);
   fc->inputs.push_back(fc_weights);
   fc->inputs.push_back(fc_bias);
@@ -121,10 +121,10 @@ TEST(OpVersionTest, MinimumVersionForMixedOpVersions) {
   // FullyConnected op with kShuffled4x16Int8 weight format is introduced from
   // '1.10.0'.
   std::unique_ptr fc(new FullyConnectedOperator());
-  const string fc_input = "fc_input";
-  const string fc_weights = "fc_weights";
-  const string fc_bias = "fc_bias";
-  const string fc_output = "fc_output";
+  const std::string fc_input = "fc_input";
+  const std::string fc_weights = "fc_weights";
+  const std::string fc_bias = "fc_bias";
+  const std::string fc_output = "fc_output";
   fc->inputs.push_back(fc_input);
   fc->inputs.push_back(fc_weights);
   fc->inputs.push_back(fc_bias);
diff --git a/tensorflow/lite/toco/tflite/operator.cc b/tensorflow/lite/toco/tflite/operator.cc
index fee10a19787..be539cf6054 100644
--- a/tensorflow/lite/toco/tflite/operator.cc
+++ b/tensorflow/lite/toco/tflite/operator.cc
@@ -238,7 +238,7 @@ class SpaceToBatchND
                    TocoOperator* op) const override {}
 
   int GetVersion(const OperatorSignature& op_signature) const override {
-    const string& input_name = op_signature.op->inputs[0];
+    const std::string& input_name = op_signature.op->inputs[0];
     const Array& input_array = op_signature.model->GetArray(input_name);
     ::tflite::OpSignature op_sig =
         GetVersioningOpSig(builtin_op(), op_signature);
@@ -268,8 +268,8 @@ class Sub : public BuiltinOperatorinputs[0];
-    const string& input2_name = op_signature.op->inputs[1];
+    const std::string& input1_name = op_signature.op->inputs[0];
+    const std::string& input2_name = op_signature.op->inputs[1];
     const Array& input1_array = op_signature.model->GetArray(input1_name);
     const Array& input2_array = op_signature.model->GetArray(input2_name);
     ::tflite::OpSignature op_sig =
@@ -305,8 +305,8 @@ class Div : public BuiltinOperatorinputs[0];
-    const string& input2_name = op_signature.op->inputs[1];
+    const std::string& input1_name = op_signature.op->inputs[0];
+    const std::string& input2_name = op_signature.op->inputs[1];
     const Array& input1_array = op_signature.model->GetArray(input1_name);
     const Array& input2_array = op_signature.model->GetArray(input2_name);
     ::tflite::OpSignature op_sig =
@@ -339,7 +339,7 @@ class BatchToSpaceND
                    TocoOperator* op) const override {}
 
   int GetVersion(const OperatorSignature& op_signature) const override {
-    const string& input_name = op_signature.op->inputs[0];
+    const std::string& input_name = op_signature.op->inputs[0];
     const Array& input_array = op_signature.model->GetArray(input_name);
     ::tflite::OpSignature op_sig =
         GetVersioningOpSig(builtin_op(), op_signature);
@@ -662,9 +662,9 @@ class Mul : public BuiltinOperatorinputs[0];
-    const string& input2_name = op_signature.op->inputs[1];
-    const string& output_name = op_signature.op->outputs[0];
+    const std::string& input1_name = op_signature.op->inputs[0];
+    const std::string& input2_name = op_signature.op->inputs[1];
+    const std::string& output_name = op_signature.op->outputs[0];
     const Array& input1_array = op_signature.model->GetArray(input1_name);
     const Array& input2_array = op_signature.model->GetArray(input2_name);
     const Array& output_array = op_signature.model->GetArray(output_name);
@@ -1440,7 +1440,7 @@ class Unpack : public BuiltinOperatorinputs[0];
+    const std::string& input_name = op_signature.op->inputs[0];
     const Array& input_array = op_signature.model->GetArray(input_name);
     // If the op take int8/uint8 input, it is version 2.
     if (input_array.data_type == ArrayDataType::kInt8 ||
@@ -1577,7 +1577,7 @@ class Where : public BuiltinOperator WriteFlexOpOptions(
-    const string& tensorflow_node_def) {
+    const std::string& tensorflow_node_def) {
   auto fbb = absl::make_unique();
 
   ::tensorflow::NodeDef node_def;
@@ -1597,7 +1597,7 @@ std::unique_ptr WriteFlexOpOptions(
 
 class TensorFlowUnsupported : public BaseOperator {
  public:
-  TensorFlowUnsupported(const string& name, OperatorType type,
+  TensorFlowUnsupported(const std::string& name, OperatorType type,
                         bool enable_select_tf_ops)
       : BaseOperator(name, type), enable_select_tf_ops_(enable_select_tf_ops) {}
 
@@ -1676,7 +1676,7 @@ class TensorFlowUnsupported : public BaseOperator {
         case tensorflow::AttrValue::kList:
           if (attr.list().s_size() > 0) {
             auto start = fbb->StartVector(key);
-            for (const string& v : attr.list().s()) {
+            for (const std::string& v : attr.list().s()) {
               fbb->Add(v);
             }
             fbb->EndVector(start, /*typed=*/true, /*fixed=*/false);
@@ -1736,10 +1736,11 @@ class TensorFlowUnsupported : public BaseOperator {
           break;
         case flexbuffers::FBT_BOOL:
           (*attr)[key].set_b(value.AsBool());
-          if (string(key) == "_output_quantized") {
+          if (std::string(key) == "_output_quantized") {
             op->quantized = value.AsBool();
           }
-          if (string(key) == "_support_output_type_float_in_quantized_op") {
+          if (std::string(key) ==
+              "_support_output_type_float_in_quantized_op") {
             op->support_output_type_float_in_quantized_op = value.AsBool();
           }
           break;
@@ -2095,9 +2096,9 @@ std::map> BuildOperatorByTypeMap(
   return result;
 }
 
-std::map> BuildOperatorByNameMap(
+std::map> BuildOperatorByNameMap(
     bool enable_select_tf_ops) {
-  std::map> result;
+  std::map> result;
 
   std::vector> ops =
       BuildOperatorList(enable_select_tf_ops);
@@ -2109,7 +2110,7 @@ std::map> BuildOperatorByNameMap(
 }
 
 bool ShouldExportAsFlexOp(bool enable_select_tf_ops,
-                          const string& tensorflow_op_name) {
+                          const std::string& tensorflow_op_name) {
   // If Flex ops aren't allow at all, simply return false.
   if (!enable_select_tf_ops) {
     return false;
diff --git a/tensorflow/lite/toco/tflite/operator.h b/tensorflow/lite/toco/tflite/operator.h
index 19d92145e0c..fb79b97f46e 100644
--- a/tensorflow/lite/toco/tflite/operator.h
+++ b/tensorflow/lite/toco/tflite/operator.h
@@ -30,7 +30,7 @@ class BaseOperator;
 // Return a map contained all know TF Lite Operators, keyed by their names.
 // TODO(ycling): The pattern to propagate parameters (e.g. enable_select_tf_ops)
 // is ugly here. Consider refactoring.
-std::map> BuildOperatorByNameMap(
+std::map> BuildOperatorByNameMap(
     bool enable_select_tf_ops = false);
 
 // Return a map contained all know TF Lite Operators, keyed by the type of
@@ -41,7 +41,7 @@ std::map> BuildOperatorByTypeMap(
 // Write the custom option FlexBuffer with a serialized TensorFlow NodeDef
 // for a Flex op.
 std::unique_ptr WriteFlexOpOptions(
-    const string& tensorflow_node_def);
+    const std::string& tensorflow_node_def);
 
 // These are the flatbuffer types for custom and builtin options.
 using CustomOptions = flatbuffers::Vector;
@@ -71,11 +71,11 @@ struct Options {
 class BaseOperator {
  public:
   // Build an operator with the given TF Lite name and tf.mini type.
-  BaseOperator(const string& name, OperatorType type)
+  BaseOperator(const std::string& name, OperatorType type)
       : name_(name), type_(type) {}
   virtual ~BaseOperator() = default;
 
-  string name() const { return name_; }
+  std::string name() const { return name_; }
   OperatorType type() const { return type_; }
 
   // Given a tf.mini operator, create the corresponding flatbuffer options and
@@ -111,7 +111,7 @@ class BaseOperator {
   }
 
  private:
-  string name_;
+  std::string name_;
   OperatorType type_;
 };
 
@@ -123,7 +123,7 @@ class BaseOperator {
 // Helper function to determine if a unsupported TensorFlow op should be
 // exported as an Flex op or a regular custom op.
 bool ShouldExportAsFlexOp(bool enable_select_tf_ops,
-                          const string& tensorflow_op_name);
+                          const std::string& tensorflow_op_name);
 
 }  // namespace tflite
 
diff --git a/tensorflow/lite/toco/tflite/operator_test.cc b/tensorflow/lite/toco/tflite/operator_test.cc
index a4fe01e4afd..cb466fef079 100644
--- a/tensorflow/lite/toco/tflite/operator_test.cc
+++ b/tensorflow/lite/toco/tflite/operator_test.cc
@@ -30,8 +30,8 @@ namespace {
 class OperatorTest : public ::testing::Test {
  protected:
   // Return the operator for the given name and type.
-  const BaseOperator& GetOperator(const string& name, OperatorType type) {
-    using OpsByName = std::map>;
+  const BaseOperator& GetOperator(const std::string& name, OperatorType type) {
+    using OpsByName = std::map>;
     using OpsByType = std::map>;
 
     static auto* by_name = new OpsByName(BuildOperatorByNameMap());
@@ -86,7 +86,7 @@ class OperatorTest : public ::testing::Test {
   // Verify serialization and deserialization of simple operators (those
   // that don't have any configuration parameters).
   template 
-  void CheckSimpleOperator(const string& name, OperatorType type) {
+  void CheckSimpleOperator(const std::string& name, OperatorType type) {
     Options options;
     auto output_toco_op =
         SerializeAndDeserialize(GetOperator(name, type), T(), &options);
@@ -99,7 +99,7 @@ class OperatorTest : public ::testing::Test {
   }
 
   template 
-  void CheckReducerOperator(const string& name, OperatorType type) {
+  void CheckReducerOperator(const std::string& name, OperatorType type) {
     T op;
 
     op.keep_dims = false;
diff --git a/tensorflow/lite/toco/tflite/types.cc b/tensorflow/lite/toco/tflite/types.cc
index 96cad557baf..9d4ab8434d1 100644
--- a/tensorflow/lite/toco/tflite/types.cc
+++ b/tensorflow/lite/toco/tflite/types.cc
@@ -25,7 +25,7 @@ DataBuffer::FlatBufferOffset CopyStringToBuffer(
     const Array& array, flatbuffers::FlatBufferBuilder* builder) {
   const auto& src_data = array.GetBuffer().data;
   ::tflite::DynamicBuffer dyn_buffer;
-  for (const string& str : src_data) {
+  for (const std::string& str : src_data) {
     dyn_buffer.AddString(str.c_str(), str.length());
   }
   char* tensor_buffer;
@@ -58,12 +58,12 @@ DataBuffer::FlatBufferOffset CopyBuffer(
 
 void CopyStringFromBuffer(const ::tflite::Buffer& buffer, Array* array) {
   auto* src_data = reinterpret_cast(buffer.data()->data());
-  std::vector* dst_data =
+  std::vector* dst_data =
       &array->GetMutableBuffer().data;
   int32_t num_strings = ::tflite::GetStringCount(src_data);
   for (int i = 0; i < num_strings; i++) {
     ::tflite::StringRef str_ref = ::tflite::GetString(src_data, i);
-    string this_str(str_ref.str, str_ref.len);
+    std::string this_str(str_ref.str, str_ref.len);
     dst_data->push_back(this_str);
   }
 }

From 642ad434d8315561ef9cc02cc9157436fe9c0f72 Mon Sep 17 00:00:00 2001
From: Jingyue Wu 
Date: Fri, 19 Jun 2020 09:50:56 -0700
Subject: [PATCH 0618/1390] Fix exports_files.

cl/317237033 replaced cwise_op_neg.cc with cwise_op_neg_1.cc and
cwise_op_neg_2.cc.

PiperOrigin-RevId: 317325461
Change-Id: Ib44ff36474b7e55d9e84ff737ea82b9dac46b9f9
---
 tensorflow/core/kernels/BUILD | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 0b7a092033b..e2ff5aed283 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -8802,7 +8802,8 @@ exports_files([
     "cwise_op_mod.cc",
     "cwise_op_mul_1.cc",
     "cwise_op_mul_2.cc",
-    "cwise_op_neg.cc",
+    "cwise_op_neg_1.cc",
+    "cwise_op_neg_2.cc",
     "cwise_op_not_equal_to_1.cc",
     "cwise_op_not_equal_to_2.cc",
     "cwise_op_round.cc",

From c3cc3c40b08a37535a281e1d7a5fd7d3d802aac6 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Fri, 19 Jun 2020 09:55:49 -0700
Subject: [PATCH 0619/1390] Move fuzzers for TF ops to own subdir. Trim some
 dependencies.

This duplicates some of the BUILD dependency tree to go around the need to link huge bottleneck dependencies (such as `//tensorflow/core:framework`). Until TF can use `cc_shared_library` in a stable way (and all support in Bazel exists), we will need to use the duplicated tree for fuzzing.

PiperOrigin-RevId: 317326319
Change-Id: I1493e3ae7340298971fe15bd3702b63657f9bf9f
---
 tensorflow/core/framework/BUILD               |   1 +
 tensorflow/security/fuzzing/BUILD             |  14 --
 tensorflow/security/fuzzing/op_fuzzing/BUILD  |  39 +++++
 .../fuzzing/op_fuzzing/fuzz_session.h         | 156 ++++++++++++++++++
 .../fuzzing/{ => op_fuzzing}/identity_fuzz.cc |   2 +-
 5 files changed, 197 insertions(+), 15 deletions(-)
 create mode 100644 tensorflow/security/fuzzing/op_fuzzing/BUILD
 create mode 100644 tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h
 rename tensorflow/security/fuzzing/{ => op_fuzzing}/identity_fuzz.cc (95%)

diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 52f15dcb5c2..d47c74a629d 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -719,6 +719,7 @@ tf_cuda_library(
     visibility = [
         "//tensorflow/core:__pkg__",
         "//tensorflow/core/util:__pkg__",
+        "//tensorflow/security/fuzzing:__subpackages__",
     ],
     deps = [
         ":allocation_description_proto_cc",
diff --git a/tensorflow/security/fuzzing/BUILD b/tensorflow/security/fuzzing/BUILD
index 9b5aeec2d36..6b6c8275275 100644
--- a/tensorflow/security/fuzzing/BUILD
+++ b/tensorflow/security/fuzzing/BUILD
@@ -18,17 +18,3 @@ tf_fuzz_target(
         "//tensorflow/core/platform:status",
     ],
 )
-
-# A trivial fuzzer with no pre-specified corpus.
-# TODO(mihaimaruseac): Move fuzz_session and the op fuzzers to a subdirectory
-tf_fuzz_target(
-    name = "identity_fuzz",
-    srcs = ["identity_fuzz.cc"],
-    deps = [
-        "//tensorflow/cc:cc_ops",
-        "//tensorflow/core/kernels/fuzzing:fuzz_session",
-        # Needed only to transitiviely link dependencies
-        "//tensorflow/cc:scope",
-        "//tensorflow/core:core_cpu",
-    ],
-)
diff --git a/tensorflow/security/fuzzing/op_fuzzing/BUILD b/tensorflow/security/fuzzing/op_fuzzing/BUILD
new file mode 100644
index 00000000000..aacd2f16cc4
--- /dev/null
+++ b/tensorflow/security/fuzzing/op_fuzzing/BUILD
@@ -0,0 +1,39 @@
+# Fuzzing TensorFlow ops..
+# Most ops have a similar set of dependencies and a similar fuzzing
+# infrastructure. Hence, we gather everything in one single place.
+# Note that these fuzzers cover a large part of TF, they are not granular.
+
+load(
+    "//tensorflow/security/fuzzing:tf_fuzzing.bzl",
+    "tf_fuzz_target",
+)
+
+package(
+    licenses = ["notice"],  # Apache 2.0
+)
+
+# Since all ops need to have a graph created before being fuzzed, we define
+# this header-only library to handle the needed plumbing.
+cc_library(
+    name = "fuzz_session",
+    hdrs = ["fuzz_session.h"],
+    deps = [
+        "//tensorflow/cc:scope",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core/common_runtime:direct_session_internal",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/platform:status",
+    ],
+)
+
+# A trivial fuzzer with no pre-specified corpus.
+tf_fuzz_target(
+    name = "identity_fuzz",
+    srcs = ["identity_fuzz.cc"],
+    deps = [
+        ":fuzz_session",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/core/kernels:array",
+    ],
+)
diff --git a/tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h b/tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h
new file mode 100644
index 00000000000..575212b3b86
--- /dev/null
+++ b/tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h
@@ -0,0 +1,156 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_
+#define TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_
+
+#include 
+#include 
+#include 
+#include 
+
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+// Standard invoking function macro to dispatch to a fuzzer class.
+#define STANDARD_TF_FUZZ_FUNCTION(FuzzerClass)                              \
+  extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) { \
+    static FuzzerClass* fuzzer = new FuzzerClass();                         \
+    return fuzzer->Fuzz(data, size);                                        \
+  }
+
+// Standard builder for hooking one placeholder to one op.
+#define SINGLE_INPUT_OP_BUILDER(dtype, opName)                          \
+  void BuildGraph(const Scope& scope) override {                        \
+    auto op_node =                                                      \
+        tensorflow::ops::Placeholder(scope.WithOpName("input"), dtype); \
+    (void)tensorflow::ops::opName(scope.WithOpName("output"), op_node); \
+  }
+
+namespace tensorflow {
+namespace fuzzing {
+
+// Create a TensorFlow session using a specific GraphDef created
+// by BuildGraph(), and make it available for fuzzing.
+// Users must override BuildGraph and FuzzImpl to specify
+// (1) which operations are being fuzzed; and
+// (2) How to translate the uint8_t* buffer from the fuzzer
+//     to a Tensor or Tensors that are semantically appropriate
+//     for the op under test.
+// For the simple cases of testing a single op that takes a single
+// input Tensor, use the SINGLE_INPUT_OP_BUILDER(dtype, opName) macro in place
+// of defining BuildGraphDef.
+//
+// Typical use:
+// class FooFuzzer : public FuzzSession {
+//   SINGLE_INPUT_OP_BUILDER(DT_INT8, Identity);
+//   void FuzzImpl(const uint8_t* data, size_t size) {
+//      ... convert data and size to a Tensor, pass it to:
+//      RunInputs({{"input", input_tensor}});
+//
+class FuzzSession {
+ public:
+  FuzzSession() : initialized_(false) {}
+  virtual ~FuzzSession() {}
+
+  // Constructs a Graph using the supplied Scope.
+  // By convention, the graph should have inputs named "input1", ...
+  // "inputN", and one output node, named "output".
+  // Users of FuzzSession should override this method to create their graph.
+  virtual void BuildGraph(const Scope& scope) = 0;
+
+  // Implements the logic that converts an opaque byte buffer
+  // from the fuzzer to Tensor inputs to the graph.  Users must override.
+  virtual void FuzzImpl(const uint8_t* data, size_t size) = 0;
+
+  // Initializes the FuzzSession.  Not safe for multithreading.
+  // Separate init function because the call to virtual BuildGraphDef
+  // can't be put into the constructor.
+  Status InitIfNeeded() {
+    if (initialized_) {
+      return Status::OK();
+    }
+    initialized_ = true;
+
+    Scope root = Scope::DisabledShapeInferenceScope().ExitOnError();
+    SessionOptions options;
+    session_ = std::unique_ptr(NewSession(options));
+
+    BuildGraph(root);
+
+    GraphDef graph_def;
+    TF_CHECK_OK(root.ToGraphDef(&graph_def));
+
+    Status status = session_->Create(graph_def);
+    if (!status.ok()) {
+      // This is FATAL, because this code is designed to fuzz an op
+      // within a session.  Failure to create the session means we
+      // can't send any data to the op.
+      LOG(FATAL) << "Could not create session: " << status.error_message();
+    }
+    return status;
+  }
+
+  // Runs the TF session by pulling on the "output" node, attaching
+  // the supplied input_tensor to the input node(s), and discarding
+  // any returned output.
+  // Note: We are ignoring Status from Run here since fuzzers don't need to
+  // check it (as that will slow them down and printing/logging is useless).
+  void RunInputs(const std::vector >& inputs) {
+    RunInputsWithStatus(inputs).IgnoreError();
+  }
+
+  // Same as RunInputs but don't ignore status
+  Status RunInputsWithStatus(
+      const std::vector >& inputs) {
+    return session_->Run(inputs, {}, {"output"}, nullptr);
+  }
+
+  // Dispatches to FuzzImpl;  small amount of sugar to keep the code
+  // of the per-op fuzzers tiny.
+  int Fuzz(const uint8_t* data, size_t size) {
+    Status status = InitIfNeeded();
+    TF_CHECK_OK(status) << "Fuzzer graph initialization failed: "
+                        << status.error_message();
+    // No return value from fuzzing:  Success is defined as "did not
+    // crash".  The actual application results are irrelevant.
+    FuzzImpl(data, size);
+    return 0;
+  }
+
+ private:
+  bool initialized_;
+  std::unique_ptr session_;
+};
+
+// A specialized fuzz implementation for ops that take
+// a single string.  Caller must still define the op
+// to plumb by overriding BuildGraph or using
+// a plumbing macro.
+class FuzzStringInputOp : public FuzzSession {
+  void FuzzImpl(const uint8_t* data, size_t size) final {
+    Tensor input_tensor(tensorflow::DT_STRING, TensorShape({}));
+    input_tensor.scalar()() =
+        string(reinterpret_cast(data), size);
+    RunInputs({{"input", input_tensor}});
+  }
+};
+
+}  // end namespace fuzzing
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_SECURITY_FUZZING_OP_FUZZING_FUZZ_SESSION_H_
diff --git a/tensorflow/security/fuzzing/identity_fuzz.cc b/tensorflow/security/fuzzing/op_fuzzing/identity_fuzz.cc
similarity index 95%
rename from tensorflow/security/fuzzing/identity_fuzz.cc
rename to tensorflow/security/fuzzing/op_fuzzing/identity_fuzz.cc
index 4c1049d381b..a63c35b45e2 100644
--- a/tensorflow/security/fuzzing/identity_fuzz.cc
+++ b/tensorflow/security/fuzzing/op_fuzzing/identity_fuzz.cc
@@ -14,7 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/cc/ops/standard_ops.h"
-#include "tensorflow/core/kernels/fuzzing/fuzz_session.h"
+#include "tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h"
 
 namespace tensorflow {
 namespace fuzzing {

From 8e305b9cec35ccf9821c2fd2a82194e328d98704 Mon Sep 17 00:00:00 2001
From: Denisa Roberts 
Date: Wed, 6 May 2020 11:29:37 -0400
Subject: [PATCH 0620/1390] Add Qr Grad for wide matrices

---
 tensorflow/python/kernel_tests/qr_op_test.py |  5 +--
 tensorflow/python/ops/linalg_grad.py         | 47 +++++++++++++++-----
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/tensorflow/python/kernel_tests/qr_op_test.py b/tensorflow/python/kernel_tests/qr_op_test.py
index d5337c183a6..0c291dbd940 100644
--- a/tensorflow/python/kernel_tests/qr_op_test.py
+++ b/tensorflow/python/kernel_tests/qr_op_test.py
@@ -278,14 +278,13 @@ if __name__ == "__main__":
                                     use_static_shape))
 
   # TODO(pfau): Get working with complex types.
-  # TODO(pfau): Get working with full_matrices when rows != cols
-  # TODO(pfau): Get working when rows < cols
+  # TODO(pfau): Get working with full_matrices when rows > cols
   # TODO(pfau): Get working with shapeholders (dynamic shapes)
   for full_matrices in False, True:
     for dtype in np.float32, np.float64:
       for rows in 1, 2, 5, 10:
         for cols in 1, 2, 5, 10:
-          if rows == cols or (not full_matrices and rows > cols):
+          if rows <= cols or (not full_matrices and rows > cols):
             for batch_dims in [(), (3,)] + [(3, 2)] * (max(rows, cols) < 10):
               shape = batch_dims + (rows, cols)
               name = "%s_%s_full_%s" % (dtype.__name__,
diff --git a/tensorflow/python/ops/linalg_grad.py b/tensorflow/python/ops/linalg_grad.py
index 437e28e7e6b..5ec372430ba 100644
--- a/tensorflow/python/ops/linalg_grad.py
+++ b/tensorflow/python/ops/linalg_grad.py
@@ -493,15 +493,10 @@ def _QrGrad(op, dq, dr):
   if (r.shape.ndims is None or r.shape.as_list()[-2] is None or
       r.shape.as_list()[-1] is None):
     raise NotImplementedError("QrGrad not implemented with dynamic shapes.")
-  if r.shape.dims[-2].value != r.shape.dims[-1].value:
+  if (r.shape.dims[-2].value > r.shape.dims[-1].value and
+      q.shape.dims[-2].value == q.shape.dims[-1].value):
     raise NotImplementedError("QrGrad not implemented when ncols > nrows "
-                              "or full_matrices is true and ncols != nrows.")
-
-  qdq = math_ops.matmul(q, dq, adjoint_a=True)
-  qdq_ = qdq - _linalg.adjoint(qdq)
-  rdr = math_ops.matmul(r, dr, adjoint_b=True)
-  rdr_ = rdr - _linalg.adjoint(rdr)
-  tril = array_ops.matrix_band_part(qdq_ + rdr_, -1, 0)
+                              "and full_matrices is true.")
 
   def _TriangularSolve(x, r):
     """Equiv to matmul(x, adjoint(matrix_inverse(r))) if r is upper-tri."""
@@ -509,9 +504,39 @@ def _QrGrad(op, dq, dr):
         linalg_ops.matrix_triangular_solve(
             r, _linalg.adjoint(x), lower=False, adjoint=False))
 
-  grad_a = math_ops.matmul(q, dr + _TriangularSolve(tril, r))
-  grad_b = _TriangularSolve(dq - math_ops.matmul(q, qdq), r)
-  return grad_a + grad_b
+  def _QrGradSquareAndDeepMatrices(q, r, dq, dr):
+    """Gradient for matrix orders num_rows >= num_cols
+    and full_matrices is false.
+    """
+    qdq = math_ops.matmul(q, dq, adjoint_a=True)
+    qdq_ = qdq - _linalg.adjoint(qdq)
+    rdr = math_ops.matmul(r, dr, adjoint_b=True)
+    rdr_ = rdr - _linalg.adjoint(rdr)
+    tril = array_ops.matrix_band_part(qdq_ + rdr_, -1, 0)
+
+    grad_a = math_ops.matmul(q, dr + _TriangularSolve(tril, r))
+    grad_b = _TriangularSolve(dq - math_ops.matmul(q, qdq), r)
+    return grad_a + grad_b
+
+  num_rows, num_cols = q.shape.dims[-2].value, r.shape.dims[-1]
+
+  if num_rows >= num_cols:
+    return _QrGradSquareAndDeepMatrices(q, r, dq, dr)
+
+  # Partition a = [x, y], r = [u, v] and reduce to the square case
+  a = op.inputs[0]
+  y = a[..., :, num_rows:]
+  u = r[..., :, :num_rows]
+  dv = dr[..., :, num_rows:]
+  du = dr[..., :, :num_rows]
+  dy = math_ops.matmul(q, dv)
+  dx = _QrGradSquareAndDeepMatrices(q,
+                                    u,
+                                    dq + math_ops.matmul(y,
+                                                         dv,
+                                                         adjoint_b=True),
+                                    du)
+  return array_ops.concat([dx, dy], axis=-1)
 
 
 @ops.RegisterGradient("MatrixSolve")

From 390b259dfc8407484533c1cb61ca3515ed5166e8 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Fri, 19 Jun 2020 10:16:47 -0700
Subject: [PATCH 0621/1390] Fix unused variable issue with fuzzing methods

---
 tensorflow/security/fuzzing/status_group_fuzz.cc | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index bc80cd72bc9..989e1c9d1cb 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -54,9 +54,10 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     }
   }
 
-  sg.as_summary_status();
-  sg.as_concatenated_status();
-  sg.AttachLogMessages();
+  // Ignore warnings that these values are unused
+  sg.as_summary_status().IgnoreError();;
+  sg.as_concatenated_status().IgnoreError();;
+  sg.AttachLogMessages().IgnoreError();;
 
   return 0;
 }

From 40b9713f6459abd043248c55b1b06ebf60712961 Mon Sep 17 00:00:00 2001
From: Gabriel Rasskin 
Date: Fri, 19 Jun 2020 10:17:22 -0700
Subject: [PATCH 0622/1390] Doubling syntax

---
 tensorflow/security/fuzzing/status_group_fuzz.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index 989e1c9d1cb..002785734bb 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -55,9 +55,9 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   }
 
   // Ignore warnings that these values are unused
-  sg.as_summary_status().IgnoreError();;
-  sg.as_concatenated_status().IgnoreError();;
-  sg.AttachLogMessages().IgnoreError();;
+  sg.as_summary_status().IgnoreError();
+  sg.as_concatenated_status().IgnoreError();
+  sg.AttachLogMessages().IgnoreError();
 
   return 0;
 }

From 4d751f9da4914d6e1fc7aafe2f5f27e0f96830d6 Mon Sep 17 00:00:00 2001
From: Katherine Wu 
Date: Fri, 19 Jun 2020 10:12:17 -0700
Subject: [PATCH 0623/1390] Roll forward of cl/316247127: Make sure compiled
 metrics are accessible after loading from H5 or SavedModel.

PiperOrigin-RevId: 317329395
Change-Id: I33578515f36aa0ba227e75bda52966d493a4bebb
---
 .../python/keras/engine/compile_utils.py      |  14 +--
 tensorflow/python/keras/engine/functional.py  |  10 +-
 tensorflow/python/keras/engine/training.py    |   7 +-
 tensorflow/python/keras/metrics.py            |   5 +-
 tensorflow/python/keras/saving/hdf5_format.py |   1 +
 .../python/keras/saving/hdf5_format_test.py   | 114 +++++++-----------
 .../python/keras/saving/saved_model/load.py   |   1 +
 .../python/keras/saving/saving_utils.py       |  14 +++
 8 files changed, 79 insertions(+), 87 deletions(-)

diff --git a/tensorflow/python/keras/engine/compile_utils.py b/tensorflow/python/keras/engine/compile_utils.py
index 3858639f024..ba7ce624090 100644
--- a/tensorflow/python/keras/engine/compile_utils.py
+++ b/tensorflow/python/keras/engine/compile_utils.py
@@ -37,7 +37,7 @@ class Container(object):
   def __init__(self, output_names=None):
     self._output_names = output_names
 
-  def _build(self, y_pred):
+  def build(self, y_pred):
     if self._output_names is None:
       # In Subclass API, output names like 'output_1' are used for
       # `Metric` names.
@@ -131,9 +131,9 @@ class LossesContainer(Container):
     ]
     return [self._loss_metric] + per_output_metrics
 
-  def _build(self, y_pred):
+  def build(self, y_pred):
     """One-time setup of loss objects."""
-    super(LossesContainer, self)._build(y_pred)
+    super(LossesContainer, self).build(y_pred)
 
     self._losses = self._maybe_broadcast_to_outputs(y_pred, self._losses)
     self._losses = self._conform_to_outputs(y_pred, self._losses)
@@ -184,7 +184,7 @@ class LossesContainer(Container):
     sample_weight = self._conform_to_outputs(y_pred, sample_weight)
 
     if not self._built:
-      self._build(y_pred)
+      self.build(y_pred)
 
     y_pred = nest.flatten(y_pred)
     y_true = nest.flatten(y_true)
@@ -295,9 +295,9 @@ class MetricsContainer(Container):
       return []
     return self._metrics_in_order
 
-  def _build(self, y_pred, y_true):
+  def build(self, y_pred, y_true):
     """One-time setup of metric objects."""
-    super(MetricsContainer, self)._build(y_pred)
+    super(MetricsContainer, self).build(y_pred)
 
     self._metrics = self._maybe_broadcast_to_outputs(y_pred, self._metrics)
     self._metrics = self._conform_to_outputs(y_pred, self._metrics)
@@ -385,7 +385,7 @@ class MetricsContainer(Container):
     sample_weight = self._conform_to_outputs(y_pred, sample_weight)
 
     if not self._built:
-      self._build(y_pred, y_true)
+      self.build(y_pred, y_true)
 
     y_pred = nest.flatten(y_pred)
     y_true = nest.flatten(y_true) if y_true is not None else []
diff --git a/tensorflow/python/keras/engine/functional.py b/tensorflow/python/keras/engine/functional.py
index 0ef4840b651..0612d70044d 100644
--- a/tensorflow/python/keras/engine/functional.py
+++ b/tensorflow/python/keras/engine/functional.py
@@ -1007,10 +1007,12 @@ def _map_subgraph_network(inputs, outputs):
 
 def _should_skip_first_node(layer):
   """Returns True if the first layer node should not be saved or loaded."""
-  # Networks start with a pre-existing node linking their input to output.
-  # For a sequential model, it is first created with _is_graph_network = False,
-  # we have to keep the _is_graph_network check here.
-  return isinstance(layer, Functional) and layer._is_graph_network
+  # Networks that are constructed with an Input layer/shape start with a
+  # pre-existing node linking their input to output. This node is excluded from
+  # the network config.
+  return (isinstance(layer, Functional) and
+          # Filter out Sequential models without an input shape.
+          isinstance(layer._layers[0], input_layer_module.InputLayer))
 
 
 def _deserialize_keras_tensors(kwargs, layer_map):
diff --git a/tensorflow/python/keras/engine/training.py b/tensorflow/python/keras/engine/training.py
index ccd184a8bc4..a0ebec4f95e 100644
--- a/tensorflow/python/keras/engine/training.py
+++ b/tensorflow/python/keras/engine/training.py
@@ -436,7 +436,6 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
                            'Instead, in order to instantiate and build your '
                            'model, `call` your model on real tensor data (of '
                            'the correct dtype).')
-
     super(Model, self).build(input_shape)
 
   def call(self, inputs, training=None, mask=None):
@@ -2417,6 +2416,12 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
 
     self._saved_model_inputs_spec = specs
 
+    # Store the input shapes
+    if (self.__class__.__name__ == 'Sequential' and
+        self._build_input_shape is None):
+      self._build_input_shape = nest.map_structure(
+          lambda x: None if x is None else x.shape, specs)
+
   def _assert_weights_created(self):
     """Asserts that all the weights for the model have been created.
 
diff --git a/tensorflow/python/keras/metrics.py b/tensorflow/python/keras/metrics.py
index a67755b9333..7f40423595b 100644
--- a/tensorflow/python/keras/metrics.py
+++ b/tensorflow/python/keras/metrics.py
@@ -630,10 +630,9 @@ class MeanMetricWrapper(Mean):
   def from_config(cls, config):
     # Note that while MeanMetricWrapper itself isn't public, objects of this
     # class may be created and added to the model by calling model.compile.
+    fn = config.pop('fn', None)
     if cls is MeanMetricWrapper:
-      fn = get(config.pop('fn'))
-      return cls(fn, **config)
-
+      return cls(get(fn), **config)
     return super(MeanMetricWrapper, cls).from_config(config)
 
 
diff --git a/tensorflow/python/keras/saving/hdf5_format.py b/tensorflow/python/keras/saving/hdf5_format.py
index 800d609fe99..01a5e12e4c6 100644
--- a/tensorflow/python/keras/saving/hdf5_format.py
+++ b/tensorflow/python/keras/saving/hdf5_format.py
@@ -192,6 +192,7 @@ def load_model_from_hdf5(filepath, custom_objects=None, compile=True):  # pylint
       # Compile model.
       model.compile(**saving_utils.compile_args_from_training_config(
           training_config, custom_objects))
+      saving_utils.try_build_compiled_arguments(model)
 
       # Set optimizer weights.
       if 'optimizer_weights' in f:
diff --git a/tensorflow/python/keras/saving/hdf5_format_test.py b/tensorflow/python/keras/saving/hdf5_format_test.py
index 757385a25ea..b079bf8cac8 100644
--- a/tensorflow/python/keras/saving/hdf5_format_test.py
+++ b/tensorflow/python/keras/saving/hdf5_format_test.py
@@ -26,10 +26,12 @@ from absl.testing import parameterized
 import numpy as np
 
 from tensorflow.python import keras
+from tensorflow.python import tf2
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
 from tensorflow.python.keras import combinations
 from tensorflow.python.keras import keras_parameterized
 from tensorflow.python.keras import optimizers
@@ -368,48 +370,54 @@ class TestWeightSavingAndLoading(test.TestCase, parameterized.TestCase):
 
 
 @keras_parameterized.run_with_all_saved_model_formats
-class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
+class TestWholeModelSaving(keras_parameterized.TestCase):
 
   def _save_model_dir(self, dirname='saved_model'):
     temp_dir = self.get_temp_dir()
     self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True)
     return os.path.join(temp_dir, dirname)
 
-  def _assert_same_weights(self, model, loaded_model,
-                           original_optimizer_has_iterations_variable=True):
-    """Checks that the loaded weighs are the same as the original weights.
+  def _assert_same_weights_and_metrics(self, model, loaded_model):
+    """Checks that the loaded weights and metrics are the same as the original.
 
     Args:
       model: original model
       loaded_model: loaded model
-      original_optimizer_has_iterations_variable: If the original optimizer
-        uses an iterations variable. The loaded model will have a v2
-        optimizer, which always contains an iterations variable. So when
-        comparing the weights, the first variable in the loaded optimizer
-        weights may need to be ignored.
     """
     self.assertAllClose(model.weights, loaded_model.weights)
+
     if loaded_model.optimizer:
       if testing_utils.get_save_format() == 'tf':
         # TODO(b/153110928): Keras TF format doesn't restore optimizer weights
         # currently.
         return
-      if original_optimizer_has_iterations_variable:
-        self.assertAllClose(model.optimizer.weights,
-                            loaded_model.optimizer.weights)
-      else:
-        self.assertAllClose(model.optimizer.weights,
-                            loaded_model.optimizer.weights[1:])
+      self.assertAllClose(model.optimizer.weights,
+                          loaded_model.optimizer.weights)
 
-  def test_sequential_model_saving(self):
+    # In V1/Graph mode, the model isn't built, so the metrics are not loaded
+    # immediately (requires model to be called on some data before building
+    # metrics).
+    check_metrics = tf2.enabled() and context.executing_eagerly()
+
+    if check_metrics:
+      self.assertAllEqual([m.name for m in model.metrics],
+                          [m.name for m in loaded_model.metrics])
+
+  @keras_parameterized.run_with_all_model_types
+  @keras_parameterized.run_all_keras_modes
+  def test_save_and_load(self):
     saved_model_dir = self._save_model_dir()
     save_format = testing_utils.get_save_format()
 
+    if save_format == 'h5' and testing_utils.get_model_type() == 'subclass':
+      return  # HDF5 format currently does not allow saving classed models.
+
     with self.cached_session():
-      model = keras.models.Sequential()
-      model.add(keras.layers.Dense(2, input_shape=(3,)))
-      model.add(keras.layers.RepeatVector(3))
-      model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
+      model = testing_utils.get_model_from_layers(
+          [keras.layers.Dense(2),
+           keras.layers.RepeatVector(3),
+           keras.layers.TimeDistributed(keras.layers.Dense(3))],
+          input_shape=(3,))
       model.compile(
           loss=keras.losses.MSE,
           optimizer=keras.optimizer_v2.rmsprop.RMSprop(lr=0.0001),
@@ -432,43 +440,35 @@ class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
       out = model.predict(x)
       keras.models.save_model(model, saved_model_dir, save_format=save_format)
 
-      new_model = keras.models.load_model(saved_model_dir)
-      self._assert_same_weights(model, new_model)
+      loaded_model = keras.models.load_model(saved_model_dir)
+      self._assert_same_weights_and_metrics(model, loaded_model)
 
-      out2 = new_model.predict(x)
+      out2 = loaded_model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
 
-      # test that new updates are the same with both models
-      model.train_on_batch(x, y)
-      new_model.train_on_batch(x, y)
-
       eval_out = model.evaluate(x, y)
-      eval_out2 = new_model.evaluate(x, y)
+      eval_out2 = loaded_model.evaluate(x, y)
       self.assertArrayNear(eval_out, eval_out2, 0.001)
 
-      out = model.predict(x)
-      out2 = new_model.predict(x)
-      # The model has been trained on two batches. So the tolerance is larger.
-      self.assertAllClose(out, out2, atol=0.01)
-
+  @test_util.run_in_graph_and_eager_modes
   def test_sequential_model_saving_without_input_shape(self):
     saved_model_dir = self._save_model_dir()
     save_format = testing_utils.get_save_format()
-    with ops.Graph().as_default(), self.cached_session():
+    with self.cached_session():
       model = keras.models.Sequential()
       model.add(keras.layers.Dense(2))
       model.add(keras.layers.RepeatVector(3))
       model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
       model.compile(
           loss=keras.losses.MSE,
-          optimizer=keras.optimizers.RMSprop(lr=0.0001),
+          optimizer='rmsprop',
           metrics=[
               keras.metrics.categorical_accuracy,
-              keras.metrics.CategoricalAccuracy()
+              keras.metrics.CategoricalAccuracy(name='cat_acc')
           ],
           weighted_metrics=[
               keras.metrics.categorical_accuracy,
-              keras.metrics.CategoricalAccuracy()
+              keras.metrics.CategoricalAccuracy(name='cat_acc2')
           ],
           sample_weight_mode='temporal')
       x = np.random.random((1, 3))
@@ -479,12 +479,13 @@ class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
       model.save(saved_model_dir, save_format=save_format)
 
       new_model = keras.models.load_model(saved_model_dir)
-      self._assert_same_weights(
-          model, new_model, original_optimizer_has_iterations_variable=False)
+
+      self._assert_same_weights_and_metrics(model, new_model)
 
       out2 = new_model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
 
+  @test_util.run_in_graph_and_eager_modes
   def test_sequential_model_saving_without_compile(self):
     saved_model_dir = self._save_model_dir()
     save_format = testing_utils.get_save_format()
@@ -501,7 +502,7 @@ class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
       keras.models.save_model(model, saved_model_dir, save_format=save_format)
 
       new_model = keras.models.load_model(saved_model_dir)
-      self._assert_same_weights(model, new_model)
+      self._assert_same_weights_and_metrics(model, new_model)
 
       out2 = new_model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
@@ -535,42 +536,11 @@ class TestWholeModelSaving(test.TestCase, parameterized.TestCase):
           saved_model_dir,
           custom_objects={'CustomOp': CustomOp,
                           'custom_loss': custom_loss})
-      self._assert_same_weights(model, new_model)
+      self._assert_same_weights_and_metrics(model, new_model)
 
       out2 = new_model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
 
-  def test_functional_model_saving(self):
-    saved_model_dir = self._save_model_dir()
-    save_format = testing_utils.get_save_format()
-    with ops.Graph().as_default(), self.cached_session():
-      inputs = keras.layers.Input(shape=(3,))
-      x = keras.layers.Dense(2)(inputs)
-      output = keras.layers.Dense(3)(x)
-
-      model = keras.models.Model(inputs, output)
-      model.compile(
-          loss=keras.losses.MSE,
-          optimizer=keras.optimizers.RMSprop(lr=0.0001),
-          metrics=[
-              keras.metrics.categorical_accuracy,
-              keras.metrics.CategoricalAccuracy()
-          ],
-          weighted_metrics=[
-              keras.metrics.categorical_accuracy,
-              keras.metrics.CategoricalAccuracy()
-          ])
-      x = np.random.random((1, 3))
-      y = np.random.random((1, 3))
-      model.train_on_batch(x, y)
-
-      out = model.predict(x)
-      keras.models.save_model(model, saved_model_dir, save_format=save_format)
-      model = keras.models.load_model(saved_model_dir)
-
-      out2 = model.predict(x)
-      self.assertAllClose(out, out2, atol=1e-05)
-
   def test_saving_without_compilation(self):
     saved_model_dir = self._save_model_dir()
     save_format = testing_utils.get_save_format()
diff --git a/tensorflow/python/keras/saving/saved_model/load.py b/tensorflow/python/keras/saving/saved_model/load.py
index a378c1b98e7..0b55e30c27b 100644
--- a/tensorflow/python/keras/saving/saved_model/load.py
+++ b/tensorflow/python/keras/saving/saved_model/load.py
@@ -129,6 +129,7 @@ def load(path, compile=True, options=None):  # pylint: disable=redefined-builtin
     if training_config is not None:
       model.compile(**saving_utils.compile_args_from_training_config(
           training_config))
+      saving_utils.try_build_compiled_arguments(model)
     else:
       logging.warning('No training configuration found in save file, so the '
                       'model was *not* compiled. Compile it manually.')
diff --git a/tensorflow/python/keras/saving/saving_utils.py b/tensorflow/python/keras/saving/saving_utils.py
index 3c9c33531bf..9fdf81cae2a 100644
--- a/tensorflow/python/keras/saving/saving_utils.py
+++ b/tensorflow/python/keras/saving/saving_utils.py
@@ -27,6 +27,7 @@ from tensorflow.python.keras import losses
 from tensorflow.python.keras import optimizers
 from tensorflow.python.keras.engine import base_layer_utils
 from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.keras.utils import version_utils
 from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import nest
@@ -307,3 +308,16 @@ def _enforce_names_consistency(specs):
   if name_inconsistency:
     specs = nest.map_structure(_clear_name, specs)
   return specs
+
+
+def try_build_compiled_arguments(model):
+  if (not version_utils.is_v1_layer_or_model(model) and
+      model.outputs is not None):
+    try:
+      model.compiled_loss.build(model.outputs)
+      model.compiled_metrics.build(model.outputs, model.outputs)
+    except:  # pylint: disable=bare-except
+      logging.warning(
+          'Compiled the loaded model, but the compiled metrics have yet to '
+          'be built. `model.compile_metrics` will be empty until you train '
+          'or evaluate the model.')

From f1f5ed68595a56357d92985466d7e0687b23303e Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Fri, 19 Jun 2020 17:22:51 +0000
Subject: [PATCH 0624/1390] Update
 tensorflow/security/fuzzing/status_group_fuzz.cc

---
 tensorflow/security/fuzzing/status_group_fuzz.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/security/fuzzing/status_group_fuzz.cc b/tensorflow/security/fuzzing/status_group_fuzz.cc
index 002785734bb..a560766410a 100644
--- a/tensorflow/security/fuzzing/status_group_fuzz.cc
+++ b/tensorflow/security/fuzzing/status_group_fuzz.cc
@@ -57,7 +57,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   // Ignore warnings that these values are unused
   sg.as_summary_status().IgnoreError();
   sg.as_concatenated_status().IgnoreError();
-  sg.AttachLogMessages().IgnoreError();
+  sg.AttachLogMessages();
 
   return 0;
 }

From d00691f7aaa954bfc9194f62e729f8b41de899af Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Fri, 19 Jun 2020 10:15:24 -0700
Subject: [PATCH 0625/1390] Fork the keras related mirrored_strategy_test to
 keras/distribute.

PiperOrigin-RevId: 317329998
Change-Id: I7dc55499cb0409129729696c286862a6b6d574aa
---
 tensorflow/python/distribute/BUILD            |  1 -
 .../distribute/mirrored_strategy_test.py      | 47 ----------
 tensorflow/python/keras/distribute/BUILD      | 23 +++++
 .../distribute/mirrored_strategy_test.py      | 89 +++++++++++++++++++
 4 files changed, 112 insertions(+), 48 deletions(-)
 create mode 100644 tensorflow/python/keras/distribute/mirrored_strategy_test.py

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 0062705126f..38c5550be16 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -1486,7 +1486,6 @@ cuda_py_test(
         "//tensorflow/python/autograph/core:test_lib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:test",
-        "//tensorflow/python/keras/layers",
     ],
 )
 
diff --git a/tensorflow/python/distribute/mirrored_strategy_test.py b/tensorflow/python/distribute/mirrored_strategy_test.py
index 950b6f2446b..e6414b2704a 100644
--- a/tensorflow/python/distribute/mirrored_strategy_test.py
+++ b/tensorflow/python/distribute/mirrored_strategy_test.py
@@ -22,7 +22,6 @@ import json
 import sys
 
 from absl.testing import parameterized
-import numpy as np
 
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python import tf2
@@ -50,16 +49,12 @@ from tensorflow.python.framework import func_graph
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_util
-from tensorflow.python.keras.engine import training as keras_training
-from tensorflow.python.keras.layers import core as keras_core
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import gradients
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables
-from tensorflow.python.training import gradient_descent
-from tensorflow.python.training import optimizer as optimizer_lib
 from tensorflow.python.training import server_lib
 
 
@@ -988,22 +983,6 @@ class MockModel(object):
     return x
 
 
-class MiniModel(keras_training.Model):
-  """Minimal model for mnist.
-
-  Useful for testing and debugging on slow TPU simulators.
-  """
-
-  def __init__(self):
-    super(MiniModel, self).__init__(name="")
-    self.fc = keras_core.Dense(1, name="fc", kernel_initializer="ones",
-                               bias_initializer="ones")
-
-  def call(self, inputs, training=True):
-    inputs = array_ops.ones([1, 10])
-    return self.fc(inputs)
-
-
 @combinations.generate(
     combinations.combine(
         distribution=[
@@ -1116,32 +1095,6 @@ class MirroredStrategyDefunTest(test.TestCase):
     expected_result = values.PerReplica((5.0 * 1.25, 3.0 * 1.25))
     self._call_and_check(distribution, fn1, [factors], expected_result, [fn1])
 
-  def testTrain(self, distribution):
-    with distribution.scope():
-      mock_model = MiniModel()
-      mock_model.call = function.defun(mock_model.call)
-
-      def loss_fn(ctx):
-        del ctx
-        return mock_model(array_ops.ones([1, 10]))
-
-      gradients_fn = backprop.implicit_grad(loss_fn)
-      gradients_fn = optimizer_lib.get_filtered_grad_fn(gradients_fn)
-      grads_and_vars = distribution.extended.call_for_each_replica(
-          gradients_fn, args=(None,))
-
-      optimizer = gradient_descent.GradientDescentOptimizer(0.25)
-      update_ops = optimizer._distributed_apply(distribution, grads_and_vars)  # pylint: disable=protected-access
-
-      if not context.executing_eagerly():
-        self.evaluate(variables.global_variables_initializer())
-        self.evaluate(update_ops)
-
-      updated_var_values = self.evaluate(mock_model.variables)
-      # All variables start at 1.0 and get two updates of 0.25.
-      self.assertAllEqual(0.5 * np.ones([10, 1]), updated_var_values[0])
-      self.assertAllEqual([0.5], updated_var_values[1])
-
 
 @combinations.generate(
     combinations.combine(
diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
index ddf274f299f..247e655621c 100644
--- a/tensorflow/python/keras/distribute/BUILD
+++ b/tensorflow/python/keras/distribute/BUILD
@@ -324,6 +324,29 @@ cuda_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "mirrored_strategy_test",
+    srcs = ["mirrored_strategy_test.py"],
+    python_version = "PY3",
+    tags = [
+        "multi_and_single_gpu",
+        "no_windows_gpu",  # TODO(b/130551176)
+    ],
+    deps = [
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:training_lib",
+        "//tensorflow/python:variables",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:strategy_combinations",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/keras/engine",
+        "//tensorflow/python/keras/layers:core",
+    ],
+)
+
 cuda_py_test(
     name = "multi_worker_test",
     srcs = ["multi_worker_test.py"],
diff --git a/tensorflow/python/keras/distribute/mirrored_strategy_test.py b/tensorflow/python/keras/distribute/mirrored_strategy_test.py
new file mode 100644
index 00000000000..2844af8cc3a
--- /dev/null
+++ b/tensorflow/python/keras/distribute/mirrored_strategy_test.py
@@ -0,0 +1,89 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for MirroredStrategy."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import function
+from tensorflow.python.eager import test
+from tensorflow.python.keras.engine import training as keras_training
+from tensorflow.python.keras.layers import core as keras_core
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import optimizer as optimizer_lib
+
+
+class MiniModel(keras_training.Model):
+  """Minimal model for mnist.
+
+  Useful for testing and debugging on slow TPU simulators.
+  """
+
+  def __init__(self):
+    super(MiniModel, self).__init__(name="")
+    self.fc = keras_core.Dense(1, name="fc", kernel_initializer="ones",
+                               bias_initializer="ones")
+
+  def call(self, inputs, training=True):
+    inputs = array_ops.ones([1, 10])
+    return self.fc(inputs)
+
+
+@combinations.generate(
+    combinations.combine(
+        distribution=[
+            strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+        ],
+        mode=["graph", "eager"]))
+class MirroredStrategyDefunTest(test.TestCase):
+
+  def testTrain(self, distribution):
+    with distribution.scope():
+      mock_model = MiniModel()
+      mock_model.call = function.defun(mock_model.call)
+
+      def loss_fn(ctx):
+        del ctx
+        return mock_model(array_ops.ones([1, 10]))
+
+      gradients_fn = backprop.implicit_grad(loss_fn)
+      gradients_fn = optimizer_lib.get_filtered_grad_fn(gradients_fn)
+      grads_and_vars = distribution.extended.call_for_each_replica(
+          gradients_fn, args=(None,))
+
+      optimizer = gradient_descent.GradientDescentOptimizer(0.25)
+      update_ops = optimizer._distributed_apply(distribution, grads_and_vars)  # pylint: disable=protected-access
+
+      if not context.executing_eagerly():
+        self.evaluate(variables.global_variables_initializer())
+        self.evaluate(update_ops)
+
+      updated_var_values = self.evaluate(mock_model.variables)
+      # All variables start at 1.0 and get two updates of 0.25.
+      self.assertAllEqual(0.5 * np.ones([10, 1]), updated_var_values[0])
+      self.assertAllEqual([0.5], updated_var_values[1])
+
+
+if __name__ == "__main__":
+  test.main()

From 57f9d638c00083e864aaa7e4f8114c0dd3ba479c Mon Sep 17 00:00:00 2001
From: Bixia Zheng 
Date: Fri, 19 Jun 2020 10:17:18 -0700
Subject: [PATCH 0626/1390] [TF:TRT] Add a prefix to the warning messages from
 TF-TRT.

Add LOG_WARNING_WITH_PREFIX to common/utils.h. Replace the use of LOG(WARNING)
with this new macro.

PiperOrigin-RevId: 317330336
Change-Id: Ife0aa0347dd72f6eb0f8805af4d46a7d4cb099ea
---
 tensorflow/compiler/tf2tensorrt/BUILD         | 14 ++++++++
 .../compiler/tf2tensorrt/common/utils.h       | 35 +++++++++++++++++++
 .../tf2tensorrt/convert/convert_graph.cc      | 30 +++++++++-------
 .../tf2tensorrt/convert/convert_nodes.cc      | 27 +++++++-------
 .../tf2tensorrt/kernels/trt_engine_op.cc      | 25 +++++++------
 .../compiler/tf2tensorrt/segment/segment.cc   | 10 +++---
 .../compiler/tf2tensorrt/utils/py_utils.cc    |  8 +++--
 .../compiler/tf2tensorrt/utils/trt_logger.cc  |  3 +-
 8 files changed, 108 insertions(+), 44 deletions(-)
 create mode 100644 tensorflow/compiler/tf2tensorrt/common/utils.h

diff --git a/tensorflow/compiler/tf2tensorrt/BUILD b/tensorflow/compiler/tf2tensorrt/BUILD
index 4a8599e29f6..368cb5af2ed 100644
--- a/tensorflow/compiler/tf2tensorrt/BUILD
+++ b/tensorflow/compiler/tf2tensorrt/BUILD
@@ -79,6 +79,15 @@ tf_cuda_cc_test(
     ]),
 )
 
+cc_library(
+    name = "common_utils",
+    hdrs = ["common/utils.h"],
+    copts = tf_copts(),
+    deps = [
+        "//tensorflow/core/platform:logging",
+    ] + if_tensorrt([":tensorrt_lib"]),
+)
+
 cc_library(
     name = "trt_op_kernels",
     srcs = [
@@ -95,6 +104,7 @@ cc_library(
         ":trt_plugins",
         ":trt_resources",
         ":utils",
+        ":common_utils",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@local_config_cuda//cuda:cuda_headers",
@@ -240,6 +250,7 @@ tf_cuda_library(
     hdrs = ["utils/trt_logger.h"],
     visibility = ["//visibility:public"],
     deps = [
+        ":common_utils",
         ":logger_registry",
         "//tensorflow/core:lib_proto_parsing",
     ] + if_tensorrt([":tensorrt_lib"]),
@@ -375,6 +386,7 @@ tf_cuda_library(
         "convert/trt_optimization_pass.h",
     ],
     deps = [
+        ":common_utils",
         ":logger_registry",
         ":segment",
         ":trt_allocator",
@@ -488,6 +500,7 @@ cc_library(
     ],
     copts = tf_copts(),
     deps = [
+        ":common_utils",
         "//tensorflow/core:graph",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
@@ -575,6 +588,7 @@ cc_library(
     hdrs = ["utils/py_utils.h"],
     copts = tf_copts(),
     deps = if_tensorrt([
+        ":common_utils",
         ":tensorrt_lib",
         "//tensorflow/stream_executor/platform:dso_loader",
     ]),
diff --git a/tensorflow/compiler/tf2tensorrt/common/utils.h b/tensorflow/compiler/tf2tensorrt/common/utils.h
new file mode 100644
index 00000000000..9ab0145e1ec
--- /dev/null
+++ b/tensorflow/compiler/tf2tensorrt/common/utils.h
@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+#define LOG_WARNING_WITH_PREFIX LOG(WARNING) << "TF-TRT Warning: "
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif
+#endif
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
index 414d27477bc..1c51d51f1c9 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include 
 
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
@@ -276,8 +277,9 @@ Status GetEngineInfo(const Graph* g,
   if (segment_devices.size() == 1) {
     info->device = *segment_devices.begin();
   } else if (segment_devices.size() > 1) {
-    LOG(WARNING) << "Detected multiple (" << segment_devices.size()
-                 << ") devices for the segment. Picking first one to continue.";
+    LOG_WARNING_WITH_PREFIX
+        << "Detected multiple (" << segment_devices.size()
+        << ") devices for the segment. Picking first one to continue.";
     info->device = *segment_devices.begin();
   } else {
     TfGpuId tf_gpu_id;
@@ -663,7 +665,7 @@ std::pair GetDeviceAndAllocator(const ConversionParams& params,
       StrAppend(&msg, engine.device, "': ");
       for (auto d : devices) StrAppend(&msg, d->name(), ", ");
       StrAppend(&msg, ". Will get the allocator from first one.");
-      LOG(WARNING) << msg;
+      LOG_WARNING_WITH_PREFIX << msg;
     }
     AllocatorAttributes alloc_attr;
     cuda_device_id = devices[0]->tensorflow_gpu_device_info()->gpu_id;
@@ -671,8 +673,8 @@ std::pair GetDeviceAndAllocator(const ConversionParams& params,
     VLOG(1) << "Using allocator " << dev_allocator->Name()
             << " and cuda_device_id " << cuda_device_id;
   } else {
-    LOG(WARNING) << "Cluster is set but device '" << engine.device
-                 << "' is not found in the cluster";
+    LOG_WARNING_WITH_PREFIX << "Cluster is set but device '" << engine.device
+                            << "' is not found in the cluster";
   }
   return std::make_pair(cuda_device_id, dev_allocator);
 }
@@ -770,8 +772,8 @@ Status ConvertAfterShapes(const ConversionParams& params) {
     Status status = GetEngineInfo(&graph, static_graph_properties, curr_segment,
                                   node_map, reverse_topo_order, &curr_engine);
     if (!status.ok()) {
-      LOG(WARNING) << "Failed to get engine info for segment " << t << ": "
-                   << status;
+      LOG_WARNING_WITH_PREFIX << "Failed to get engine info for segment " << t
+                              << ": " << status;
       continue;
     }
     curr_engine.precision_mode = params.precision_mode;
@@ -784,8 +786,9 @@ Status ConvertAfterShapes(const ConversionParams& params) {
                                             &graph, curr_engine.engine_name);
 
     if (!status.ok()) {
-      LOG(WARNING) << "Failed to register segment graphdef to the library " << t
-                   << ": " << status;
+      LOG_WARNING_WITH_PREFIX
+          << "Failed to register segment graphdef to the library " << t << ": "
+          << status;
       continue;
     }
 
@@ -836,7 +839,8 @@ Status ConvertAfterShapes(const ConversionParams& params) {
       alloc.reset(new TRTDeviceAllocator(device_alloc.second));
     } else {
       // Setting allocator as nullptr should get revert to the cudamalloc
-      LOG(WARNING) << "Can't identify the cuda device. Running on device 0 ";
+      LOG_WARNING_WITH_PREFIX
+          << "Can't identify the cuda device. Running on device 0 ";
     }
     cudaSetDevice(cuda_device_id);
     auto status =
@@ -850,9 +854,9 @@ Status ConvertAfterShapes(const ConversionParams& params) {
       LOG(INFO) << "Replaced " << msg << ".";
     } else {
       // Graph is not modified.
-      LOG(WARNING) << "Cannot replace " << msg
-                   << " reason: " << status.error_message()
-                   << " (keeping original segment).";
+      LOG_WARNING_WITH_PREFIX << "Cannot replace " << msg
+                              << " reason: " << status.error_message()
+                              << " (keeping original segment).";
     }
     if (VLOG_IS_ON(1)) {
       msg = "Segment consists of nodes: ";
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
index 28b27959afc..96cec556942 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h"
@@ -1214,15 +1215,16 @@ static void InitializeTrtPlugins(nvinfer1::ILogger* trt_logger) {
   nvinfer1::IPluginCreator* const* trt_plugin_creator_list =
       getPluginRegistry()->getPluginCreatorList(&num_trt_plugins);
   if (!trt_plugin_creator_list) {
-    LOG(WARNING) << "Can not find any TensorRT plugins in registry.";
+    LOG_WARNING_WITH_PREFIX << "Can not find any TensorRT plugins in registry.";
   } else {
     VLOG(1) << "Found the following " << num_trt_plugins
             << " TensorRT plugins in registry:";
     for (int i = 0; i < num_trt_plugins; ++i) {
       if (!trt_plugin_creator_list[i]) {
-        LOG(WARNING) << "TensorRT plugin at index " << i
-                     << " is not accessible (null pointer returned by "
-                        "getPluginCreatorList for this plugin)";
+        LOG_WARNING_WITH_PREFIX
+            << "TensorRT plugin at index " << i
+            << " is not accessible (null pointer returned by "
+               "getPluginCreatorList for this plugin)";
       } else {
         VLOG(1) << "  " << trt_plugin_creator_list[i]->getPluginName();
       }
@@ -1827,9 +1829,9 @@ void Converter::MaybeApplyQuantizationRanges() {
       // are tensors which are created internally by TF-TRT. The ranges for
       // these unnamed ITensors are always inferred from user provided ranges,
       // thus there will also be a warning for the range(s) the user missed.
-      LOG(WARNING) << "Quantization range was not found for "
-                   << tensor->getName() << ". "
-                   << "Setting invalid quantization range.";
+      LOG_WARNING_WITH_PREFIX << "Quantization range was not found for "
+                              << tensor->getName() << ". "
+                              << "Setting invalid quantization range.";
       // Set the range to something unusable so the engine will fail if it
       // tries to actually use the tensor's range.
       tensor->setDynamicRange(0, 0);
@@ -4898,10 +4900,11 @@ Status ConvertFusedBatchNorm(OpConverterParams* params) {
     // Trying to use batchnorm in training mode is a very common problem.
     // Because the error message will only be printed in VLOG(1) by the
     // segmenter, we issue a special warning so that users will actually see it.
-    LOG(WARNING) << node_def.op() << " only supports is_training=false. If you "
-                 << "are using Keras, please call "
-                 << "keras.backend.set_learning_phase(0) before constructing "
-                 << "your model. At " << node_def.name();
+    LOG_WARNING_WITH_PREFIX
+        << node_def.op() << " only supports is_training=false. If you "
+        << "are using Keras, please call "
+        << "keras.backend.set_learning_phase(0) before constructing "
+        << "your model. At " << node_def.name();
     return errors::Unimplemented(node_def.op(),
                                  " only supports is_training=false, at ",
                                  node_def.name());
@@ -6039,7 +6042,7 @@ Status ConvertGraphDefToEngine(
         const string error_message =
             StrCat("Validation failed for ", node_name, " and input slot ",
                    slot_number, ": ", status.error_message());
-        LOG(WARNING) << error_message;
+        LOG_WARNING_WITH_PREFIX << error_message;
         return Status(status.code(), error_message);
       }
       VLOG(2) << "Adding engine input tensor " << node_name << " with shape "
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index d9b8e198f4f..ac4a331041d 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
@@ -613,8 +614,8 @@ void TRTEngineOp::ComputeAsync(OpKernelContext* ctx,
   }
   Status stat = ExecuteTrtEngine(ctx, engine_context, trt_context_idx);
   if (!stat.ok()) {
-    LOG(WARNING) << "Failed to execute engine: " << stat
-                 << " Retrying with native segment for " << name();
+    LOG_WARNING_WITH_PREFIX << "Failed to execute engine: " << stat
+                            << " Retrying with native segment for " << name();
     // Release any outputs that are allocated, ExecuteNativeSegment will
     // re-allocate them and fail if they are currently allocated.
     for (int i = 0; i < ctx->num_outputs(); i++) {
@@ -727,9 +728,9 @@ StatusOr> TRTEngineOp::BuildEngine(
       calibrator, &engine, use_calibration, use_implicit_batch_, nullptr,
       &cache_resource->profiles_);
   if (!status.ok()) {
-    LOG(WARNING) << "Engine creation for " << name() << " failed. "
-                 << "The native segment will be used instead. "
-                 << "Reason: " << status;
+    LOG_WARNING_WITH_PREFIX << "Engine creation for " << name() << " failed. "
+                            << "The native segment will be used instead. "
+                            << "Reason: " << status;
     // Store an empty engine in the cache for these input shapes so we don't try
     // to build the same failing engine again.
     cache_resource->cache_.emplace(input_concrete_shapes,
@@ -791,8 +792,9 @@ StatusOr> TRTEngineOp::GetEngine(
               FunctionDefToGraphDef(func_handle_, lib, &segment_graph_def_);
         }
         if (!status.ok()) {
-          LOG(WARNING) << "Getting segment graph for " << name() << " failed. "
-                       << "Reason: " << status;
+          LOG_WARNING_WITH_PREFIX << "Getting segment graph for " << name()
+                                  << " failed. "
+                                  << "Reason: " << status;
         }
       }
       auto result = BuildEngine(input_concrete_shapes, batch_size,
@@ -851,10 +853,11 @@ StatusOr> TRTEngineOp::GetEngine(
   // If cache does not have a compatible engine then create a new engine.
   if (engine_contexts == nullptr) {
     if (!allow_build_at_runtime_) {
-      LOG(WARNING) << "Found no engine in cache matching input shapes. "
-                   << "Not building a new engine because "
-                   << "allow_build_at_runtime=False. "
-                   << "The native segment will be used instead.";
+      LOG_WARNING_WITH_PREFIX
+          << "Found no engine in cache matching input shapes. "
+          << "Not building a new engine because "
+          << "allow_build_at_runtime=False. "
+          << "The native segment will be used instead.";
       // Store an empty engine in the cache for these input shapes so we don't
       // try to build the same failing engine again.
       cache.emplace(input_concrete_shapes, absl::make_unique());
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
index 749335f1b09..32e30006f58 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/segment/union_find.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/graph/algorithm.h"
@@ -748,9 +749,10 @@ Status SegmentGraph(const Graph* tf_graph,
         exclude_node(status.error_message());
       } else if (tftrt_op_blacklist.count(node->tf_node()->type_string())) {
         // WARNING verbosity since the user explicitly requests this behavior.
-        LOG(WARNING) << "Blacklisted as TF-TRT candidate, "
-                     << "(Op type: " << node->tf_node()->type_string() << "), "
-                     << "(Op name: " << node->name() << ")";
+        LOG_WARNING_WITH_PREFIX
+            << "Blacklisted as TF-TRT candidate, "
+            << "(Op type: " << node->tf_node()->type_string() << "), "
+            << "(Op name: " << node->name() << ")";
         exclude_node("Blacklisted with the env var TF_TRT_OP_BLACKLIST");
       } else {
         VLOG(2) << "Accepted as a TF-TRT candidate, "
@@ -1038,7 +1040,7 @@ Status SegmentGraph(const Graph* tf_graph,
       for (const auto& dev : dev_itr->second) {
         StrAppend(&s, dev, ", ");
       }
-      LOG(WARNING) << s;
+      LOG_WARNING_WITH_PREFIX << s;
     }
 
     segments->emplace_back(segment_nodes);
diff --git a/tensorflow/compiler/tf2tensorrt/utils/py_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/py_utils.cc
index 885f58cd70c..a8e24aa8983 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/py_utils.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/py_utils.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2tensorrt/utils/py_utils.h"
 
 #if GOOGLE_CUDA && GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/stream_executor/platform/dso_loader.h"
 #include "third_party/tensorrt/NvInfer.h"
 #endif
@@ -27,9 +28,10 @@ bool IsGoogleTensorRTEnabled() {
 #if GOOGLE_CUDA && GOOGLE_TENSORRT
   auto handle_or = se::internal::DsoLoader::TryDlopenTensorRTLibraries();
   if (!handle_or.ok()) {
-    LOG(WARNING) << "Cannot dlopen some TensorRT libraries. If you would like "
-                    "to use Nvidia GPU with TensorRT, please make sure the "
-                    "missing libraries mentioned above are installed properly.";
+    LOG_WARNING_WITH_PREFIX
+        << "Cannot dlopen some TensorRT libraries. If you would like "
+           "to use Nvidia GPU with TensorRT, please make sure the "
+           "missing libraries mentioned above are installed properly.";
     return false;
   } else {
     return true;
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
index 6bb6f1f9dd8..193687ebc8c 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
+#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 #include "tensorflow/core/platform/logging.h"
 
@@ -35,7 +36,7 @@ void Logger::log(Severity severity, const char* msg) {
       break;
     }
     case Severity::kWARNING: {
-      LOG(WARNING) << name_ << " " << msg;
+      LOG_WARNING_WITH_PREFIX << name_ << " " << msg;
       break;
     }
     case Severity::kERROR: {

From ef1cabc7a8fb05f4c66b21aa2409a74e22635e45 Mon Sep 17 00:00:00 2001
From: Thomas Joerg 
Date: Fri, 19 Jun 2020 10:18:32 -0700
Subject: [PATCH 0627/1390] [XLA:GPU] Split reduce ops with large but
 non-consecutive reduction dimensions.

PiperOrigin-RevId: 317330616
Change-Id: Icdcf320b233479c2f74c5b40ee4c8d9a73a6088a
---
 tensorflow/compiler/xla/service/gpu/BUILD     |  28 ++++
 .../compiler/xla/service/gpu/gpu_compiler.cc  |   2 +
 .../xla/service/gpu/reduction_splitter.cc     | 117 +++++++++++++++
 .../xla/service/gpu/reduction_splitter.h      |  49 ++++++
 .../service/gpu/reduction_splitter_test.cc    | 140 ++++++++++++++++++
 .../reduction_degenerate_dim_remover_test.cc  |   1 +
 .../tests/reduction_layout_normalizer_test.cc |   1 +
 7 files changed, 338 insertions(+)
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter.h
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc

diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 0eb82128159..472d2117a2c 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -1174,6 +1174,7 @@ cc_library(
         ":reduction_degenerate_dim_remover",
         ":reduction_dimension_grouper",
         ":reduction_layout_normalizer",
+        ":reduction_splitter",
         ":stream_assignment",
         ":stream_executor_util",
         ":target_constants",
@@ -1819,6 +1820,33 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "reduction_splitter",
+    srcs = ["reduction_splitter.cc"],
+    hdrs = ["reduction_splitter.h"],
+    deps = [
+        ":ir_emission_utils",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service:hlo_pass",
+    ],
+)
+
+tf_cc_test(
+    name = "reduction_splitter_test",
+    srcs = ["reduction_splitter_test.cc"],
+    deps = [
+        ":reduction_splitter",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:test",
+        "//tensorflow/compiler/xla:test_helpers",
+        "//tensorflow/compiler/xla/service:hlo_matchers",
+        "//tensorflow/compiler/xla/service:hlo_parser",
+        "//tensorflow/compiler/xla/tests:hlo_test_base",
+        "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+    ],
+)
+
 cc_library(
     name = "reduction_layout_normalizer",
     srcs = ["reduction_layout_normalizer.cc"],
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
index cddbee92874..156cb112285 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
@@ -65,6 +65,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/gpu/reduction_degenerate_dim_remover.h"
 #include "tensorflow/compiler/xla/service/gpu/reduction_dimension_grouper.h"
 #include "tensorflow/compiler/xla/service/gpu/reduction_layout_normalizer.h"
+#include "tensorflow/compiler/xla/service/gpu/reduction_splitter.h"
 #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h"
 #include "tensorflow/compiler/xla/service/gpu/stream_executor_util.h"
 #include "tensorflow/compiler/xla/service/gpu/target_constants.h"
@@ -371,6 +372,7 @@ Status GpuCompiler::OptimizeHloPostLayoutAssignment(
   pipeline.AddPass();
   pipeline.AddPass();
   pipeline.AddPass();
+  pipeline.AddPass>();
 
   // The LayoutAssignment pass may leave behind kCopy instructions which are
   // duplicate or NOPs, so remove them with algebraic simplification and CSE.
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc b/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
new file mode 100644
index 00000000000..b68213ec35f
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
@@ -0,0 +1,117 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/gpu/reduction_splitter.h"
+
+#include 
+
+#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
+#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
+#include "tensorflow/compiler/xla/service/hlo_instructions.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+
+namespace xla {
+namespace gpu {
+
+class ReductionSplitterVisitor : public DfsHloRewriteVisitor {
+ public:
+  Status HandleReduce(HloInstruction *reduce) override {
+    VLOG(4) << "Input: " << reduce->ToString();
+
+    // Reductions with contiguous dimensions are lowered to efficient code. No
+    // need to split such ops.
+    if (IsReductionFromOrToContiguousDimensions(*reduce)) {
+      return Status::OK();
+    }
+    if (reduce->dimensions().size() < 2) {
+      return Status::OK();
+    }
+    if (!reduce->shape().IsArray()) {
+      // TODO(cheshire): Handle variadic reduction.
+      return Status::OK();
+    }
+
+    HloInstruction *operand = reduce->mutable_operand(0);
+    const Shape &shape = operand->shape();
+    CHECK(shape == LayoutUtil::GetWithDefaultLayout(shape))
+        << "Default layout should be enforced on reduction operand";
+    // Verify that contiguous dimensions have been grouped by the
+    // ReductionDimensionGrouper pass.
+    for (int64 i = 0; i < reduce->dimensions().size(); ++i) {
+      for (int64 j = i + 1; j < reduce->dimensions().size(); ++j) {
+        CHECK(abs(reduce->dimensions(i) - reduce->dimensions(j)) > 1)
+            << "Reduction dimensions must not be consecutive";
+      }
+    }
+
+    // The reduce op has non-contiguous dimensions. Look for the dimension with
+    // the largest shape dimension. Reducing along this dimension first will
+    // reduce the output size most effectively.
+    int64 max_shape_dim = 0;
+    int64 max_reduce_dim = 0;
+    const auto &input_shape = reduce->operand(0)->shape();
+    for (int64 i = 0; i < reduce->dimensions().size(); ++i) {
+      if (input_shape.dimensions(reduce->dimensions(i)) > max_shape_dim) {
+        max_reduce_dim = reduce->dimensions(i);
+        max_shape_dim = input_shape.dimensions(max_reduce_dim);
+      }
+    }
+    // TODO(tjoerg): Run microbenchmarks to tune this threshold.
+    if (max_shape_dim < 128) {
+      return Status::OK();
+    }
+
+    // Split the reduction into a pre-reduction and a final reduction.
+    VLOG(3) << "Splitting reduction " << reduce->name() << " at dimension "
+            << max_reduce_dim;
+    std::vector pre_reduce_dims;
+    pre_reduce_dims.push_back(max_reduce_dim);
+    std::vector pre_reduce_shape_dims(input_shape.dimensions().begin(),
+                                             input_shape.dimensions().end());
+    pre_reduce_shape_dims.erase(pre_reduce_shape_dims.begin() + max_reduce_dim);
+    Shape pre_reduce_shape = ShapeUtil::MakeShape(
+        reduce->shape().element_type(), pre_reduce_shape_dims);
+    std::unique_ptr pre_reduce = HloInstruction::CreateReduce(
+        pre_reduce_shape, reduce->mutable_operand(0),
+        reduce->mutable_operand(1), pre_reduce_dims, reduce->to_apply());
+    pre_reduce->set_metadata(reduce->metadata());
+
+    std::vector final_reduce_dims(reduce->dimensions().begin(),
+                                         reduce->dimensions().end());
+    final_reduce_dims.erase(
+        std::remove(final_reduce_dims.begin(), final_reduce_dims.end(),
+                    max_reduce_dim),
+        final_reduce_dims.end());
+    for (int64 i = 0; i < final_reduce_dims.size(); ++i) {
+      if (final_reduce_dims[i] > max_reduce_dim) {
+        final_reduce_dims[i]--;
+      }
+    }
+    std::unique_ptr final_reduce = HloInstruction::CreateReduce(
+        reduce->shape(),
+        reduce->parent()->AddInstruction(std::move(pre_reduce)),
+        reduce->mutable_operand(1), final_reduce_dims, reduce->to_apply());
+    return ReplaceWithNewInstruction(reduce, std::move(final_reduce));
+  }
+};
+
+StatusOr ReductionSplitter::Run(HloModule *module) {
+  TF_ASSIGN_OR_RETURN(bool changed,
+                      ReductionSplitterVisitor().RunOnModule(module));
+  return changed;
+}
+
+}  // namespace gpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter.h b/tensorflow/compiler/xla/service/gpu/reduction_splitter.h
new file mode 100644
index 00000000000..f161b579eb8
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter.h
@@ -0,0 +1,49 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
+
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
+
+namespace xla {
+namespace gpu {
+
+// Splits a reduce op into two consecutive reduce ops if
+// * the reduce dimensions are not contiguous and
+// * at least one reduce dimension is large (i.e. corresponds to a large input
+//   shape dimension).
+//
+// Reductions with non-contiguous dimensions are emitted as simple element-wise
+// loops. This is inefficient when reducing large input shape dimensions.
+// Splitting such reductions allows using more efficient reduction emitters.
+//
+// This pass splits reduce ops into two consecutive reduce ops. Run it to a
+// fixpoint to split reduce ops along multiple large dimensions.
+//
+// Precondition: ReductionDimensionGrouper has been run and adjacent reduce
+// dimentsions have been grouped. Reduction layouts have been normalized.
+
+class ReductionSplitter : public HloModulePass {
+ public:
+  absl::string_view name() const override { return "reduction-splitter"; }
+
+  StatusOr Run(HloModule* module) override;
+};
+
+}  // namespace gpu
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc b/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc
new file mode 100644
index 00000000000..1be55b84204
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc
@@ -0,0 +1,140 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/gpu/reduction_splitter.h"
+
+#include "tensorflow/compiler/xla/service/hlo_matchers.h"
+#include "tensorflow/compiler/xla/service/hlo_parser.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/test_helpers.h"
+#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+
+namespace xla {
+namespace gpu {
+namespace {
+
+namespace op = xla::testing::opcode_matchers;
+
+class ReductionSplitterTest : public HloTestBase {};
+
+TEST_F(ReductionSplitterTest, SplitReductionAtDimensionTwo) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f16[6,16,512,64]{3,2,1,0} parameter(0)
+    transpose.1781 = f16[6,512,16,64]{3,1,2,0} transpose(param_0), dimensions={0,2,1,3}
+    convert.6986 = f32[6,512,16,64]{3,1,2,0} convert(transpose.1781)
+    bitcast.2136 = f32[6,16,512,64]{3,2,1,0} bitcast(convert.6986)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[16,64]{1,0} reduce(bitcast.2136, constant_11111), dimensions={0,2}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  ASSERT_TRUE(ReductionSplitter().Run(module.get()).ValueOrDie());
+  SCOPED_TRACE(module->ToString());
+  const HloInstruction* root_reduction =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(root_reduction, op::Reduce(op::Reduce(), op::Constant()));
+
+  auto* pre_reduction = root_reduction->operand(0);
+  EXPECT_THAT(pre_reduction->dimensions(), std::vector({2}));
+  EXPECT_THAT(pre_reduction->shape(), ShapeUtil::MakeShape(F32, {6, 16, 64}));
+  EXPECT_THAT(root_reduction->dimensions(), std::vector({0}));
+  EXPECT_THAT(root_reduction->shape(), ShapeUtil::MakeShape(F32, {16, 64}));
+}
+
+TEST_F(ReductionSplitterTest, SplitReductionAtDimensionZero) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[1024,16,512,64,128]{4,3,2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[16,64]{1,0} reduce(param_0, constant_11111), dimensions={2,0,4}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  ASSERT_TRUE(ReductionSplitter().Run(module.get()).ValueOrDie());
+  SCOPED_TRACE(module->ToString());
+  const HloInstruction* root_reduction =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(root_reduction, op::Reduce(op::Reduce(), op::Constant()));
+
+  auto* pre_reduction = root_reduction->operand(0);
+  EXPECT_THAT(pre_reduction->dimensions(), std::vector({0}));
+  EXPECT_THAT(pre_reduction->shape(),
+              ShapeUtil::MakeShape(F32, {16, 512, 64, 128}));
+  EXPECT_THAT(root_reduction->dimensions(), std::vector({1, 3}));
+  EXPECT_THAT(root_reduction->shape(), ShapeUtil::MakeShape(F32, {16, 64}));
+}
+
+TEST_F(ReductionSplitterTest, DontSplitReductionWithSmallDimensions) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[8,1024,8]{2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[1024]{0} reduce(param_0, constant_11111), dimensions={2,0}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  EXPECT_FALSE(ReductionSplitter().Run(module.get()).ValueOrDie());
+}
+
+TEST_F(ReductionSplitterTest, DontSplitReductionsWithContiguousDimensions) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[128,128,64,128]{3,2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    // The dimenstions to keep (1 and 2) are contiguous.
+    ROOT reduce.982 = f32[128,64]{1,0} reduce(param_0, constant_11111), dimensions={3,0}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  EXPECT_FALSE(ReductionSplitter().Run(module.get()).ValueOrDie());
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc
index 2c5e704d7c2..92f558ee98d 100644
--- a/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/tests/reduction_degenerate_dim_remover_test.cc
@@ -37,6 +37,7 @@ class ReductionDegenerateDimRemoverTest : public GpuCodegenTest {
     DebugOptions debug_options = GpuCodegenTest::GetDebugOptionsForTest();
     debug_options.add_xla_disable_hlo_passes("reduction-layout-normalizer");
     debug_options.add_xla_disable_hlo_passes("reduction-dimension-grouper");
+    debug_options.add_xla_disable_hlo_passes("reduction-splitter");
     debug_options.add_xla_disable_hlo_passes("gpu-tree-reduction-rewriter");
     return debug_options;
   }
diff --git a/tensorflow/compiler/xla/service/gpu/tests/reduction_layout_normalizer_test.cc b/tensorflow/compiler/xla/service/gpu/tests/reduction_layout_normalizer_test.cc
index d06385480e5..b65c2842320 100644
--- a/tensorflow/compiler/xla/service/gpu/tests/reduction_layout_normalizer_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/tests/reduction_layout_normalizer_test.cc
@@ -33,6 +33,7 @@ class ReductionLayoutNormalizerTest : public GpuCodegenTest {
   DebugOptions GetDebugOptionsForTest() override {
     DebugOptions debug_options = GpuCodegenTest::GetDebugOptionsForTest();
     debug_options.add_xla_disable_hlo_passes("reduction-dimension-grouper");
+    debug_options.add_xla_disable_hlo_passes("reduction-splitter");
     debug_options.add_xla_disable_hlo_passes("layout-assignment");
     debug_options.add_xla_disable_hlo_passes("gpu-tree-reduction-rewriter");
     return debug_options;

From e0780ef031fc27f4f2a71d745004d859a711c90a Mon Sep 17 00:00:00 2001
From: Reed Wanderman-Milne 
Date: Fri, 19 Jun 2020 10:21:14 -0700
Subject: [PATCH 0628/1390] Fix fp16 FusedBatchNorm CPU crash if batch
 dimension is 1.

The GPU kernel outputs NaNs for the variance in this case, which is also incorrect, but better than crashing.

PiperOrigin-RevId: 317331280
Change-Id: Iea4e5a3337625796c50244e51d7ccb4b89f4c3e4
---
 tensorflow/core/kernels/redux_functor.h       |  6 ++--
 .../python/ops/nn_fused_batchnorm_test.py     | 28 +++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/redux_functor.h b/tensorflow/core/kernels/redux_functor.h
index 30038c62dbd..e07fa5364f1 100644
--- a/tensorflow/core/kernels/redux_functor.h
+++ b/tensorflow/core/kernels/redux_functor.h
@@ -57,7 +57,8 @@ struct ReduceOuterDimensions {
     if (1 == outer_dim) {
       // Nothing to do but passing input to output.
       output->template flat() =
-          input.template flat().reshape(output_dims);
+          input.template flat().template cast().reshape(
+              output_dims);
       return;
     }
 
@@ -226,7 +227,8 @@ struct ReduceMiddleDimensions {
     if ((1 == inner_dim * outer_dim)) {
       // Nothing to do.
       output->template flat() =
-          input.template flat().reshape(output_dims);
+          input.template flat().template cast().reshape(
+              output_dims);
       return;
     } else if (1 == inner_dim) {
       // Equivalent to ReduceOuterDimensions.
diff --git a/tensorflow/python/ops/nn_fused_batchnorm_test.py b/tensorflow/python/ops/nn_fused_batchnorm_test.py
index 5497325f6c0..1742a919216 100644
--- a/tensorflow/python/ops/nn_fused_batchnorm_test.py
+++ b/tensorflow/python/ops/nn_fused_batchnorm_test.py
@@ -375,9 +375,10 @@ class BatchNormalizationTest(test.TestCase):
     self.assertLess(err_grad_x_2, err_tolerance)
     self.assertLess(err_grad_scale, err_tolerance)
 
-  def _runtests(self, x_shape, is_training, gradient_test=False):
+  def _runtests(self, x_shape, is_training, gradient_test=False,
+                cpu_only=False):
     use_gpu_vals = [False]
-    if test.is_gpu_available(cuda_only=True):
+    if test.is_gpu_available(cuda_only=True) and not cpu_only:
       use_gpu_vals += [True]
     factors = [1.0, 0.6]
     for dtype in [np.float16, np.float32]:
@@ -438,6 +439,11 @@ class BatchNormalizationTest(test.TestCase):
     x_shape = [0, 131, 127, 6]
     self._runtests(x_shape, False)
 
+  def testInferenceShape6(self):
+    x_shape = [1, 1, 1, 1]
+    # GPU kernel doesn't properly handle case where non-channel dimensions are 1
+    self._runtests(x_shape, False, cpu_only=True)
+
   def testTrainingShape1(self):
     x_shape = [1, 1, 6, 1]
     self._runtests(x_shape, True)
@@ -459,6 +465,11 @@ class BatchNormalizationTest(test.TestCase):
     x_shape = [0, 131, 127, 6]
     self._runtests(x_shape, True)
 
+  def testTrainingShape6(self):
+    x_shape = [1, 1, 1, 1]
+    # GPU kernel doesn't properly handle case where non-channel dimensions are 1
+    self._runtests(x_shape, True, cpu_only=True)
+
   @test_util.run_deprecated_v1
   def testBatchNormGradInferenceShape1(self):
     x_shape = [1, 1, 6, 1]
@@ -485,6 +496,13 @@ class BatchNormalizationTest(test.TestCase):
     x_shape = [0, 7, 11, 4]
     self._runtests(x_shape, is_training=False, gradient_test=True)
 
+  @test_util.run_deprecated_v1
+  def testBatchNormGradInferenceShape6(self):
+    x_shape = [1, 1, 1, 1]
+    # GPU kernel doesn't properly handle case where non-channel dimensions are 1
+    self._runtests(x_shape, is_training=False, gradient_test=True,
+                   cpu_only=True)
+
   @test_util.run_deprecated_v1
   def testBatchNormGradTrainingShape1(self):
     x_shape = [1, 1, 6, 1]
@@ -511,6 +529,12 @@ class BatchNormalizationTest(test.TestCase):
     x_shape = [0, 7, 11, 4]
     self._runtests(x_shape, is_training=True, gradient_test=True)
 
+  @test_util.run_deprecated_v1
+  def testBatchNormGradTrainingShape6(self):
+    x_shape = [1, 1, 1, 1]
+    # GPU kernel doesn't properly handle case where non-channel dimensions are 1
+    self._runtests(x_shape, is_training=True, gradient_test=True, cpu_only=True)
+
   def _testBatchNormGradGrad(self, config):
     shape = config['shape']
     err_tolerance = config['err_tolerance']

From 94f241379a351da7cb7a180967915beec48d14ef Mon Sep 17 00:00:00 2001
From: Vo Van Nghia 
Date: Fri, 19 Jun 2020 20:42:39 +0700
Subject: [PATCH 0629/1390] Add gcs_filesystem_test

---
 .../experimental/filesystem/plugins/gcs/BUILD | 24 +++++++++-
 .../filesystem/plugins/gcs/gcs_filesystem.cc  | 10 ++++-
 .../plugins/gcs/gcs_filesystem_test.cc        | 45 +++++++++++++++++++
 3 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc

diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
index 43f3a507f0b..3a65824cd7c 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
@@ -1,5 +1,5 @@
 # Experimental gcs filesystem plugin.
-load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object")
+load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
 
 package(
     licenses = ["notice"],  # Apache 2.0
@@ -42,3 +42,25 @@ cc_library(
         "//tensorflow/c:env",
     ],
 )
+
+tf_cc_test(
+    name = "gcs_filesystem_test",
+    srcs = [
+        "gcs_filesystem.cc",
+        "gcs_filesystem_test.cc",
+    ],
+    local_defines = ["TF_GCS_FILESYSTEM_TEST"],
+    tags = [
+        "manual",
+    ],
+    deps = [
+        "gcs_helper",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c/experimental/filesystem:filesystem_interface",
+        "@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
+        "@com_google_absl//absl/strings",
+    ] + [
+        "//tensorflow/c:tf_status_helper",
+        "@com_google_googletest//:gtest",
+    ],
+)
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
index 24d85f359ef..f459710ddc6 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
@@ -24,6 +24,14 @@ limitations under the License.
 #include "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
 #include "tensorflow/c/tf_status.h"
 
+#ifdef TF_GCS_FILESYSTEM_TEST
+// For testing purpose, we expose some functions.
+#define TF_STATIC
+#else
+// Otherwise, we don't expose any symbol.
+#define TF_STATIC static
+#endif
+
 // Implementation of a filesystem for GCS environments.
 // This filesystem will support `gs://` URI schemes.
 namespace gcs = google::cloud::storage;
@@ -122,7 +130,7 @@ namespace tf_read_only_memory_region {
 namespace tf_gcs_filesystem {
 
 // TODO(vnvo2409): Add lazy-loading and customizing parameters.
-static void Init(TF_Filesystem* filesystem, TF_Status* status) {
+TF_STATIC void Init(TF_Filesystem* filesystem, TF_Status* status) {
   google::cloud::StatusOr client =
       gcs::Client::CreateDefaultClient();
   if (!client) {
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc
new file mode 100644
index 00000000000..eb0fbfc33f7
--- /dev/null
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc
@@ -0,0 +1,45 @@
+#include "gtest/gtest.h"
+#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
+#include "tensorflow/c/tf_status_helper.h"
+
+#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x))
+
+// Forward declaration
+namespace tf_gcs_filesystem {
+void Init(TF_Filesystem* filesystem, TF_Status* status);
+}
+
+namespace tensorflow {
+namespace {
+
+class GCSFilesystemTest : public ::testing::Test {
+ public:
+  void SetUp() override {
+    status = TF_NewStatus();
+    filesystem = new TF_Filesystem;
+    tf_gcs_filesystem::Init(filesystem, status);
+    ASSERT_TF_OK(status) << "Can not initialize filesystem. "
+                         << TF_Message(status);
+  }
+  void TearDown() override {
+    TF_DeleteStatus(status);
+    // TODO(vnvo2409): Add filesystem cleanup
+    delete filesystem;
+  }
+
+ protected:
+  TF_Filesystem* filesystem;
+  TF_Status* status;
+};
+
+// We have to add this test here because there must be at least one test.
+// This test will be removed in the future.
+TEST_F(GCSFilesystemTest, TestInit) { ASSERT_TF_OK(status); }
+
+}  // namespace
+}  // namespace tensorflow
+
+GTEST_API_ int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

From fbf407383c93774d10bd7c45cd66788a070b0e07 Mon Sep 17 00:00:00 2001
From: Nick Kreeger 
Date: Fri, 19 Jun 2020 10:40:30 -0700
Subject: [PATCH 0630/1390] Reduce the size of TfLiteTensor for the TF Micro
 runtime.

This change uses the existing micro-specific build flag (TF_LITE_STATIC_MEMORY) to reduce the size of TfLiteTensor. In this build setting, only the minimum number of fields required for preparing and initializing a model in TFLM are used. This build define is opt-in only for internal builds and continues to be enabled by default in Makefile builds./

All TFLM internal targets can be built with this flag by adding '--copt=-DTF_LITE_STATIC_MEMORY'.

This change reduces the sizeof(TfLiteTensor) to 64 bytes (64bit systems) down from 112 bytes (64 bit systems).

TfLiteTensor struct reduced by 1.75x (~43% reduction)
Tail allocation reduced by: 2,592kb (~12.5% reduction)
Total allocation reduced by: 2,592kb (~12% reduction)

Optimized results from memory_arena_threshold_test:
Keyword Model:
--------------
[RecordingMicroAllocator] Arena allocation total 18448 bytes
[RecordingMicroAllocator] Arena allocation head 672 bytes
[RecordingMicroAllocator] Arena allocation tail 17776 bytes
[RecordingMicroAllocator] 'TfLiteTensor struct' used 3456 bytes with alignment overhead (requested 3456 bytes for 54 tensors)
[RecordingMicroAllocator] 'TfLiteTensor quantization data' used 1728 bytes with alignment overhead (requested 1728 bytes for 108 allocations)
[RecordingMicroAllocator] 'TfLiteTensor variable buffer data' used 10240 bytes with alignment overhead (requested 10240 bytes for 7 allocations)
[RecordingMicroAllocator] 'NodeAndRegistration struct' used 1200 bytes with alignment overhead (requested 1200 bytes for 15 NodeAndRegistration structs)
[RecordingMicroAllocator] 'Operator runtime data' used 148 bytes with alignment overhead (requested 148 bytes for 13 OpData structs)

Test Conv Model:
----------------
[RecordingMicroAllocator] Arena allocation total 10960 bytes
[RecordingMicroAllocator] Arena allocation head 7744 bytes
[RecordingMicroAllocator] Arena allocation tail 3216 bytes
[RecordingMicroAllocator] 'TfLiteTensor struct' used 960 bytes with alignment overhead (requested 960 bytes for 15 tensors)
[RecordingMicroAllocator] 'TfLiteTensor quantization data' used 768 bytes with alignment overhead (requested 752 bytes for 24 allocations)
[RecordingMicroAllocator] 'TfLiteTensor variable buffer data' used 0 bytes with alignment overhead (requested 0 bytes for 0 allocations)
[RecordingMicroAllocator] 'NodeAndRegistration struct' used 560 bytes with alignment overhead (requested 560 bytes for 7 NodeAndRegistration structs)
[RecordingMicroAllocator] 'Operator runtime data' used 136 bytes with alignment overhead (requested 136 bytes for 5 OpData structs)

PiperOrigin-RevId: 317335359
Change-Id: Ic3d4d2c3e62249f072ece8f621f9ef94eaa28589
---
 tensorflow/lite/c/common.h                    | 46 +++++++++++++++++++
 .../lite/micro/memory_arena_threshold_test.cc | 19 ++++++--
 .../lite/micro/micro_interpreter_test.cc      |  2 +-
 tensorflow/lite/micro/tools/make/Makefile     |  2 +
 .../benchmark/experimental/c/c_api_types.h    | 46 +++++++++++++++++++
 5 files changed, 111 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/c/common.h b/tensorflow/lite/c/common.h
index 15823784d12..9093e5d50ad 100644
--- a/tensorflow/lite/c/common.h
+++ b/tensorflow/lite/c/common.h
@@ -375,6 +375,7 @@ typedef struct TfLiteSparsity {
 
 // An tensor in the interpreter system which is a wrapper around a buffer of
 // data including a dimensionality (or NULL if not currently defined).
+#ifndef TF_LITE_STATIC_MEMORY
 typedef struct TfLiteTensor {
   // The data type specification for data stored in `data`. This affects
   // what member of `data` union should be used.
@@ -439,6 +440,51 @@ typedef struct TfLiteTensor {
   // `dims_signature` contains [1, -1, -1, 3]).
   const TfLiteIntArray* dims_signature;
 } TfLiteTensor;
+#else
+// Specific reduced TfLiteTensor struct for TF Micro runtime. This struct
+// contains only the minimum fields required to initialize and prepare a micro
+// inference graph. The fields in this struct have been ordered from
+// largest-to-smallest for optimal struct sizeof.
+//
+// NOTE: This flag is opt-in only at compile time.
+typedef struct TfLiteTensor {
+  // TODO(b/155784997): Consider consolidating these quantization fields:
+  // Quantization information. Replaces params field above.
+  TfLiteQuantization quantization;
+
+  // Quantization information.
+  TfLiteQuantizationParams params;
+
+  // A union of data pointers. The appropriate type should be used for a typed
+  // tensor based on `type`.
+  TfLitePtrUnion data;
+
+  // A pointer to a structure representing the dimensionality interpretation
+  // that the buffer should have. NOTE: the product of elements of `dims`
+  // and the element datatype size should be equal to `bytes` below.
+  TfLiteIntArray* dims;
+
+  // The number of bytes required to store the data of this Tensor. I.e.
+  // (bytes of each element) * dims[0] * ... * dims[n-1].  For example, if
+  // type is kTfLiteFloat32 and dims = {3, 2} then
+  // bytes = sizeof(float) * 3 * 2 = 4 * 3 * 2 = 24.
+  size_t bytes;
+
+  // The data type specification for data stored in `data`. This affects
+  // what member of `data` union should be used.
+  TfLiteType type;
+
+  // How memory is mapped
+  //  kTfLiteMmapRo: Memory mapped read only.
+  //  i.e. weights
+  //  kTfLiteArenaRw: Arena allocated read write memory
+  //  (i.e. temporaries, outputs).
+  TfLiteAllocationType allocation_type;
+
+  // True if the tensor is a variable.
+  bool is_variable;
+} TfLiteTensor;
+#endif  // TF_LITE_STATIC_MEMORY
 
 #ifndef TF_LITE_STATIC_MEMORY
 // Free data memory of tensor `t`.
diff --git a/tensorflow/lite/micro/memory_arena_threshold_test.cc b/tensorflow/lite/micro/memory_arena_threshold_test.cc
index 4f49b57112a..b45de85a21b 100644
--- a/tensorflow/lite/micro/memory_arena_threshold_test.cc
+++ b/tensorflow/lite/micro/memory_arena_threshold_test.cc
@@ -41,9 +41,17 @@ constexpr int kKeywordModelNodeAndRegistrationCount = 15;
 
 // NOTE: These values are measured on x86-64:
 // TODO(b/158651472): Consider auditing these values on non-64 bit systems.
+//
+// Run this test with '--copt=-DTF_LITE_MICRO_OPTIMIZED_RUNTIME' to get
+// optimized memory runtime values:
+#ifdef TF_LITE_STATIC_MEMORY
+constexpr int kKeywordModelTotalSize = 18448;
+constexpr int kKeywordModelTailSize = 17776;
+#else
 constexpr int kKeywordModelTotalSize = 21040;
-constexpr int kKeywordModelHeadSize = 672;
 constexpr int kKeywordModelTailSize = 20368;
+#endif
+constexpr int kKeywordModelHeadSize = 672;
 constexpr int kKeywordModelTfLiteTensorVariableBufferDataSize = 10240;
 constexpr int kKeywordModelTfLiteTensorQuantizationDataSize = 1728;
 constexpr int kKeywordModelOpRuntimeDataSize = 148;
@@ -56,9 +64,14 @@ constexpr int kTestConvModelNodeAndRegistrationCount = 7;
 
 // NOTE: These values are measured on x86-64:
 // TODO(b/158651472): Consider auditing these values on non-64 bit systems.
+#ifdef TF_LITE_STATIC_MEMORY
+constexpr int kTestConvModelTotalSize = 10960;
+constexpr int kTestConvModelTailSize = 3216;
+#else
 constexpr int kTestConvModelTotalSize = 11680;
-constexpr int kTestConvModelHeadSize = 7744;
 constexpr int kTestConvModelTailSize = 3936;
+#endif
+constexpr int kTestConvModelHeadSize = 7744;
 constexpr int kTestConvModelTfLiteTensorQuantizationDataSize = 768;
 constexpr int kTestConvModelOpRuntimeDataSize = 136;
 
@@ -81,7 +94,7 @@ void EnsureAllocatedSizeThreshold(const char* allocation_type, size_t actual,
     TF_LITE_MICRO_EXPECT_NEAR(actual, expected, kAllocationThreshold);
     if (actual != expected) {
       TF_LITE_REPORT_ERROR(micro_test::reporter,
-                           "%s threshold failed: %ld != %ld", allocation_type,
+                           "%s threshold failed: %d != %d", allocation_type,
                            actual, expected);
     }
   } else {
diff --git a/tensorflow/lite/micro/micro_interpreter_test.cc b/tensorflow/lite/micro/micro_interpreter_test.cc
index f54c212b573..c577d8cb513 100644
--- a/tensorflow/lite/micro/micro_interpreter_test.cc
+++ b/tensorflow/lite/micro/micro_interpreter_test.cc
@@ -284,7 +284,7 @@ TF_LITE_MICRO_TEST(TestIncompleteInitializationAllocationsWithSmallArena) {
 
   tflite::testing::MockOpResolver mock_resolver;
   // 1kb is too small for the ComplexMockModel:
-  constexpr size_t allocator_buffer_size = 1048;
+  constexpr size_t allocator_buffer_size = 500;
   uint8_t allocator_buffer[allocator_buffer_size];
 
   tflite::RecordingMicroAllocator* allocator =
diff --git a/tensorflow/lite/micro/tools/make/Makefile b/tensorflow/lite/micro/tools/make/Makefile
index 8b6cba06a0b..a75c59b05c9 100644
--- a/tensorflow/lite/micro/tools/make/Makefile
+++ b/tensorflow/lite/micro/tools/make/Makefile
@@ -75,6 +75,8 @@ TEST_SCRIPT := tensorflow/lite/micro/testing/test_linux_binary.sh
 MICROLITE_LIBS := -lm
 
 # TODO(b/150240249): Add in -fno-rtti once that works for the Xtensa toolchain.
+# TODO(b/159155203): Consider TF_LITE_STATIC_MEMORY to align more with the fact
+# this flag is for an optimized micro runtime.
 CXXFLAGS := -std=c++11 -DTF_LITE_STATIC_MEMORY
 CCFLAGS  := -std=c11   -DTF_LITE_STATIC_MEMORY
 ARFLAGS := -r
diff --git a/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h b/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
index 15823784d12..9093e5d50ad 100644
--- a/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
+++ b/tensorflow/lite/tools/benchmark/experimental/c/c_api_types.h
@@ -375,6 +375,7 @@ typedef struct TfLiteSparsity {
 
 // An tensor in the interpreter system which is a wrapper around a buffer of
 // data including a dimensionality (or NULL if not currently defined).
+#ifndef TF_LITE_STATIC_MEMORY
 typedef struct TfLiteTensor {
   // The data type specification for data stored in `data`. This affects
   // what member of `data` union should be used.
@@ -439,6 +440,51 @@ typedef struct TfLiteTensor {
   // `dims_signature` contains [1, -1, -1, 3]).
   const TfLiteIntArray* dims_signature;
 } TfLiteTensor;
+#else
+// Specific reduced TfLiteTensor struct for TF Micro runtime. This struct
+// contains only the minimum fields required to initialize and prepare a micro
+// inference graph. The fields in this struct have been ordered from
+// largest-to-smallest for optimal struct sizeof.
+//
+// NOTE: This flag is opt-in only at compile time.
+typedef struct TfLiteTensor {
+  // TODO(b/155784997): Consider consolidating these quantization fields:
+  // Quantization information. Replaces params field above.
+  TfLiteQuantization quantization;
+
+  // Quantization information.
+  TfLiteQuantizationParams params;
+
+  // A union of data pointers. The appropriate type should be used for a typed
+  // tensor based on `type`.
+  TfLitePtrUnion data;
+
+  // A pointer to a structure representing the dimensionality interpretation
+  // that the buffer should have. NOTE: the product of elements of `dims`
+  // and the element datatype size should be equal to `bytes` below.
+  TfLiteIntArray* dims;
+
+  // The number of bytes required to store the data of this Tensor. I.e.
+  // (bytes of each element) * dims[0] * ... * dims[n-1].  For example, if
+  // type is kTfLiteFloat32 and dims = {3, 2} then
+  // bytes = sizeof(float) * 3 * 2 = 4 * 3 * 2 = 24.
+  size_t bytes;
+
+  // The data type specification for data stored in `data`. This affects
+  // what member of `data` union should be used.
+  TfLiteType type;
+
+  // How memory is mapped
+  //  kTfLiteMmapRo: Memory mapped read only.
+  //  i.e. weights
+  //  kTfLiteArenaRw: Arena allocated read write memory
+  //  (i.e. temporaries, outputs).
+  TfLiteAllocationType allocation_type;
+
+  // True if the tensor is a variable.
+  bool is_variable;
+} TfLiteTensor;
+#endif  // TF_LITE_STATIC_MEMORY
 
 #ifndef TF_LITE_STATIC_MEMORY
 // Free data memory of tensor `t`.

From 1332b9365e403baf448358465e6ffd3e5368e614 Mon Sep 17 00:00:00 2001
From: Rahul Joshi 
Date: Fri, 19 Jun 2020 10:42:32 -0700
Subject: [PATCH 0631/1390] [NFC] Adopt FuncOp::isPublic() and friends

PiperOrigin-RevId: 317335755
Change-Id: I6f1b25798150bb3a255d4572831d3be4747d28c7
---
 tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc      | 4 ++--
 .../mlir/tensorflow/transforms/stack_ops_decomposition.cc     | 2 +-
 .../tensorflow/transforms/tensor_array_ops_decomposition.cc   | 2 +-
 .../tensorflow/transforms/tensor_list_ops_decomposition.cc    | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index d59532fef65..ef248379d2e 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -231,12 +231,12 @@ static LogicalResult VerifySavedModelModule(
   for (auto func : module.getOps()) {
     const bool is_exported = IsExported(func);
 
-    if (is_exported && func.getVisibility() != FuncOp::Visibility::Public) {
+    if (is_exported && !func.isPublic()) {
       return func.emitError()
              << "exported function @" << func.getName() << " should be public";
     }
 
-    if (!is_exported && func.getVisibility() == FuncOp::Visibility::Public) {
+    if (!is_exported && func.isPublic()) {
       return func.emitError() << "non-exported function @" << func.getName()
                               << " should be private";
     }
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/stack_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/stack_ops_decomposition.cc
index c349c2b4c3e..734a7d04a86 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/stack_ops_decomposition.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/stack_ops_decomposition.cc
@@ -343,7 +343,7 @@ LogicalResult HandlePartitionedCallOp(
   }
   llvm::SmallDenseMap callee_map;
   FuncOp lowered_callee = callee;
-  if (callee.getVisibility() != SymbolTable::Visibility::Private) {
+  if (!callee.isPrivate()) {
     // Clone non-private callee in case of signature change.
     lowered_callee = callee.clone();
     lowered_callee.setVisibility(SymbolTable::Visibility::Private);
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_array_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_array_ops_decomposition.cc
index cfeb2b1f031..a9e1243714e 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_array_ops_decomposition.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_array_ops_decomposition.cc
@@ -759,7 +759,7 @@ LogicalResult HandlePartitionedCallOp(
     return it->getSecond().accumulate_on_write;
   };
   FuncOp lowered_callee = callee;
-  if (callee.getVisibility() != SymbolTable::Visibility::Private) {
+  if (!callee.isPrivate()) {
     // Clone non-private callee in case of signature change.
     lowered_callee = callee.clone();
     lowered_callee.setVisibility(SymbolTable::Visibility::Private);
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc
index 9733bfe2290..b118ab6c6c9 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc
@@ -322,7 +322,7 @@ LogicalResult HandlePartitionedCallOp(
   // Rewrite the callee.
   llvm::SmallDenseMap callee_map;
   FuncOp lowered_callee = callee;
-  if (callee.getVisibility() != SymbolTable::Visibility::Private) {
+  if (!callee.isPrivate()) {
     // Clone non-private callee in case of signature change.
     lowered_callee = callee.clone();
     lowered_callee.setVisibility(SymbolTable::Visibility::Private);

From 89de554b69d9b74896445a41250943d97e3b9e77 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Fri, 19 Jun 2020 10:43:39 -0700
Subject: [PATCH 0632/1390] Fork keras related mirror_variable_test to
 keras/distribute.

PiperOrigin-RevId: 317335959
Change-Id: Ie34b8ddecdd137926959868157ab7ade7c72e1b9
---
 tensorflow/python/distribute/BUILD            |   1 -
 .../distribute/mirrored_variable_test.py      |  34 ------
 tensorflow/python/keras/distribute/BUILD      |  22 ++++
 .../distribute/mirrored_variable_test.py      | 106 ++++++++++++++++++
 4 files changed, 128 insertions(+), 35 deletions(-)
 create mode 100644 tensorflow/python/keras/distribute/mirrored_variable_test.py

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 38c5550be16..f0f3766afe1 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -1511,7 +1511,6 @@ cuda_py_test(
         "//tensorflow/python:variable_scope",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:test",
-        "//tensorflow/python/keras/layers",
     ],
 )
 
diff --git a/tensorflow/python/distribute/mirrored_variable_test.py b/tensorflow/python/distribute/mirrored_variable_test.py
index 6623422b45f..df32a6babea 100644
--- a/tensorflow/python/distribute/mirrored_variable_test.py
+++ b/tensorflow/python/distribute/mirrored_variable_test.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.distribute import collective_all_reduce_strategy
 from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import distribution_strategy_context as ds_context
@@ -32,7 +31,6 @@ from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import func_graph
 from tensorflow.python.framework import ops
-from tensorflow.python.keras.layers import core
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import rnn
@@ -208,38 +206,6 @@ class MirroredVariableCreationTest(test.TestCase):
       # The resulting mirrored variable will use the name from the first device.
       self.assertEqual("foo_0:0", result.name)
 
-  def testWithLayers(self, distribution):
-
-    def model_fn(features):
-
-      layer1 = core.Dense(1)
-      layer1(features)
-      layer2 = core.Dense(1)
-      layer2(features)
-      # We rely on names and orders to make sure replica references the same
-      # MirroredVariable. Uniquifying names may involve global states,
-      # merge_call switches threads so we need to test things work after
-      # merge_call.
-      ds_context.get_replica_context().merge_call(lambda _: _)
-      layer3 = core.Dense(1)
-      layer3(features)
-      return [(layer1.kernel, layer1.bias), (layer2.kernel, layer2.bias),
-              (layer3.kernel, layer3.bias)]
-
-    iterator = distribution.make_input_fn_iterator(
-        lambda _: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10))
-    self.evaluate(iterator.initializer)
-    features = iterator.get_next()
-
-    with distribution.scope():
-      result = distribution.extended.call_for_each_replica(
-          model_fn, args=(features,))
-      for kernel, bias in result:
-        self.assertIsInstance(kernel, values.MirroredVariable)
-        self.assertAllDifferent(distribution.experimental_local_results(kernel))
-        self.assertIsInstance(bias, values.MirroredVariable)
-        self.assertAllDifferent(distribution.experimental_local_results(kernel))
-
   def testWithVariableAndVariableScope(self, distribution):
 
     def model_fn():
diff --git a/tensorflow/python/keras/distribute/BUILD b/tensorflow/python/keras/distribute/BUILD
index 247e655621c..c6a8f2c5f91 100644
--- a/tensorflow/python/keras/distribute/BUILD
+++ b/tensorflow/python/keras/distribute/BUILD
@@ -347,6 +347,28 @@ cuda_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "mirrored_variable_test",
+    srcs = ["mirrored_variable_test.py"],
+    python_version = "PY3",
+    tags = [
+        "guitar",
+        "multi_and_single_gpu",
+    ],
+    deps = [
+        "//tensorflow/python:config",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:strategy_combinations",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/keras/layers:core",
+    ],
+)
+
 cuda_py_test(
     name = "multi_worker_test",
     srcs = ["multi_worker_test.py"],
diff --git a/tensorflow/python/keras/distribute/mirrored_variable_test.py b/tensorflow/python/keras/distribute/mirrored_variable_test.py
new file mode 100644
index 00000000000..0edfa4806f2
--- /dev/null
+++ b/tensorflow/python/keras/distribute/mirrored_variable_test.py
@@ -0,0 +1,106 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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 MirroredVariable in MirroredStrategy and MultiWorkerMirroredStrategy."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribution_strategy_context as ds_context
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import values
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.keras.layers import core
+
+
+def _mimic_two_cpus():
+  cpus = config.list_physical_devices("CPU")
+
+  config.set_logical_device_configuration(cpus[0], [
+      context.LogicalDeviceConfiguration(),
+      context.LogicalDeviceConfiguration(),
+  ])
+
+
+@combinations.generate(
+    combinations.combine(
+        distribution=[
+            strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+            combinations.NamedDistribution(
+                "Collective2CPUs",
+                # pylint: disable=g-long-lambda
+                lambda: collective_all_reduce_strategy.
+                CollectiveAllReduceStrategy._from_local_devices((
+                    "/device:CPU:0", "/device:CPU:1")),
+                required_gpus=0)
+        ],
+        mode=["graph", "eager"]))
+class MirroredVariableCreationTest(test.TestCase):
+  """Base class that tests mirrored variable creator.
+
+  Currently it assumes all strategy objects have two replicas.
+  """
+
+  @classmethod
+  def setUpClass(cls):
+    _mimic_two_cpus()
+
+  def assertAllDifferent(self, objs):
+    for i in range(len(objs)):
+      for j in range(len(objs)):
+        if i == j:
+          continue
+        self.assertIsNot(objs[i], objs[j])
+
+  def testWithLayers(self, distribution):
+
+    def model_fn(features):
+
+      layer1 = core.Dense(1)
+      layer1(features)
+      layer2 = core.Dense(1)
+      layer2(features)
+      # We rely on names and orders to make sure replica references the same
+      # MirroredVariable. Uniquifying names may involve global states,
+      # merge_call switches threads so we need to test things work after
+      # merge_call.
+      ds_context.get_replica_context().merge_call(lambda _: _)
+      layer3 = core.Dense(1)
+      layer3(features)
+      return [(layer1.kernel, layer1.bias), (layer2.kernel, layer2.bias),
+              (layer3.kernel, layer3.bias)]
+
+    iterator = distribution.make_input_fn_iterator(
+        lambda _: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10))
+    self.evaluate(iterator.initializer)
+    features = iterator.get_next()
+
+    with distribution.scope():
+      result = distribution.extended.call_for_each_replica(
+          model_fn, args=(features,))
+      for kernel, bias in result:
+        self.assertIsInstance(kernel, values.MirroredVariable)
+        self.assertAllDifferent(distribution.experimental_local_results(kernel))
+        self.assertIsInstance(bias, values.MirroredVariable)
+        self.assertAllDifferent(distribution.experimental_local_results(kernel))
+
+
+if __name__ == "__main__":
+  test.main()

From 8088eddf203220799cd51ede0142329fa4665d3c Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Fri, 19 Jun 2020 10:58:32 -0700
Subject: [PATCH 0633/1390] Let tensorflow op take precedence when doing
 "ndarray  tensor"

Also add a few more interop tests.

PiperOrigin-RevId: 317339113
Change-Id: Ic28fab7abefea681e1e8d840b8e4cf4f98b63f1e
---
 tensorflow/python/ops/numpy_ops/np_arrays.py  |  3 +-
 .../python/ops/numpy_ops/np_interop_test.py   | 50 ++++++++++++++++++
 .../python/ops/numpy_ops/np_math_ops.py       | 52 +++++++++++--------
 3 files changed, 82 insertions(+), 23 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index 88bf4e7499a..65e8273375f 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -262,7 +262,8 @@ class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
     """
     return np.asarray(self.data, dtype)
 
-  __array_priority__ = 110
+  # NOTE: we currently prefer interop with TF to allow TF to take precedence.
+  __array_priority__ = 90
 
   def __index__(self):
     """Returns a python scalar.
diff --git a/tensorflow/python/ops/numpy_ops/np_interop_test.py b/tensorflow/python/ops/numpy_ops/np_interop_test.py
index 052949dff9d..f52d3dae78b 100644
--- a/tensorflow/python/ops/numpy_ops/np_interop_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_interop_test.py
@@ -19,12 +19,18 @@ from __future__ import division
 from __future__ import print_function
 
 
+import numpy as onp
+
+
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.numpy_ops import np_array_ops
 from tensorflow.python.ops.numpy_ops import np_arrays
+from tensorflow.python.ops.numpy_ops import np_math_ops
 from tensorflow.python.platform import test
 
 
@@ -88,6 +94,50 @@ class InteropTest(test.TestCase):
     self.assertEqual(10000, fn()[0])
     self.assertEqual(10000, def_function.function(fn)()[0])
 
+  def testTensorTFNPArrayInterop(self):
+    arr = np_array_ops.asarray(0.)
+    t = constant_op.constant(10.)
+
+    arr_plus_t = arr + t
+    t_plus_arr = t + arr
+
+    self.assertIsInstance(arr_plus_t, ops.Tensor)
+    self.assertIsInstance(t_plus_arr, ops.Tensor)
+    self.assertEqual(10., arr_plus_t.numpy())
+    self.assertEqual(10., t_plus_arr.numpy())
+
+  def testTensorTFNPOp(self):
+    t = constant_op.constant(10.)
+
+    sq = np_math_ops.square(t)
+    self.assertIsInstance(sq, np_arrays.ndarray)
+    self.assertEqual(100., sq)
+
+  def testTFNPArrayTFOpInterop(self):
+    arr = np_array_ops.asarray(10.)
+
+    # TODO(nareshmodi): Test more ops.
+    sq = math_ops.square(arr)
+    self.assertIsInstance(sq, ops.Tensor)
+    self.assertEqual(100., sq.numpy())
+
+  def testTFNPArrayNPOpInterop(self):
+    arr = np_array_ops.asarray([10.])
+
+    # TODO(nareshmodi): Test more ops.
+    sq = onp.square(arr)
+    self.assertIsInstance(sq, onp.ndarray)
+    self.assertEqual(100., sq[0])
+
+    # TODO(nareshmodi): Fails since the autopacking code doesn't use
+    # nest.flatten.
+#   def testAutopacking(self):
+#     arr1 = np_array_ops.asarray(1.)
+#     arr2 = np_array_ops.asarray(2.)
+#     arr3 = np_array_ops.asarray(3.)
+#     t = ops.convert_to_tensor_v2([arr1, arr2, arr3])
+
+#     self.assertEqual(t.numpy(), [1., 2., 3.])
 
 if __name__ == '__main__':
   ops.enable_eager_execution()
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index abfd9087ffd..361bfb50dec 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -917,29 +917,37 @@ def diff(a, n=1, axis=-1):  # pylint: disable=missing-function-docstring
   return _scalar(f, a)
 
 
-def _flip_args(f):
+def _wrap(f, reverse=False):
+  """Wraps binary ops so they can be added as operator overloads on ndarray."""
 
   def _f(a, b):
-    return f(b, a)
+    if reverse:
+      a, b = b, a
+
+    if getattr(b, '__array_priority__',
+               0) > np_arrays.ndarray.__array_priority__:
+      return NotImplemented
+
+    return f(a, b)
 
   return _f
 
 
 setattr(np_arrays.ndarray, '__abs__', absolute)
-setattr(np_arrays.ndarray, '__floordiv__', floor_divide)
-setattr(np_arrays.ndarray, '__rfloordiv__', _flip_args(floor_divide))
-setattr(np_arrays.ndarray, '__mod__', mod)
-setattr(np_arrays.ndarray, '__rmod__', _flip_args(mod))
-setattr(np_arrays.ndarray, '__add__', add)
-setattr(np_arrays.ndarray, '__radd__', _flip_args(add))
-setattr(np_arrays.ndarray, '__sub__', subtract)
-setattr(np_arrays.ndarray, '__rsub__', _flip_args(subtract))
-setattr(np_arrays.ndarray, '__mul__', multiply)
-setattr(np_arrays.ndarray, '__rmul__', _flip_args(multiply))
-setattr(np_arrays.ndarray, '__pow__', power)
-setattr(np_arrays.ndarray, '__rpow__', _flip_args(power))
-setattr(np_arrays.ndarray, '__truediv__', true_divide)
-setattr(np_arrays.ndarray, '__rtruediv__', _flip_args(true_divide))
+setattr(np_arrays.ndarray, '__floordiv__', _wrap(floor_divide))
+setattr(np_arrays.ndarray, '__rfloordiv__', _wrap(floor_divide, True))
+setattr(np_arrays.ndarray, '__mod__', _wrap(mod))
+setattr(np_arrays.ndarray, '__rmod__', _wrap(mod, True))
+setattr(np_arrays.ndarray, '__add__', _wrap(add))
+setattr(np_arrays.ndarray, '__radd__', _wrap(add, True))
+setattr(np_arrays.ndarray, '__sub__', _wrap(subtract))
+setattr(np_arrays.ndarray, '__rsub__', _wrap(subtract, True))
+setattr(np_arrays.ndarray, '__mul__', _wrap(multiply))
+setattr(np_arrays.ndarray, '__rmul__', _wrap(multiply, True))
+setattr(np_arrays.ndarray, '__pow__', _wrap(power))
+setattr(np_arrays.ndarray, '__rpow__', _wrap(power, True))
+setattr(np_arrays.ndarray, '__truediv__', _wrap(true_divide))
+setattr(np_arrays.ndarray, '__rtruediv__', _wrap(true_divide, True))
 
 
 def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):
@@ -1031,12 +1039,12 @@ def logical_not(x):
 
 
 setattr(np_arrays.ndarray, '__invert__', logical_not)
-setattr(np_arrays.ndarray, '__lt__', less)
-setattr(np_arrays.ndarray, '__le__', less_equal)
-setattr(np_arrays.ndarray, '__gt__', greater)
-setattr(np_arrays.ndarray, '__ge__', greater_equal)
-setattr(np_arrays.ndarray, '__eq__', equal)
-setattr(np_arrays.ndarray, '__ne__', not_equal)
+setattr(np_arrays.ndarray, '__lt__', _wrap(less))
+setattr(np_arrays.ndarray, '__le__', _wrap(less_equal))
+setattr(np_arrays.ndarray, '__gt__', _wrap(greater))
+setattr(np_arrays.ndarray, '__ge__', _wrap(greater_equal))
+setattr(np_arrays.ndarray, '__eq__', _wrap(equal))
+setattr(np_arrays.ndarray, '__ne__', _wrap(not_equal))
 
 
 @np_utils.np_doc(np.linspace)

From 9152edc1f09e7420771628fd55ffef4b35b1cbc3 Mon Sep 17 00:00:00 2001
From: Bixia Zheng 
Date: Fri, 19 Jun 2020 11:04:49 -0700
Subject: [PATCH 0634/1390] [TF:TRT] Add flag
 TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION.

The default value of the flag is True. When the flag value is false, the
bridge will report an error when the native segment of a TRTEngineOp is
executed.

Add test cases.

PiperOrigin-RevId: 317340632
Change-Id: Iacded09b38e63442bbd93076a079d385fb8a77e6
---
 .../tf2tensorrt/kernels/trt_engine_op.cc      | 34 ++++++++++++++++---
 .../compiler/tensorrt/trt_convert_test.py     | 33 +++++++++++++++++-
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index ac4a331041d..98d199ca9ab 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -45,6 +45,7 @@ limitations under the License.
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
 #if GOOGLE_CUDA
@@ -521,6 +522,17 @@ Status TRTEngineOp::VerifyInputShapes(
   return Status::OK();
 }
 
+static bool AllowEngineNativeSegmentExecution() {
+  bool value;
+  Status status =
+      ReadBoolFromEnvVar("TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION",
+                         /*default_value=*/true, &value);
+  if (!status.ok()) {
+    LOG(ERROR) << status;
+  }
+  return value;
+}
+
 void TRTEngineOp::ComputeAsync(OpKernelContext* ctx,
                                AsyncOpKernel::DoneCallback done) {
   auto helper = new AsyncHelper(done);
@@ -605,17 +617,31 @@ void TRTEngineOp::ComputeAsync(OpKernelContext* ctx,
 
   EngineContext* engine_context = status.ValueOrDie().first;
   int trt_context_idx = status.ValueOrDie().second;
+  auto may_execute_native_segment = [&] {
+    if (!AllowEngineNativeSegmentExecution()) {
+      ctx->CtxFailure(
+          errors::Aborted("User disallowed engine native segment execution"));
+      return false;
+    }
+    return true;
+  };
   if (!engine_context->cuda_engine) {
-    VLOG(1) << "Engine retrieval for input shapes: "
-            << TensorShapeUtils::ShapeListString(input_concrete_shapes)
-            << " failed. Running native segment for " << name();
-    ExecuteNativeSegment(ctx, helper);
+    LOG_WARNING_WITH_PREFIX
+        << "Engine retrieval for input shapes: "
+        << TensorShapeUtils::ShapeListString(input_concrete_shapes)
+        << " failed. Running native segment for " << name();
+    if (may_execute_native_segment()) {
+      ExecuteNativeSegment(ctx, helper);
+    }
     return;
   }
   Status stat = ExecuteTrtEngine(ctx, engine_context, trt_context_idx);
   if (!stat.ok()) {
     LOG_WARNING_WITH_PREFIX << "Failed to execute engine: " << stat
                             << " Retrying with native segment for " << name();
+    if (!may_execute_native_segment()) {
+      return;
+    }
     // Release any outputs that are allocated, ExecuteNativeSegment will
     // re-allocate them and fail if they are currently allocated.
     for (int i = 0; i < ctx->num_outputs(); i++) {
diff --git a/tensorflow/python/compiler/tensorrt/trt_convert_test.py b/tensorflow/python/compiler/tensorrt/trt_convert_test.py
index df21e93f836..05ff6fcaebe 100644
--- a/tensorflow/python/compiler/tensorrt/trt_convert_test.py
+++ b/tensorflow/python/compiler/tensorrt/trt_convert_test.py
@@ -439,6 +439,7 @@ class TrtConvertTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       self,
       input_saved_model_dir,
       input_saved_model_signature_key=_SAVED_MODEL_SIGNATURE_KEY,
+      max_workspace_size_bytes=10 << 20,  # Use a smaller workspace.
       precision_mode=trt_convert.TrtPrecisionMode.FP32,
       is_dynamic_op=True,
       maximum_cached_engines=2):
@@ -446,7 +447,7 @@ class TrtConvertTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         input_saved_model_dir=input_saved_model_dir,
         input_saved_model_signature_key=input_saved_model_signature_key,
         conversion_params=trt_convert.DEFAULT_TRT_CONVERSION_PARAMS._replace(
-            max_workspace_size_bytes=10 << 20,  # Use a smaller workspace.
+            max_workspace_size_bytes=max_workspace_size_bytes,
             precision_mode=precision_mode,
             is_dynamic_op=is_dynamic_op,
             maximum_cached_engines=maximum_cached_engines))
@@ -924,6 +925,36 @@ class TrtConvertTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         # to fall back to TF function.
         self._TestRun(sess, 2)
 
+  @test_util.run_v2_only
+  def testTrtGraphConverter_AllowEngineNativeSegmentExecution(self):
+    if not is_tensorrt_enabled():
+      return
+
+    np_input1, np_input2 = self._RandomInput([4, 1, 1])
+
+    # Create a model and save it.
+    input_saved_model_dir = self.mkdtemp()
+    root = self._GetModelForV2()
+    save.save(root, input_saved_model_dir,
+              {_SAVED_MODEL_SIGNATURE_KEY: root.run})
+
+    def _InputFn():
+      yield np_input1, np_input2
+
+    # Run TRT conversion and request an unreasonably large workspace.
+    converter = self._CreateConverterV2(
+        input_saved_model_dir, max_workspace_size_bytes=10 << 40)
+    converter.convert()
+
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
+    with self.assertRaisesRegex(
+        errors.AbortedError,
+        r"User disallowed engine native segment execution"):
+      converter.build(input_fn=_InputFn)
+
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+    converter.build(input_fn=_InputFn)
+
   @test_util.run_v2_only
   def testBackwardCompatibility(self):
     """Load and execute a model that was saved in TF2.0."""

From 1129f21360b24fd957a248b86ddefa311d9f5658 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Fri, 19 Jun 2020 11:16:53 -0700
Subject: [PATCH 0635/1390] Consistently call LSTM's gate biases in variable
 names _gate_bias.

PiperOrigin-RevId: 317343158
Change-Id: I385ddbad6c1283b84574b2ec0b523ce9f88a4cd3
---
 .../kernels/bidirectional_sequence_lstm.cc    | 20 ++--
 .../bidirectional_sequence_lstm_test.cc       | 48 +++++-----
 tensorflow/lite/kernels/lstm.cc               | 56 ++++++-----
 tensorflow/lite/kernels/lstm_eval.cc          | 64 +++++++------
 tensorflow/lite/kernels/lstm_eval.h           |  8 +-
 tensorflow/lite/kernels/lstm_eval_test.cc     | 93 ++++++++++---------
 tensorflow/lite/kernels/lstm_test.cc          | 70 +++++++-------
 .../lite/kernels/optional_tensor_test.cc      |  8 +-
 .../kernels/unidirectional_sequence_lstm.cc   | 18 ++--
 .../unidirectional_sequence_lstm_test.cc      | 32 +++----
 .../calibration/builtin_logging_ops/lstm.cc   | 13 +--
 11 files changed, 224 insertions(+), 206 deletions(-)

diff --git a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
index 439fc94afad..fd60fe573ef 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
@@ -318,11 +318,11 @@ TfLiteStatus CheckLstmTensorDimensionsAndTypes(
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell);
   TF_LITE_ENSURE_TYPES_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
 
-  const TfLiteTensor* cell_bias =
+  const TfLiteTensor* cell_gate_bias =
       GetInput(context, node, cell_gate_bias_tensor);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell);
-  TF_LITE_ENSURE_EQ(context, cell_bias->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->size, 1);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->data[0], n_cell);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->type, kTfLiteFloat32);
 
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, output_gate_bias_tensor);
@@ -886,7 +886,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       GetOptionalInputTensor(context, node, kFwInputGateBiasTensor);
   const TfLiteTensor* fw_forget_gate_bias =
       GetInput(context, node, kFwForgetGateBiasTensor);
-  const TfLiteTensor* fw_cell_bias =
+  const TfLiteTensor* fw_cell_gate_bias =
       GetInput(context, node, kFwCellGateBiasTensor);
   const TfLiteTensor* fw_output_gate_bias =
       GetInput(context, node, kFwOutputGateBiasTensor);
@@ -934,7 +934,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       GetOptionalInputTensor(context, node, kBwInputGateBiasTensor);
   const TfLiteTensor* bw_forget_gate_bias =
       GetInput(context, node, kBwForgetGateBiasTensor);
-  const TfLiteTensor* bw_cell_bias =
+  const TfLiteTensor* bw_cell_gate_bias =
       GetInput(context, node, kBwCellGateBiasTensor);
   const TfLiteTensor* bw_output_gate_bias =
       GetInput(context, node, kBwOutputGateBiasTensor);
@@ -1029,7 +1029,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*output_layer_norm_coefficients=*/nullptr, real_aux_input,
           fw_aux_input_to_input_weights, fw_aux_input_to_forget_weights,
           fw_aux_input_to_cell_weights, fw_aux_input_to_output_weights,
-          fw_input_gate_bias, fw_forget_gate_bias, fw_cell_bias,
+          fw_input_gate_bias, fw_forget_gate_bias, fw_cell_gate_bias,
           fw_output_gate_bias, fw_projection_weights, fw_projection_bias,
           &lstm_params,
           /*forward_sequence=*/true, time_major, /*output_offset=*/0,
@@ -1049,7 +1049,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*output_layer_norm_coefficients=*/nullptr, real_aux_input,
           bw_aux_input_to_input_weights, bw_aux_input_to_forget_weights,
           bw_aux_input_to_cell_weights, bw_aux_input_to_output_weights,
-          bw_input_gate_bias, bw_forget_gate_bias, bw_cell_bias,
+          bw_input_gate_bias, bw_forget_gate_bias, bw_cell_gate_bias,
           bw_output_gate_bias, bw_projection_weights, bw_projection_bias,
           &lstm_params,
           /*forward_sequence=*/false, time_major, bw_output_offset,
@@ -1099,7 +1099,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*output_layer_norm_coefficients=*/nullptr, real_aux_input,
           fw_aux_input_to_input_weights, fw_aux_input_to_forget_weights,
           fw_aux_input_to_cell_weights, fw_aux_input_to_output_weights,
-          fw_input_gate_bias, fw_forget_gate_bias, fw_cell_bias,
+          fw_input_gate_bias, fw_forget_gate_bias, fw_cell_gate_bias,
           fw_output_gate_bias, fw_projection_weights, fw_projection_bias,
           &lstm_params,
           /*forward_sequence=*/true, /*time_major=*/true, /*output_offset=*/0,
@@ -1125,7 +1125,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*output_layer_norm_coefficients=*/nullptr, real_aux_input,
           bw_aux_input_to_input_weights, bw_aux_input_to_forget_weights,
           bw_aux_input_to_cell_weights, bw_aux_input_to_output_weights,
-          bw_input_gate_bias, bw_forget_gate_bias, bw_cell_bias,
+          bw_input_gate_bias, bw_forget_gate_bias, bw_cell_gate_bias,
           bw_output_gate_bias, bw_projection_weights, bw_projection_bias,
           &lstm_params,
           /*forward_sequence=*/false, /*time_major=*/true, bw_output_offset,
diff --git a/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc b/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
index 3a52de130e3..778751aa04b 100644
--- a/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
+++ b/tensorflow/lite/kernels/bidirectional_sequence_lstm_test.cc
@@ -89,7 +89,7 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
       fw_input_gate_bias_ = AddInput(TensorType_FLOAT32);
     }
     fw_forget_gate_bias_ = AddInput(TensorType_FLOAT32);
-    fw_cell_bias_ = AddInput(TensorType_FLOAT32);
+    fw_cell_gate_bias_ = AddInput(TensorType_FLOAT32);
     fw_output_gate_bias_ = AddInput(TensorType_FLOAT32);
 
     if (use_projection_weights) {
@@ -144,7 +144,7 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
       bw_input_gate_bias_ = AddInput(TensorType_FLOAT32);
     }
     bw_forget_gate_bias_ = AddInput(TensorType_FLOAT32);
-    bw_cell_bias_ = AddInput(TensorType_FLOAT32);
+    bw_cell_gate_bias_ = AddInput(TensorType_FLOAT32);
     bw_output_gate_bias_ = AddInput(TensorType_FLOAT32);
 
     if (use_projection_weights) {
@@ -288,8 +288,8 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
   }
 
   void SetCellBias(const std::vector& f) {
-    PopulateTensor(fw_cell_bias_, f);
-    PopulateTensor(bw_cell_bias_, f);
+    PopulateTensor(fw_cell_gate_bias_, f);
+    PopulateTensor(bw_cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(const std::vector& f) {
@@ -364,7 +364,7 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
 
   int fw_input_gate_bias_;
   int fw_forget_gate_bias_;
-  int fw_cell_bias_;
+  int fw_cell_gate_bias_;
   int fw_output_gate_bias_;
 
   int fw_projection_weights_;
@@ -386,7 +386,7 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
 
   int bw_input_gate_bias_;
   int bw_forget_gate_bias_;
-  int bw_cell_bias_;
+  int bw_cell_gate_bias_;
   int bw_output_gate_bias_;
 
   int bw_projection_weights_;
@@ -467,7 +467,7 @@ TEST_P(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClipping) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -490,7 +490,7 @@ TEST_P(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClipping) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -633,7 +633,7 @@ TEST_P(LSTMOpTest, BlackBoxTestMergedOutput) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -656,7 +656,7 @@ TEST_P(LSTMOpTest, BlackBoxTestMergedOutput) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -796,7 +796,7 @@ TEST(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClippingReverse) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -819,7 +819,7 @@ TEST(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClippingReverse) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -956,7 +956,7 @@ TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) {
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -978,7 +978,7 @@ TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) {
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -1107,7 +1107,7 @@ TEST(LSTMOpTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -1129,7 +1129,7 @@ TEST(LSTMOpTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -1258,7 +1258,7 @@ TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClipping) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1280,7 +1280,7 @@ TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClipping) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1961,7 +1961,7 @@ TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClippingBatchMajor) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1983,7 +1983,7 @@ TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClippingBatchMajor) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2667,7 +2667,7 @@ TEST_P(LSTMOpTest, BlackBoxTestWithAuxInputZeroAuxWeight) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -2690,7 +2690,7 @@ TEST_P(LSTMOpTest, BlackBoxTestWithAuxInputZeroAuxWeight) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -2841,7 +2841,7 @@ TEST_P(LSTMOpTest, BlackBoxTestWithAuxInput) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -2864,7 +2864,7 @@ TEST_P(LSTMOpTest, BlackBoxTestWithAuxInput) {
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index aa6a112a022..b941f2237ca 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -407,7 +407,8 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
       GetOptionalInputTensor(context, node, kInputGateBiasTensor);
   const TfLiteTensor* forget_gate_bias =
       GetInput(context, node, kForgetGateBiasTensor);
-  const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
+  const TfLiteTensor* cell_gate_bias =
+      GetInput(context, node, kCellGateBiasTensor);
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, kOutputGateBiasTensor);
 
@@ -446,10 +447,10 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   int16_t* layer_norm_forget_weight_ptr = nullptr;
   int16_t* layer_norm_cell_weight_ptr = nullptr;
   int16_t* layer_norm_output_weight_ptr = nullptr;
-  int32_t* input_bias_ptr = nullptr;
-  int32_t* forget_bias_ptr = nullptr;
-  int32_t* cell_bias_ptr = nullptr;
-  int32_t* output_bias_ptr = nullptr;
+  int32_t* input_gate_bias_ptr = nullptr;
+  int32_t* forget_gate_bias_ptr = nullptr;
+  int32_t* cell_gate_bias_ptr = nullptr;
+  int32_t* output_gate_bias_ptr = nullptr;
   int32_t* proj_bias_ptr = nullptr;
   int16_t* cell_ptr = nullptr;
   int8_t* output_state_ptr = nullptr;
@@ -497,7 +498,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   if (!use_cifg) {
     input_to_input_weight_ptr = input_to_input_weights->data.int8;
     recurrent_to_input_weight_ptr = recurrent_to_input_weights->data.int8;
-    input_bias_ptr = input_gate_bias->data.i32;
+    input_gate_bias_ptr = input_gate_bias->data.i32;
     input_to_input_weight_scale = input_to_input_weights->params.scale;
     recurrent_to_input_weight_scale = recurrent_to_input_weights->params.scale;
   }
@@ -547,9 +548,9 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   recurrent_to_cell_weight_scale = recurrent_to_cell_weights->params.scale;
   recurrent_to_output_weight_ptr = recurrent_to_output_weights->data.int8;
   recurrent_to_output_weight_scale = recurrent_to_output_weights->params.scale;
-  forget_bias_ptr = forget_gate_bias->data.i32;
-  cell_bias_ptr = cell_bias->data.i32;
-  output_bias_ptr = output_gate_bias->data.i32;
+  forget_gate_bias_ptr = forget_gate_bias->data.i32;
+  cell_gate_bias_ptr = cell_gate_bias->data.i32;
+  output_gate_bias_ptr = output_gate_bias->data.i32;
   output_state_ptr = output_state->data.int8;
   cell_ptr = cell_state->data.i16;
   input_scale = input->params.scale;
@@ -875,13 +876,14 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
     TF_LITE_ENSURE_TYPES_EQ(context, forget_gate_bias->type, kTfLiteFloat32);
   }
 
-  const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell);
+  const TfLiteTensor* cell_gate_bias =
+      GetInput(context, node, kCellGateBiasTensor);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->size, 1);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->data[0], n_cell);
   if (is_integer) {
-    TF_LITE_ENSURE_TYPES_EQ(context, cell_bias->type, kTfLiteInt32);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_gate_bias->type, kTfLiteInt32);
   } else {
-    TF_LITE_ENSURE_TYPES_EQ(context, cell_bias->type, kTfLiteFloat32);
+    TF_LITE_ENSURE_TYPES_EQ(context, cell_gate_bias->type, kTfLiteFloat32);
   }
 
   const TfLiteTensor* output_gate_bias =
@@ -1526,7 +1528,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       GetOptionalInputTensor(context, node, kInputGateBiasTensor);
   const TfLiteTensor* forget_gate_bias =
       GetInput(context, node, kForgetGateBiasTensor);
-  const TfLiteTensor* cell_bias = GetInput(context, node, kCellGateBiasTensor);
+  const TfLiteTensor* cell_gate_bias =
+      GetInput(context, node, kCellGateBiasTensor);
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, kOutputGateBiasTensor);
 
@@ -1560,8 +1563,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*aux_input_to_forget_weights=*/nullptr,
           /*aux_input_to_cell_weights=*/nullptr,
           /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
-          forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
-          projection_bias, params, /*forward_sequence=*/true,
+          forget_gate_bias, cell_gate_bias, output_gate_bias,
+          projection_weights, projection_bias, params,
+          /*forward_sequence=*/true,
           /*time_major=*/true,
           /*output_offset=*/0, scratch_buffer, output_state, cell_state,
           output);
@@ -1603,8 +1607,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
             /*aux_input_to_forget_weights=*/nullptr,
             /*aux_input_to_cell_weights=*/nullptr,
             /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
-            forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
-            projection_bias, params, /*forward_sequence=*/true,
+            forget_gate_bias, cell_gate_bias, output_gate_bias,
+            projection_weights, projection_bias, params,
+            /*forward_sequence=*/true,
             /*time_major=*/true, /*output_offset=*/0, scratch_buffer,
             scaling_factors, prod_scaling_factors, recovered_cell_weights,
             input_quantized,
@@ -1631,10 +1636,11 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               cell_to_output_weights, input_layer_norm_coefficients,
               forget_layer_norm_coefficients, cell_layer_norm_coefficients,
               output_layer_norm_coefficients, input_gate_bias, forget_gate_bias,
-              cell_bias, output_gate_bias, projection_weights, projection_bias,
-              params, &op_data->integer_lstm_param, output_state, cell_state,
-              output, scratch0, scratch1, scratch2, scratch3, scratch4,
-              scratch5, CpuBackendContext::GetFromContext(context));
+              cell_gate_bias, output_gate_bias, projection_weights,
+              projection_bias, params, &op_data->integer_lstm_param,
+              output_state, cell_state, output, scratch0, scratch1, scratch2,
+              scratch3, scratch4, scratch5,
+              CpuBackendContext::GetFromContext(context));
         } else {
           TfLiteTensor* scratch0 = GetTemporary(context, node, /*index=*/0);
           TfLiteTensor* scratch1 = GetTemporary(context, node, /*index=*/1);
@@ -1653,8 +1659,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
               cell_to_output_weights, input_layer_norm_coefficients,
               forget_layer_norm_coefficients, cell_layer_norm_coefficients,
               output_layer_norm_coefficients, input_gate_bias, forget_gate_bias,
-              cell_bias, output_gate_bias, projection_weights, projection_bias,
-              params, output_state, cell_state, output,
+              cell_gate_bias, output_gate_bias, projection_weights,
+              projection_bias, params, output_state, cell_state, output,
               &op_data->integer_lstm_param, scratch0, scratch1, scratch2,
               scratch3, scratch4, scratch5, scratch6, scratch7);
           return kTfLiteOk;
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index f45d46762bf..9bdbfa9d48d 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -942,10 +942,10 @@ inline void LstmStepHybrid(
 //   effective_proj_scale_b                  - optional
 //
 // Gate biases of size 'n_cell':
-//   input_bias_ptr                 - optional
-//   forget_bias_ptr
+//   input_gate_bias_ptr                 - optional
+//   forget_gate_bias_ptr
 //   cell_gate_bias_ptr
-//   output_bias_ptr
+//   output_gate_bias_ptr
 //
 // Layer norm coefficients of size 'n_cell', representing diagonal matrices.
 //   layer_norm_input_weight_ptr    - optional
@@ -1031,8 +1031,8 @@ inline void LstmStepInteger(
     int32_t layer_norm_cell_scale_b,
     const int16_t* layer_norm_output_weight_ptr,
     int32_t layer_norm_output_scale_a, int32_t layer_norm_output_scale_b,
-    const int32_t* input_bias_ptr, const int32_t* forget_bias_ptr,
-    const int32_t* cell_gate_bias_ptr, const int32_t* output_bias_ptr,
+    const int32_t* input_gate_bias_ptr, const int32_t* forget_gate_bias_ptr,
+    const int32_t* cell_gate_bias_ptr, const int32_t* output_gate_bias_ptr,
     int16_t quantized_cell_clip, int8_t quantized_proj_clip, int32_t cell_scale,
     int32_t input_variance_guard, int32_t forget_variance_guard,
     int32_t cell_variance_guard, int32_t output_variance_guard,
@@ -1098,7 +1098,7 @@ inline void LstmStepInteger(
 
   if (use_layer_norm) {
     tensor_utils::ApplyLayerNorm(
-        scratch_1_ptr, layer_norm_forget_weight_ptr, forget_bias_ptr,
+        scratch_1_ptr, layer_norm_forget_weight_ptr, forget_gate_bias_ptr,
         layer_norm_forget_scale_a, layer_norm_forget_scale_b,
         forget_variance_guard, n_batch, n_cell, scratch_1_ptr);
   }
@@ -1149,7 +1149,7 @@ inline void LstmStepInteger(
 
     if (use_layer_norm) {
       tensor_utils::ApplyLayerNorm(
-          scratch_0_ptr, layer_norm_input_weight_ptr, input_bias_ptr,
+          scratch_0_ptr, layer_norm_input_weight_ptr, input_gate_bias_ptr,
           layer_norm_input_scale_a, layer_norm_input_scale_b,
           input_variance_guard, n_batch, n_cell, scratch_0_ptr);
     }
@@ -1190,7 +1190,7 @@ inline void LstmStepInteger(
 
   if (use_layer_norm) {
     tensor_utils::ApplyLayerNorm(
-        scratch_3_ptr, layer_norm_output_weight_ptr, output_bias_ptr,
+        scratch_3_ptr, layer_norm_output_weight_ptr, output_gate_bias_ptr,
         layer_norm_output_scale_a, layer_norm_output_scale_b,
         output_variance_guard, n_batch, n_cell, scratch_3_ptr);
   }
@@ -1268,10 +1268,10 @@ inline void LstmStepInteger(
 //   effective_proj_scale_b                  - optional
 //
 // Gate biases of size 'n_cell':
-//   input_bias_ptr                 - optional
-//   forget_bias_ptr
+//   input_gate_bias_ptr                 - optional
+//   forget_gate_bias_ptr
 //   cell_gate_bias_ptr
-//   output_bias_ptr
+//   output_gate_bias_ptr
 //
 // Layer norm coefficients of size 'n_cell', representing diagonal matrices.
 //   layer_norm_input_weight_ptr    - optional
@@ -1358,8 +1358,8 @@ void LstmStepInteger(
     int32_t layer_norm_cell_scale_b,
     const int16_t* layer_norm_output_weight_ptr,
     int32_t layer_norm_output_scale_a, int32_t layer_norm_output_scale_b,
-    const int32_t* input_bias_ptr, const int32_t* forget_bias_ptr,
-    const int32_t* cell_gate_bias_ptr, const int32_t* output_bias_ptr,
+    const int32_t* input_gate_bias_ptr, const int32_t* forget_gate_bias_ptr,
+    const int32_t* cell_gate_bias_ptr, const int32_t* output_gate_bias_ptr,
     const int32_t* proj_bias_ptr, const TfLiteLSTMParams* params,
     const int32_t* intermediate_scale_a, const int32_t* intermediate_scale_b,
     const int32_t* intermediate_zp, int16_t quantized_cell_clip,
@@ -1391,7 +1391,8 @@ void LstmStepInteger(
   // Forget gate layer norm.
   tensor_utils::ApplyLayerNormFloat(
       scratch2, layer_norm_forget_weight_ptr, layer_norm_forget_scale_a,
-      layer_norm_forget_scale_b, forget_bias_ptr, n_batch, n_cell, scratch2);
+      layer_norm_forget_scale_b, forget_gate_bias_ptr, n_batch, n_cell,
+      scratch2);
 
   // Forget gate sigmoid.
   tensor_utils::ApplySigmoidFloat(scratch2, n_batch, n_cell, scratch2);
@@ -1444,7 +1445,8 @@ void LstmStepInteger(
   // Output gate with layer norm.
   tensor_utils::ApplyLayerNormFloat(
       scratch4, layer_norm_output_weight_ptr, layer_norm_output_scale_a,
-      layer_norm_output_scale_b, output_bias_ptr, n_batch, n_cell, scratch4);
+      layer_norm_output_scale_b, output_gate_bias_ptr, n_batch, n_cell,
+      scratch4);
 
   // Output gate sigmoid.
   tensor_utils::ApplySigmoidFloat(scratch4, n_batch, n_cell, scratch4);
@@ -1512,7 +1514,7 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* aux_input_to_cell_weights,
     const TfLiteTensor* aux_input_to_output_weights,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
     int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
@@ -1595,7 +1597,7 @@ TfLiteStatus EvalFloat(
           GetTensorData(output_layer_norm_coefficients),
           GetTensorData(input_gate_bias),
           GetTensorData(forget_gate_bias),
-          GetTensorData(cell_bias),
+          GetTensorData(cell_gate_bias),
           GetTensorData(output_gate_bias),
           GetTensorData(projection_weights),
           GetTensorData(projection_bias), params, n_batch, n_cell,
@@ -1656,7 +1658,7 @@ TfLiteStatus EvalFloat(
             GetTensorData(output_layer_norm_coefficients),
             GetTensorData(input_gate_bias),
             GetTensorData(forget_gate_bias),
-            GetTensorData(cell_bias),
+            GetTensorData(cell_gate_bias),
             GetTensorData(output_gate_bias),
             GetTensorData(projection_weights),
             GetTensorData(projection_bias), params, /*n_batch=*/1,
@@ -1693,7 +1695,7 @@ TfLiteStatus EvalHybrid(
     const TfLiteTensor* aux_input_to_cell_weights,
     const TfLiteTensor* aux_input_to_output_weights,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
     int output_offset, TfLiteTensor* scratch_buffer,
@@ -1802,7 +1804,7 @@ TfLiteStatus EvalHybrid(
           GetTensorData(output_layer_norm_coefficients),
           GetTensorData(input_gate_bias),
           GetTensorData(forget_gate_bias),
-          GetTensorData(cell_bias),
+          GetTensorData(cell_gate_bias),
           GetTensorData(output_gate_bias),
           GetTensorData(projection_weights),
           GetTensorScale(projection_weights),
@@ -1888,7 +1890,7 @@ TfLiteStatus EvalHybrid(
             GetTensorData(output_layer_norm_coefficients),
             GetTensorData(input_gate_bias),
             GetTensorData(forget_gate_bias),
-            GetTensorData(cell_bias),
+            GetTensorData(cell_gate_bias),
             GetTensorData(output_gate_bias),
             GetTensorData(projection_weights),
             GetTensorScale(projection_weights),
@@ -1930,7 +1932,7 @@ TfLiteStatus EvalInteger8x8_16(
     const TfLiteTensor* cell_layer_norm_coefficients,
     const TfLiteTensor* output_layer_norm_coefficients,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
@@ -2020,7 +2022,7 @@ TfLiteStatus EvalInteger8x8_16(
         integer_lstm_param->layer_norm_output_scale_b,
         GetTensorData(input_gate_bias),
         GetTensorData(forget_gate_bias),
-        GetTensorData(cell_bias),
+        GetTensorData(cell_gate_bias),
         GetTensorData(output_gate_bias),
         integer_lstm_param->quantized_cell_clip,
         integer_lstm_param->quantized_proj_clip, integer_lstm_param->cell_scale,
@@ -2065,7 +2067,7 @@ TfLiteStatus EvalInteger8x8_8(
     const TfLiteTensor* cell_layer_norm_coefficients,
     const TfLiteTensor* output_layer_norm_coefficients,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, TfLiteTensor* output_state,
     TfLiteTensor* cell_state, TfLiteTensor* output,
@@ -2120,10 +2122,12 @@ TfLiteStatus EvalInteger8x8_8(
       GetTensorData(cell_layer_norm_coefficients);
   const int16_t* layer_norm_output_weight_ptr =
       GetTensorData(output_layer_norm_coefficients);
-  const int32_t* input_bias_ptr = GetTensorData(input_gate_bias);
-  const int32_t* forget_bias_ptr = GetTensorData(forget_gate_bias);
-  const int32_t* cell_gate_bias_ptr = GetTensorData(cell_bias);
-  const int32_t* output_bias_ptr = GetTensorData(output_gate_bias);
+  const int32_t* input_gate_bias_ptr = GetTensorData(input_gate_bias);
+  const int32_t* forget_gate_bias_ptr =
+      GetTensorData(forget_gate_bias);
+  const int32_t* cell_gate_bias_ptr = GetTensorData(cell_gate_bias);
+  const int32_t* output_gate_bias_ptr =
+      GetTensorData(output_gate_bias);
   const int32_t* proj_bias_ptr = GetTensorData(projection_bias);
   int16_t* cell_ptr = GetTensorData(cell_state);
   int8_t* output_state_ptr = GetTensorData(output_state);
@@ -2209,8 +2213,8 @@ TfLiteStatus EvalInteger8x8_8(
         integer_lstm_param->layer_norm_output_scale_a,
         integer_lstm_param->layer_norm_output_scale_b,
 
-        input_bias_ptr, forget_bias_ptr, cell_gate_bias_ptr, output_bias_ptr,
-        proj_bias_ptr,
+        input_gate_bias_ptr, forget_gate_bias_ptr, cell_gate_bias_ptr,
+        output_gate_bias_ptr, proj_bias_ptr,
 
         params, integer_lstm_param->intermediate_scale_a,
         integer_lstm_param->intermediate_scale_b,
diff --git a/tensorflow/lite/kernels/lstm_eval.h b/tensorflow/lite/kernels/lstm_eval.h
index 3c9b4bccf42..9b3bd0c54ec 100644
--- a/tensorflow/lite/kernels/lstm_eval.h
+++ b/tensorflow/lite/kernels/lstm_eval.h
@@ -117,7 +117,7 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* aux_input_to_cell_weights,
     const TfLiteTensor* aux_input_to_output_weights,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
     int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
@@ -145,7 +145,7 @@ TfLiteStatus EvalHybrid(
     const TfLiteTensor* aux_input_to_cell_weights,
     const TfLiteTensor* aux_input_to_output_weights,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
     int output_offset, TfLiteTensor* scratch_buffer,
@@ -174,7 +174,7 @@ TfLiteStatus EvalInteger8x8_16(
     const TfLiteTensor* cell_layer_norm_coefficients,
     const TfLiteTensor* output_layer_norm_coefficients,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params,
     const lstm_eval::IntegerLstmParameter* integer_lstm_param,
@@ -200,7 +200,7 @@ TfLiteStatus EvalInteger8x8_8(
     const TfLiteTensor* cell_layer_norm_coefficients,
     const TfLiteTensor* output_layer_norm_coefficients,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, TfLiteTensor* output_state,
     TfLiteTensor* cell_state, TfLiteTensor* output,
diff --git a/tensorflow/lite/kernels/lstm_eval_test.cc b/tensorflow/lite/kernels/lstm_eval_test.cc
index baf2e5e83df..78459117859 100644
--- a/tensorflow/lite/kernels/lstm_eval_test.cc
+++ b/tensorflow/lite/kernels/lstm_eval_test.cc
@@ -113,10 +113,10 @@ class BaseLstmParam {
     TfLiteIntArrayFree(layer_norm_forget_tensor_.dims);
     TfLiteIntArrayFree(layer_norm_cell_tensor_.dims);
     TfLiteIntArrayFree(layer_norm_output_tensor_.dims);
-    TfLiteIntArrayFree(input_bias_tensor_.dims);
-    TfLiteIntArrayFree(forget_bias_tensor_.dims);
-    TfLiteIntArrayFree(cell_bias_tensor_.dims);
-    TfLiteIntArrayFree(output_bias_tensor_.dims);
+    TfLiteIntArrayFree(input_gate_bias_tensor_.dims);
+    TfLiteIntArrayFree(forget_gate_bias_tensor_.dims);
+    TfLiteIntArrayFree(cell_gate_bias_tensor_.dims);
+    TfLiteIntArrayFree(output_gate_bias_tensor_.dims);
     TfLiteIntArrayFree(projection_tensor_.dims);
     TfLiteIntArrayFree(projection_bias_tensor_.dims);
     TfLiteIntArrayFree(activation_tensor_.dims);
@@ -275,17 +275,17 @@ class BaseLstmParam {
   std::vector layer_norm_output_size_ = {n_cell_};
   TfLiteTensor layer_norm_output_tensor_;
 
-  std::vector input_bias_size_ = {n_cell_};
-  TfLiteTensor input_bias_tensor_;
+  std::vector input_gate_bias_size_ = {n_cell_};
+  TfLiteTensor input_gate_bias_tensor_;
 
-  std::vector forget_bias_size_ = {n_cell_};
-  TfLiteTensor forget_bias_tensor_;
+  std::vector forget_gate_bias_size_ = {n_cell_};
+  TfLiteTensor forget_gate_bias_tensor_;
 
-  std::vector cell_bias_size_ = {n_cell_};
-  TfLiteTensor cell_bias_tensor_;
+  std::vector cell_gate_bias_size_ = {n_cell_};
+  TfLiteTensor cell_gate_bias_tensor_;
 
-  std::vector output_bias_size_ = {n_cell_};
-  TfLiteTensor output_bias_tensor_;
+  std::vector output_gate_bias_size_ = {n_cell_};
+  TfLiteTensor output_gate_bias_tensor_;
 
   // projection_weights.
   std::vector projection_ = {
@@ -350,24 +350,28 @@ class QuantizedLstmParam : public BaseLstmParam {
     return &layer_norm_output_tensor_;
   }
   TfLiteTensor* GetInputBias() {
-    PackWeightToTensor(&input_bias_tensor_, input_bias_, input_bias_size_);
-    input_bias_tensor_.data.i32 = input_bias_.data();
-    return &input_bias_tensor_;
+    PackWeightToTensor(&input_gate_bias_tensor_, input_gate_bias_,
+                       input_gate_bias_size_);
+    input_gate_bias_tensor_.data.i32 = input_gate_bias_.data();
+    return &input_gate_bias_tensor_;
   }
   TfLiteTensor* GetForgetBias() {
-    PackWeightToTensor(&forget_bias_tensor_, forget_bias_, forget_bias_size_);
-    forget_bias_tensor_.data.i32 = forget_bias_.data();
-    return &forget_bias_tensor_;
+    PackWeightToTensor(&forget_gate_bias_tensor_, forget_gate_bias_,
+                       forget_gate_bias_size_);
+    forget_gate_bias_tensor_.data.i32 = forget_gate_bias_.data();
+    return &forget_gate_bias_tensor_;
   }
   TfLiteTensor* GetCellBias() {
-    PackWeightToTensor(&cell_bias_tensor_, cell_bias_, cell_bias_size_);
-    cell_bias_tensor_.data.i32 = cell_bias_.data();
-    return &cell_bias_tensor_;
+    PackWeightToTensor(&cell_gate_bias_tensor_, cell_gate_bias_,
+                       cell_gate_bias_size_);
+    cell_gate_bias_tensor_.data.i32 = cell_gate_bias_.data();
+    return &cell_gate_bias_tensor_;
   }
   TfLiteTensor* GetOutputBias() {
-    PackWeightToTensor(&output_bias_tensor_, output_bias_, output_bias_size_);
-    output_bias_tensor_.data.i32 = output_bias_.data();
-    return &output_bias_tensor_;
+    PackWeightToTensor(&output_gate_bias_tensor_, output_gate_bias_,
+                       output_gate_bias_size_);
+    output_gate_bias_tensor_.data.i32 = output_gate_bias_.data();
+    return &output_gate_bias_tensor_;
   }
   TfLiteTensor* GetProjectionBias() {
     PackWeightToTensor(&projection_bias_tensor_, projection_bias_,
@@ -539,22 +543,22 @@ class QuantizedLstmParam : public BaseLstmParam {
   };
 
   // input_gate_bias.
-  std::vector input_bias_ = {
+  std::vector input_gate_bias_ = {
       16, 4, 5, 6, 1, 1, 3, 4, -5, 6,  //
   };
 
   // forget_gate_bias.
-  std::vector forget_bias_ = {
+  std::vector forget_gate_bias_ = {
       16, 4, 5, 6, 1, 1, 3, 4, -5, 6,  //
   };
 
-  // cell_bias.
-  std::vector cell_bias_ = {
+  // cell_gate_bias.
+  std::vector cell_gate_bias_ = {
       16, 4, 5, 6, 1, 1, 3, 4, -5, 6,  //
   };
 
   // output_gate_bias.
-  std::vector output_bias_ = {
+  std::vector output_gate_bias_ = {
       16, 4, 5, 6, 1, 1, 3, 4, -5, 6,  //
   };
 
@@ -711,27 +715,28 @@ class HybridLstmParam : public BaseLstmParam {
     return &accum_scratch_tensor_;
   }
   TfLiteTensor* GetInputBias() {
-    PackWeightToTensor(&input_bias_tensor_, input_float_bias_,
-                       input_bias_size_);
-    input_bias_tensor_.data.f = input_float_bias_.data();
-    return &input_bias_tensor_;
+    PackWeightToTensor(&input_gate_bias_tensor_, input_float_bias_,
+                       input_gate_bias_size_);
+    input_gate_bias_tensor_.data.f = input_float_bias_.data();
+    return &input_gate_bias_tensor_;
   }
   TfLiteTensor* GetForgetBias() {
-    PackWeightToTensor(&forget_bias_tensor_, forget_float_bias_,
-                       forget_bias_size_);
-    forget_bias_tensor_.data.f = forget_float_bias_.data();
-    return &forget_bias_tensor_;
+    PackWeightToTensor(&forget_gate_bias_tensor_, forget_float_bias_,
+                       forget_gate_bias_size_);
+    forget_gate_bias_tensor_.data.f = forget_float_bias_.data();
+    return &forget_gate_bias_tensor_;
   }
   TfLiteTensor* GetCellBias() {
-    PackWeightToTensor(&cell_bias_tensor_, cell_float_bias_, cell_bias_size_);
-    cell_bias_tensor_.data.f = cell_float_bias_.data();
-    return &cell_bias_tensor_;
+    PackWeightToTensor(&cell_gate_bias_tensor_, cell_float_bias_,
+                       cell_gate_bias_size_);
+    cell_gate_bias_tensor_.data.f = cell_float_bias_.data();
+    return &cell_gate_bias_tensor_;
   }
   TfLiteTensor* GetOutputBias() {
-    PackWeightToTensor(&output_bias_tensor_, output_float_bias_,
-                       output_bias_size_);
-    output_bias_tensor_.data.f = output_float_bias_.data();
-    return &output_bias_tensor_;
+    PackWeightToTensor(&output_gate_bias_tensor_, output_float_bias_,
+                       output_gate_bias_size_);
+    output_gate_bias_tensor_.data.f = output_float_bias_.data();
+    return &output_gate_bias_tensor_;
   }
   TfLiteTensor* GetProjectionBias() {
     PackWeightToTensor(&projection_bias_tensor_, projection_float_bias_,
diff --git a/tensorflow/lite/kernels/lstm_test.cc b/tensorflow/lite/kernels/lstm_test.cc
index f8594f9adf0..a9023dce371 100644
--- a/tensorflow/lite/kernels/lstm_test.cc
+++ b/tensorflow/lite/kernels/lstm_test.cc
@@ -89,7 +89,7 @@ class LSTMOpModel : public SingleOpModel {
       input_gate_bias_ = AddInput(TensorType_FLOAT32);
     }
     forget_gate_bias_ = AddInput(TensorType_FLOAT32);
-    cell_bias_ = AddInput(TensorType_FLOAT32);
+    cell_gate_bias_ = AddInput(TensorType_FLOAT32);
     output_gate_bias_ = AddInput(TensorType_FLOAT32);
 
     if (use_projection_weights) {
@@ -211,7 +211,7 @@ class LSTMOpModel : public SingleOpModel {
   }
 
   void SetCellBias(const std::vector& f) {
-    PopulateTensor(cell_bias_, f);
+    PopulateTensor(cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(const std::vector& f) {
@@ -261,7 +261,7 @@ class LSTMOpModel : public SingleOpModel {
 
   int input_gate_bias_;
   int forget_gate_bias_;
-  int cell_bias_;
+  int cell_gate_bias_;
   int output_gate_bias_;
 
   int projection_weights_;
@@ -498,7 +498,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingLstmTest, LstmBlackBoxTest) {
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -545,7 +545,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingOmittedLayerNormLstmTest,
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -601,7 +601,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingLstmTest,
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -652,7 +652,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingLstmInt8Test,
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -743,7 +743,7 @@ TEST_F(CifgNoPeepholeNoProjectionNoClippingLstmTest, LstmBlackBoxTest) {
 
                        {0},       // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -791,7 +791,7 @@ TEST_P(CifgNoPeepholeNoProjectionNoClippingLstmTest,
 
                        {0},       // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -840,7 +840,7 @@ TEST_P(CifgNoPeepholeNoProjectionNoClippingLstmInt8Test,
 
                        {0},       // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -1481,7 +1481,7 @@ TEST_F(NoCifgPeepholeProjectionNoClippingLstmTest, LstmBlackBoxTest) {
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {n_output, n_cell},  // projection_weight tensor
@@ -1528,7 +1528,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLstmTest,
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {n_output, n_cell},  // projection_weight tensor
@@ -1577,7 +1577,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLstmInt8Test,
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {n_output, n_cell},  // projection_weight tensor
@@ -1689,7 +1689,7 @@ TEST_F(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1760,7 +1760,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1833,7 +1833,7 @@ TEST_P(NoCifgPeepholeProjectionNoClippingLayerNormLstmInt8Test,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1947,7 +1947,7 @@ TEST_F(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2018,7 +2018,7 @@ TEST_P(CifgPeepholeProjectionNoClippingLayerNormLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2090,7 +2090,7 @@ TEST_P(CifgPeepholeProjectionNoClippingLayerNormLstmInt8Test,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2195,8 +2195,8 @@ class LSTMIntegerOpModel : public SingleOpModel {
     }
     forget_gate_bias_ = AddInput({TensorType_INT32, input_shapes[13],
                                   ranges[13].first, ranges[13].second});
-    cell_bias_ = AddInput({TensorType_INT32, input_shapes[14], ranges[14].first,
-                           ranges[14].second});
+    cell_gate_bias_ = AddInput({TensorType_INT32, input_shapes[14],
+                                ranges[14].first, ranges[14].second});
     output_gate_bias_ = AddInput({TensorType_INT32, input_shapes[15],
                                   ranges[15].first, ranges[15].second});
 
@@ -2330,7 +2330,7 @@ class LSTMIntegerOpModel : public SingleOpModel {
   }
 
   void SetCellBias(const std::vector& f) {
-    QuantizeAndPopulate(cell_bias_, f);
+    QuantizeAndPopulate(cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(const std::vector& f) {
@@ -2379,7 +2379,7 @@ class LSTMIntegerOpModel : public SingleOpModel {
 
   int input_gate_bias_;
   int forget_gate_bias_;
-  int cell_bias_;
+  int cell_gate_bias_;
   int output_gate_bias_;
 
   int projection_weights_;
@@ -2473,7 +2473,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionNoPeephole) {
 
       {n_cell},  // input_gate_bias tensor
       {n_cell},  // forget_gate_bias tensor
-      {n_cell},  // cell_bias tensor
+      {n_cell},  // cell_gate_bias tensor
       {n_cell},  // output_gate_bias tensor
 
       {n_output, n_cell},  // projection_weight tensor
@@ -2507,7 +2507,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionNoPeephole) {
 
       {-100, 100},  // input_gate_bias tensor
       {-100, 100},  // forget_gate_bias tensor
-      {-100, 100},  // cell_bias tensor
+      {-100, 100},  // cell_gate_bias tensor
       {-100, 100},  // output_gate_bias tensor
 
       {-0.5, 0.5},  // projection_weight tensor
@@ -2675,7 +2675,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionYesPeephole) {
 
       {n_cell},  // input_gate_bias tensor
       {n_cell},  // forget_gate_bias tensor
-      {n_cell},  // cell_bias tensor
+      {n_cell},  // cell_gate_bias tensor
       {n_cell},  // output_gate_bias tensor
 
       {n_output, n_cell},  // projection_weight tensor
@@ -2709,7 +2709,7 @@ TEST(LSTMIntegerOpModel, NoCifgYesLayerNormNoYesProjectionYesPeephole) {
 
       {-100, 100},  // input_gate_bias tensor
       {-100, 80},   // forget_gate_bias tensor
-      {-100, 100},  // cell_bias tensor
+      {-100, 100},  // cell_gate_bias tensor
       {-100, 100},  // output_gate_bias tensor
 
       {-0.5, 0.5},  // projection_weight tensor
@@ -2869,8 +2869,8 @@ class LSTMIntegerOpModel8x8_8 : public SingleOpModel {
     }
     forget_gate_bias_ = AddInput({TensorType_INT32, input_shapes[13],
                                   ranges[13].first, ranges[13].second});
-    cell_bias_ = AddInput({TensorType_INT32, input_shapes[14], ranges[14].first,
-                           ranges[14].second});
+    cell_gate_bias_ = AddInput({TensorType_INT32, input_shapes[14],
+                                ranges[14].first, ranges[14].second});
     output_gate_bias_ = AddInput({TensorType_INT32, input_shapes[15],
                                   ranges[15].first, ranges[15].second});
 
@@ -3004,7 +3004,7 @@ class LSTMIntegerOpModel8x8_8 : public SingleOpModel {
   }
 
   void SetCellBias(const std::vector& f) {
-    QuantizeAndPopulate(cell_bias_, f);
+    QuantizeAndPopulate(cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(const std::vector& f) {
@@ -3053,7 +3053,7 @@ class LSTMIntegerOpModel8x8_8 : public SingleOpModel {
 
   int input_gate_bias_;
   int forget_gate_bias_;
-  int cell_bias_;
+  int cell_gate_bias_;
   int output_gate_bias_;
 
   int projection_weights_;
@@ -3148,7 +3148,7 @@ TEST(LSTMIntegerOpModel8x8_8, CifgYesLayerNormNoYesProjectionNoPeephole) {
 
       {0},       // input_gate_bias tensor
       {n_cell},  // forget_gate_bias tensor
-      {n_cell},  // cell_bias tensor
+      {n_cell},  // cell_gate_bias tensor
       {n_cell},  // output_gate_bias tensor
 
       {n_output, n_cell},  // projection_weight tensor
@@ -3182,7 +3182,7 @@ TEST(LSTMIntegerOpModel8x8_8, CifgYesLayerNormNoYesProjectionNoPeephole) {
 
       {-100, 100},  // input_gate_bias tensor
       {-100, 100},  // forget_gate_bias tensor
-      {-100, 100},  // cell_bias tensor
+      {-100, 100},  // cell_gate_bias tensor
       {-100, 100},  // output_gate_bias tensor
 
       {-0.5, 0.5},  // projection_weight tensor
@@ -3303,7 +3303,7 @@ TEST(LSTMOpModel, InvalidTypeTest) {
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
@@ -3338,7 +3338,7 @@ TEST(LSTMOpModel, InvalidTypeTest) {
 
                        {n_cell},  // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
diff --git a/tensorflow/lite/kernels/optional_tensor_test.cc b/tensorflow/lite/kernels/optional_tensor_test.cc
index 26d619276aa..9e83c74da8d 100644
--- a/tensorflow/lite/kernels/optional_tensor_test.cc
+++ b/tensorflow/lite/kernels/optional_tensor_test.cc
@@ -78,7 +78,7 @@ class LSTMOpModel : public SingleOpModel {
       input_gate_bias_ = AddInput(TensorType_FLOAT32);
     }
     forget_gate_bias_ = AddInput(TensorType_FLOAT32);
-    cell_bias_ = AddInput(TensorType_FLOAT32);
+    cell_gate_bias_ = AddInput(TensorType_FLOAT32);
     output_gate_bias_ = AddInput(TensorType_FLOAT32);
 
     if (use_projection_weights) {
@@ -161,7 +161,7 @@ class LSTMOpModel : public SingleOpModel {
   }
 
   void SetCellBias(std::initializer_list f) {
-    PopulateTensor(cell_bias_, f);
+    PopulateTensor(cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(std::initializer_list f) {
@@ -209,7 +209,7 @@ class LSTMOpModel : public SingleOpModel {
 
   int input_gate_bias_;
   int forget_gate_bias_;
-  int cell_bias_;
+  int cell_gate_bias_;
   int output_gate_bias_;
 
   int projection_weights_;
@@ -256,7 +256,7 @@ TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) {
 
                        {0},       // input_gate_bias tensor
                        {n_cell},  // forget_gate_bias tensor
-                       {n_cell},  // cell_bias tensor
+                       {n_cell},  // cell_gate_bias tensor
                        {n_cell},  // output_gate_bias tensor
 
                        {0, 0},  // projection_weight tensor
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
index f1c0f9d42a6..0b2cba72369 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
@@ -179,10 +179,10 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->size, 1);
   TF_LITE_ENSURE_EQ(context, forget_gate_bias->dims->data[0], n_cell);
 
-  const TfLiteTensor* cell_bias =
+  const TfLiteTensor* cell_gate_bias =
       GetInput(context, node, lstm::full::kCellGateBiasTensor);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->size, 1);
-  TF_LITE_ENSURE_EQ(context, cell_bias->dims->data[0], n_cell);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->size, 1);
+  TF_LITE_ENSURE_EQ(context, cell_gate_bias->dims->data[0], n_cell);
 
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, lstm::full::kOutputGateBiasTensor);
@@ -546,7 +546,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       GetOptionalInputTensor(context, node, lstm::full::kInputGateBiasTensor);
   const TfLiteTensor* forget_gate_bias =
       GetInput(context, node, lstm::full::kForgetGateBiasTensor);
-  const TfLiteTensor* cell_bias =
+  const TfLiteTensor* cell_gate_bias =
       GetInput(context, node, lstm::full::kCellGateBiasTensor);
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, lstm::full::kOutputGateBiasTensor);
@@ -611,8 +611,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*aux_input_to_forget_weights=*/nullptr,
           /*aux_input_to_cell_weights=*/nullptr,
           /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
-          forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
-          projection_bias, &lstm_params, /*forward_sequence=*/true, time_major,
+          forget_gate_bias, cell_gate_bias, output_gate_bias,
+          projection_weights, projection_bias, &lstm_params,
+          /*forward_sequence=*/true, time_major,
           /*output_offset=*/0, scratch_buffer, output_state, cell_state,
           output);
     }
@@ -648,8 +649,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
           /*aux_input_to_forget_weights=*/nullptr,
           /*aux_input_to_cell_weights=*/nullptr,
           /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
-          forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
-          projection_bias, &lstm_params, /*forward_sequence=*/true, time_major,
+          forget_gate_bias, cell_gate_bias, output_gate_bias,
+          projection_weights, projection_bias, &lstm_params,
+          /*forward_sequence=*/true, time_major,
           /*output_offset=*/0, scratch_buffer, scaling_factors,
           prod_scaling_factors, recovered_cell_weights, input_quantized,
           /*aux_input_quantized=*/nullptr, output_state_quantized,
diff --git a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
index ec20d76ae2e..74584ec9e85 100644
--- a/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
+++ b/tensorflow/lite/kernels/unidirectional_sequence_lstm_test.cc
@@ -85,7 +85,7 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
       input_gate_bias_ = AddInput(TensorType_FLOAT32);
     }
     forget_gate_bias_ = AddInput(TensorType_FLOAT32);
-    cell_bias_ = AddInput(TensorType_FLOAT32);
+    cell_gate_bias_ = AddInput(TensorType_FLOAT32);
     output_gate_bias_ = AddInput(TensorType_FLOAT32);
 
     if (use_projection_weights) {
@@ -187,7 +187,7 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
   }
 
   void SetCellBias(const std::vector& f) {
-    PopulateTensor(cell_bias_, f);
+    PopulateTensor(cell_gate_bias_, f);
   }
 
   void SetOutputGateBias(const std::vector& f) {
@@ -249,7 +249,7 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
 
   int input_gate_bias_;
   int forget_gate_bias_;
-  int cell_bias_;
+  int cell_gate_bias_;
   int output_gate_bias_;
 
   int projection_weights_;
@@ -530,7 +530,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -592,7 +592,7 @@ TEST_F(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -658,7 +658,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -721,7 +721,7 @@ TEST_P(NoCifgNoPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -833,7 +833,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -894,7 +894,7 @@ TEST_P(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -957,7 +957,7 @@ TEST_P(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -1619,7 +1619,7 @@ TEST_F(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1688,7 +1688,7 @@ TEST_P(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -1759,7 +1759,7 @@ TEST_P(NoCifgPeepholeProjectionClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2430,7 +2430,7 @@ TEST_F(NoCifgPeepholeProjectionAndBiasClippingUnidirectionalLstmTest,
 
           {n_cell},  // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {n_output, n_cell},  // projection_weight tensor
@@ -2636,7 +2636,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingLayerNormUnidirectionalLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
@@ -2707,7 +2707,7 @@ TEST_F(CifgPeepholeNoProjectionNoClippingUnidirectionalLstmTest,
 
           {0},       // input_gate_bias tensor
           {n_cell},  // forget_gate_bias tensor
-          {n_cell},  // cell_bias tensor
+          {n_cell},  // cell_gate_bias tensor
           {n_cell},  // output_gate_bias tensor
 
           {0, 0},  // projection_weight tensor
diff --git a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
index 88ea7c1d591..50138442c25 100644
--- a/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
+++ b/tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc
@@ -299,7 +299,7 @@ TfLiteStatus EvalFloat(
     const TfLiteTensor* aux_input_to_cell_weights,
     const TfLiteTensor* aux_input_to_output_weights,
     const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
-    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
     const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
     const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
     int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
@@ -384,7 +384,7 @@ TfLiteStatus EvalFloat(
           GetTensorData(output_layer_norm_coefficients),
           GetTensorData(input_gate_bias),
           GetTensorData(forget_gate_bias),
-          GetTensorData(cell_bias),
+          GetTensorData(cell_gate_bias),
           GetTensorData(output_gate_bias),
           GetTensorData(projection_weights),
           GetTensorData(projection_bias), params, n_batch, n_cell,
@@ -446,7 +446,7 @@ TfLiteStatus EvalFloat(
             GetTensorData(output_layer_norm_coefficients),
             GetTensorData(input_gate_bias),
             GetTensorData(forget_gate_bias),
-            GetTensorData(cell_bias),
+            GetTensorData(cell_gate_bias),
             GetTensorData(output_gate_bias),
             GetTensorData(projection_weights),
             GetTensorData(projection_bias), params, /*n_batch=*/1,
@@ -527,7 +527,7 @@ TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
       context, node, ops::builtin::lstm::full::kInputGateBiasTensor);
   const TfLiteTensor* forget_gate_bias =
       GetInput(context, node, ops::builtin::lstm::full::kForgetGateBiasTensor);
-  const TfLiteTensor* cell_bias =
+  const TfLiteTensor* cell_gate_bias =
       GetInput(context, node, ops::builtin::lstm::full::kCellGateBiasTensor);
   const TfLiteTensor* output_gate_bias =
       GetInput(context, node, ops::builtin::lstm::full::kOutputGateBiasTensor);
@@ -570,8 +570,9 @@ TfLiteStatus lstm_eval(TfLiteContext* context, TfLiteNode* node, Logger* logger,
           /*aux_input_to_forget_weights=*/nullptr,
           /*aux_input_to_cell_weights=*/nullptr,
           /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
-          forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
-          projection_bias, params, /*forward_sequence=*/true,
+          forget_gate_bias, cell_gate_bias, output_gate_bias,
+          projection_weights, projection_bias, params,
+          /*forward_sequence=*/true,
           /*time_major=*/true,
           /*output_offset=*/0, scratch_buffer, output_state, cell_state, output,
           logger, intermediate_tensor_indexes, error_reporter);

From 7cd6c3115badaf79fff2b8809cb4a0b49a5f9c7c Mon Sep 17 00:00:00 2001
From: Denisa Roberts 
Date: Fri, 19 Jun 2020 14:26:51 -0400
Subject: [PATCH 0636/1390] Allow gradient access to QR input

---
 tensorflow/python/eager/pywrap_gradient_exclusions.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
index 7da45e36118..7e9f0b16334 100644
--- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc
+++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
@@ -50,7 +50,7 @@ auto OpGradientInfoInit(const T &a) {
 
 absl::optional> OpGradientUnusedInputIndices(
     const tensorflow::string &op_name) {
-  static std::array a = {{
+  static std::array a = {{
       {"Acosh"},
       {"AllToAll", 1, {0}},
       {"ApproximateEqual"},
@@ -222,7 +222,6 @@ absl::optional> OpGradientUnusedInputIndices(
       {"PlaceholderWithDefault"},
       {"PopulationCount"},
       {"PreventGradient"},
-      {"Qr"},
       {"QuantizeAndDequantize"},
       {"QuantizeAndDequantizeV2"},
       {"QuantizeAndDequantizeV3"},

From 9aea03e98d2ef5416dff09f721d6d519ea83e74f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 11:39:56 -0700
Subject: [PATCH 0637/1390] Add StatefulRandom op to TF MLIR.

PiperOrigin-RevId: 317347849
Change-Id: I19b7b1b3157d065cc556cce7aba72b79ed9b776f
---
 .../mlir/tensorflow/ir/tf_generated_ops.td    | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index 3b1f3eec699..dcef99e6971 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -9190,6 +9190,32 @@ def TF_StackV2Op : TF_Op<"StackV2", []> {
   );
 }
 
+def TF_StatelessRandomUniformOp : TF_Op<"StatelessRandomUniform", [NoSideEffect]> {
+  let summary = [{
+Outputs deterministic pseudorandom random values from a uniform distribution.
+  }];
+
+  let description = [{
+The generated values follow a uniform distribution in the range `[0, 1)`. The
+lower bound 0 is included in the range, while the upper bound 1 is excluded.
+
+The outputs are a deterministic function of `shape` and `seed`.
+  }];
+
+  let arguments = (ins
+    TF_I32OrI64Tensor:$shape,
+    TF_I32OrI64Tensor:$seed
+  );
+
+  let results = (outs
+    TF_FpTensor:$output
+  );
+
+  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
+  TF_DerivedOperandTypeAttr Tseed = TF_DerivedOperandTypeAttr<1>;
+  TF_DerivedResultTypeAttr dtype = TF_DerivedResultTypeAttr<0>;
+}
+
 def TF_StopGradientOp : TF_Op<"StopGradient", [NoSideEffect, TF_AllTypesMatch<["input", "output"]>]> {
   let summary = "Stops gradient computation.";
 

From 02d312b25edc2d8999640d7fbe2a833c2eaab0d8 Mon Sep 17 00:00:00 2001
From: Ran Chen 
Date: Fri, 19 Jun 2020 11:44:48 -0700
Subject: [PATCH 0638/1390] *** Reason for rollback ***

break pip test

PiperOrigin-RevId: 317348830
Change-Id: Ia6968b2cd3bf001f6afbc4487234d88123d3d0a8
---
 .../python/distribute/multi_process_runner.py | 205 ++----------------
 .../distribute/multi_process_runner_test.py   |  69 +-----
 2 files changed, 25 insertions(+), 249 deletions(-)

diff --git a/tensorflow/python/distribute/multi_process_runner.py b/tensorflow/python/distribute/multi_process_runner.py
index ce36287a9da..8699e59b410 100644
--- a/tensorflow/python/distribute/multi_process_runner.py
+++ b/tensorflow/python/distribute/multi_process_runner.py
@@ -67,7 +67,8 @@ except ImportError:
 # exception stack trace info is stored in exc_info to pass on to parent process
 # to be re-raised.
 _ProcessStatusInfo = collections.namedtuple(
-    '_ProcessStatusInfo', ['is_successful', 'exc_info', 'return_value'])
+    '_ProcessStatusInfo',
+    ['task_type', 'is_successful', 'exc_info', 'return_value'])
 
 # Information returned from a successful MultiProcessRunner run.
 MultiProcessRunnerResult = collections.namedtuple('MultiProcessRunnerResult',
@@ -123,7 +124,6 @@ class MultiProcessRunner(object):
                stream_stdout=True,
                list_stdout=False,
                use_dill_for_args=True,
-               daemon=False,
                args=None,
                kwargs=None):
     """Creates a multi-process runner.
@@ -157,7 +157,6 @@ class MultiProcessRunner(object):
       use_dill_for_args: Whether to use dill to pickle `args` and `kwargs`. dill
         can pickle more objects, but doesn't work with types in
         `multiprocessing` library like `Mutex`.
-      daemon: Whether to start processes as daemons.
       args: Positional arguments to be sent to functions run on processes.
       kwargs: Keyword arguments to be sent to functions run on processes.
 
@@ -189,7 +188,6 @@ class MultiProcessRunner(object):
     self._list_stdout = list_stdout
     self._dependence_on_chief = True
     self._use_dill_for_args = use_dill_for_args
-    self._daemon = daemon
     self._args = args or ()
     self._kwargs = kwargs or {}
 
@@ -270,8 +268,7 @@ class MultiProcessRunner(object):
         test_env=test_env,
         target=_ProcFunc(),
         args=(resources, test_env, proc_func, args, kwargs,
-              self._use_dill_for_args),
-        daemon=self._daemon)
+              self._use_dill_for_args))
     p.start()
     self._processes[(task_type, task_id)] = p
     self._outstanding_subprocess_count += 1
@@ -571,6 +568,7 @@ class _ProcFunc(object):
         time.sleep(0.1)
     self._resources.process_status_queue.put(
         _ProcessStatusInfo(
+            task_type=task_type,
             is_successful=True,
             exc_info=None,
             return_value=None))
@@ -630,9 +628,17 @@ class _ProcFunc(object):
     if test_env.v2_enabled:
       v2_compat.enable_v2_behavior()
 
-    with self._runtime_mode(test_env.executing_eagerly):
-      info = _run_contained(proc_func, args, kwargs)
-      self._resources.process_status_queue.put(info)
+    try:
+      with self._runtime_mode(test_env.executing_eagerly):
+        return_value = proc_func(*args, **kwargs)
+        is_successful = True
+        exc_info = None
+
+    except Exception:  # pylint: disable=broad-except
+      # Capture all exceptions to be reported to parent process.
+      return_value = None
+      is_successful = False
+      exc_info = sys.exc_info()
 
       # Re-raise the exception in addition to reporting it to the parent
       # process, so that even if `--test_timeout` flag is set and the
@@ -641,181 +647,16 @@ class _ProcFunc(object):
       # instead of silently suppressing the error due to early bazel
       # timeout. Raising an error in the subprocess produces stack trace in
       # the log, but the program continues running.
-      if not info.is_successful:
-        six.reraise(*info.exc_info)
+      raise
 
-      self._close_streaming()
-
-
-class MultiProcessPoolRunner(object):
-  """A utility class to start a process pool to simulate a cluster.
-
-  It's similar to MultiProcessRunner, but uses a pool of processes to avoid the
-  expensive initialization cost of Tensorflow.
-  """
-
-  def __init__(self, cluster_spec, initializer=None):
-    """Creates a multi-process pool runner.
-
-    Args:
-      cluster_spec: Dict for cluster spec. The following is an example of
-        cluster with three workers.
-        {"worker": ["worker0.example.com:2222",
-                    "worker1.example.com:2222",
-                    "worker2.example.com:2222"]}
-      initializer: a callable to called at the startup of worker processes.
-
-    Raises:
-      RuntimeError: if `multi_process_runner.test_main()` is not called.
-      ValueError: if there are more than one chief in the `cluster_spec`.
-    """
-    self._cluster_spec = cluster_spec
-    self._initializer = initializer
-    self._conn = {}
-    self._runner = None
-
-  def __del__(self):
-    self._reset()
-
-  def _reset(self):
-    for conn in self._conn.values():
-      conn.close()
-    self._conn = {}
-    if self._runner is not None:
-      self._runner.join()
-      self._runner = None
-
-  def _start(self):
-    """Starts the worker pool."""
-    # We need different arguments for different processes so we're passing a
-    # no-op proc_func here and use start_single_process instead.
-    #
-    # We also need to start the process pool as daemon, so that they don't block
-    # the program from exiting. Note that __del__ may not get called when
-    # there's an exception. The user may also store a pool runner in a global
-    # object to share across test cases
-    self._runner = MultiProcessRunner(
-        proc_func=lambda: None,
-        cluster_spec=self._cluster_spec,
-        use_dill_for_args=False,
-        daemon=True)
-    if self._initializer:
-      initializer = dill.dumps(self._initializer, dill.HIGHEST_PROTOCOL)
-    else:
-      initializer = None
-    for task_type, addresses in self._cluster_spec.items():
-      for task_id, _ in enumerate(addresses):
-        conn1, conn2 = multiprocessing.Pipe(duplex=True)
-        self._conn[(task_type, task_id)] = conn1
-        self._runner.start_single_process(
-            task_type,
-            task_id,
-            proc_func=_pool_runner_worker,
-            args=(initializer, conn2))
-
-  def run(self, proc_func, args=None, kwargs=None):
-    """Runs `proc_func` with `args` and `kwargs` on all jobs.
-
-    Args:
-      proc_func: The function to be run.
-      args: Optional positional arguments to be supplied in `proc_func`.
-      kwargs: Optional keyword arguments to be supplied in `proc_func`.
-
-    Returns:
-      A list of return values.
-    """
-    if self._runner is None:
-      self._start()
-
-    # Since we start the processes as daemon they're going to be killed by
-    # SIGTERM when the program exits. We only turn on streaming during run() to
-    # avoid printing the stacktrace caused by the SIGTERM.
-    self._runner._stream_stdout = True  # pylint: disable=protected-access
-
-    try:
-      proc_func = dill.dumps(proc_func, dill.HIGHEST_PROTOCOL)
-      for conn in self._conn.values():
-        conn.send((proc_func, args or [], kwargs or {}))
-
-      process_statuses = []
-      for (task_type, task_id), conn in self._conn.items():
-        logging.info('Waiting for the result from %s-%d', task_type, task_id)
-        try:
-          process_statuses.append(conn.recv())
-        except EOFError:
-          # This shouldn't happen due to exceptions in proc_func. This usually
-          # means bugs in the runner.
-          self._reset()
-          raise RuntimeError('Unexpected EOF. Worker process may have died. '
-                             'Please report a bug')
-
-      return_values = []
-      for process_status in process_statuses:
-        assert isinstance(process_status, _ProcessStatusInfo)
-        if not process_status.is_successful:
-          six.reraise(*process_status.exc_info)
-        if process_status.return_value is not None:
-          return_values.append(process_status.return_value)
-
-      return return_values
     finally:
-      self._runner._stream_stdout = False  # pylint: disable=protected-access
-
-
-def _pool_runner_worker(initializer, conn):
-  """Function that runs on the workers in a pool.
-
-  It listens for callables to run and returns the result until `conn` is closed.
-  It captures the exceptions during executing the callable and return it through
-  `conn`.
-
-  Args:
-    initializer: A callable to execute during startup.
-    conn: A multiprocessing.Connection object to listen for tasks and send
-      results.
-  """
-  if initializer:
-    initializer = dill.loads(initializer)
-    initializer()
-  while True:
-    try:
-      proc_func, args, kwargs = conn.recv()
-    except EOFError:
-      break
-    proc_func = dill.loads(proc_func)
-    info = _run_contained(proc_func, args, kwargs)
-    sys.stdout.flush()
-    sys.stderr.flush()
-    conn.send(info)
-
-
-def _run_contained(proc_func, args, kwargs):
-  """Runs `proc_func` with `args` and `kwargs`.
-
-  The function returns _ProcessStatusInfo which captures the return value and
-  the exception.
-
-  Args:
-    proc_func: The function to be run.
-    args: Optional positional arguments to be supplied in `proc_func`.
-    kwargs: Optional keyword arguments to be supplied in `proc_func`.
-
-  Returns:
-    a _ProcessStatusInfo.
-  """
-  try:
-    return_value = proc_func(*args, **kwargs)
-    is_successful = True
-    exc_info = None
-  except Exception:  # pylint: disable=broad-except
-    return_value = None
-    is_successful = False
-    exc_info = sys.exc_info()
-  finally:
-    return _ProcessStatusInfo(  # pylint: disable=lost-exception
-        is_successful=is_successful,
-        exc_info=exc_info,
-        return_value=return_value)
+      info = _ProcessStatusInfo(
+          task_type=test_env.task_type,
+          is_successful=is_successful,
+          exc_info=exc_info,
+          return_value=return_value)
+      self._resources.process_status_queue.put(info)
+      self._close_streaming()
 
 
 class SubprocessTimeoutError(RuntimeError):
diff --git a/tensorflow/python/distribute/multi_process_runner_test.py b/tensorflow/python/distribute/multi_process_runner_test.py
index d76ef5a5a3c..aeba43b6b7c 100644
--- a/tensorflow/python/distribute/multi_process_runner_test.py
+++ b/tensorflow/python/distribute/multi_process_runner_test.py
@@ -22,8 +22,6 @@ import json
 import os
 import threading
 import time
-import unittest
-
 from absl import logging
 
 from tensorflow.python.distribute import multi_process_runner
@@ -47,7 +45,7 @@ def proc_func_that_adds_simple_return_data():
   return 'dummy_data'
 
 
-def proc_func_that_returns_args_and_kwargs(*args, **kwargs):
+def proc_func_that_return_args_and_kwargs(*args, **kwargs):
   return list(args) + list(kwargs.items())
 
 
@@ -55,20 +53,6 @@ def proc_func_with_barrier():
   return multi_process_runner.barrier()
 
 
-def proc_func_that_returns_pid():
-  return os.getpid()
-
-
-V = None
-
-
-def proc_func_that_sets_global(val):
-  global V
-  old_val = V
-  V = val
-  return old_val
-
-
 class MultiProcessRunnerTest(test.TestCase):
 
   def _worker_idx(self):
@@ -111,7 +95,7 @@ class MultiProcessRunnerTest(test.TestCase):
 
   def test_multi_process_runner_args_passed_correctly(self):
     return_value = multi_process_runner.run(
-        proc_func_that_returns_args_and_kwargs,
+        proc_func_that_return_args_and_kwargs,
         multi_worker_test_base.create_cluster_spec(num_workers=1),
         args=('a', 'b'),
         kwargs={
@@ -341,54 +325,5 @@ class MultiProcessRunnerTest(test.TestCase):
                 for line in list_to_assert))
 
 
-class MultiProcessPoolRunnerTest(test.TestCase):
-
-  def test_same_process_across_runs(self):
-    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
-    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
-    pid = runner.run(proc_func_that_returns_pid)
-    for _ in range(3):
-      self.assertAllEqual(runner.run(proc_func_that_returns_pid), pid)
-
-  def test_exceptions_in_sub_process(self):
-    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
-    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
-    pid = runner.run(proc_func_that_returns_pid)
-    with self.assertRaisesRegexp(ValueError, 'This is an error.'):
-      runner.run(proc_func_that_errors)
-    self.assertAllEqual(runner.run(proc_func_that_returns_pid), pid)
-
-  def test_tf_config(self):
-    cluster_spec = multi_worker_test_base.create_cluster_spec(
-        has_chief=True, num_workers=2)
-    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
-    result = runner.run(proc_func_that_adds_task_type_in_return_data)
-
-    job_count_dict = {'worker': 2, 'chief': 1}
-    for data in result:
-      job_count_dict[data] -= 1
-
-    self.assertEqual(job_count_dict['worker'], 0)
-    self.assertEqual(job_count_dict['chief'], 0)
-
-  @unittest.expectedFailure
-  def test_exception_in_main_process(self):
-    # When there's an exception in the main process, __del__() is not called.
-    # This test is to verify MultiProcessPoolRunner can cope with __del__() not
-    # being called.
-    cluster_spec = multi_worker_test_base.create_cluster_spec(
-        has_chief=True, num_workers=2)
-    runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec)
-    runner.run(proc_func_that_returns_pid)
-    raise ValueError('failure')
-
-  def test_initializer(self):
-    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
-    runner = multi_process_runner.MultiProcessPoolRunner(
-        cluster_spec, initializer=lambda: proc_func_that_sets_global(1))
-    result = runner.run(proc_func_that_sets_global, args=(2,))
-    self.assertAllEqual(result, [1, 1])
-
-
 if __name__ == '__main__':
   multi_process_runner.test_main()

From 2417a15bf186cde76b601a53324ebbc2c9193124 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 11:46:26 -0700
Subject: [PATCH 0639/1390] Clean up post fusion using DCE.

PiperOrigin-RevId: 317349120
Change-Id: I479d9967323d86e924315d2b1302bafd01ed4151
---
 tensorflow/compiler/xla/service/BUILD                  | 1 +
 tensorflow/compiler/xla/service/multi_output_fusion.cc | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index acd35cbc153..2fd457e8e47 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -1679,6 +1679,7 @@ cc_library(
     hdrs = ["multi_output_fusion.h"],
     deps = [
         ":hlo",
+        ":hlo_dce",
         ":hlo_pass",
         ":hlo_reachability",
         "//tensorflow/compiler/xla:debug_options_flags",
diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc
index b95b27d6291..a21cec538d1 100644
--- a/tensorflow/compiler/xla/service/multi_output_fusion.cc
+++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "absl/container/flat_hash_set.h"
 #include "tensorflow/compiler/xla/debug_options_flags.h"
+#include "tensorflow/compiler/xla/service/hlo_dce.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
 #include "tensorflow/compiler/xla/service/hlo_reachability.h"
@@ -126,6 +127,10 @@ StatusOr MultiOutputFusion::Run(HloModule* module) {
   candidates_index_.clear();
   all_fusion_candidates_.clear();
   reachability_.reset();
+  if (changed) {
+    HloDCE dce;
+    TF_RETURN_IF_ERROR(dce.Run(module).status());
+  }
   return changed;
 }
 

From f805153a25b00d12072bd728e91bb1621bfcf1b1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 11:48:22 -0700
Subject: [PATCH 0640/1390] Fix regression that omitted -DNDEBUG in opt builds
 after toolchain refactoring.

Not adding -DNDEBUG leads to enabling assertions which makes LLVM often 10x slower.

PiperOrigin-RevId: 317349500
Change-Id: I52df6ab5013ad5a02101dd96f18220e054c3e94c
---
 third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
index 082ed950b04..eb320a94201 100644
--- a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
+++ b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
@@ -583,7 +583,11 @@ def _features(cpu, compiler, ctx):
                     ),
                 ],
             ),
-            feature(name = "opt"),
+            feature(name = "disable-assertions"),
+            feature(
+                name = "opt",
+                implies = ["disable-assertions"],
+            ),
             feature(name = "fastbuild"),
             feature(name = "dbg"),
             feature(name = "supports_dynamic_linker", enabled = True),

From f53e1aac6576b3f5e575f870da1fbdb286094444 Mon Sep 17 00:00:00 2001
From: Lu Wang 
Date: Fri, 19 Jun 2020 11:50:15 -0700
Subject: [PATCH 0641/1390] Stamp the version number for all metadata
 components simultaneously

PiperOrigin-RevId: 317349859
Change-Id: Ica912c3fb310889185c026e6d73ce4c69a9f0505
---
 .../lite/experimental/support/metadata/BUILD  | 24 ++++++++++-
 .../support/metadata/build_defs.bzl           | 43 +++++++++++++++++++
 .../experimental/support/metadata/cc/BUILD    | 13 +++++-
 .../metadata/cc/metadata_parser.h.template    | 28 ++++++++++++
 .../support/metadata/cc/test/BUILD            |  9 ++++
 .../metadata/cc/test/metadata_parser_test.cc  | 33 ++++++++++++++
 .../metadata/cc/test/metadata_version_test.cc |  2 +-
 .../experimental/support/metadata/java/BUILD  |  8 +++-
 .../support/metadata/MetadataExtractor.java   |  8 +---
 .../lite/support/metadata/MetadataParser.java | 27 ++++++++++++
 .../metadata/metadata_parser.py.template      | 26 +++++++++++
 .../support/metadata/metadata_parser_test.py  | 38 ++++++++++++++++
 .../support/metadata/metadata_schema.fbs      | 29 ++++++++-----
 13 files changed, 267 insertions(+), 21 deletions(-)
 create mode 100644 tensorflow/lite/experimental/support/metadata/build_defs.bzl
 create mode 100644 tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
 create mode 100644 tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
 create mode 100644 tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
 create mode 100644 tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
 create mode 100644 tensorflow/lite/experimental/support/metadata/metadata_parser_test.py

diff --git a/tensorflow/lite/experimental/support/metadata/BUILD b/tensorflow/lite/experimental/support/metadata/BUILD
index 4621c8c55d2..ba410d914c7 100644
--- a/tensorflow/lite/experimental/support/metadata/BUILD
+++ b/tensorflow/lite/experimental/support/metadata/BUILD
@@ -1,5 +1,6 @@
 load("//tensorflow:tensorflow.bzl", "py_test")
 load("@flatbuffers//:build_defs.bzl", "flatbuffer_android_library", "flatbuffer_cc_library", "flatbuffer_java_library", "flatbuffer_py_library")
+load("//tensorflow/lite/experimental/support/metadata:build_defs.bzl", "stamp_metadata_parser_version")
 
 package(
     default_visibility = [
@@ -51,9 +52,19 @@ flatbuffer_android_library(
     custom_package = "org.tensorflow.lite.support.metadata.schema",
 )
 
+# TODO(b/157813075): move the metadata python library to metadata/python/ when migrating to the new repo.
+stamp_metadata_parser_version(
+    name = "metadata_parser_py",
+    srcs = ["metadata_parser.py.template"],
+    outs = ["metadata_parser.py"],
+)
+
 py_library(
     name = "metadata",
-    srcs = ["metadata.py"],
+    srcs = [
+        "metadata.py",
+        ":metadata_parser_py",
+    ],
     data = [
         "//tensorflow/lite/experimental/support/metadata:metadata_schema.fbs",
     ],
@@ -89,3 +100,14 @@ py_test(
         "@six_archive//:six",
     ],
 )
+
+py_test(
+    name = "metadata_parser_test",
+    srcs = ["metadata_parser_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY2AND3",
+    deps = [
+        ":metadata",
+        "//tensorflow/python:client_testlib",
+    ],
+)
diff --git a/tensorflow/lite/experimental/support/metadata/build_defs.bzl b/tensorflow/lite/experimental/support/metadata/build_defs.bzl
new file mode 100644
index 00000000000..3ea945770e0
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/build_defs.bzl
@@ -0,0 +1,43 @@
+"""Build rules to generate metadata schema versions."""
+
+METADATA_SCHEMA_FILE = "//tensorflow/lite/experimental/support/metadata:metadata_schema.fbs"
+
+def stamp_metadata_parser_version(
+        name,
+        srcs,
+        outs):
+    """Stamps the latest metadata parser version into the srcs files.
+
+    Replaces all the occurrences of "{LATEST_METADATA_PARSER_VERSION}" in the
+    srcs files with the metadata schema version extracted from
+    METADATA_SCHEMA_FILE and then outputs the generated file into outs,
+    respectively. The number of srcs files needs to match the number of outs
+    files.
+
+    Args:
+        name: Rule name. (required)
+        srcs: List of source files. (required)
+        outs: List of output files. (required)
+    """
+    if len(srcs) != len(outs):
+        fail(("The number of srcs files (%d) does not match that of the outs" +
+              " files (%d).") %
+             (len(srcs), len(outs)))
+
+    for i in range(0, len(srcs)):
+        native.genrule(
+            name = "%s_file%d" % (name, i),
+            srcs = [srcs[i]],
+            outs = [outs[i]],
+            tools = [METADATA_SCHEMA_FILE],
+            # Gets the metadata schema version from the file, and stamps it
+            # into the srcs file.
+            cmd = "version=$$(sed -n -e '/Schema Semantic version/ s/.*\\: *//p' $(location %s));" %
+                  METADATA_SCHEMA_FILE +
+                  'sed "s/{LATEST_METADATA_PARSER_VERSION}/$$version/" $< > $@',
+        )
+
+    native.filegroup(
+        name = name,
+        srcs = outs,
+    )
diff --git a/tensorflow/lite/experimental/support/metadata/cc/BUILD b/tensorflow/lite/experimental/support/metadata/cc/BUILD
index 832e2edb56d..8febc7a2237 100644
--- a/tensorflow/lite/experimental/support/metadata/cc/BUILD
+++ b/tensorflow/lite/experimental/support/metadata/cc/BUILD
@@ -1,12 +1,23 @@
+load("//tensorflow/lite/experimental/support/metadata:build_defs.bzl", "stamp_metadata_parser_version")
+
 package(
     default_visibility = ["//tensorflow/lite/experimental/support:users"],
     licenses = ["notice"],  # Apache 2.0
 )
 
+stamp_metadata_parser_version(
+    name = "metadata_parser_h",
+    srcs = ["metadata_parser.h.template"],
+    outs = ["metadata_parser.h"],
+)
+
 cc_library(
     name = "metadata_version",
     srcs = ["metadata_version.cc"],
-    hdrs = ["metadata_version.h"],
+    hdrs = [
+        "metadata_version.h",
+        ":metadata_parser_h",
+    ],
     deps = [
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/experimental/support/metadata:metadata_schema_cc",
diff --git a/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template b/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
new file mode 100644
index 00000000000..dfb62d0de81
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
+
+namespace tflite {
+namespace metadata {
+
+// The version of the metadata parser that this metadata versioning library is
+// depending on.
+inline constexpr char kMatadataParserVersion[] = "{LATEST_METADATA_PARSER_VERSION}";
+
+}  // namespace metadata
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/BUILD b/tensorflow/lite/experimental/support/metadata/cc/test/BUILD
index fd829124c73..f9d78567d70 100644
--- a/tensorflow/lite/experimental/support/metadata/cc/test/BUILD
+++ b/tensorflow/lite/experimental/support/metadata/cc/test/BUILD
@@ -13,3 +13,12 @@ cc_test(
         "@flatbuffers",
     ],
 )
+
+cc_test(
+    name = "metadata_parser_test",
+    srcs = ["metadata_parser_test.cc"],
+    deps = [
+        "//tensorflow/lite/experimental/support/metadata/cc:metadata_version",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
new file mode 100644
index 00000000000..af7b8791fe8
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h"
+
+#include 
+#include 
+
+namespace tflite {
+namespace metadata {
+namespace {
+
+using ::testing::MatchesRegex;
+
+TEST(MetadataParserTest, MatadataParserVersionIsWellFormed) {
+  // Validates that the version is well-formed (x.y.z).
+  EXPECT_THAT(kMatadataParserVersion, MatchesRegex("[0-9]+\\.[0-9]+\\.[0-9]+"));
+}
+
+}  // namespace
+}  // namespace metadata
+}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc
index 02ecfdbd232..03f4d3bf28b 100644
--- a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc
+++ b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc
@@ -44,7 +44,7 @@ TEST(MetadataVersionTest,
                                             builder.GetSize(), &min_version),
             kTfLiteOk);
   // Validates that the version is well-formed (x.y.z).
-  EXPECT_THAT(min_version, MatchesRegex("[0-9]*\\.[0-9]*\\.[0-9]"));
+  EXPECT_THAT(min_version, MatchesRegex("[0-9]+\\.[0-9]+\\.[0-9]+"));
 }
 
 TEST(MetadataVersionTest,
diff --git a/tensorflow/lite/experimental/support/metadata/java/BUILD b/tensorflow/lite/experimental/support/metadata/java/BUILD
index c208752ae24..00d10bcca56 100644
--- a/tensorflow/lite/experimental/support/metadata/java/BUILD
+++ b/tensorflow/lite/experimental/support/metadata/java/BUILD
@@ -9,9 +9,13 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+METADATA_SRCS = glob(
+    ["src/java/org/tensorflow/lite/support/metadata/**/*.java"],
+)
+
 android_library(
     name = "tensorflow-lite-support-metadata",
-    srcs = glob(["src/java/org/tensorflow/lite/support/metadata/**/*.java"]),
+    srcs = METADATA_SRCS,
     manifest = "AndroidManifest.xml",
     deps = [
         "//tensorflow/lite/experimental/support/metadata:metadata_schema_fbs_android",
@@ -22,7 +26,7 @@ android_library(
 
 java_library(
     name = "tensorflow-lite-support-metadata-lib",
-    srcs = glob(["src/java/org/tensorflow/lite/support/metadata/**/*.java"]),
+    srcs = METADATA_SRCS,
     javacopts = JAVACOPTS,
     resource_jars = [
         "//tensorflow/lite/experimental/support/metadata:libmetadata_schema_java.jar",
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java
index 9da5b59cf46..9bf5ae93138 100644
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java
+++ b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java
@@ -52,10 +52,6 @@ import org.tensorflow.lite.support.metadata.schema.TensorMetadata;
  * MetadataExtractor} omits subgraph index as an input in its methods.
  */
 public class MetadataExtractor {
-  // TODO(b/156539454): remove the hardcode versioning number and populate the version through
-  // genrule.
-  /** The version of the metadata parser that this {@link MetadataExtractor} library depends on. */
-  public static final String METADATA_PARSER_VERSION = "1.0.1";
 
   /** The helper class to load metadata from TFLite model FlatBuffer. */
   private final ModelInfo modelInfo;
@@ -85,7 +81,7 @@ public class MetadataExtractor {
         System.err.printf(
             " Some fields in the metadata belong to a future schema. The minimum parser"
                 + " version required is %s, but the version of the current metadata parser is %s",
-            metadataInfo.getMininumParserVersion(), METADATA_PARSER_VERSION);
+            metadataInfo.getMininumParserVersion(), MetadataParser.VERSION);
       }
 
       checkArgument(
@@ -290,7 +286,7 @@ public class MetadataExtractor {
     if (minVersion == null) {
       return true;
     }
-    return compareVersions(minVersion, METADATA_PARSER_VERSION) <= 0;
+    return compareVersions(minVersion, MetadataParser.VERSION) <= 0;
   }
 
   /**
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
new file mode 100644
index 00000000000..195a330462b
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
@@ -0,0 +1,27 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+package org.tensorflow.lite.support.metadata;
+
+/** Information about the metadata parser that this metadata extractor library is depending on. */
+public final class MetadataParser {
+  /**
+   * The version of the metadata parser that this metadata extractor library is depending on. The
+   * value should match the value of "Schema Semantic version" in metadata_schema.fbs.
+   */
+  public static final String VERSION = "1.0.1";
+
+  private MetadataParser() {}
+}
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template b/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
new file mode 100644
index 00000000000..a41ac06969c
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
@@ -0,0 +1,26 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Information about the metadata parser that this python library depends on."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+class MetadataParser(object):
+  """Information about the metadata parser."""
+
+  # The version of the metadata parser.
+  VERSION = "{LATEST_METADATA_PARSER_VERSION}"
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py b/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py
new file mode 100644
index 00000000000..3b1d19278cd
--- /dev/null
+++ b/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py
@@ -0,0 +1,38 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for tensorflow.lite.experimental.support.metadata.metadata_parser."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import re
+
+from tensorflow.lite.experimental.support.metadata import metadata_parser
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class MetadataParserTest(test_util.TensorFlowTestCase):
+
+  def test_version_wellFormedSemanticVersion(self):
+    # Validates that the version is well-formed (x.y.z).
+    self.assertTrue(
+        re.match('[0-9]+\\.[0-9]+\\.[0-9]+',
+                 metadata_parser.MetadataParser.VERSION))
+
+
+if __name__ == '__main__':
+  test.main()
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs b/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs
index 53c26b3e079..a88225f1960 100644
--- a/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs
+++ b/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs
@@ -29,18 +29,31 @@ namespace tflite;
 // generate the model interface. It is recommended to fill in at least those
 // enties to boost the codegen performance.
 
-// LINT.IfChange
-
-// The Metadata schema is versioned by the Semantic versioning number, which
-// tracks the schema changes according to the Semantic versioning rules.
+// The Metadata schema is versioned by the Semantic versioning number, such as
+// MAJOR.MINOR.PATCH. It tracks the schema changes according to the rules below:
+//  * Bump up the MAJOR number when making potentially backwards incompatible
+//    changes. It must be incremented if the new changes break the backwards
+//    compatibility. It may also include minor and patch level changes as
+//    needed. The true backwards compatibility is indicated by the file
+//    identifier.
+//  * Bump up the MINOR number when making backwards compatible updates for
+//    major features, such as supporting new content types or adding new
+//    processing units.
+//  * Bump up the PATCH number when making small backwards compatible changes,
+//    such as adding a new fields or deprecating certain fields (not deleting
+//    them).
 //
 // ModelMetadata.min_parser_version indicates the minimum necessary metadata
 // parser version to fully understand all fields in a given metadata flatbuffer.
 //
-// New fields and types will have associated comments with the schema version for
-// which they were added.
+// New fields and types will have associated comments with the schema version
+// for which they were added.
 //
+// LINT.IfChange
 // Schema Semantic version: 1.0.1
+// LINT.ThenChange(//tensorflow/lite/experimental/\
+//.    support/metadata/java/src/java/org/tensorflow/lite/support/metadata/\
+//.    MetadataParser.java)
 
 // This indicates the flatbuffer compatibility. The number will bump up when a
 // break change is applied to the schema, such as removing fields or adding new
@@ -53,10 +66,6 @@ file_identifier "M001";
 // File extension of any written files.
 file_extension "tflitemeta";
 
-// LINT.ThenChange(//tensorflow/lite/experimental/\
-//     /support/metadata/java/src/java/org/tensorflow/lite/support/metadata/\
-//     MetadataExtractor.java)
-
 // LINT.IfChange
 enum AssociatedFileType : byte {
   UNKNOWN = 0,

From a8456eae42648e81f8a8f75c793ad53473bebc70 Mon Sep 17 00:00:00 2001
From: Yunxing Dai 
Date: Fri, 19 Jun 2020 11:59:51 -0700
Subject: [PATCH 0642/1390] Check the value shape of set dimension size in
 shape inference.

PiperOrigin-RevId: 317351738
Change-Id: Ia185e7745753711ca9ebf657522ef7422c9696ca
---
 tensorflow/compiler/xla/client/xla_builder.cc |  6 +++--
 .../compiler/xla/service/hlo_verifier.cc      |  3 ++-
 .../compiler/xla/service/shape_inference.cc   |  7 +++++-
 .../compiler/xla/service/shape_inference.h    |  5 ++++-
 .../xla/service/shape_inference_test.cc       | 22 +++++++++++++++++++
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc
index 03ae23ea18b..56e9aba6112 100644
--- a/tensorflow/compiler/xla/client/xla_builder.cc
+++ b/tensorflow/compiler/xla/client/xla_builder.cc
@@ -2762,9 +2762,11 @@ XlaOp XlaBuilder::SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension) {
   return ReportErrorOrReturn([&]() -> StatusOr {
     HloInstructionProto instr;
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+    TF_ASSIGN_OR_RETURN(const Shape* val_shape, GetShapePtr(val));
 
-    TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferSetDimensionSizeShape(
-                                         *operand_shape, dimension));
+    TF_ASSIGN_OR_RETURN(Shape shape,
+                        ShapeInference::InferSetDimensionSizeShape(
+                            *operand_shape, *val_shape, dimension));
     // Setting an op's dynamic dimension to the static size is a noop.
     TF_ASSIGN_OR_RETURN(const HloInstructionProto* val_proto,
                         LookUpInstruction(val));
diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc
index 4661b8fd9e3..d8baebd6fdd 100644
--- a/tensorflow/compiler/xla/service/hlo_verifier.cc
+++ b/tensorflow/compiler/xla/service/hlo_verifier.cc
@@ -1123,7 +1123,8 @@ Status ShapeVerifier::HandleGetDimensionSize(HloInstruction* get_size) {
 Status ShapeVerifier::HandleSetDimensionSize(HloInstruction* set_size) {
   return CheckShape(set_size,
                     ShapeInference::InferSetDimensionSizeShape(
-                        set_size->operand(0)->shape(), set_size->dimension()));
+                        set_size->operand(0)->shape(),
+                        set_size->operand(1)->shape(), set_size->dimension()));
 }
 
 Status ShapeVerifier::CheckShape(const HloInstruction* instruction,
diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc
index 75a80747c1d..bb4a38ded1e 100644
--- a/tensorflow/compiler/xla/service/shape_inference.cc
+++ b/tensorflow/compiler/xla/service/shape_inference.cc
@@ -2248,12 +2248,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation,
 }
 
 /* static */ StatusOr ShapeInference::InferSetDimensionSizeShape(
-    const Shape& shape, int64 dimension) {
+    const Shape& shape, const Shape& val_shape, int64 dimension) {
   if (dimension < 0 || dimension >= shape.rank()) {
     return InvalidArgument("SetDimensionSize dimension out of bounds: %d.",
                            dimension);
   }
 
+  if (val_shape.rank() != 0 || val_shape.element_type() != S32) {
+    return InvalidArgument(
+        "SetDimensionSize's value has to be S32 scalar, got %s",
+        val_shape.ToString());
+  }
   // TODO(b/119580730): Remove this restriction when very large dimension size
   // is needed.
   if (shape.dimensions(dimension) > std::numeric_limits::max()) {
diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h
index 2cb5930d098..d47d96ab52d 100644
--- a/tensorflow/compiler/xla/service/shape_inference.h
+++ b/tensorflow/compiler/xla/service/shape_inference.h
@@ -303,10 +303,13 @@ class ShapeInference {
       const Shape& updates_shape, const ProgramShape& to_apply_shape,
       const ScatterDimensionNumbers& scatter_dim_numbers);
 
+  // Helper that validates the given input shape to GetDimensionSize.
   static StatusOr InferGetDimensionSizeShape(const Shape& shape,
                                                     int64 dimension);
 
-  static StatusOr InferSetDimensionSizeShape(const Shape& shape,
+  // Helper that validates the given input shape to SetDimensionSize.
+  static StatusOr InferSetDimensionSizeShape(const Shape& operand_shape,
+                                                    const Shape& val_shape,
                                                     int64 dimension);
 
   // Helper function for creating a Window proto from user-supplied data.
diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc
index b5ecf6e583e..916d3ab15c8 100644
--- a/tensorflow/compiler/xla/service/shape_inference_test.cc
+++ b/tensorflow/compiler/xla/service/shape_inference_test.cc
@@ -1365,6 +1365,28 @@ TEST_F(ShapeInferenceTest, DotWithTwoContractingDimsPasses) {
   EXPECT_TRUE(ShapeUtil::Equal(inferred_status.ValueOrDie(), output_shape));
 }
 
+TEST_F(ShapeInferenceTest, ErrorSetDimensionSize) {
+  Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3});
+  Shape val_shape = ShapeUtil::MakeShape(S32, {1});
+  auto inferred_status = ShapeInference::InferSetDimensionSizeShape(
+      arg_shape, val_shape, /*dimension=*/0);
+
+  EXPECT_FALSE(inferred_status.ok());
+  EXPECT_THAT(inferred_status.status().error_message(),
+              HasSubstr("value has to be S32 scalar"));
+}
+
+TEST_F(ShapeInferenceTest, ErrorSetDimensionSizeWrongType) {
+  Shape arg_shape = ShapeUtil::MakeShape(F32, {5, 3});
+  Shape val_shape = ShapeUtil::MakeShape(U32, {});
+  auto inferred_status = ShapeInference::InferSetDimensionSizeShape(
+      arg_shape, val_shape, /*dimension=*/0);
+
+  EXPECT_FALSE(inferred_status.ok());
+  EXPECT_THAT(inferred_status.status().error_message(),
+              HasSubstr("value has to be S32 scalar"));
+}
+
 // BatchMatMul with different batch dimension sizes fails.
 TEST_F(ShapeInferenceTest, DotWithMismatchedBatchDimSizesFails) {
   Shape lhs_shape = ShapeUtil::MakeShape(F32, {2, 11, 3});

From 2229ae89c927b46355a15e8af22365d24afc25bf Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Fri, 19 Jun 2020 12:06:36 -0700
Subject: [PATCH 0643/1390] Use group_id as step_id.

PiperOrigin-RevId: 317353238
Change-Id: If52b2b4872c92d3f65af8f6ce1651e8c6da7dae7
---
 tensorflow/core/profiler/convert/BUILD        |  1 +
 .../convert/xplane_to_memory_profile.cc       | 59 +++++++------------
 .../convert/xplane_to_memory_profile.h        |  1 +
 .../convert/xplane_to_memory_profile_test.cc  |  3 +-
 .../profiler/protobuf/memory_profile.proto    |  4 +-
 5 files changed, 25 insertions(+), 43 deletions(-)

diff --git a/tensorflow/core/profiler/convert/BUILD b/tensorflow/core/profiler/convert/BUILD
index abf0176bf6f..06594b1aeaf 100644
--- a/tensorflow/core/profiler/convert/BUILD
+++ b/tensorflow/core/profiler/convert/BUILD
@@ -525,6 +525,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core/profiler/protobuf:memory_profile_proto_cc",
         "//tensorflow/core/profiler/protobuf:xplane_proto_cc",
+        "//tensorflow/core/profiler/utils:group_events",
         "//tensorflow/core/profiler/utils:xplane_builder",
         "//tensorflow/core/profiler/utils:xplane_schema",
         "//tensorflow/core/profiler/utils:xplane_test_utils",
diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc b/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
index d039ca8da32..d7104c2bbf5 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile.cc
@@ -42,6 +42,8 @@ namespace profiler {
 
 namespace {
 
+constexpr int64 kInvalidStepId = -1;
+
 // Index of the time-sorted memory_profile_snapshots list, and the
 // MemoryActivityMetadata proto it contains.
 using IndexMetaPair = std::pair;
@@ -63,7 +65,7 @@ struct ActivityMetadata {
   int64 allocation_bytes = 0;
   uint64 address = 0;
   absl::string_view tf_op_name;
-  int64 step_id = -1;
+  int64 step_id = kInvalidStepId;
   absl::string_view region_type;
   int64 data_type = 0;
   absl::string_view tensor_shape;
@@ -129,7 +131,6 @@ void UpdateProfileSummary(const AggregationStats& stats, int64 time_offset_ps,
 MemoryProfile GenerateMemoryProfile(const XPlane* host_trace) {
   XPlaneVisitor plane = CreateTfXPlaneVisitor(host_trace);
   MemoryProfile memory_profile;
-  auto* step_count = memory_profile.mutable_step_count();
   // Iterate over all XEvents in the XPlane, and add the XStats to a new
   // MemoryProfileSnapshot if the EventType is kMemoryAllocation or
   // kMemoryDeallocation.
@@ -181,9 +182,8 @@ MemoryProfile GenerateMemoryProfile(const XPlane* host_trace) {
           case StatType::kTfOp:
             metadata.tf_op_name = stat.StrOrRefValue();
             break;
-          case StatType::kStepId:
+          case StatType::kGroupId:
             metadata.step_id = stat.IntValue();
-            if (metadata.step_id != 0) (*step_count)[metadata.step_id]++;
             break;
           case StatType::kRegionType:
             metadata.region_type = stat.StrOrRefValue();
@@ -214,40 +214,21 @@ MemoryProfile GenerateMemoryProfile(const XPlane* host_trace) {
   return memory_profile;
 }
 
-// Sequentialize step ids for the memory profile.
-void UpdateStepId(const tensorflow::protobuf::Map<
-                      tensorflow::protobuf_int64 /*orig_step_id*/,
-                      tensorflow::protobuf_int64 /*count*/>& step_count,
-                  PerAllocatorMemoryProfile* memory_profile) {
-  // Map from original random step id to sequential step id.
-  absl::flat_hash_map step_map;
-  constexpr int kUnknownStep = -2;
-  constexpr double kStepFilterRatio = 0.1;  // Magic number for filtering.
-  tensorflow::protobuf_int64 max_step_count = 0;
-  for (const auto& step_and_count : step_count) {
-    max_step_count = std::max(max_step_count, step_and_count.second);
-  }
-  // Filter out noisy and incomplete original step ids.
-  for (const auto& step_and_count : step_count) {
-    if (static_cast(step_and_count.second) / max_step_count >
-        kStepFilterRatio) {
-      step_map[step_and_count.first] = kUnknownStep;
-    }
-  }
-
-  // Update the step ids in memory_profile for this allocator.
-  int64 step_id = -1;
+// Fix invalid step ids of snapshots at the beginning/end of the profile or at
+// the step boundaries. The snapshots with invalid step ids at the beginning get
+// 0 for their step ids. Those at the step boundaries or at the end get the
+// previous snapshot's step id + 1.
+void UpdateStepId(PerAllocatorMemoryProfile* memory_profile) {
+  int64 last_valid_step_id = -1;
+  // Snapshots are already sorted in time.
   for (auto& snapshot : *memory_profile->mutable_memory_profile_snapshots()) {
     DCHECK(snapshot.has_activity_metadata());
-    // Convert the random step id to sequential step id.
-    int64 orig_step_id = snapshot.activity_metadata().step_id();
-    if (step_map.contains(orig_step_id) &&
-        step_map[orig_step_id] == kUnknownStep) {
-      step_map[orig_step_id] = ++step_id;
+    if (snapshot.mutable_activity_metadata()->step_id() == kInvalidStepId) {
+      snapshot.mutable_activity_metadata()->set_step_id(last_valid_step_id + 1);
+    } else {
+      last_valid_step_id = snapshot.mutable_activity_metadata()->step_id();
     }
-    snapshot.mutable_activity_metadata()->set_step_id(step_id);
   }
-  VLOG(2) << "Max sequential step id in profile: " << step_id;
 }
 
 // Update the MemoryActivityMetadata for each deallocation event by copying from
@@ -481,14 +462,14 @@ void ProcessMemoryProfileProto(int64 max_num_snapshots,
       return a.time_offset_ps() < b.time_offset_ps();
     });
 
-    UpdateStepId(memory_profile->step_count(), allocator_memory_profile);
+    UpdateStepId(allocator_memory_profile);
     UpdateDeallocation(allocator_memory_profile);
 
-    int64 peak_bytes_profile = allocator_memory_profile->profile_summary()
-                                   .peak_stats()
-                                   .peak_bytes_in_use();
     int64 peak_step_id =
-        GetPeakMemoryStep(peak_bytes_profile, allocator_memory_profile);
+        GetPeakMemoryStep(allocator_memory_profile->profile_summary()
+                              .peak_stats()
+                              .peak_bytes_in_use(),
+                          allocator_memory_profile);
     ProcessActiveAllocations(peak_step_id, allocator_memory_profile);
     SampleSnapshots(max_num_snapshots, snapshots);
   }
diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile.h b/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
index 873ac800aa5..6eddaeeec71 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile.h
@@ -25,6 +25,7 @@ namespace profiler {
 
 // Process the host threads XPlane and generate MemoryProfile result; at most
 // max_num_snapshots will be displayed on the UI.
+// REQUIRED: host_plane should have been grouped by calling GroupTfEvents().
 MemoryProfile ConvertXPlaneToMemoryProfile(const XPlane& host_plane,
                                            int64 max_num_snapshots = 1000);
 
diff --git a/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc b/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
index 5ddcbcfc75d..c334318dcfe 100644
--- a/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_memory_profile_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/protobuf/memory_profile.pb.h"
 #include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+#include "tensorflow/core/profiler/utils/group_events.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
 #include "tensorflow/core/profiler/utils/xplane_test_utils.h"
@@ -84,11 +85,11 @@ TEST(ConvertXPlaneToMemoryProfile, OneAllocatorMultiActivitiesTest) {
                 {StatType::kRegionType, "temp"},
                 {StatType::kTensorShapes, "[1, 2]"}});
 
+  tensorflow::profiler::GroupTfEvents(&space, nullptr);
   MemoryProfile memory_profile = ConvertXPlaneToMemoryProfile(*host_plane);
   EXPECT_EQ(memory_profile.memory_profile_per_allocator().size(), 1);
   EXPECT_EQ(memory_profile.num_hosts(), 1);
   EXPECT_EQ(memory_profile.memory_ids_size(), 1);
-  EXPECT_EQ(memory_profile.step_count().size(), 1);
   EXPECT_EQ(memory_profile.memory_profile_per_allocator().begin()->first,
             "GPU_0_bfc");
   const auto& allocator_memory_profile =
diff --git a/tensorflow/core/profiler/protobuf/memory_profile.proto b/tensorflow/core/profiler/protobuf/memory_profile.proto
index 7a5272c60b2..4d492a56255 100644
--- a/tensorflow/core/profiler/protobuf/memory_profile.proto
+++ b/tensorflow/core/profiler/protobuf/memory_profile.proto
@@ -122,7 +122,5 @@ message MemoryProfile {
   // Ids for profiled memory allocators, used to populate memory selection list
   // at front end.
   repeated string memory_ids = 3;
-  // Map of original random int64 step id to the count of memory activity events
-  // assigned with it.
-  map step_count = 4;
+  reserved 4;
 }

From 2ec0214b48878b94763c4f41e095b4579d78b58f Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen 
Date: Fri, 19 Jun 2020 12:13:11 -0700
Subject: [PATCH 0644/1390] Internal rollback for saved model importer.

PiperOrigin-RevId: 317354577
Change-Id: I6fdc05921b2f1d85b34b1751ed00f65525d45ad3
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |   4 +-
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  25 ---
 .../mlir/tensorflow/ir/tf_saved_model_ops.td  |  24 ---
 .../tests/tf_saved_model/common_v1.py         |   1 -
 .../tests/tf_saved_model/hash_table_v1.py     |  92 -----------
 .../tensorflow/tests/tf_saved_model_ops.mlir  |   5 -
 .../tests/tf_saved_model_ops_invalid.mlir     |  33 ----
 .../mlir/tensorflow/translate/import_model.cc | 149 +++++-------------
 8 files changed, 43 insertions(+), 290 deletions(-)
 delete mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 7c0d427e87b..bcabb13d301 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -661,9 +661,7 @@ cc_library(
         ":tensorflow_types",
         ":translate_utils",
         "//tensorflow/cc/saved_model:bundle_v2",
-        "//tensorflow/cc/saved_model:constants",
         "//tensorflow/cc/saved_model:loader_lite",
-        "//tensorflow/cc/saved_model:loader_util",
         "//tensorflow/compiler/jit:shape_inference_helpers",
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/tf2xla:functionalize_control_flow",
@@ -675,7 +673,6 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/utils:transitive_fanin",
-        "//tensorflow/core/platform:protobuf_internal",
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
@@ -685,6 +682,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:StandardOps",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index ef248379d2e..7db0eed7713 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -76,23 +76,6 @@ static LogicalResult Verify(GlobalTensorOp global_tensor) {
   return success();
 }
 
-static LogicalResult Verify(SessionInitializerOp session_initializer) {
-  mlir::SymbolTable symbol_table(
-      session_initializer.getParentOfType());
-
-  auto init_func_op =
-      symbol_table.lookup(session_initializer.initializer());
-  if (!init_func_op)
-    return session_initializer.emitOpError()
-           << "the initializer function does not exist";
-
-  if (!init_func_op.getType().getResults().empty())
-    return session_initializer.emitOpError()
-           << "the initializer function should have no output";
-
-  return success();
-}
-
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc.inc"
 
@@ -249,14 +232,6 @@ static LogicalResult VerifySavedModelModule(
       }
     }
   }
-
-  auto session_initializers = module.getOps();
-  if (std::distance(session_initializers.begin(), session_initializers.end()) >
-      1) {
-    return (*++session_initializers.begin()).emitError()
-           << "there must be no more than one session_initializer op";
-  }
-
   SymbolTable symbol_table(module);
   auto symbol_uses = SymbolTable::getSymbolUses(&module.getBodyRegion());
   if (!symbol_uses.hasValue()) {
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
index 497f4d90cb9..4431a160edf 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
@@ -128,28 +128,4 @@ def TfSavedModel_GlobalTensorOp : TfSavedModel_Op<"global_tensor"> {
   let verifier = [{ return Verify(*this); }];
 }
 
-def TfSavedModel_SessionInitializerOp: TfSavedModel_Op<"session_initializer"> {
-  let summary = "Initializes TensorFlow session state.";
-  let description = [{
-    Represents a session initializer function initializes TensorFlow session
-    state. It is used to initialize resources in the saved model before calling
-    any exported functions. There must be no more than one session initializer
-    in a saved model.
-
-    The `initializer` represents the initialization function. The function have
-    no output and this function should be only called once.
-
-    This is used, for example, to initialize hash tables stored in resources and
-    accessed by resource name (rather than as resource handles or bound inputs
-    which is how `global_tensor`s are referenced)
-  }];
-
-  let arguments = (ins
-    FlatSymbolRefAttr:$initializer
-  );
-
-
-  let verifier = [{ return Verify(*this); }];
-}
-
 #endif // SAVED_MODEL_DIALECT
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
index 51ccbeb1fbd..7171f63bb05 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
@@ -84,7 +84,6 @@ def do_test(signature_def_map, show_debug_info=False):
     builder.add_meta_graph_and_variables(
         sess, [tf.saved_model.tag_constants.SERVING],
         signature_def_map,
-        main_op=tf.tables_initializer(),
         strip_default_attrs=True)
     builder.save()
 
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
deleted file mode 100644
index 64847434b82..00000000000
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
+++ /dev/null
@@ -1,92 +0,0 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-
-# RUN: %p/hash_table_v1 | FileCheck %s
-
-# pylint: disable=missing-docstring,line-too-long
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow.compat.v1 as tf
-from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
-
-# Verify that the tf.versions attribute exists. It is difficult to enforce
-# contents, since the version numbers change over time. The conversion logic
-# itself is verified in the common graphdef converter, so here just assert
-# it is being invoked.
-# CHECK: module
-# CHECK-SAME: tf.versions
-# CHECK-SAME: bad_consumers
-# CHECK-SAME: min_consumer
-# CHECK-SAME: producer
-
-# CHECK: "tf_saved_model.session_initializer"() {initializer = [[init:@.*]]} : () -> ()
-# CHECK: "tf_saved_model.global_tensor"()
-
-# CHECK:      func {{@[a-zA-Z_0-9]+}}(
-# CHECK-SAME: [[ARG0:%.*]]: tensor
-# CHECK-SAME: [[ARG1:%.*]]: tensor ()
-
   // Representation for constants: (immutable) global tensor.
   // CHECK: tf_saved_model.global_tensor
   "tf_saved_model.global_tensor"() {
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index f04e1a60b36..7287fcf66c8 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -261,39 +261,6 @@ module attributes {tf_saved_model.semantics} {
 
 // -----
 
-module attributes {tf_saved_model.semantics} {
-
-  // expected-error@+1 {{the initializer function does not exist}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  // expected-error@+1 {{the initializer function should have no output}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
-    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
-    return %0 : tensor<1xf32>
-  }
-}
-
-// -----
-
-module attributes {tf_saved_model.semantics} {
-
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  // expected-error@+1 {{there must be no more than one session_initializer op}}
-  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
-  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
-    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
-    return %0 : tensor<1xf32>
-  }
-}
-
-// -----
-
 module attributes {tf_saved_model.semantics} {
 
   // expected-error@+1 {{exported function @f should be public}}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 3cff4217215..820d0ce31fb 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -60,8 +60,6 @@ limitations under the License.
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
-#include "tensorflow/cc/saved_model/constants.h"
-#include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
@@ -101,7 +99,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/platform/protobuf_internal.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
@@ -119,7 +116,6 @@ using mlir::NamedAttrList;
 using mlir::TensorType;
 using mlir::TF::VarHandleOp;
 using mlir::tf_saved_model::GlobalTensorOp;
-using mlir::tf_saved_model::SessionInitializerOp;
 using stream_executor::port::StatusOr;
 
 namespace {
@@ -2959,13 +2955,6 @@ void SortSavedModelModule(mlir::ModuleOp module) {
     named_global_tensor.global_tensor.getOperation()->moveBefore(
         &module.getBody()->front());
   }
-
-  auto initializers = module.getOps();
-  if (!initializers.empty()) {
-    (*initializers.begin())
-        .getOperation()
-        ->moveBefore(&module.getBody()->front());
-  }
 }
 
 Status CreateSavedModelIR(
@@ -3252,29 +3241,17 @@ class SavedModelSignatureDefImporter {
                                  absl::Span exported_names,
                                  mlir::MLIRContext* context)
       : bundle_(bundle),
-        flib_def_(OpRegistry::Global(), graph_def().library()),
-        debug_info_(),
         exported_names_(exported_names),
-        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {
-    // debug_info might not be loaded with loader_lite.
-    if (bundle_.debug_info != nullptr) debug_info_ = *bundle_.debug_info;
-  }
+        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {}
 
   // Converts the SavedModel to the SavedModel dialect. Creates an MLIR function
   // for each signature.
   StatusOr ConvertSignatures();
-  Status ConvertSignature(const std::string& sig_def_key,
-                          const SignatureDef& signature_def);
-
-  // Converts the initialization graph in the SavedModel to an MLIR function.
-  Status ConvertInitializer();
-
-  // Converts a graph with feeds and fetches to an MLIR function.
-  StatusOr ConvertGraph(
-      const std::string& name,
-      const std::vector>& inputs,
-      const std::vector>& outputs,
-      const std::vector control_outputs);
+  Status ConvertSignature(const GraphDef& graphdef,
+                          const std::string& sig_def_key,
+                          const SignatureDef& signature_def,
+                          const GraphDebugInfo& debug_info,
+                          const FunctionLibraryDefinition& flib_def);
 
   // Creates GlobalTensorOp for each variable and moves each VarHandle op to
   // the enclosing function's arguments.
@@ -3296,62 +3273,18 @@ class SavedModelSignatureDefImporter {
   GraphImportConfig::InputArrays ParseInputArrays(
       const std::vector>& inputs);
 
-  const GraphDef& graph_def() const {
-    return bundle_.meta_graph_def.graph_def();
-  }
-  const FunctionLibraryDefinition& flib_def() const { return flib_def_; }
-  const GraphDebugInfo& debug_info() const { return debug_info_; }
-
   const SavedModelBundle& bundle_;
-  FunctionLibraryDefinition flib_def_;
-  GraphDebugInfo debug_info_;
   absl::Span exported_names_;
   mlir::OwningModuleRef module_;
 };
 
-Status SavedModelSignatureDefImporter::ConvertInitializer() {
-  std::vector asset_file_defs;
-  TF_RETURN_IF_ERROR(
-      internal::GetAssetFileDefs(bundle_.meta_graph_def, &asset_file_defs));
-
-  if (!asset_file_defs.empty())
-    return errors::Unimplemented(
-        absl::StrCat("Assets are not supported in signaturedef importer"));
-
-  std::string init_node_name;
-  TF_RETURN_IF_ERROR(
-      internal::GetInitOp("", bundle_.meta_graph_def, &init_node_name));
-
-  if (init_node_name.empty()) return Status::OK();
-
-  TF_ASSIGN_OR_RETURN(auto sub_module,
-                      ConvertGraph(init_node_name, {}, {}, {init_node_name}));
-
-  mlir::SymbolTable symbol_table(*sub_module);
-
-  auto init_func_op = symbol_table.lookup(init_node_name);
-
-  init_func_op.removeAttr("tf.entry_function");
-
-  mlir::OpBuilder builder(module_->getBodyRegion());
-
-  builder.create(
-      module_->getLoc(), builder.getSymbolRefAttr(init_func_op.getName()));
-
-  // Move the converted functions to top level MLIR module.
-  auto* block = module_->getBody();
-  auto* sub_block = sub_module->getBody();
-  block->getOperations().splice(
-      mlir::Block::iterator(block->getTerminator()), sub_block->getOperations(),
-      sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
-
-  return Status::OK();
-}
-
 StatusOr
 SavedModelSignatureDefImporter::ConvertSignatures() {
   const auto& signatures = bundle_.GetSignatures();
-  PopulateTfVersions(module_.get(), graph_def().versions());
+  const auto& graphdef = bundle_.meta_graph_def.graph_def();
+  PopulateTfVersions(module_.get(), graphdef.versions());
+
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graphdef.library());
 
   // debug_info might not be loaded with loader_lite.
   GraphDebugInfo debug_info;
@@ -3374,10 +3307,9 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
       continue;
     }
 
-    TF_RETURN_IF_ERROR(ConvertSignature(sig_def_key, signature_def));
+    TF_RETURN_IF_ERROR(ConvertSignature(graphdef, sig_def_key, signature_def,
+                                        debug_info, flib_def));
   }
-
-  TF_RETURN_IF_ERROR(ConvertInitializer());
   TF_RETURN_IF_ERROR(LiftVariables());
 
   mlir::OpBuilder builder(module_->getBodyRegion());
@@ -3388,32 +3320,10 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
   return std::move(module_);
 }
 
-StatusOr SavedModelSignatureDefImporter::ConvertGraph(
-    const std::string& name,
-    const std::vector>& inputs,
-    const std::vector>& outputs,
-    const std::vector control_outputs) {
-  GraphImportConfig specs;
-  specs.prune_unused_nodes = true;
-  specs.inputs = ParseInputArrays(inputs);
-  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
-  specs.control_outputs = control_outputs;
-
-  // Convert sub-graphdef to sub-graph.
-  GraphConstructorOptions options;
-  options.allow_internal_ops = true;
-  options.add_default_attributes = true;
-  Graph graph(OpRegistry::Global());
-
-  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, graph_def(), &graph));
-
-  // Convert sub-graph to MLIR module.true
-  return GraphDefImporter::Convert(module_->getContext(), graph, debug_info(),
-                                   flib_def(), specs, name);
-}
-
 Status SavedModelSignatureDefImporter::ConvertSignature(
-    const std::string& sig_def_key, const SignatureDef& signature_def) {
+    const GraphDef& graphdef, const std::string& sig_def_key,
+    const SignatureDef& signature_def, const GraphDebugInfo& debug_info,
+    const FunctionLibraryDefinition& flib_def) {
   // Create local vectors for the input and output and sort them to be
   // deterministic. We don't want anyone to really depend on the order, client
   // should lookup argument/result mapping by attribute name.
@@ -3429,9 +3339,34 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
     return lhs.first.size() < rhs.first.size() || lhs.first > rhs.first;
   });
 
+  GraphImportConfig specs;
+  specs.prune_unused_nodes = true;
+  specs.inputs = ParseInputArrays(inputs);
+  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
+
+  // Remove unused nodes and create sub-graphdef.
+  GraphDef sub_graph_def;
+  TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(
+      graphdef, &sub_graph_def,
+      /*terminal_nodes=*/{specs.outputs.begin(), specs.outputs.end()}));
+
+  // Set the function library definitions in the pruned graphdef.
+  *sub_graph_def.mutable_library() = flib_def.ToProto();
+
+  // Convert sub-graphdef to sub-graph.
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.add_default_attributes = true;
+  Graph sub_graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(
+      ConvertGraphDefToGraph(options, sub_graph_def, &sub_graph));
+
   // Convert sub-graph to MLIR module.
-  TF_ASSIGN_OR_RETURN(auto sub_module,
-                      ConvertGraph(sig_def_key, inputs, outputs, {}));
+  TF_ASSIGN_OR_RETURN(
+      auto sub_module,
+      GraphDefImporter::Convert(module_->getContext(), sub_graph, debug_info,
+                                flib_def, specs, sig_def_key));
   mlir::OpBuilder builder(sub_module->getBodyRegion());
 
   // Find the FuncOp which corresponds to current SignatureDef.

From 08a99df84299b4c55f40800c148f55773f3d6198 Mon Sep 17 00:00:00 2001
From: Jonah Kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Fri, 19 Jun 2020 12:34:16 -0700
Subject: [PATCH 0645/1390] Implement format_dtype as intended.

---
 tensorflow/python/keras/utils/vis_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/utils/vis_utils.py b/tensorflow/python/keras/utils/vis_utils.py
index 21de8014c2a..49407d92ce3 100644
--- a/tensorflow/python/keras/utils/vis_utils.py
+++ b/tensorflow/python/keras/utils/vis_utils.py
@@ -197,7 +197,7 @@ def model_to_dot(model,
           else:
             return str(dtype)
           
-      label = '%s|%s' % (label, dtype)
+      label = '%s|%s' % (label, format_dtype(dtype))
 
     # Rebuild the label as a table including input/output shapes.
     if show_shapes:

From eeb6ccf2cf72a758db2ef41c925b03998229b3a8 Mon Sep 17 00:00:00 2001
From: Jonah Kohn <51345541+jonah-kohn@users.noreply.github.com>
Date: Fri, 19 Jun 2020 12:48:52 -0700
Subject: [PATCH 0646/1390] Correctly access layer attributes.

---
 tensorflow/python/keras/utils/vis_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/utils/vis_utils.py b/tensorflow/python/keras/utils/vis_utils.py
index 49407d92ce3..aa2a4c978cb 100644
--- a/tensorflow/python/keras/utils/vis_utils.py
+++ b/tensorflow/python/keras/utils/vis_utils.py
@@ -197,7 +197,7 @@ def model_to_dot(model,
           else:
             return str(dtype)
           
-      label = '%s|%s' % (label, format_dtype(dtype))
+      label = '%s|%s' % (label, format_dtype(layer.dtype))
 
     # Rebuild the label as a table including input/output shapes.
     if show_shapes:

From 0869ff0af5cacc9d9ea85f6cfc0027166b964f4e Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Fri, 19 Jun 2020 12:50:42 -0700
Subject: [PATCH 0647/1390] Unabbreviate "proj" to "projection" a few places,
 to be consistent over the codebase.

PiperOrigin-RevId: 317361865
Change-Id: I72fedd0bda16a5668fb7b880128f36d2595f0042
---
 tensorflow/lite/kernels/lstm.cc      | 20 +++++++--------
 tensorflow/lite/kernels/lstm_eval.cc | 38 +++++++++++++++-------------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index b941f2237ca..803fbba4eae 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -182,7 +182,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   float input_to_output_weight_scale = default_scale;
   float recurrent_to_output_weight_scale = default_scale;
   float cell_to_output_weight_scale = default_scale;
-  float proj_weight_scale = default_scale;
+  float projection_weight_scale = default_scale;
   float layer_norm_input_scale = default_scale;
   float layer_norm_forget_scale = default_scale;
   float layer_norm_cell_scale = default_scale;
@@ -229,7 +229,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
   }
 
   if (use_projection) {
-    proj_weight_scale = projection_weights->params.scale;
+    projection_weight_scale = projection_weights->params.scale;
   }
   output_state_scale = output_state->params.scale;
 
@@ -276,7 +276,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_16(
       std::pow(2, -15) / intermediate_scale[4] * std::pow(2, -15);
 
   effective_proj_scale =
-      proj_weight_scale * intermediate_scale[4] / output_state_scale;
+      projection_weight_scale * intermediate_scale[4] / output_state_scale;
 
   if (use_peephole) {
     if (!use_cifg) {
@@ -442,7 +442,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   int8_t* input_to_output_weight_ptr = nullptr;
   int8_t* recurrent_to_output_weight_ptr = nullptr;
   int8_t* cell_to_output_weight_ptr = nullptr;
-  int8_t* proj_weight_ptr = nullptr;
+  int8_t* projection_weight_ptr = nullptr;
   int16_t* layer_norm_input_weight_ptr = nullptr;
   int16_t* layer_norm_forget_weight_ptr = nullptr;
   int16_t* layer_norm_cell_weight_ptr = nullptr;
@@ -451,7 +451,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   int32_t* forget_gate_bias_ptr = nullptr;
   int32_t* cell_gate_bias_ptr = nullptr;
   int32_t* output_gate_bias_ptr = nullptr;
-  int32_t* proj_bias_ptr = nullptr;
+  int32_t* projection_bias_ptr = nullptr;
   int16_t* cell_ptr = nullptr;
   int8_t* output_state_ptr = nullptr;
 
@@ -469,7 +469,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   float input_to_output_weight_scale = default_scale;
   float recurrent_to_output_weight_scale = default_scale;
   float cell_to_output_weight_scale = default_scale;
-  float proj_weight_scale = default_scale;
+  float projection_weight_scale = default_scale;
   float layer_norm_input_scale = default_scale;
   float layer_norm_forget_scale = default_scale;
   float layer_norm_cell_scale = default_scale;
@@ -528,10 +528,10 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
   }
 
   if (use_projection) {
-    proj_weight_ptr = projection_weights->data.int8;
-    proj_weight_scale = projection_weights->params.scale;
+    projection_weight_ptr = projection_weights->data.int8;
+    projection_weight_scale = projection_weights->params.scale;
     if (projection_bias) {
-      proj_bias_ptr = projection_bias->data.i32;
+      projection_bias_ptr = projection_bias->data.i32;
     }
   }
   output_state_scale = output_state->params.scale;
@@ -593,7 +593,7 @@ TfLiteStatus PopulateQuantizedLstmParams8x8_8(
                                         output_state_scale /
                                         intermediate_scale[11];
   effective_proj_scale =
-      proj_weight_scale * std::pow(2, -15) / output_state_scale;
+      projection_weight_scale * std::pow(2, -15) / output_state_scale;
 
   if (use_peephole) {
     if (!use_cifg) {
diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 9bdbfa9d48d..f38fdc95f3e 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -919,7 +919,7 @@ inline void LstmStepHybrid(
 //   cell_to_output_weights              - optional
 //
 // Quantized projection weights of size 'n_output * n_cell'
-//   proj_weight_ptr                     - optional
+//   projection_weight_ptr                     - optional
 //
 // Weight scales (scalars) for each of the weights above.
 //   effective_input_to_input_scale_a    - optional
@@ -1019,10 +1019,10 @@ inline void LstmStepInteger(
     int32_t effective_cell_to_forget_scale_b,
     const int16_t* cell_to_output_weight_ptr,
     int32_t effective_cell_to_output_scale_a,
-    int32_t effective_cell_to_output_scale_b, const int8_t* proj_weight_ptr,
-    int32_t effective_proj_scale_a, int32_t effective_proj_scale_b,
-    int32_t hidden_zp, int32_t effective_hidden_scale_a,
-    int32_t effective_hidden_scale_b,
+    int32_t effective_cell_to_output_scale_b,
+    const int8_t* projection_weight_ptr, int32_t effective_proj_scale_a,
+    int32_t effective_proj_scale_b, int32_t hidden_zp,
+    int32_t effective_hidden_scale_a, int32_t effective_hidden_scale_b,
     const int16_t* layer_norm_input_weight_ptr,
     int32_t layer_norm_input_scale_a, int32_t layer_norm_input_scale_b,
     const int16_t* layer_norm_forget_weight_ptr,
@@ -1055,7 +1055,7 @@ inline void LstmStepInteger(
   const bool use_cifg = (input_to_input_weight_ptr == nullptr);
   const bool use_peephole = (cell_to_output_weight_ptr != nullptr);
   const bool use_layer_norm = (layer_norm_forget_weight_ptr != nullptr);
-  const bool use_projection = (proj_weight_ptr != nullptr);
+  const bool use_projection = (projection_weight_ptr != nullptr);
 
   // Check for nullptrs.
   TFLITE_DCHECK(input_to_forget_effective_bias);
@@ -1208,7 +1208,7 @@ inline void LstmStepInteger(
   if (use_projection) {
     std::fill_n(output_ptr, n_batch * n_output, 0);
     tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-        scratch_4_ptr, projection_effective_bias, proj_weight_ptr,
+        scratch_4_ptr, projection_effective_bias, projection_weight_ptr,
         effective_proj_scale_a, effective_proj_scale_b, n_batch, n_cell,
         n_output, output_state_zp, scratch_5_ptr, output_ptr, context);
     if (quantized_proj_clip > 0) {
@@ -1245,7 +1245,7 @@ inline void LstmStepInteger(
 //   cell_to_output_weights              - optional
 //
 // Quantized projection weights of size 'n_output * n_cell'
-//   proj_weight_ptr                     - optional
+//   projection_weight_ptr                     - optional
 //
 // Weight scales (scalars) for each of the weights above.
 //   effective_input_to_input_scale_a    - optional
@@ -1348,9 +1348,9 @@ void LstmStepInteger(
     int32_t effective_cell_to_forget_scale_b,
     const int8_t* cell_to_output_weight_ptr,
     int32_t effective_cell_to_output_scale_a,
-    int32_t effective_cell_to_output_scale_b, const int8_t* proj_weight_ptr,
-    int32_t effective_proj_scale_a, int32_t effective_proj_scale_b,
-    const int16_t* layer_norm_input_weight_ptr,
+    int32_t effective_cell_to_output_scale_b,
+    const int8_t* projection_weight_ptr, int32_t effective_proj_scale_a,
+    int32_t effective_proj_scale_b, const int16_t* layer_norm_input_weight_ptr,
     int32_t layer_norm_input_scale_a, int32_t layer_norm_input_scale_b,
     const int16_t* layer_norm_forget_weight_ptr,
     int32_t layer_norm_forget_scale_a, int32_t layer_norm_forget_scale_b,
@@ -1360,7 +1360,7 @@ void LstmStepInteger(
     int32_t layer_norm_output_scale_a, int32_t layer_norm_output_scale_b,
     const int32_t* input_gate_bias_ptr, const int32_t* forget_gate_bias_ptr,
     const int32_t* cell_gate_bias_ptr, const int32_t* output_gate_bias_ptr,
-    const int32_t* proj_bias_ptr, const TfLiteLSTMParams* params,
+    const int32_t* projection_bias_ptr, const TfLiteLSTMParams* params,
     const int32_t* intermediate_scale_a, const int32_t* intermediate_scale_b,
     const int32_t* intermediate_zp, int16_t quantized_cell_clip,
     int8_t quantized_proj_clip, int n_batch, int n_cell, int n_input,
@@ -1476,8 +1476,9 @@ void LstmStepInteger(
 
   // Projection.
   tensor_utils::MatrixBatchVectorMultiply(
-      scratch3, proj_weight_ptr, effective_proj_scale_a, effective_proj_scale_b,
-      proj_bias_ptr, n_batch, n_cell, n_output, output_state_zp, output_ptr);
+      scratch3, projection_weight_ptr, effective_proj_scale_a,
+      effective_proj_scale_b, projection_bias_ptr, n_batch, n_cell, n_output,
+      output_state_zp, output_ptr);
 
   // Projection clipping.
   if (quantized_proj_clip > 0) {
@@ -2113,7 +2114,8 @@ TfLiteStatus EvalInteger8x8_8(
       GetTensorData(recurrent_to_output_weights);
   const int8_t* cell_to_output_weight_ptr =
       GetTensorData(cell_to_output_weights);
-  const int8_t* proj_weight_ptr = GetTensorData(projection_weights);
+  const int8_t* projection_weight_ptr =
+      GetTensorData(projection_weights);
   const int16_t* layer_norm_input_weight_ptr =
       GetTensorData(input_layer_norm_coefficients);
   const int16_t* layer_norm_forget_weight_ptr =
@@ -2128,7 +2130,7 @@ TfLiteStatus EvalInteger8x8_8(
   const int32_t* cell_gate_bias_ptr = GetTensorData(cell_gate_bias);
   const int32_t* output_gate_bias_ptr =
       GetTensorData(output_gate_bias);
-  const int32_t* proj_bias_ptr = GetTensorData(projection_bias);
+  const int32_t* projection_bias_ptr = GetTensorData(projection_bias);
   int16_t* cell_ptr = GetTensorData(cell_state);
   int8_t* output_state_ptr = GetTensorData(output_state);
   int8_t* output_ptr = nullptr;
@@ -2195,7 +2197,7 @@ TfLiteStatus EvalInteger8x8_8(
         integer_lstm_param->effective_cell_to_output_scale_a,
         integer_lstm_param->effective_cell_to_output_scale_b,
 
-        proj_weight_ptr, integer_lstm_param->effective_proj_scale_a,
+        projection_weight_ptr, integer_lstm_param->effective_proj_scale_a,
         integer_lstm_param->effective_proj_scale_b,
 
         layer_norm_input_weight_ptr,
@@ -2214,7 +2216,7 @@ TfLiteStatus EvalInteger8x8_8(
         integer_lstm_param->layer_norm_output_scale_b,
 
         input_gate_bias_ptr, forget_gate_bias_ptr, cell_gate_bias_ptr,
-        output_gate_bias_ptr, proj_bias_ptr,
+        output_gate_bias_ptr, projection_bias_ptr,
 
         params, integer_lstm_param->intermediate_scale_a,
         integer_lstm_param->intermediate_scale_b,

From bd98ba765aac3e20c9db89131aa3016964749133 Mon Sep 17 00:00:00 2001
From: Haoyu Zhang 
Date: Fri, 19 Jun 2020 13:00:32 -0700
Subject: [PATCH 0648/1390] Fix cancellation race condition in
 BaseRendezvousMgr::RegisterCall

PiperOrigin-RevId: 317363743
Change-Id: Ide89dd360a9885b5e8f67b12f362cbce8cb85d80
---
 .../base_rendezvous_mgr.cc                    | 67 +++++++++----------
 .../distributed_runtime/base_rendezvous_mgr.h | 16 ++---
 .../rpc/rpc_rendezvous_mgr.cc                 |  1 +
 3 files changed, 37 insertions(+), 47 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/base_rendezvous_mgr.cc b/tensorflow/core/distributed_runtime/base_rendezvous_mgr.cc
index 7849e094cb9..4b398e4ecef 100644
--- a/tensorflow/core/distributed_runtime/base_rendezvous_mgr.cc
+++ b/tensorflow/core/distributed_runtime/base_rendezvous_mgr.cc
@@ -139,7 +139,7 @@ Status BaseRemoteRendezvous::Initialize(WorkerSession* session) {
   CHECK_NE(session, nullptr) << "session must not be null!";
   std::vector deferred_calls;
   {
-    mutex_lock l(init_mu_);
+    mutex_lock l(mu_);
     if (session_ != nullptr) {
       if (session_->worker_name() == session->worker_name()) {
         VLOG(1) << "Skipping rendezvous re-initialization.";
@@ -161,12 +161,12 @@ Status BaseRemoteRendezvous::Initialize(WorkerSession* session) {
 }
 
 WorkerSession* BaseRemoteRendezvous::session() {
-  tf_shared_lock l(init_mu_);
+  tf_shared_lock l(mu_);
   return session_;
 }
 
 bool BaseRemoteRendezvous::is_initialized() {
-  tf_shared_lock l(init_mu_);
+  tf_shared_lock l(mu_);
   return is_initialized_locked();
 }
 
@@ -176,7 +176,7 @@ Status BaseRemoteRendezvous::Send(const Rendezvous::ParsedKey& parsed,
   VLOG(1) << "BaseRemoteRendezvous Send " << this << " " << parsed.FullKey();
   WorkerSession* sess = nullptr;
   {
-    tf_shared_lock l(init_mu_);
+    tf_shared_lock l(mu_);
     if (!status_.ok()) return status_;
     DCHECK(is_initialized_locked());
     sess = session_;
@@ -198,7 +198,7 @@ Status BaseRemoteRendezvous::ValidateDevices(const ParsedKey& parsed,
   // (e.g. calling session())
   WorkerSession* sess = nullptr;
   {
-    tf_shared_lock l(init_mu_);
+    tf_shared_lock l(mu_);
     if (!status_.ok()) return status_;
     if (!is_initialized_locked()) {
       return errors::Internal("ValidateDevices called before initialization.");
@@ -345,7 +345,7 @@ void BaseRemoteRendezvous::RecvLocalAsync(const ParsedKey& parsed,
   // Test whether the rendezvous is initialized using a shared lock, to avoid
   // the need for exclusive access in the common case.
   if (TF_PREDICT_FALSE(!is_initialized())) {
-    mutex_lock l(init_mu_);
+    mutex_lock l(mu_);
     if (!is_initialized_locked()) {
       // RecvLocalAsync can be called (due to an incoming RecvTensor RPC from a
       // remote worker) before the RunStep (or PartialRunStep) RPC from the
@@ -386,8 +386,7 @@ void BaseRemoteRendezvous::StartAbort(const Status& s) {
   local_->StartAbort(derived_status);
   {
     // Aborts all active RecvTensor calls.
-    mutex_lock l(init_mu_);
-    mutex_lock l2(active_mu_);
+    mutex_lock l(mu_);
     if (status_.ok()) {
       status_ = derived_status;
       for (auto& entry : active_) {
@@ -402,42 +401,36 @@ void BaseRemoteRendezvous::StartAbort(const Status& s) {
 void BaseRemoteRendezvous::RegisterCall(BaseRecvTensorCall* call,
                                         const Rendezvous::Args& args) {
   CancellationManager* cm = args.cancellation_manager;
-  Status captured_status;
-  {
-    tf_shared_lock l(init_mu_);
-    if (!status_.ok()) {
-      captured_status = status_;
-    }
-  }
-  if (!captured_status.ok()) {
-    call->StartAbort(captured_status);
-    return;
-  }
-
   bool already_cancelled = false;
   InactiveCallback callback = [] {};
-  if (cm != nullptr) {
-    auto token = cm->get_cancellation_token();
-    already_cancelled = !cm->RegisterCallback(token, [this, call] {
-      {
-        tf_shared_lock l(active_mu_);
-        if (active_.find(call) == active_.end()) return;
-      }
+  {
+    mutex_lock l(mu_);
+    if (!status_.ok()) {
+      call->StartAbort(status_);
+      return;
+    }
+    if (cm != nullptr) {
+      auto token = cm->get_cancellation_token();
+      already_cancelled = !cm->RegisterCallback(token, [this, call] {
+        {
+          mutex_lock l(mu_);
+          if (active_.find(call) == active_.end()) return;
+          call->StartAbort(
+              errors::Cancelled("RecvFromRemoteAsync is cancelled."));
+        }
+      });
+      callback = [cm, token] { cm->TryDeregisterCallback(token); };
+    }
+    if (already_cancelled) {
       call->StartAbort(errors::Cancelled("RecvFromRemoteAsync is cancelled."));
-    });
-    callback = [cm, token] { cm->TryDeregisterCallback(token); };
-  }
-
-  if (already_cancelled) {
-    call->StartAbort(errors::Cancelled("RecvFromRemoteAsync is cancelled."));
-  } else {
-    mutex_lock l(active_mu_);
-    CHECK(active_.emplace(call, callback).second);
+    } else {
+      CHECK(active_.emplace(call, callback).second);
+    }
   }
 }
 
 void BaseRemoteRendezvous::DeregisterCall(BaseRecvTensorCall* call) {
-  mutex_lock l(active_mu_);
+  mutex_lock l(mu_);
   auto it = active_.find(call);
   if (it != active_.end()) {
     // Deregister the cancellation callback, if one was registered.
diff --git a/tensorflow/core/distributed_runtime/base_rendezvous_mgr.h b/tensorflow/core/distributed_runtime/base_rendezvous_mgr.h
index afa0f74ea2c..63409a31549 100644
--- a/tensorflow/core/distributed_runtime/base_rendezvous_mgr.h
+++ b/tensorflow/core/distributed_runtime/base_rendezvous_mgr.h
@@ -174,14 +174,12 @@ class BaseRemoteRendezvous : public RemoteRendezvous {
  private:
   Rendezvous* local_;  // Owns a Ref on this object.
 
-  // Guards mutable state that is read-mostly after this rendezvous is
-  // initialized.
-  mutable mutex init_mu_;
+  mutable mutex mu_;
 
   // Status given by StartAbort() if any.
-  Status status_ TF_GUARDED_BY(init_mu_);
+  Status status_ TF_GUARDED_BY(mu_);
 
-  WorkerSession* session_ TF_GUARDED_BY(init_mu_);  // Not owned.
+  WorkerSession* session_ TF_GUARDED_BY(mu_);  // Not owned.
 
   // Data structures to handle calls when partially initialized.
   struct DeferredCall {
@@ -190,16 +188,14 @@ class BaseRemoteRendezvous : public RemoteRendezvous {
 
     DeferredCall(const ParsedKey& parsed, DoneCallback done);
   };
-  std::vector deferred_calls_ TF_GUARDED_BY(init_mu_);
+  std::vector deferred_calls_ TF_GUARDED_BY(mu_);
 
   typedef std::function InactiveCallback;
 
-  // Active outstanding RecvTensor calls.
-  mutex active_mu_;
   std::unordered_map active_
-      TF_GUARDED_BY(active_mu_);
+      TF_GUARDED_BY(mu_);
 
-  bool is_initialized_locked() TF_SHARED_LOCKS_REQUIRED(init_mu_) {
+  bool is_initialized_locked() TF_SHARED_LOCKS_REQUIRED(mu_) {
     return session_ != nullptr;
   }
 
diff --git a/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc b/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc
index 512c17fcfcf..89fe6ced725 100644
--- a/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc
+++ b/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc
@@ -282,6 +282,7 @@ void RpcRemoteRendezvous::RecvFromRemoteAsync(
     // callback.
     call->ReleaseWorker(sess->worker_cache());
     call->done()(call->status(), Args(), Args(), Tensor(), false);
+    DeregisterCall(call);
     get_call_freelist()->Release(call);
     return;
   }

From 5ed030734c63118c3fe09470a1d2da6af204a0eb Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Fri, 19 Jun 2020 13:03:24 -0700
Subject: [PATCH 0649/1390] [tf.data service] Avoid holding locks during RPC
 calls.

This CL fixes a deadlock where a worker holds its lock while making an RPC to the master, and the master holds its lock while making an RPC to the worker. The RPCs require locks to serve, so we end up deadlocked.

We can avoid this by never holding a lock while performing RPCs. This CL modifies the master locking to release the lock when making the `ProcessTask` RPC to the worker.

This change shouldn't affect any functionality - it should only reduce the scope of some locking.

PiperOrigin-RevId: 317364346
Change-Id: I21e5ed8cdaced1192a89ffda4f8f93418e5dc4a5
---
 tensorflow/core/data/service/master_impl.cc | 131 ++++++++++++--------
 tensorflow/core/data/service/master_impl.h  |  19 ++-
 2 files changed, 94 insertions(+), 56 deletions(-)

diff --git a/tensorflow/core/data/service/master_impl.cc b/tensorflow/core/data/service/master_impl.cc
index 37a884d540e..5c7917b4154 100644
--- a/tensorflow/core/data/service/master_impl.cc
+++ b/tensorflow/core/data/service/master_impl.cc
@@ -61,7 +61,8 @@ Status DataServiceMasterImpl::RegisterWorker(
   VLOG(3) << "Received register worker request";
   mutex_lock l(mu_);
   int64 worker_id = next_worker_id_++;
-  workers_.emplace_back(worker_id, request->worker_address());
+  workers_.push_back(
+      std::make_shared(worker_id, request->worker_address()));
   response->set_worker_id(worker_id);
 
   // Allocate tasks to the worker.
@@ -70,17 +71,18 @@ Status DataServiceMasterImpl::RegisterWorker(
     if (job->finished()) {
       continue;
     }
-    int64 task_id = CreateTask(job.get(), request->worker_address());
+    const Task& task = CreateTaskLocked(job.get(), request->worker_address());
 
     TaskDef* task_def = response->add_tasks();
     *task_def->mutable_dataset() =
         datasets_by_id_[job->dataset_id()]->dataset_def();
     task_def->set_dataset_id(job->dataset_id());
     task_def->set_job_id(job->job_id());
-    task_def->set_task_id(task_id);
+    task_def->set_task_id(task.task_id());
   }
 
-  VLOG(1) << "Registered worker " << workers_.back().DebugString();
+  VLOG(1) << "Registered worker at address " << request->worker_address()
+          << " with id " << worker_id;
   return Status::OK();
 }
 
@@ -145,7 +147,6 @@ Status DataServiceMasterImpl::CreateJob(const CreateJobRequest* request,
   VLOG(3) << "Received create job request for dataset id "
           << request->dataset_id();
   ProcessingMode processing_mode = ProcessingMode(request->processing_mode());
-  mutex_lock l(mu_);
   int64 job_id;
   TF_RETURN_IF_ERROR(CreateJob(request->dataset_id(), processing_mode,
                                absl::optional(), &job_id));
@@ -161,25 +162,30 @@ Status DataServiceMasterImpl::GetOrCreateJob(
   VLOG(3) << "Received get or create job request for dataset id "
           << request->dataset_id() << " with name " << request->job_name()
           << " and index " << request->job_name_index();
-  mutex_lock l(mu_);
   NamedJobKey key(request->job_name(), request->job_name_index());
   ProcessingMode requested_processing_mode =
       ProcessingMode(request->processing_mode());
-  std::shared_ptr* job = gtl::FindOrNull(named_jobs_, key);
-  if (job != nullptr) {
-    TF_RETURN_IF_ERROR(ValidateMatchingJob(**job, requested_processing_mode,
-                                           request->dataset_id()));
-    int64 job_id = (*job)->job_id();
-    response->set_job_id(job_id);
-    VLOG(3) << "Found existing job for name=" << request->job_name()
-            << ", index=" << request->job_name_index()
-            << ". job_id: " << job_id;
-    return Status::OK();
+  {
+    mutex_lock l(mu_);
+    std::shared_ptr* job = gtl::FindOrNull(named_jobs_, key);
+    if (job != nullptr) {
+      TF_RETURN_IF_ERROR(ValidateMatchingJob(**job, requested_processing_mode,
+                                             request->dataset_id()));
+      int64 job_id = (*job)->job_id();
+      response->set_job_id(job_id);
+      VLOG(3) << "Found existing job for name=" << request->job_name()
+              << ", index=" << request->job_name_index()
+              << ". job_id: " << job_id;
+      return Status::OK();
+    }
   }
   int64 job_id;
   TF_RETURN_IF_ERROR(CreateJob(request->dataset_id(), requested_processing_mode,
                                request->job_name(), &job_id));
-  named_jobs_[key] = jobs_[job_id];
+  {
+    mutex_lock l(mu_);
+    named_jobs_[key] = jobs_[job_id];
+  }
   response->set_job_id(job_id);
   VLOG(3) << "Created job " << job_id << " for dataset "
           << request->dataset_id() << " and name " << request->job_name();
@@ -211,8 +217,7 @@ Status DataServiceMasterImpl::ValidateMatchingJob(
 Status DataServiceMasterImpl::CreateJob(int64 dataset_id,
                                         ProcessingMode processing_mode,
                                         absl::optional job_name,
-                                        int64* out_job_id)
-    EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+                                        int64* out_job_id) LOCKS_EXCLUDED(mu_) {
   switch (processing_mode) {
     case ProcessingMode::PARALLEL_EPOCHS:
       break;
@@ -225,41 +230,64 @@ Status DataServiceMasterImpl::CreateJob(int64 dataset_id,
                                    ProcessingModeToString(processing_mode),
                                    " not recognized");
   }
-  if (!datasets_by_id_.contains(dataset_id)) {
-    return errors::NotFound("Dataset id: <", dataset_id, "> not found.");
+  std::shared_ptr job;
+  std::vector> workers;
+  {
+    mutex_lock l(mu_);
+    if (!datasets_by_id_.contains(dataset_id)) {
+      return errors::NotFound("Dataset id: <", dataset_id, "> not found.");
+    }
+
+    int64 job_id = next_job_id_++;
+    DCHECK(!jobs_.contains(job_id));
+    job = std::make_shared(job_id, dataset_id, processing_mode, job_name);
+    jobs_[job_id] = job;
+
+    // Copy workers_ so that we can iterate through the workers without holding
+    // the lock. When a new worker is added in `RegisterWorker`, we iterate
+    // through the jobs in `jobs_` and give it a task for each job. So even if a
+    // new worker is registered after we release the lock, because this job has
+    // been added to `jobs_`, it will still receive a task for this job.
+    workers = workers_;
+    const Dataset& dataset = *datasets_by_id_[dataset_id];
+    if (VLOG_IS_ON(1)) {
+      VLOG(1) << "Sending tasks to workers for job " << job->job_id()
+              << ". Dataset id: " << dataset_id
+              << ". Dataset fingerprint: " << dataset.fingerprint()
+              << ". Dataset definition size: "
+              << datasets_by_id_[dataset_id]->dataset_def().ByteSizeLong();
+    }
   }
 
-  int64 job_id = next_job_id_++;
-  DCHECK(!jobs_.contains(job_id));
-  auto job =
-      std::make_shared(job_id, dataset_id, processing_mode, job_name);
-  jobs_[job_id] = job;
-
-  for (auto& worker : workers_) {
-    int64 task_id = CreateTask(job.get(), worker.address());
-
-    // TODO(aaudibert): perform these calls asynchronously.
-    // TODO(aaudibert): clean up in case some calls succeed, but later calls
-    // fail
-    TF_RETURN_IF_ERROR(AllocateTaskToWorker(tasks_.at(task_id), &worker));
+  for (auto& worker : workers) {
+    const Task& task = CreateTask(job.get(), worker->address());
+    Status s = AllocateTaskToWorker(task, worker.get());
+    if (!s.ok()) {
+      LOG(WARNING) << "Failed to allocate task with id " << task.task_id()
+                   << " to worker at address " << worker->address() << ": "
+                   << s.error_message();
+    }
   }
+  VLOG(1) << "Done sending tasks to workers for job " << job->job_id();
 
-  *out_job_id = job_id;
+  *out_job_id = job->job_id();
   return Status::OK();
 }
 
-int64 DataServiceMasterImpl::CreateTask(Job* job,
-                                        const std::string& worker_address)
-    EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+const DataServiceMasterImpl::Task& DataServiceMasterImpl::CreateTask(
+    Job* job, const std::string& worker_address) LOCKS_EXCLUDED(mu_) {
+  mutex_lock l(mu_);
+  return CreateTaskLocked(job, worker_address);
+}
+
+const DataServiceMasterImpl::Task& DataServiceMasterImpl::CreateTaskLocked(
+    Job* job, const std::string& worker_address) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
   int64 task_id = next_task_id_++;
   DCHECK(!tasks_.contains(task_id));
-  auto result =
-      tasks_.emplace(std::piecewise_construct, std::forward_as_tuple(task_id),
-                     std::forward_as_tuple(task_id, job->job_id(),
-                                           job->dataset_id(), worker_address));
+  tasks_.insert({task_id, Task(task_id, job->job_id(), job->dataset_id(),
+                               worker_address)});
   job->add_task_id(task_id);
-  DCHECK(result.second);
-  return task_id;
+  return tasks_.at(task_id);
 }
 
 Status DataServiceMasterImpl::EnsureWorkerStubInitialized(Worker* worker) {
@@ -273,14 +301,17 @@ Status DataServiceMasterImpl::EnsureWorkerStubInitialized(Worker* worker) {
 
 Status DataServiceMasterImpl::AllocateTaskToWorker(const Task& task,
                                                    Worker* worker)
-    EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    LOCKS_EXCLUDED(mu_) {
   TF_RETURN_IF_ERROR(EnsureWorkerStubInitialized(worker));
   grpc::ClientContext client_ctx;
   ProcessTaskRequest req;
   req.mutable_task()->set_dataset_id(task.dataset_id());
-  DCHECK(datasets_by_id_.contains(task.dataset_id()));
-  *req.mutable_task()->mutable_dataset() =
-      datasets_by_id_.at(task.dataset_id())->dataset_def();
+  {
+    mutex_lock l(mu_);
+    DCHECK(datasets_by_id_.contains(task.dataset_id()));
+    *req.mutable_task()->mutable_dataset() =
+        datasets_by_id_.at(task.dataset_id())->dataset_def();
+  }
   req.mutable_task()->set_task_id(task.task_id());
   ProcessTaskResponse resp;
   grpc::Status s = worker->stub()->ProcessTask(&client_ctx, req, &resp);
@@ -321,8 +352,8 @@ Status DataServiceMasterImpl::GetWorkers(const GetWorkersRequest* request,
   VLOG(3) << "Enter GetWorkers";
   for (auto& worker : workers_) {
     WorkerInfo* info = response->add_workers();
-    info->set_address(worker.address());
-    info->set_id(worker.worker_id());
+    info->set_address(worker->address());
+    info->set_id(worker->worker_id());
   }
   VLOG(3) << "Returning list of " << workers_.size()
           << " workers from GetWorkers";
diff --git a/tensorflow/core/data/service/master_impl.h b/tensorflow/core/data/service/master_impl.h
index 0dc049a389c..67df2613118 100644
--- a/tensorflow/core/data/service/master_impl.h
+++ b/tensorflow/core/data/service/master_impl.h
@@ -177,16 +177,23 @@ class DataServiceMasterImpl {
   };
 
   // Registers a dataset with the given fingerprint, returning a new dataset id.
-  int64 RegisterDataset(uint64 fingerprint, const DatasetDef& dataset);
+  int64 RegisterDataset(uint64 fingerprint, const DatasetDef& dataset)
+      EXCLUSIVE_LOCKS_REQUIRED(mu_);
   // Initializes a workers stub, if it hasn't been initialized already.
   Status EnsureWorkerStubInitialized(Worker* worker);
   // Instructs a worker to begin processing a task.
-  Status AllocateTaskToWorker(const Task& task_id, Worker* worker);
+  Status AllocateTaskToWorker(const Task& task_id, Worker* worker)
+      LOCKS_EXCLUDED(mu_);
   // Creates a job and stores its job_id in `*job_id`.
   Status CreateJob(int64 dataset_id, ProcessingMode processing_mode,
-                   absl::optional job_name, int64* out_job_id);
-  // Creates a new task for a job, returning the new task's id.
-  int64 CreateTask(Job* job, const std::string& worker_address);
+                   absl::optional job_name, int64* out_job_id)
+      LOCKS_EXCLUDED(mu_);
+  // Creates a new task for a job, returning a reference to the task.
+  const Task& CreateTask(Job* job, const std::string& worker_address)
+      LOCKS_EXCLUDED(mu_);
+  // Same as `CreateTask`, but expects that the master lock is already held.
+  const Task& CreateTaskLocked(Job* job, const std::string& worker_address)
+      EXCLUSIVE_LOCKS_REQUIRED(mu_);
   // Validates that an existing job matches the given processing_mode and
   // dataset_id, returning an error status describing any difference.
   Status ValidateMatchingJob(const Job& job, ProcessingMode processing_mode,
@@ -202,7 +209,7 @@ class DataServiceMasterImpl {
   int64 next_task_id_ TF_GUARDED_BY(mu_) = 0;
 
   // Registered workers.
-  std::vector workers_ TF_GUARDED_BY(mu_);
+  std::vector> workers_ TF_GUARDED_BY(mu_);
   // Registered datasets, keyed by dataset ids.
   absl::flat_hash_map> datasets_by_id_
       TF_GUARDED_BY(mu_);

From 3ae2cf96105c32a71db7796b901e3c8ae5d8f71a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 13:05:16 -0700
Subject: [PATCH 0650/1390] Suppress Clang error: unused variable 'row_base'
 when building with MSAN instrumentation

PiperOrigin-RevId: 317364738
Change-Id: Ic55fdfa9520f341d254aa42333da67ae5d340680
---
 tensorflow/core/kernels/eigen_contraction_kernel.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/kernels/eigen_contraction_kernel.h b/tensorflow/core/kernels/eigen_contraction_kernel.h
index 7233020c1c1..ef4b9dbc012 100644
--- a/tensorflow/core/kernels/eigen_contraction_kernel.h
+++ b/tensorflow/core/kernels/eigen_contraction_kernel.h
@@ -171,6 +171,7 @@ struct mkldnn_gemm_kernel
Date: Fri, 19 Jun 2020 13:06:52 -0700
Subject: [PATCH 0651/1390] Remove `run_deprecated_v1` annotations from
 collective ops tests.

PiperOrigin-RevId: 317365063
Change-Id: Ibf13ad8629947becd40038d41ee213d3466b6292
---
 tensorflow/python/ops/collective_ops_test.py | 380 ++++++++++---------
 1 file changed, 203 insertions(+), 177 deletions(-)

diff --git a/tensorflow/python/ops/collective_ops_test.py b/tensorflow/python/ops/collective_ops_test.py
index 9727593a1c5..8e3a95d7dbf 100644
--- a/tensorflow/python/ops/collective_ops_test.py
+++ b/tensorflow/python/ops/collective_ops_test.py
@@ -104,39 +104,42 @@ class CollectiveOpTest(test.TestCase):
     for i in range(group_size * num_instances):
       self.assertAllClose(results[i], expected, rtol=1e-5, atol=1e-5)
 
-  @test_util.run_deprecated_v1
   def testCollectiveReduce(self):
-    self._testCollectiveReduce(
-        inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
-                [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
-        expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
-        set_graph_key=True)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=True)
 
-  @test_util.run_deprecated_v1
   def testCollectiveAutoGraphKey(self):
-    self._testCollectiveReduce(
-        inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
-                [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
-        expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
-        set_graph_key=False)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=False)
 
-  @test_util.run_deprecated_v1
   def testFp16Reduce(self):
-    self._testCollectiveReduce(
-        inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
-                [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
-        expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
-        set_graph_key=True,
-        fp16=True)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=True,
+          fp16=True)
 
-  @test_util.run_deprecated_v1
   def testCollectiveMultipleConcurrentReduce(self):
-    self._testMultipleConcurrentCollectiveReduce(
-        [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
-        [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
-        [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testMultipleConcurrentCollectiveReduce(
+          [0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+          [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3],
+          [0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2])
 
-  @test_util.run_deprecated_v1
   def testCollectiveTimeoutV1(self):
     timeout = 4.5
     kwargs = dict(
@@ -145,14 +148,17 @@ class CollectiveOpTest(test.TestCase):
         set_graph_key=True,
         timeout=timeout)
 
-    self._testCollectiveReduce(**kwargs)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(**kwargs)
 
     start_time = time.time()
-    with self.assertRaisesRegex(
-        errors.DeadlineExceededError,
-        'Collective has timed out waiting for other workers'):
-      self._testCollectiveReduce(
-          reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(
+          errors.DeadlineExceededError,
+          'Collective has timed out waiting for other workers'):
+        self._testCollectiveReduce(
+            reported_group_size=len(kwargs['inputs']) + 1, **kwargs)
     elapsed = time.time() - start_time
     self.assertAllGreaterEqual(elapsed, timeout)
 
@@ -199,17 +205,18 @@ class CollectiveOpTest(test.TestCase):
     elapsed = time.time() - start_time
     self.assertAllGreaterEqual(elapsed, timeout)
 
-  @test_util.run_deprecated_v1
   def testNcclHintFallbackToRingReduce(self):
     """Tests that setting `communication_hint=nccl` works on non-GPU builds."""
     if kernels.get_registered_kernels_for_op('NcclAllReduce'):
       self.skipTest('Run only on non-GPU environments')
-    self._testCollectiveReduce(
-        inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
-                [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
-        expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
-        set_graph_key=False,
-        communication_hint='nccl')
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1],
+                  [0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3]],
+          expected=[0.2, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2],
+          set_graph_key=False,
+          communication_hint='nccl')
 
   def _testWhile(self, num_vars, num_iterations, key_base):
     group_size = 2
@@ -262,15 +269,16 @@ class CollectiveOpTest(test.TestCase):
           [((1 << (num_iterations + v)) * 1.) for v in range(num_vars)]
           for _ in range(group_size)])
 
-  @test_util.run_deprecated_v1
   def testSimpleWhile(self):
-    self._testWhile(num_vars=1, num_iterations=4, key_base=20)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testWhile(num_vars=1, num_iterations=4, key_base=20)
 
-  @test_util.run_deprecated_v1
   def testWhileMultipleAllReduce(self):
-    self._testWhile(num_vars=2, num_iterations=4, key_base=20)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testWhile(num_vars=2, num_iterations=4, key_base=20)
 
-  @test_util.run_deprecated_v1
   def testWhileWithScopedAllocator(self):
     group_size = 2
     group_key = 1
@@ -284,47 +292,52 @@ class CollectiveOpTest(test.TestCase):
     del rewrite_options.scoped_allocator_opts.enable_op[:]
     rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
 
-    with self.session(config=config) as sess:
-      run_ops = []
-      for i in range(group_size):
-        with ops.device('CPU:%d' % i):
-          constant = constant_op.constant(0.)
-          cond = lambda i: math_ops.less(i, 10.)
-          body = lambda i: math_ops.add(i, 1.)
-          input0 = control_flow_ops.while_loop(cond, body, [constant])
-          input1 = math_ops.add(constant, 5)
-          colred0 = collective_ops.all_reduce(input0, group_size, group_key,
-                                              instance_key0, 'Add', 'Id')
-          colred1 = collective_ops.all_reduce(input1, group_size, group_key,
-                                              instance_key1, 'Add', 'Id')
-          run_ops.append(math_ops.add_n([colred0, colred1]))
-      results = sess.run(run_ops)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(config=config) as sess:
+        run_ops = []
+        for i in range(group_size):
+          with ops.device('CPU:%d' % i):
+            constant = constant_op.constant(0.)
+            cond = lambda i: math_ops.less(i, 10.)
+            body = lambda i: math_ops.add(i, 1.)
+            input0 = control_flow_ops.while_loop(cond, body, [constant])
+            input1 = math_ops.add(constant, 5)
+            colred0 = collective_ops.all_reduce(input0, group_size, group_key,
+                                                instance_key0, 'Add', 'Id')
+            colred1 = collective_ops.all_reduce(input1, group_size, group_key,
+                                                instance_key1, 'Add', 'Id')
+            run_ops.append(math_ops.add_n([colred0, colred1]))
+        results = sess.run(run_ops)
       self.assertEqual(results, [30., 30.])
 
-  @test_util.run_deprecated_v1
   def testCollectiveReduceScalar(self):
-    self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
-                               set_graph_key=True)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(inputs=[0.1, 0.3], expected=0.2,
+                                 set_graph_key=True)
 
-  @test_util.run_deprecated_v1
   def testCollectiveReduceMaximum(self):
-    self._testCollectiveReduce(
-        inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
-        expected=[10., 20., 30., 40., 50.],
-        set_graph_key=True,
-        instance_key=30,
-        merge_op='Max',
-        final_op='Id')
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
+          expected=[10., 20., 30., 40., 50.],
+          set_graph_key=True,
+          instance_key=30,
+          merge_op='Max',
+          final_op='Id')
 
-  @test_util.run_deprecated_v1
   def testCollectiveReduceMinimum(self):
-    self._testCollectiveReduce(
-        inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
-        expected=[1., 2., 3., 4., 5.],
-        set_graph_key=True,
-        instance_key=40,
-        merge_op='Min',
-        final_op='Id')
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveReduce(
+          inputs=[[1., 20., 3., 40., 5.], [10., 2., 30., 4., 50.]],
+          expected=[1., 2., 3., 4., 5.],
+          set_graph_key=True,
+          instance_key=40,
+          merge_op='Min',
+          final_op='Id')
 
   def _testCollectiveBroadcast(self, in_val):
     group_key = 1
@@ -345,13 +358,15 @@ class CollectiveOpTest(test.TestCase):
     self.assertAllClose(results[0], in_val, rtol=1e-5, atol=1e-5)
     self.assertAllClose(results[1], in_val, rtol=1e-5, atol=1e-5)
 
-  @test_util.run_deprecated_v1
   def testCollectiveBroadcast(self):
-    self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveBroadcast([0.1, 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1])
 
-  @test_util.run_deprecated_v1
   def testCollectiveBroadcastBool(self):
-    self._testCollectiveBroadcast([True, False])
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveBroadcast([True, False])
 
   def _testCollectiveGather(self, t0, t1, expected, set_graph_key):
     group_key = 1
@@ -371,94 +386,101 @@ class CollectiveOpTest(test.TestCase):
     self.assertAllClose(results[0], expected, rtol=1e-5, atol=1e-5)
     self.assertAllClose(results[1], expected, rtol=1e-5, atol=1e-5)
 
-  @test_util.run_deprecated_v1
   def testCollectiveGather(self):
-    self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
-                               [10, 11, 12, 13, 14, 15, 16, 17],
-                               [0, 1, 2, 3, 4, 5, 6, 7,
-                                10, 11, 12, 13, 14, 15, 16, 17],
-                               True)
-    self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
-                               [[10, 11, 12, 13], [14, 15, 16, 17]],
-                               [[0, 1, 2, 3], [4, 5, 6, 7],
-                                [10, 11, 12, 13], [14, 15, 16, 17]],
-                               True)
-    self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
-                               [[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
-                               [[[0, 1], [2, 3]], [[4, 5], [6, 7]],
-                                [[10, 11], [12, 13]], [[14, 15], [16, 17]]],
-                               True)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      self._testCollectiveGather([0, 1, 2, 3, 4, 5, 6, 7],
+                                 [10, 11, 12, 13, 14, 15, 16, 17],
+                                 [0, 1, 2, 3, 4, 5, 6, 7,
+                                  10, 11, 12, 13, 14, 15, 16, 17],
+                                 True)
+      self._testCollectiveGather([[0, 1, 2, 3], [4, 5, 6, 7]],
+                                 [[10, 11, 12, 13], [14, 15, 16, 17]],
+                                 [[0, 1, 2, 3], [4, 5, 6, 7],
+                                  [10, 11, 12, 13], [14, 15, 16, 17]],
+                                 True)
+      self._testCollectiveGather([[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
+                                 [[[10, 11], [12, 13]], [[14, 15], [16, 17]]],
+                                 [[[0, 1], [2, 3]], [[4, 5], [6, 7]],
+                                  [[10, 11], [12, 13]], [[14, 15], [16, 17]]],
+                                 True)
 
-  @test_util.run_deprecated_v1
   def testCollectiveGatherShapeMismatch(self):
     group_key = 1
     instance_key = 1
     t0 = [1, 2, 3, 4]
     t1 = [5, 6, 7, 8]
     t2 = [9, 10]
-    with self.session(
-        config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
-      with ops.device('/CPU:0'):
-        in0 = constant_op.constant(t0)
-        c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
-      with ops.device('/CPU:1'):
-        in1 = constant_op.constant(t1)
-        in2 = constant_op.constant(t2)
-        c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
-        c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
-      run_options = config_pb2.RunOptions()
-      run_options.experimental.collective_graph_key = 1
-      sess.run([c0, c1], options=run_options)
-      with self.assertRaisesRegexp(errors.InvalidArgumentError,
-                                   'Shape mismatch'):
-        sess.run([c0, c2], options=run_options)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = constant_op.constant(t0)
+          c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
+        with ops.device('/CPU:1'):
+          in1 = constant_op.constant(t1)
+          in2 = constant_op.constant(t2)
+          c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
+          c2 = collective_ops.all_gather(in2, 2, group_key, instance_key)
+        run_options = config_pb2.RunOptions()
+        run_options.experimental.collective_graph_key = 1
+        sess.run([c0, c1], options=run_options)
+        with self.assertRaisesRegexp(errors.InvalidArgumentError,
+                                     'Shape mismatch'):
+          sess.run([c0, c2], options=run_options)
 
-  @test_util.run_deprecated_v1
   def testCollectiveGatherShapeMismatchAcrossDevices(self):
     group_key = 1
     instance_key = 1
     t0 = [1, 2, 3, 4]
     t1 = [5, 6]
-    with self.session(
-        config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
-      with ops.device('/CPU:0'):
-        in0 = constant_op.constant(t0)
-        c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
-      with ops.device('/CPU:1'):
-        in1 = constant_op.constant(t1)
-        c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
-      run_options = config_pb2.RunOptions()
-      run_options.experimental.collective_graph_key = 1
-      with self.assertRaisesRegexp(errors.InvalidArgumentError,
-                                   'Shape mismatch'):
-        sess.run([c0, c1], options=run_options)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(device_count={'CPU': 2})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = constant_op.constant(t0)
+          c0 = collective_ops.all_gather(in0, 2, group_key, instance_key)
+        with ops.device('/CPU:1'):
+          in1 = constant_op.constant(t1)
+          c1 = collective_ops.all_gather(in1, 2, group_key, instance_key)
+        run_options = config_pb2.RunOptions()
+        run_options.experimental.collective_graph_key = 1
+        with self.assertRaisesRegexp(errors.InvalidArgumentError,
+                                     'Shape mismatch'):
+          sess.run([c0, c1], options=run_options)
 
-  @test_util.run_deprecated_v1
   def testCollectiveGatherPolymorphicShape(self):
     t0 = [0, 1, 2, 3, 4, 5, 6, 7]
     t1 = [10, 11, 12, 13, 14, 15, 16, 17]
     group_size = 2
     group_key = 1
     instance_key = 123
-    with self.session(
-        config=config_pb2.ConfigProto(
-            device_count={'CPU': group_size})) as sess:
-      with ops.device('/CPU:0'):
-        in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
-        c0 = collective_ops.all_gather(in0, group_size, group_key, instance_key)
-      with ops.device('/CPU:1'):
-        in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
-        c1 = collective_ops.all_gather(in1, group_size, group_key, instance_key)
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(
+          config=config_pb2.ConfigProto(
+              device_count={'CPU': group_size})) as sess:
+        with ops.device('/CPU:0'):
+          in0 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
+          c0 = collective_ops.all_gather(in0, group_size, group_key,
+                                         instance_key)
+        with ops.device('/CPU:1'):
+          in1 = array_ops.placeholder(dtype=dtypes.int32, shape=[None])
+          c1 = collective_ops.all_gather(in1, group_size, group_key,
+                                         instance_key)
 
-      results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
-      expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
-      self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
-      self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
+        results = sess.run([c0, c1], feed_dict={in0: t0, in1: t1})
+        results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
 
-      results_ = sess.run([c0, c1], feed_dict={in0: t0[1:], in1: t1[1:]})
-      expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
-      self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
-      self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
+    expected_output = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 15, 16, 17]
+    self.assertAllClose(results[0], expected_output, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results[1], expected_output, rtol=1e-5, atol=1e-5)
+
+    expected_output_ = [1, 2, 3, 4, 5, 6, 7, 11, 12, 13, 14, 15, 16, 17]
+    self.assertAllClose(results_[0], expected_output_, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(results_[1], expected_output_, rtol=1e-5, atol=1e-5)
 
   @test_util.run_v2_only
   def testCollectiveGroupSizeMismatch(self):
@@ -492,8 +514,17 @@ class CollectiveOpTest(test.TestCase):
                                  'but that group has size'):
       run_all_reduce()
 
-  @test_util.run_deprecated_v1
+  @test_util.run_v2_only
   def testCollectiveTensorsHaveNoDeviceSpecified(self):
+    context._reset_context()
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    context.ensure_initialized()
+
     group_size = 2
     group_key = 1
     instance_key = 1
@@ -517,20 +548,12 @@ class CollectiveOpTest(test.TestCase):
 
       return results
 
-    with self.session(config=config_pb2.ConfigProto(
-        device_count={'CPU': 2})) as sess:
-      with ops.device('/CPU:0'):
-        in0 = constant_op.constant(1)
-      with ops.device('/CPU:1'):
-        in1 = constant_op.constant(3)
-
-      result_op = fn([in0, in1])
-
-      run_options = config_pb2.RunOptions()
-      run_options.experimental.collective_graph_key = 1
-      result = sess.run(result_op, options=run_options)
-
-      self.assertAllClose(result, [2, 2])
+    with ops.device('/CPU:0'):
+      in0 = constant_op.constant(1)
+    with ops.device('/CPU:1'):
+      in1 = constant_op.constant(3)
+    result = fn([in0, in1])
+    self.assertAllClose(result, [2, 2])
 
   @test_util.run_v2_only
   def testCollectiveGroupSizeOne(self):
@@ -548,7 +571,6 @@ class CollectiveOpTest(test.TestCase):
         in_tensor, group_size, group_key, instance_key)
     self.assertAllEqual(in_value, gathered_tensor.numpy())
 
-  @test_util.run_deprecated_v1
   def testConstantWithScopedAllocator(self):
     group_size = 2
     group_key = 1
@@ -565,21 +587,25 @@ class CollectiveOpTest(test.TestCase):
     del rewrite_options.scoped_allocator_opts.enable_op[:]
     rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce')
 
-    with self.session(config=cfg) as sess:
-      run_ops = []
-      for i in range(group_size):
-        with ops.device('CPU:%d' % i):
-          constant = constant_op.constant(i + 1.)
-          input_tensor1 = array_ops.identity(constant)
-          input_tensor2 = array_ops.identity(constant)
-          reduced_tensor1 = collective_ops.all_reduce(
-              input_tensor1, group_size, group_key, instance_key1, 'Add', 'Id')
-          reduced_tensor2 = collective_ops.all_reduce(
-              input_tensor2, group_size, group_key, instance_key2, 'Add', 'Id')
-          run_ops.append(array_ops.identity(reduced_tensor1))
-          run_ops.append(array_ops.identity(reduced_tensor2))
-      results = sess.run(run_ops)
-      self.assertEqual(results, [3., 3., 3., 3.])
+    # Tests that execute collectives need to be enclosed in graph or tf.function
+    with ops.Graph().as_default():
+      with self.session(config=cfg) as sess:
+        run_ops = []
+        for i in range(group_size):
+          with ops.device('CPU:%d' % i):
+            constant = constant_op.constant(i + 1.)
+            input_tensor1 = array_ops.identity(constant)
+            input_tensor2 = array_ops.identity(constant)
+            reduced_tensor1 = collective_ops.all_reduce(
+                input_tensor1, group_size, group_key, instance_key1, 'Add',
+                'Id')
+            reduced_tensor2 = collective_ops.all_reduce(
+                input_tensor2, group_size, group_key, instance_key2, 'Add',
+                'Id')
+            run_ops.append(array_ops.identity(reduced_tensor1))
+            run_ops.append(array_ops.identity(reduced_tensor2))
+        results = sess.run(run_ops)
+    self.assertEqual(results, [3., 3., 3., 3.])
 
   @test_util.run_v2_only
   def testMultipleGroups(self):

From 72d30dfb8bc58be931604f853bd161a11b7c9fcc Mon Sep 17 00:00:00 2001
From: Haoyu Zhang 
Date: Fri, 19 Jun 2020 13:19:39 -0700
Subject: [PATCH 0652/1390] Raise error type corresponding to status code
 instead of generic RuntimeError.

PiperOrigin-RevId: 317367352
Change-Id: I35378b88a33269ac225632ae848398b819c694a1
---
 tensorflow/python/eager/pywrap_tensor.cc | 32 ++++++++++++++++--------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/tensorflow/python/eager/pywrap_tensor.cc b/tensorflow/python/eager/pywrap_tensor.cc
index 031545531f1..0789eab6270 100644
--- a/tensorflow/python/eager/pywrap_tensor.cc
+++ b/tensorflow/python/eager/pywrap_tensor.cc
@@ -180,6 +180,15 @@ int ConvertDeviceName(PyObject* obj, const char** dst) {
   return 1;
 }
 
+void RaiseExceptionTypeFromTFStatus(TF_Status* status) {
+  TF_Code code = TF_GetCode(status);
+  PyObject* exception = tensorflow::PyExceptionRegistry::Lookup(code);
+  PyErr_SetObject(exception,
+                  pybind11::make_tuple(pybind11::none(), pybind11::none(),
+                                       TF_Message(status))
+                      .ptr());
+}
+
 }  // namespace
 
 namespace tensorflow {
@@ -305,13 +314,7 @@ TFE_TensorHandle* ConvertToEagerTensorUncached(TFE_Context* ctx,
                                                     device_name, status.get()));
     const TF_Code code = TF_GetCode(status.get());
     if (code != TF_OK) {
-      // Instead of raising a generic RuntimeError, raise an exception type
-      // based on the status error code.
-      PyObject* exception = PyExceptionRegistry::Lookup(code);
-      PyErr_SetObject(exception,
-                      pybind11::make_tuple(pybind11::none(), pybind11::none(),
-                                           TF_Message(status.get()))
-                          .ptr());
+      RaiseExceptionTypeFromTFStatus(status.get());
       return nullptr;
     }
   }
@@ -512,7 +515,9 @@ static PyObject* EagerTensor_datatype_enum(EagerTensor* self) {
 static PyObject* EagerTensor_shape_tuple(EagerTensor* self) {
   auto handle = self->handle;
   int n = TFE_TensorHandleNumDims(handle, &self->status);
-  if (MaybeRaiseExceptionFromTFStatus(&self->status, nullptr)) {
+  TF_Code code = TF_GetCode(&self->status);
+  if (code != TF_OK) {
+    RaiseExceptionTypeFromTFStatus(&self->status);
     // Cleanup self->status before returning.
     self->status.status = tensorflow::Status::OK();
     return nullptr;
@@ -522,13 +527,18 @@ static PyObject* EagerTensor_shape_tuple(EagerTensor* self) {
   for (int i = 0; i < n; ++i) {
     PyObject* dim =
         PyLong_FromLongLong(TFE_TensorHandleDim(handle, i, &self->status));
-    if (MaybeRaiseExceptionFromTFStatus(&self->status, nullptr) ||
-        dim == nullptr || PyTuple_SetItem(shape, i, dim) != 0) {
+    code = TF_GetCode(&self->status);
+    if (code != TF_OK || dim == nullptr ||
+        PyTuple_SetItem(shape, i, dim) != 0) {
+      if (code != TF_OK) {
+        RaiseExceptionTypeFromTFStatus(&self->status);
+      } else {
+        PyErr_SetString(PyExc_RuntimeError, "Error while creating shape");
+      }
       // Cleanup self->status before returning.
       self->status.status = tensorflow::Status::OK();
       Py_DECREF(shape);
       if (dim != nullptr) Py_DECREF(dim);
-      PyErr_SetString(PyExc_RuntimeError, "Error while creating shape");
       return nullptr;
     }
   }

From 4b3576c08171aee7edc80e929d2babdd899d67b7 Mon Sep 17 00:00:00 2001
From: Frank Chen 
Date: Fri, 19 Jun 2020 14:01:00 -0700
Subject: [PATCH 0653/1390] Change quick_exit() to exit() to fix Mac OS nightly
 build problems

PiperOrigin-RevId: 317375174
Change-Id: Ia0adc40f0d952cb478118fc07189ccd0f2d9a073
---
 tensorflow/core/tpu/kernels/tpu_compile_op_common.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
index 92d1fa1337e..79556cfa544 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compile_op_common.cc
@@ -353,7 +353,7 @@ Status TpuCompileOpKernelCommon::CompileTFFunctionToHlo(
     return;
   }
 
-  std::quick_exit(42);
+  std::exit(42);
 }
 
 /* static */ Status TpuCompileOpKernelCommon::GetDynamicShapes(

From 83fe1bad15e65e8db5e546d683d0ad591f19fad7 Mon Sep 17 00:00:00 2001
From: Andrew Audibert 
Date: Fri, 19 Jun 2020 14:02:55 -0700
Subject: [PATCH 0654/1390] [tf.data service] Add test that different workers
 use independent shuffle orders.

If shuffle seeds are unspecified, shuffle order should be non-deterministically chosen on each worker.

PiperOrigin-RevId: 317375549
Change-Id: I35e32cfbbfb8558451a079875b495708347a23bf
---
 .../kernel_tests/data_service_ops_test.py     | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tensorflow/python/data/kernel_tests/data_service_ops_test.py b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
index 796ab328980..488bf97f184 100644
--- a/tensorflow/python/data/kernel_tests/data_service_ops_test.py
+++ b/tensorflow/python/data/kernel_tests/data_service_ops_test.py
@@ -31,6 +31,7 @@ from tensorflow.python.eager import def_function
 from tensorflow.python.framework import combinations
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import tensor_array_ops
@@ -78,6 +79,28 @@ class DataServiceOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
     results = [elem.numpy() for elem in ds]
     self.assertEqual(list(range(num_elements)), results)
 
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDifferentShuffleOrders(self):
+    random_seed.set_random_seed(None)
+    num_elements = 100
+    master_address = self.create_cluster(2)
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = ds.shuffle(num_elements)
+    ds = _make_distributed_dataset(ds, master_address)
+    output = [elem.numpy() for elem in ds]
+
+    # The output will be two sequences of range(num_elements)
+    # non-deterministically interleaved together. If the orders of the elements
+    # were the same, first_order and second_order computed below will be equal.
+    first_order = {}
+    second_order = {}
+    for element in output:
+      if element in first_order:
+        second_order[element] = len(second_order)
+      else:
+        first_order[element] = len(first_order)
+    self.assertNotEqual(first_order, second_order)
+
   @combinations.generate(test_base.eager_only_combinations())
   def testMultipleEpochs(self):
     num_elements = 3

From 94bf57d06c546df46b2af36cfa11dc62dd35aebb Mon Sep 17 00:00:00 2001
From: Sanjoy Das 
Date: Fri, 19 Jun 2020 14:07:16 -0700
Subject: [PATCH 0655/1390] Do not try to compile trivially dead branches in
 the Case tf2xla lowering

This is important for the upcoming DeviceIndex op which can be used to select
one of many implementations depending on the device, and some of them may not be
compilable by tf2xla.

PiperOrigin-RevId: 317376420
Change-Id: I6428df6f4da238e5d2bc3618d51c579e34454945
---
 tensorflow/compiler/tests/BUILD               | 20 +++++
 tensorflow/compiler/tests/case_test.py        | 87 +++++++++++++++++++
 tensorflow/compiler/tf2xla/kernels/BUILD      |  1 +
 tensorflow/compiler/tf2xla/kernels/case_op.cc | 46 ++++++++--
 tensorflow/compiler/tf2xla/kernels/case_op.h  | 11 ++-
 5 files changed, 156 insertions(+), 9 deletions(-)
 create mode 100644 tensorflow/compiler/tests/case_test.py

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index 034ec82de10..42353451408 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1453,6 +1453,26 @@ tf_xla_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "case_test",
+    size = "small",
+    srcs = ["case_test.py"],
+    disabled_backends = ["cpu_ondemand"],
+    python_version = "PY3",
+    tags = [
+        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
+    ],
+    use_xla_device = False,  # Uses tf.function(experimental_compile=True)
+    deps = [
+        ":xla_test",
+        "//tensorflow/compiler/tf2xla/python:xla",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:platform_test",
+        "//tensorflow/python:training",
+    ],
+)
+
 tf_xla_py_test(
     name = "gather_test",
     size = "medium",
diff --git a/tensorflow/compiler/tests/case_test.py b/tensorflow/compiler/tests/case_test.py
new file mode 100644
index 00000000000..3b2dff537da
--- /dev/null
+++ b/tensorflow/compiler/tests/case_test.py
@@ -0,0 +1,87 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for while loops in XLA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import test
+
+
+class CaseTest(xla_test.XLATestCase):
+
+  def testCaseBasic(self):
+
+    @def_function.function(experimental_compile=True)
+    def switch_case_test(branch_index):
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        return array_ops.constant(31)
+
+      def f3():
+        return array_ops.constant(-1)
+
+      return control_flow_ops.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f3)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test(array_ops.constant(0)).numpy(), 17)
+      self.assertEqual(switch_case_test(array_ops.constant(1)).numpy(), 31)
+      self.assertEqual(switch_case_test(array_ops.constant(2)).numpy(), -1)
+      self.assertEqual(switch_case_test(array_ops.constant(3)).numpy(), -1)
+
+  def testBranchIsPruned(self):
+
+    @def_function.function(experimental_compile=True)
+    def switch_case_test():
+      branch_index = array_ops.constant(0)
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        # Some operations that XLA cannot compile.
+        image_ops.decode_image(io_ops.read_file('/tmp/bmp'))
+        return array_ops.constant(31)
+
+      # This tests that we do not try to compile all branches if the branch
+      # index in trivially constant.
+      return control_flow_ops.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f2)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test().numpy(), 17)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index bfdfe38305b..bdaeeafd295 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -316,6 +316,7 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/compiler/xla:literal",
         "//tensorflow/compiler/xla/client:xla_builder",
+        "//tensorflow/compiler/xla/client/lib:constants",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
diff --git a/tensorflow/compiler/tf2xla/kernels/case_op.cc b/tensorflow/compiler/tf2xla/kernels/case_op.cc
index 1b15c09f7e3..fbd54f1ef39 100644
--- a/tensorflow/compiler/tf2xla/kernels/case_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/case_op.cc
@@ -21,13 +21,14 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_context.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/lib/constants.h"
 #include "tensorflow/compiler/xla/client/xla_builder.h"
 #include "tensorflow/core/lib/core/errors.h"
 
 namespace tensorflow {
 
 XlaCaseOp::XlaCaseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("branches", &branches_));
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("branches", &unpruned_branches_));
   OP_REQUIRES_OK(ctx, ctx->GetAttr("Tin", &input_types_));
   OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_types_));
   if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) {
@@ -41,12 +42,29 @@ XlaCaseOp::XlaCaseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
   }
 }
 
+std::pair, xla::XlaOp>
+XlaCaseOp::GetPrunedBranchesAndIndex(XlaOpKernelContext* ctx) {
+  xla::Literal branch_index_literal;
+  bool branch_index_is_constant =
+      ctx->ConstantInput(0, &branch_index_literal).ok();
+
+  if (!branch_index_is_constant) {
+    return {unpruned_branches_, ctx->Input(0)};
+  }
+
+  int32 branch_index = branch_index_literal.Get({});
+  if (branch_index < 0 || branch_index >= unpruned_branches_.size()) {
+    branch_index = unpruned_branches_.size() - 1;
+  }
+
+  std::vector pruned_branch = {unpruned_branches_[branch_index]};
+  return {pruned_branch, xla::ZerosLike(ctx->Input(0))};
+}
+
 // TODO(b/35949885): There is duplication here with the handling of the
 // while_op/if_op. Refactor the common code out/rework.
 void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
-  xla::XlaBuilder* b = ctx->builder();
-  int num_branches = branches_.size();
-  OP_REQUIRES(ctx, num_branches >= 1,
+  OP_REQUIRES(ctx, !unpruned_branches_.empty(),
               errors::InvalidArgument("Must provide at least one case branch"));
   OP_REQUIRES(ctx, input_type(0) == DT_INT32,
               errors::InvalidArgument(
@@ -55,6 +73,18 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
               errors::InvalidArgument(
                   "branch_index argument must be scalar for XLA compilation"));
 
+  xla::XlaBuilder* b = ctx->builder();
+
+  // We opportunistically prune out branches if the branch index is a
+  // compile-time constant.  This is important in the context of the DeviceIndex
+  // ops (and other such ops that may come later) since we may have a Case with
+  // trivially unselected branches that cannot be compiled into HLO.
+  std::vector branches;
+  xla::XlaOp branch_index;
+  std::tie(branches, branch_index) = GetPrunedBranchesAndIndex(ctx);
+
+  int num_branches = branches.size();
+
   VLOG(1) << "Building Case: " << input_types_.size() << " inputs";
 
   std::vector arguments(input_types_.size());
@@ -94,7 +124,7 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
     std::vector case_bodies(num_branches);
     for (int branch_idx = 0; branch_idx < num_branches; branch_idx++) {
       OP_REQUIRES_OK(ctx, FindMustBeConstNodes(
-                              ctx, branches_[branch_idx],
+                              ctx, branches[branch_idx],
                               &case_branch_must_be_const_nodes[branch_idx],
                               &case_bodies[branch_idx]));
     }
@@ -133,7 +163,7 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
   std::vector branch_results_p(num_branches);
   for (int j = 0; j < num_branches; ++j) {
     OP_REQUIRES_OK(ctx,
-                   compiler->CompileFunction(options, branches_[j], arguments,
+                   compiler->CompileFunction(options, branches[j], arguments,
                                              &branch_results[j]));
     branch_results_p[j] = &branch_results[j];
   }
@@ -171,7 +201,7 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
     for (int j = 0; j < num_branches; ++j) {
       branch_results[j] = {};
       OP_REQUIRES_OK(ctx,
-                     compiler->CompileFunction(options, branches_[j], arguments,
+                     compiler->CompileFunction(options, branches[j], arguments,
                                                &branch_results[j]));
     }
   }
@@ -277,7 +307,7 @@ void XlaCaseOp::Compile(XlaOpKernelContext* ctx) {
   auto input_tuple = xla::Tuple(b, inputs);
 
   xla::XlaOp outputs =
-      xla::Conditional(ctx->Input(0), absl::MakeSpan(result_computations),
+      xla::Conditional(branch_index, absl::MakeSpan(result_computations),
                        std::vector(num_branches, input_tuple));
   // Sets non-variable outputs.
   for (int i = 0; i < output_types_.size(); ++i) {
diff --git a/tensorflow/compiler/tf2xla/kernels/case_op.h b/tensorflow/compiler/tf2xla/kernels/case_op.h
index 4a61707864e..4d22a3db830 100644
--- a/tensorflow/compiler/tf2xla/kernels/case_op.h
+++ b/tensorflow/compiler/tf2xla/kernels/case_op.h
@@ -50,7 +50,16 @@ class XlaCaseOp : public XlaOpKernel {
  private:
   TF_DISALLOW_COPY_AND_ASSIGN(XlaCaseOp);
 
-  std::vector branches_;
+  // If the branch_index input is a constant: prunes out all but the branch
+  // corrresponding to that constant branch index, and returns that branch and
+  // the literal 0 (as the first and second component of the pair).
+  //
+  // If the branch_index input is not a constant: returns unpruned_branches_ and
+  // the branch_index input.
+  std::pair, xla::XlaOp> GetPrunedBranchesAndIndex(
+      XlaOpKernelContext* ctx);
+
+  std::vector unpruned_branches_;
   DataTypeVector input_types_;
   DataTypeVector output_types_;
   bool has_token_input_output_;

From f60f6f0c1f68d729b2e501d5a0a668466acb7cda Mon Sep 17 00:00:00 2001
From: Michael Banfield 
Date: Fri, 19 Jun 2020 14:11:21 -0700
Subject: [PATCH 0656/1390] Support uploading only new data when a file in GCS
 Filesystem is Flush()'d multiple times.

This uses the GCS compose API to avoid reuploading the entire file.

Also add some Vmodules for profiling GCS write paths.

PiperOrigin-RevId: 317377298
Change-Id: I3f36fd684c44070331ba1d9e6689efd0f74bfc0e
---
 .../core/platform/cloud/gcs_file_system.cc    |  150 +-
 .../core/platform/cloud/gcs_file_system.h     |    6 +-
 .../platform/cloud/gcs_file_system_test.cc    | 1620 ++++++++++-------
 3 files changed, 1065 insertions(+), 711 deletions(-)

diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 5d395f3d821..1bd4d86eef6 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -14,14 +14,18 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/platform/cloud/gcs_file_system.h"
+
 #include 
 #include 
+
 #include 
 #include 
 #include 
 #include 
 #include 
 #include 
+
+#include "tensorflow/core/platform/strcat.h"
 #ifdef _WIN32
 #include   // for _mktemp
 #endif
@@ -128,6 +132,15 @@ constexpr char kAllowedBucketLocations[] = "GCS_ALLOWED_BUCKET_LOCATIONS";
 // is running in and restricts to buckets in that region.
 constexpr char kDetectZoneSentinelValue[] = "auto";
 
+// How to upload new data when Flush() is called multiple times.
+// By default the entire file is reuploaded.
+constexpr char kAppendMode[] = "GCS_APPEND_MODE";
+// If GCS_APPEND_MODE=compose then instead the new data is uploaded to a
+// temporary object and composed with the original object. This is disabled by
+// default as the multiple API calls required add a risk of stranding temporary
+// objects.
+constexpr char kComposeAppend[] = "compose";
+
 Status GetTmpFilename(string* filename) {
   *filename = io::GetTempFilename("");
   return Status::OK();
@@ -379,15 +392,18 @@ class GcsWritableFile : public WritableFile {
                   GcsFileSystem* filesystem,
                   GcsFileSystem::TimeoutConfig* timeouts,
                   std::function file_cache_erase,
-                  RetryConfig retry_config)
+                  RetryConfig retry_config, bool compose_append)
       : bucket_(bucket),
         object_(object),
         filesystem_(filesystem),
         timeouts_(timeouts),
         file_cache_erase_(std::move(file_cache_erase)),
         sync_needed_(true),
-        retry_config_(retry_config) {
+        retry_config_(retry_config),
+        compose_append_(compose_append),
+        start_offset_(0) {
     // TODO: to make it safer, outfile_ should be constructed from an FD
+    VLOG(3) << "GcsWritableFile: " << GetGcsPath();
     if (GetTmpFilename(&tmp_content_filename_).ok()) {
       outfile_.open(tmp_content_filename_,
                     std::ofstream::binary | std::ofstream::app);
@@ -403,14 +419,18 @@ class GcsWritableFile : public WritableFile {
                   GcsFileSystem* filesystem, const string& tmp_content_filename,
                   GcsFileSystem::TimeoutConfig* timeouts,
                   std::function file_cache_erase,
-                  RetryConfig retry_config)
+                  RetryConfig retry_config, bool compose_append)
       : bucket_(bucket),
         object_(object),
         filesystem_(filesystem),
         timeouts_(timeouts),
         file_cache_erase_(std::move(file_cache_erase)),
         sync_needed_(true),
-        retry_config_(retry_config) {
+        retry_config_(retry_config),
+        compose_append_(compose_append),
+        start_offset_(0) {
+    VLOG(3) << "GcsWritableFile: " << GetGcsPath() << "with existing file "
+            << tmp_content_filename;
     tmp_content_filename_ = tmp_content_filename;
     outfile_.open(tmp_content_filename_,
                   std::ofstream::binary | std::ofstream::app);
@@ -423,6 +443,7 @@ class GcsWritableFile : public WritableFile {
 
   Status Append(StringPiece data) override {
     TF_RETURN_IF_ERROR(CheckWritable());
+    VLOG(3) << "Append: " << GetGcsPath() << " size " << data.length();
     sync_needed_ = true;
     outfile_ << data;
     if (!outfile_.good()) {
@@ -433,6 +454,7 @@ class GcsWritableFile : public WritableFile {
   }
 
   Status Close() override {
+    VLOG(3) << "Close:" << GetGcsPath();
     if (outfile_.is_open()) {
       Status sync_status = Sync();
       if (sync_status.ok()) {
@@ -443,18 +465,23 @@ class GcsWritableFile : public WritableFile {
     return Status::OK();
   }
 
-  Status Flush() override { return Sync(); }
+  Status Flush() override {
+    VLOG(3) << "Flush:" << GetGcsPath();
+    return Sync();
+  }
 
   Status Name(StringPiece* result) const override {
     return errors::Unimplemented("GCSWritableFile does not support Name()");
   }
 
   Status Sync() override {
+    VLOG(3) << "Sync started:" << GetGcsPath();
     TF_RETURN_IF_ERROR(CheckWritable());
     if (!sync_needed_) {
       return Status::OK();
     }
     Status status = SyncImpl();
+    VLOG(3) << "Sync finished " << GetGcsPath();
     if (status.ok()) {
       sync_needed_ = false;
     }
@@ -483,11 +510,26 @@ class GcsWritableFile : public WritableFile {
           "Could not write to the internal temporary file.");
     }
     string session_uri;
-    TF_RETURN_IF_ERROR(CreateNewUploadSession(&session_uri));
+    uint64 start_offset = 0;
+    string object_to_upload = object_;
+    bool should_compose = false;
+    if (compose_append_) {
+      start_offset = start_offset_;
+      // Only compose if the object has already been uploaded to GCS
+      should_compose = start_offset > 0;
+      if (should_compose) {
+        object_to_upload =
+            strings::StrCat(io::Dirname(object_), "/.tmpcompose/",
+                            io::Basename(object_), ".", start_offset_);
+      }
+    }
+    TF_RETURN_IF_ERROR(
+        CreateNewUploadSession(&session_uri, start_offset, object_to_upload));
     uint64 already_uploaded = 0;
     bool first_attempt = true;
     const Status upload_status = RetryingUtils::CallWithRetries(
-        [&first_attempt, &already_uploaded, &session_uri, this]() {
+        [&first_attempt, &already_uploaded, &session_uri, &start_offset,
+         this]() {
           if (!first_attempt) {
             bool completed;
             TF_RETURN_IF_ERROR(RequestUploadSessionStatus(
@@ -502,7 +544,7 @@ class GcsWritableFile : public WritableFile {
             }
           }
           first_attempt = false;
-          return UploadToSession(session_uri, already_uploaded);
+          return UploadToSession(session_uri, start_offset, already_uploaded);
         },
         retry_config_);
     if (upload_status.code() == errors::Code::NOT_FOUND) {
@@ -512,6 +554,12 @@ class GcsWritableFile : public WritableFile {
           strings::StrCat("Upload to gs://", bucket_, "/", object_,
                           " failed, caused by: ", upload_status.ToString()));
     }
+    if (upload_status.ok()) {
+      if (should_compose) {
+        TF_RETURN_IF_ERROR(AppendObject(object_to_upload));
+      }
+      TF_RETURN_IF_ERROR(GetCurrentFileSize(&start_offset_));
+    }
     return upload_status;
   }
 
@@ -534,7 +582,8 @@ class GcsWritableFile : public WritableFile {
   }
 
   /// Initiates a new resumable upload session.
-  Status CreateNewUploadSession(string* session_uri) {
+  Status CreateNewUploadSession(string* session_uri, uint64 start_offset,
+                                string object_to_upload) {
     uint64 file_size;
     TF_RETURN_IF_ERROR(GetCurrentFileSize(&file_size));
 
@@ -542,10 +591,11 @@ class GcsWritableFile : public WritableFile {
     std::unique_ptr request;
     TF_RETURN_IF_ERROR(filesystem_->CreateHttpRequest(&request));
 
-    request->SetUri(strings::StrCat(
-        kGcsUploadUriBase, "b/", bucket_,
-        "/o?uploadType=resumable&name=", request->EscapeString(object_)));
-    request->AddHeader("X-Upload-Content-Length", std::to_string(file_size));
+    request->SetUri(strings::StrCat(kGcsUploadUriBase, "b/", bucket_,
+                                    "/o?uploadType=resumable&name=",
+                                    request->EscapeString(object_to_upload)));
+    request->AddHeader("X-Upload-Content-Length",
+                       std::to_string(file_size - start_offset));
     request->SetPostEmptyBody();
     request->SetResultBuffer(&output_buffer);
     request->SetTimeouts(timeouts_->connect, timeouts_->idle,
@@ -561,6 +611,37 @@ class GcsWritableFile : public WritableFile {
     return Status::OK();
   }
 
+  /// Appends the data of append_object to the original object and deletes
+  /// append_object.
+  Status AppendObject(string append_object) {
+    VLOG(3) << "AppendObject: " << GetGcsPathWithObject(append_object) << " to "
+            << GetGcsPath();
+    std::unique_ptr request;
+    TF_RETURN_IF_ERROR(filesystem_->CreateHttpRequest(&request));
+
+    request->SetUri(strings::StrCat(kGcsUriBase, "b/", bucket_, "/o/",
+                                    request->EscapeString(object_),
+                                    "/compose"));
+
+    const string request_body =
+        strings::StrCat("{'sourceObjects': [{'name': '", object_,
+                        "'},{'name': '", append_object, "'}]}");
+    request->SetTimeouts(timeouts_->connect, timeouts_->idle,
+                         timeouts_->metadata);
+    request->AddHeader("content-type", "application/json");
+    request->SetPostFromBuffer(request_body.c_str(), request_body.size());
+    return RetryingUtils::CallWithRetries(
+        [&request, &append_object, this]() {
+          TF_RETURN_WITH_CONTEXT_IF_ERROR(request->Send(),
+                                          " when composing to ", GetGcsPath());
+          TF_RETURN_WITH_CONTEXT_IF_ERROR(
+              filesystem_->DeleteFile(GetGcsPathWithObject(append_object)),
+              " when cleaning up.");
+          return Status::OK();
+        },
+        retry_config_);
+  }
+
   /// \brief Requests status of a previously initiated upload session.
   ///
   /// If the upload has already succeeded, sets 'completed' to true.
@@ -628,7 +709,8 @@ class GcsWritableFile : public WritableFile {
     return Status::OK();
   }
 
-  Status UploadToSession(const string& session_uri, uint64 start_offset) {
+  Status UploadToSession(const string& session_uri, uint64 start_offset,
+                         uint64 already_uploaded) {
     uint64 file_size;
     TF_RETURN_IF_ERROR(GetCurrentFileSize(&file_size));
 
@@ -637,13 +719,14 @@ class GcsWritableFile : public WritableFile {
     request->SetUri(session_uri);
     if (file_size > 0) {
       request->AddHeader("Content-Range",
-                         strings::StrCat("bytes ", start_offset, "-",
-                                         file_size - 1, "/", file_size));
+                         strings::StrCat("bytes ", already_uploaded, "-",
+                                         file_size - start_offset - 1, "/",
+                                         file_size - start_offset));
     }
     request->SetTimeouts(timeouts_->connect, timeouts_->idle, timeouts_->write);
 
-    TF_RETURN_IF_ERROR(
-        request->SetPutFromFile(tmp_content_filename_, start_offset));
+    TF_RETURN_IF_ERROR(request->SetPutFromFile(
+        tmp_content_filename_, start_offset + already_uploaded));
     TF_RETURN_WITH_CONTEXT_IF_ERROR(request->Send(), " when uploading ",
                                     GetGcsPath());
     // Erase the file from the file cache on every successful write.
@@ -651,9 +734,10 @@ class GcsWritableFile : public WritableFile {
     return Status::OK();
   }
 
-  string GetGcsPath() const {
-    return strings::StrCat("gs://", bucket_, "/", object_);
+  string GetGcsPathWithObject(string object) const {
+    return strings::StrCat("gs://", bucket_, "/", object);
   }
+  string GetGcsPath() const { return GetGcsPathWithObject(object_); }
 
   string bucket_;
   string object_;
@@ -664,6 +748,8 @@ class GcsWritableFile : public WritableFile {
   std::function file_cache_erase_;
   bool sync_needed_;  // whether there is buffered data that needs to be synced
   RetryConfig retry_config_;
+  bool compose_append_;
+  uint64 start_offset_;
 };
 
 class GcsReadOnlyMemoryRegion : public ReadOnlyMemoryRegion {
@@ -849,6 +935,14 @@ GcsFileSystem::GcsFileSystem(bool make_default_cache) {
 
   GetEnvVar(kAllowedBucketLocations, SplitByCommaToLowercaseSet,
             &allowed_locations_);
+
+  StringPiece append_mode;
+  GetEnvVar(kAppendMode, StringPieceIdentity, &append_mode);
+  if (append_mode == kComposeAppend) {
+    compose_append_ = true;
+  } else {
+    compose_append_ = false;
+  }
 }
 
 GcsFileSystem::GcsFileSystem(
@@ -859,7 +953,8 @@ GcsFileSystem::GcsFileSystem(
     size_t stat_cache_max_entries, uint64 matching_paths_cache_max_age,
     size_t matching_paths_cache_max_entries, RetryConfig retry_config,
     TimeoutConfig timeouts, const std::unordered_set& allowed_locations,
-    std::pair* additional_header)
+    std::pair* additional_header,
+    bool compose_append)
     : auth_provider_(std::move(auth_provider)),
       http_request_factory_(std::move(http_request_factory)),
       zone_provider_(std::move(zone_provider)),
@@ -872,6 +967,7 @@ GcsFileSystem::GcsFileSystem(
       bucket_location_cache_(new BucketLocationCache(
           kCacheNeverExpire, kBucketLocationCacheMaxEntries)),
       allowed_locations_(allowed_locations),
+      compose_append_(compose_append),
       timeouts_(timeouts),
       retry_config_(retry_config),
       additional_header_(additional_header) {}
@@ -1056,9 +1152,10 @@ Status GcsFileSystem::NewWritableFile(const string& fname,
                                       std::unique_ptr* result) {
   string bucket, object;
   TF_RETURN_IF_ERROR(ParseGcsPath(fname, false, &bucket, &object));
-  result->reset(new GcsWritableFile(bucket, object, this, &timeouts_,
-                                    [this, fname]() { ClearFileCaches(fname); },
-                                    retry_config_));
+  result->reset(new GcsWritableFile(
+      bucket, object, this, &timeouts_,
+      [this, fname]() { ClearFileCaches(fname); }, retry_config_,
+      compose_append_));
   return Status::OK();
 }
 
@@ -1098,7 +1195,8 @@ Status GcsFileSystem::NewAppendableFile(const string& fname,
   TF_RETURN_IF_ERROR(ParseGcsPath(fname, false, &bucket, &object));
   result->reset(new GcsWritableFile(
       bucket, object, this, old_content_filename, &timeouts_,
-      [this, fname]() { ClearFileCaches(fname); }, retry_config_));
+      [this, fname]() { ClearFileCaches(fname); }, retry_config_,
+      compose_append_));
   return Status::OK();
 }
 
@@ -1629,6 +1727,7 @@ Status GcsFileSystem::RenameFile(const string& src, const string& target) {
 
 // Uses a GCS API command to copy the object and then deletes the old one.
 Status GcsFileSystem::RenameObject(const string& src, const string& target) {
+  VLOG(3) << "RenameObject: started gs://" << src << " to " << target;
   string src_bucket, src_object, target_bucket, target_object;
   TF_RETURN_IF_ERROR(ParseGcsPath(src, false, &src_bucket, &src_object));
   TF_RETURN_IF_ERROR(
@@ -1664,6 +1763,7 @@ Status GcsFileSystem::RenameObject(const string& src, const string& target) {
         "locations or storage classes is not supported.");
   }
 
+  VLOG(3) << "RenameObject: finished from: gs://" << src << " to " << target;
   // In case the delete API call failed, but the deletion actually happened
   // on the server side, we can't just retry the whole RenameFile operation
   // because the source object is already gone.
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.h b/tensorflow/core/platform/cloud/gcs_file_system.h
index d1d8aed54d4..f066cc31eb4 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.h
+++ b/tensorflow/core/platform/cloud/gcs_file_system.h
@@ -122,7 +122,8 @@ class GcsFileSystem : public FileSystem {
                 size_t matching_paths_cache_max_entries,
                 RetryConfig retry_config, TimeoutConfig timeouts,
                 const std::unordered_set& allowed_locations,
-                std::pair* additional_header);
+                std::pair* additional_header,
+                bool compose_append);
 
   Status NewRandomAccessFile(
       const string& fname, std::unique_ptr* result) override;
@@ -187,6 +188,8 @@ class GcsFileSystem : public FileSystem {
   std::unordered_set allowed_locations() const {
     return allowed_locations_;
   }
+
+  bool compose_append() const { return compose_append_; }
   string additional_header_name() const {
     return additional_header_ ? additional_header_->first : "";
   }
@@ -373,6 +376,7 @@ class GcsFileSystem : public FileSystem {
   using BucketLocationCache = ExpiringLRUCache;
   std::unique_ptr bucket_location_cache_;
   std::unordered_set allowed_locations_;
+  bool compose_append_;
 
   TimeoutConfig timeouts_;
 
diff --git a/tensorflow/core/platform/cloud/gcs_file_system_test.cc b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
index 544ddc32043..6892bd7cc26 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system_test.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/platform/cloud/http_request_fake.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/platform/strcat.h"
 #include "tensorflow/core/platform/test.h"
 
 // Undef DeleteFile macro defined in wndows.h.
@@ -73,16 +74,16 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_NoBlockCache) {
            "Range: 6-11\n"
            "Timeouts: 5 1 20\n",
            "6789")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -129,7 +130,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -177,7 +178,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_Errors) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -224,7 +225,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_ReadAtEOF) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -265,7 +266,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_CachedOutOfRange) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -316,7 +317,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_CachedNotSequential) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -357,7 +358,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_Growing) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -404,7 +405,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_Buffered_ReadBackwards) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -437,16 +438,16 @@ TEST(GcsFileSystemTest,
             "location":"US-EAST1"
           })")});
 
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsAuto,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsAuto,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -479,16 +480,16 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithLocationConstraintCaching) {
             "location":"US-EAST1"
           })")});
 
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsAuto,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsAuto,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
 
@@ -520,16 +521,16 @@ TEST(GcsFileSystemTest,
             "location":"BARFOO"
           })")});
 
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsAuto,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsAuto,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   EXPECT_EQ(tensorflow::errors::FailedPrecondition(
@@ -552,16 +553,16 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_NoBlockCache_DifferentN) {
            "Range: 3-12\n"
            "Timeouts: 5 1 20\n",
            "3456789")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -621,7 +622,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   char scratch[100];
   StringPiece result;
@@ -710,7 +711,7 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache_Flush) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   char scratch[100];
   StringPiece result;
@@ -750,17 +751,17 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache_MaxStaleness) {
                            "Range: 8-15\n"
                            "Timeouts: 5 1 20\n",
                            "89abcdef")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   8 /* block size */, 16 /* max bytes */,
-                   3600 /* max staleness */, 3600 /* stat cache max age */,
-                   0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 8 /* block size */,
+      16 /* max bytes */, 3600 /* max staleness */,
+      3600 /* stat cache max age */, 0 /* stat cache max entries */,
+      0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
   char scratch[100];
   StringPiece result;
   // There should only be two HTTP requests issued to GCS even though we iterate
@@ -830,7 +831,7 @@ TEST(GcsFileSystemTest,
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewRandomAccessFile("gs://bucket/random_access.txt", &file));
@@ -849,17 +850,17 @@ TEST(GcsFileSystemTest,
 
 TEST(GcsFileSystemTest, NewRandomAccessFile_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* read ahead bytes */, 0 /* max bytes */,
-                   0 /* max staleness */, 0 /* stat cache max age */,
-                   0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider),
+      0 /* read ahead bytes */, 0 /* max bytes */, 0 /* max staleness */,
+      0 /* stat cache max age */, 0 /* stat cache max entries */,
+      0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -883,16 +884,16 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_InconsistentRead) {
            "012")});
 
   // Set stat_cache_max_age to 1000s so that StatCache could work.
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   1e3 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 1e3 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Stat the file first so that the file stats are cached.
   FileStatistics stat;
@@ -959,7 +960,7 @@ TEST(GcsFileSystemTest, NewWritableFile) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Read from the file first, to fill the block cache.
   std::unique_ptr rfile;
@@ -1042,16 +1043,16 @@ TEST(GcsFileSystemTest, NewWritableFile_ResumeUploadSucceeds) {
                            "Timeouts: 5 1 30\n"
                            "Put body: t2\n",
                            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
@@ -1112,17 +1113,17 @@ TEST(GcsFileSystemTest, NewWritableFile_ResumeUploadSucceedsOnGetStatus) {
            "Range: 0-7\n"
            "Timeouts: 5 1 20\n",
            "01234567")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   8 /* block size */, 8 /* max bytes */,
-                   3600 /* max staleness */, 3600 /* stat cache max age */,
-                   0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 8 /* block size */,
+      8 /* max bytes */, 3600 /* max staleness */,
+      3600 /* stat cache max age */, 0 /* stat cache max entries */,
+      0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
   // Pull the file's first block into the cache. This will trigger the first
   // HTTP request to GCS.
   std::unique_ptr rfile;
@@ -1208,7 +1209,8 @@ TEST(GcsFileSystemTest, NewWritableFile_ResumeUploadAllAttemptsFail) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */,
       RetryConfig(2 /* .init_delay_time_us */), kTestTimeoutConfig,
-      *kAllowedLocationsDefault, nullptr /* gcs additional header */);
+      *kAllowedLocationsDefault, nullptr /* gcs additional header */,
+      false /* compose append */);
 
   std::unique_ptr file;
   TF_EXPECT_OK(fs.NewWritableFile("gs://bucket/path/writeable.txt", &file));
@@ -1262,16 +1264,16 @@ TEST(GcsFileSystemTest, NewWritableFile_UploadReturns410) {
                            "Timeouts: 5 1 30\n"
                            "Put body: content1,content2\n",
                            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   {
     std::unique_ptr file;
@@ -1302,16 +1304,16 @@ TEST(GcsFileSystemTest, NewWritableFile_UploadReturns410) {
 
 TEST(GcsFileSystemTest, NewWritableFile_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1375,7 +1377,7 @@ TEST(GcsFileSystemTest, NewAppendableFile) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Create an appendable file. This should read the file from GCS, and pull its
   // contents into the block cache.
@@ -1401,16 +1403,16 @@ TEST(GcsFileSystemTest, NewAppendableFile) {
 
 TEST(GcsFileSystemTest, NewAppendableFile_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr file;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1435,16 +1437,16 @@ TEST(GcsFileSystemTest, NewReadOnlyMemoryRegionFromFile) {
                            "Range: 0-",
                            content.size() - 1, "\n", "Timeouts: 5 1 20\n"),
            content)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr region;
   TF_EXPECT_OK(fs.NewReadOnlyMemoryRegionFromFile(
@@ -1456,16 +1458,16 @@ TEST(GcsFileSystemTest, NewReadOnlyMemoryRegionFromFile) {
 
 TEST(GcsFileSystemTest, NewReadOnlyMemoryRegionFromFile_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr region;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1480,16 +1482,16 @@ TEST(GcsFileSystemTest, FileExists_YesAsObject) {
       "Timeouts: 5 1 10\n",
       strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                       "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.FileExists("gs://bucket/path/file1.txt"));
 }
@@ -1510,16 +1512,16 @@ TEST(GcsFileSystemTest, FileExists_YesAsFolder) {
            "Timeouts: 5 1 10\n",
            "{\"items\": [ "
            "  { \"name\": \"path/subfolder/\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.FileExists("gs://bucket/path/subfolder"));
 }
@@ -1536,16 +1538,16 @@ TEST(GcsFileSystemTest, FileExists_YesAsBucket) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "{\"size\": \"100\"}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.FileExists("gs://bucket1"));
   TF_EXPECT_OK(fs.FileExists("gs://bucket1/"));
@@ -1566,16 +1568,16 @@ TEST(GcsFileSystemTest, FileExists_NotAsObjectOrFolder) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "{\"items\": []}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(errors::Code::NOT_FOUND,
             fs.FileExists("gs://bucket/path/file1.txt").code());
@@ -1593,16 +1595,16 @@ TEST(GcsFileSystemTest, FileExists_NotAsBucket) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
             fs.FileExists("gs://bucket2/").code());
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -1641,7 +1643,7 @@ TEST(GcsFileSystemTest, FileExists_StatCache) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // The stat cache will ensure that repeated lookups don't trigger additional
   // HTTP requests.
@@ -1668,7 +1670,7 @@ TEST(GcsFileSystemTest, FileExists_DirectoryMark) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.FileExists("gs://bucket/dir/"));
   TF_EXPECT_OK(fs.IsDirectory("gs://bucket/dir/"));
@@ -1682,16 +1684,16 @@ TEST(GcsFileSystemTest, GetChildren_NoItems) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{\"prefixes\": [\"path/subpath/\"]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1710,16 +1712,16 @@ TEST(GcsFileSystemTest, GetChildren_ThreeFiles) {
       "  { \"name\": \"path/file1.txt\" },"
       "  { \"name\": \"path/file3.txt\" }],"
       "\"prefixes\": [\"path/subpath/\"]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1739,16 +1741,16 @@ TEST(GcsFileSystemTest, GetChildren_SelfDirectoryMarker) {
       "  { \"name\": \"path/\" },"
       "  { \"name\": \"path/file3.txt\" }],"
       "\"prefixes\": [\"path/subpath/\"]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1767,16 +1769,16 @@ TEST(GcsFileSystemTest, GetChildren_ThreeFiles_NoSlash) {
       "  { \"name\": \"path/file1.txt\" },"
       "  { \"name\": \"path/file3.txt\" }],"
       "\"prefixes\": [\"path/subpath/\"]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path", &children));
@@ -1792,16 +1794,16 @@ TEST(GcsFileSystemTest, GetChildren_Root) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket-a-b-c", &children));
@@ -1817,16 +1819,16 @@ TEST(GcsFileSystemTest, GetChildren_Empty) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path/", &children));
@@ -1858,16 +1860,16 @@ TEST(GcsFileSystemTest, GetChildren_Pagination) {
            "  { \"name\": \"path/file4.txt\" },"
            "  { \"name\": \"path/file5.txt\" }]}")});
 
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector children;
   TF_EXPECT_OK(fs.GetChildren("gs://bucket/path", &children));
@@ -1885,16 +1887,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_NoWildcard) {
       "Timeouts: 5 1 10\n",
       "{\"items\": [ "
       "  { \"name\": \"path/subpath/file2.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(
@@ -1913,16 +1915,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_BucketAndWildcard) {
       "  { \"name\": \"path/file1.txt\" },"
       "  { \"name\": \"path/subpath/file2.txt\" },"
       "  { \"name\": \"path/file3.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/*/*", &result));
@@ -1942,16 +1944,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_FolderAndWildcard_Matches) {
       "  { \"name\": \"path/file1.txt\" },"
       "  { \"name\": \"path/subpath/file2.txt\" },"
       "  { \"name\": \"path/file3.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*/file2.txt", &result));
@@ -1968,16 +1970,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_SelfDirectoryMarker) {
       "{\"items\": [ "
       "  { \"name\": \"path/\" },"
       "  { \"name\": \"path/file3.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*", &result));
@@ -1993,16 +1995,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_SlashInObjectName) {
       "{\"items\": [ "
       "  { \"name\": \"path/\" },"
       "  { \"name\": \"path//foo.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*", &result));
@@ -2018,16 +2020,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_SlashInObjectNameEscaped) {
       "{\"items\": [ "
       "  { \"name\": \"path/\" },"
       "  { \"name\": \"path//foo.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/\\/*", &result));
@@ -2044,16 +2046,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_FolderAndWildcard_NoMatches) {
       "  { \"name\": \"path/file1.txt\" },"
       "  { \"name\": \"path/subpath/file2.txt\" },"
       "  { \"name\": \"path/file3.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   TF_EXPECT_OK(fs.GetMatchingPaths("gs://bucket/path/*/file3.txt", &result));
@@ -2062,16 +2064,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_FolderAndWildcard_NoMatches) {
 
 TEST(GcsFileSystemTest, GetMatchingPaths_OnlyWildcard) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::vector result;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -2096,16 +2098,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_Cache) {
            "  { \"name\": \"path/file1.txt\" },"
            "  { \"name\": \"path/subpath/file2.txt\" },"
            "  { \"name\": \"path/file3.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   3600 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 3600 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Repeated calls to fs.GetMatchingPaths on these patterns should not lead to
   // any additional HTTP requests to GCS.
@@ -2139,16 +2141,16 @@ TEST(GcsFileSystemTest, GetMatchingPaths_Cache_Flush) {
            "Timeouts: 5 1 10\n",
            "{\"items\": [ "
            "  { \"name\": \"path/subpath/file2.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   3600 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 3600 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // This loop should trigger the first HTTP request to GCS.
   for (int i = 0; i < 10; i++) {
@@ -2212,7 +2214,7 @@ TEST(GcsFileSystemTest, DeleteFile) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Do an initial read of the file to load its contents into the block cache.
   char scratch[100];
@@ -2231,16 +2233,16 @@ TEST(GcsFileSystemTest, DeleteFile) {
 
 TEST(GcsFileSystemTest, DeleteFile_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
             fs.DeleteFile("gs://bucket/").code());
@@ -2283,7 +2285,7 @@ TEST(GcsFileSystemTest, DeleteFile_StatCacheRemoved) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Stats the file first so the stat is cached.
   FileStatistics stat_before_deletion;
@@ -2304,16 +2306,16 @@ TEST(GcsFileSystemTest, DeleteDir_Empty) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.DeleteDir("gs://bucket/path/"));
 }
@@ -2333,16 +2335,16 @@ TEST(GcsFileSystemTest, DeleteDir_OnlyDirMarkerLeft) {
                            "Timeouts: 5 1 10\n"
                            "Delete: yes\n",
                            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.DeleteDir("gs://bucket/path/"));
 }
@@ -2353,16 +2355,16 @@ TEST(GcsFileSystemTest, DeleteDir_BucketOnly) {
       "name%2CnextPageToken&maxResults=2\nAuth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.DeleteDir("gs://bucket"));
 }
@@ -2375,16 +2377,16 @@ TEST(GcsFileSystemTest, DeleteDir_NonEmpty) {
       "Timeouts: 5 1 10\n",
       "{\"items\": [ "
       "  { \"name\": \"path/file1.txt\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(error::Code::FAILED_PRECONDITION,
             fs.DeleteDir("gs://bucket/path/").code());
@@ -2398,16 +2400,16 @@ TEST(GcsFileSystemTest, GetFileSize) {
       "Timeouts: 5 1 10\n",
       strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                       "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   uint64 size;
   TF_EXPECT_OK(fs.GetFileSize("gs://bucket/file.txt", &size));
@@ -2416,16 +2418,16 @@ TEST(GcsFileSystemTest, GetFileSize) {
 
 TEST(GcsFileSystemTest, GetFileSize_NoObjectName) {
   std::vector requests;
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   uint64 size;
   EXPECT_EQ(errors::Code::INVALID_ARGUMENT,
@@ -2502,16 +2504,16 @@ TEST(GcsFileSystemTest, RenameFile_Folder) {
            "Timeouts: 5 1 10\n"
            "Delete: yes\n",
            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.RenameFile("gs://bucket/path1", "gs://bucket/path2/"));
 }
@@ -2603,7 +2605,7 @@ TEST(GcsFileSystemTest, RenameFile_Object) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
   // Do an initial read of the source and destination files to load their
   // contents into the block cache.
   char scratch[100];
@@ -2684,7 +2686,7 @@ TEST(GcsFileSystemTest, RenameFile_Object_FlushTargetStatCache) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
   // Do an initial stat of the destination file to load their contents into the
   // stat cache.
   FileStatistics stat_before_renaming;
@@ -2742,16 +2744,16 @@ TEST(GcsFileSystemTest, RenameFile_Object_DeletionRetried) {
            "Timeouts: 5 1 10\n"
            "Delete: yes\n",
            "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(
       fs.RenameFile("gs://bucket/path/src.txt", "gs://bucket/path/dst.txt"));
@@ -2784,16 +2786,16 @@ TEST(GcsFileSystemTest, RenameFile_Object_Incomplete) {
            "Post: yes\n"
            "Timeouts: 5 1 10\n",
            "{\"done\": false}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(
       errors::Code::UNIMPLEMENTED,
@@ -2809,16 +2811,16 @@ TEST(GcsFileSystemTest, Stat_Object) {
       "Timeouts: 5 1 10\n",
       strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                       "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   TF_EXPECT_OK(fs.Stat("gs://bucket/file.txt", &stat));
@@ -2843,16 +2845,16 @@ TEST(GcsFileSystemTest, Stat_Folder) {
            "Timeouts: 5 1 10\n",
            "{\"items\": [ "
            "  { \"name\": \"subfolder/\" }]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   TF_EXPECT_OK(fs.Stat("gs://bucket/subfolder", &stat));
@@ -2876,16 +2878,16 @@ TEST(GcsFileSystemTest, Stat_ObjectOrFolderNotFound) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   EXPECT_EQ(error::Code::NOT_FOUND, fs.Stat("gs://bucket/path", &stat).code());
@@ -2897,16 +2899,16 @@ TEST(GcsFileSystemTest, Stat_Bucket) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   TF_EXPECT_OK(fs.Stat("gs://bucket/", &stat));
@@ -2921,16 +2923,16 @@ TEST(GcsFileSystemTest, Stat_BucketNotFound) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   EXPECT_EQ(error::Code::NOT_FOUND, fs.Stat("gs://bucket/", &stat).code());
@@ -2968,7 +2970,7 @@ TEST(GcsFileSystemTest, Stat_Cache) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
 
   // Repeated calls to fs.Stat on these paths should not lead to any additional
   // HTTP requests to GCS.
@@ -3010,7 +3012,7 @@ TEST(GcsFileSystemTest, Stat_Cache_Flush) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      nullptr /* gcs additional header */);
+      nullptr /* gcs additional header */, false /* compose append */);
   // There should be a single HTTP request to GCS for fs.Stat in this loop.
   for (int i = 0; i < 10; i++) {
     FileStatistics stat;
@@ -3038,16 +3040,16 @@ TEST(GcsFileSystemTest, Stat_FilenameEndingWithSlash) {
       "Timeouts: 5 1 10\n",
       strings::StrCat("{\"size\": \"5\",\"generation\": \"1\","
                       "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   FileStatistics stat;
   TF_EXPECT_OK(fs.Stat("gs://bucket/dir/", &stat));
@@ -3070,16 +3072,16 @@ TEST(GcsFileSystemTest, IsDirectory_NotFound) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(error::Code::NOT_FOUND,
             fs.IsDirectory("gs://bucket/file.txt").code());
@@ -3101,16 +3103,16 @@ TEST(GcsFileSystemTest, IsDirectory_NotDirectoryButObject) {
            "Timeouts: 5 1 10\n",
            strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                            "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(error::Code::FAILED_PRECONDITION,
             fs.IsDirectory("gs://bucket/file.txt").code());
@@ -3132,16 +3134,16 @@ TEST(GcsFileSystemTest, IsDirectory_Yes) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "{\"items\": [{\"name\": \"subfolder/\"}]}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.IsDirectory("gs://bucket/subfolder"));
   TF_EXPECT_OK(fs.IsDirectory("gs://bucket/subfolder/"));
@@ -3159,16 +3161,16 @@ TEST(GcsFileSystemTest, IsDirectory_Bucket) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.IsDirectory("gs://bucket"));
   TF_EXPECT_OK(fs.IsDirectory("gs://bucket/"));
@@ -3180,16 +3182,16 @@ TEST(GcsFileSystemTest, IsDirectory_BucketNotFound) {
       "Auth Token: fake_token\n"
       "Timeouts: 5 1 10\n",
       "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   EXPECT_EQ(error::Code::NOT_FOUND, fs.IsDirectory("gs://bucket/").code());
 }
@@ -3222,16 +3224,16 @@ TEST(GcsFileSystemTest, CreateDir_Folder) {
            "Timeouts: 5 1 10\n",
            strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                            "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.CreateDir("gs://bucket/subpath"));
   EXPECT_EQ(errors::AlreadyExists("gs://bucket/subpath/"),
@@ -3250,16 +3252,16 @@ TEST(GcsFileSystemTest, CreateDir_Bucket) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TF_EXPECT_OK(fs.CreateDir("gs://bucket/"));
   TF_EXPECT_OK(fs.CreateDir("gs://bucket"));
@@ -3322,16 +3324,16 @@ TEST(GcsFileSystemTest, DeleteRecursively_Ok) {
                            "Timeouts: 5 1 10\n"
                            "Delete: yes\n",
                            "")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   int64 undeleted_files, undeleted_dirs;
   TF_EXPECT_OK(fs.DeleteRecursively("gs://bucket/path", &undeleted_files,
@@ -3415,16 +3417,16 @@ TEST(GcsFileSystemTest, DeleteRecursively_DeletionErrors) {
            "Timeouts: 5 1 10\n",
            "", errors::NotFound("404"), 404)});
 
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   int64 undeleted_files, undeleted_dirs;
   TF_EXPECT_OK(fs.DeleteRecursively("gs://bucket/path", &undeleted_files,
@@ -3450,16 +3452,16 @@ TEST(GcsFileSystemTest, DeleteRecursively_NotAFolder) {
            "Auth Token: fake_token\n"
            "Timeouts: 5 1 10\n",
            "", errors::NotFound("404"), 404)});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   int64 undeleted_files, undeleted_dirs;
   EXPECT_EQ(error::Code::NOT_FOUND,
@@ -3543,7 +3545,7 @@ TEST(GcsFileSystemTest, AdditionalRequestHeaderTest) {
       0 /* stat cache max entries */, 0 /* matching paths cache max age */,
       0 /* matching paths cache max entries */, kTestRetryConfig,
       kTestTimeoutConfig, *kAllowedLocationsDefault,
-      add_header /* gcs additional header */);
+      add_header /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr request;
   TF_EXPECT_OK(fs7.CreateHttpRequest(&request));
@@ -3613,16 +3615,16 @@ TEST(GcsFileSystemTest, CreateHttpRequest) {
                            "Auth Token: fake_token\n"
                            "Header Hello: world\n",
                            "{}")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   std::unique_ptr request;
   TF_EXPECT_OK(fs.CreateHttpRequest(&request));
@@ -3676,16 +3678,16 @@ TEST(GcsFileSystemTest, Stat_StatsRecording) {
       "Timeouts: 5 1 10\n",
       strings::StrCat("{\"size\": \"1010\",\"generation\": \"1\","
                       "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TestGcsStats stats;
   fs.SetStats(&stats);
@@ -3703,16 +3705,16 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_StatsRecording) {
       "Range: 0-5\n"
       "Timeouts: 5 1 20\n",
       "012345")});
-  GcsFileSystem fs(std::unique_ptr(new FakeAuthProvider),
-                   std::unique_ptr(
-                       new FakeHttpRequestFactory(&requests)),
-                   std::unique_ptr(new FakeZoneProvider),
-                   0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
-                   0 /* stat cache max age */, 0 /* stat cache max entries */,
-                   0 /* matching paths cache max age */,
-                   0 /* matching paths cache max entries */, kTestRetryConfig,
-                   kTestTimeoutConfig, *kAllowedLocationsDefault,
-                   nullptr /* gcs additional header */);
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 0 /* block size */,
+      0 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
 
   TestGcsStats stats;
   fs.SetStats(&stats);
@@ -3732,5 +3734,253 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_StatsRecording) {
   EXPECT_EQ(6, stats.block_retrieved_bytes_transferred_);
 }
 
+TEST(GcsFileSystemTest, NewAppendableFile_MultipleFlushesWithCompose) {
+  std::vector contents(
+      {"content0,", "content1,", "content2,", "content3,"});
+  std::vector requests({
+      // Fetch the file (stats and then content)
+      new FakeHttpRequest(
+          "Uri: "
+          "https://www.googleapis.com/storage/v1/b/bucket/o/"
+          "some%2Fpath%2Fappendable?fields=size%2Cgeneration%2Cupdated\n"
+          "Auth Token: fake_token\n"
+          "Timeouts: 5 1 10\n",
+          strings::StrCat("{\"size\": \"8\",\"generation\": \"1\","
+                          "\"updated\": \"2016-04-29T23:15:24.896Z\"}")),
+      new FakeHttpRequest(
+          "Uri: "
+          "https://storage.googleapis.com/bucket/some%2Fpath%2Fappendable\n"
+          "Auth Token: fake_token\n"
+          "Range: 0-1048575\n"
+          "Timeouts: 5 1 20\n",
+          contents[0]),
+      // Upload entire file
+      new FakeHttpRequest(
+          "Uri: https://www.googleapis.com/upload/storage/v1/b/bucket/o?"
+          "uploadType=resumable&name=some%2Fpath%2Fappendable\n"
+          "Auth Token: fake_token\n"
+          "Header X-Upload-Content-Length: 18\n"
+          "Post: yes\n"
+          "Timeouts: 5 1 10\n",
+          "", {{"Location", "https://custom/upload/location"}}),
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-17/18\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[0], contents[1], "\n"),
+          ""),
+      // Upload new part to a temporary object
+      new FakeHttpRequest(
+          "Uri: "
+          "https://www.googleapis.com/upload/storage/v1/b/bucket/"
+          "o?uploadType=resumable&name=some%2Fpath%2F.tmpcompose%2Fappendable."
+          "18\n"
+          "Auth Token: fake_token\n"
+          "Header X-Upload-Content-Length: 9\n"
+          "Post: yes\n"
+          "Timeouts: 5 1 10\n",
+          "",
+          {{"Location",
+            "https://custom/upload/"
+            "location"}}),
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-8/9\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[2], "\n"),
+          ""),
+      // Compose the new part at the end of the original object.
+      new FakeHttpRequest("Uri: "
+                          "https://www.googleapis.com/storage/v1/b/bucket/o/"
+                          "some%2Fpath%2Fappendable/compose\n"
+                          "Auth Token: fake_token\n"
+                          "Timeouts: 5 1 10\n"
+                          "Header content-type: application/json\n"
+                          "Post body: {'sourceObjects': [{'name': "
+                          "'some/path/appendable'},{'name': "
+                          "'some/path/.tmpcompose/appendable.18'}]}\n",
+                          ""),
+      // Delete the temporary object.
+      new FakeHttpRequest("Uri: "
+                          "https://www.googleapis.com/storage/v1/b/bucket/o/"
+                          "some%2Fpath%2F.tmpcompose%2Fappendable.18\n"
+                          "Auth Token: fake_token\n"
+                          "Timeouts: 5 1 10\n"
+                          "Delete: yes\n",
+                          ""),
+      new FakeHttpRequest(
+          "Uri: https://www.googleapis.com/upload/storage/v1/b/bucket/o?"
+          "uploadType=resumable&name=some%2Fpath%2F.tmpcompose%2Fappendable."
+          "27\n"
+          "Auth Token: fake_token\n"
+          "Header X-Upload-Content-Length: 9\n"
+          "Post: yes\n"
+          "Timeouts: 5 1 10\n",
+          "", {{"Location", "https://custom/upload/location"}}),
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-8/9\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[3], "\n"),
+          ""),
+      new FakeHttpRequest("Uri: "
+                          "https://www.googleapis.com/storage/v1/b/bucket/o/"
+                          "some%2Fpath%2Fappendable/compose\n"
+                          "Auth Token: fake_token\n"
+                          "Timeouts: 5 1 10\n"
+                          "Header content-type: application/json\n"
+                          "Post body: {'sourceObjects': [{'name': "
+                          "'some/path/appendable'},{'name': "
+                          "'some/path/.tmpcompose/appendable.27'}]}\n",
+                          ""),
+      new FakeHttpRequest("Uri: "
+                          "https://www.googleapis.com/storage/v1/b/bucket/o/"
+                          "some%2Fpath%2F.tmpcompose%2Fappendable."
+                          "27\n"
+                          "Auth Token: fake_token\n"
+                          "Timeouts: 5 1 10\n"
+                          "Delete: yes\n",
+                          ""),
+  });
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 32 /* block size */,
+      32 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, true /* compose append */);
+
+  // Create an appendable file. This should read the file from GCS, and pull its
+  // contents into the block cache.
+  std::unique_ptr wfile;
+  TF_EXPECT_OK(
+      fs.NewAppendableFile("gs://bucket/some/path/appendable", &wfile));
+  TF_EXPECT_OK(wfile->Append(contents[1]));
+  TF_EXPECT_OK(wfile->Flush());
+  TF_EXPECT_OK(wfile->Append(contents[2]));
+  TF_EXPECT_OK(wfile->Flush());
+  TF_EXPECT_OK(wfile->Append(contents[3]));
+  TF_EXPECT_OK(wfile->Close());
+}
+
+TEST(GcsFileSystemTest, NewAppendableFile_MultipleFlushesWithoutCompose) {
+  std::vector contents(
+      {"content0,", "content1,", "content2,", "content3,"});
+  std::vector requests({
+      new FakeHttpRequest(
+          "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+          "path%2Fappendable?fields=size%2Cgeneration%2Cupdated\n"
+          "Auth Token: fake_token\n"
+          "Timeouts: 5 1 10\n",
+          strings::StrCat("{\"size\": \"8\",\"generation\": \"1\","
+                          "\"updated\": \"2016-04-29T23:15:24.896Z\"}")),
+      new FakeHttpRequest(
+          "Uri: https://storage.googleapis.com/bucket/path%2Fappendable\n"
+          "Auth Token: fake_token\n"
+          "Range: 0-1048575\n"
+          "Timeouts: 5 1 20\n",
+          contents[0]),
+      new FakeHttpRequest(
+          "Uri: https://www.googleapis.com/upload/storage/v1/b/bucket/o?"
+          "uploadType=resumable&name=path%2Fappendable\n"
+          "Auth Token: fake_token\n"
+          "Header X-Upload-Content-Length: 18\n"
+          "Post: yes\n"
+          "Timeouts: 5 1 10\n",
+          "", {{"Location", "https://custom/upload/location"}}),
+      // Uploads entire file.
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-17/18\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[0], contents[1], "\n"),
+          ""),
+      new FakeHttpRequest("Uri: "
+                          "https://www.googleapis.com/upload/storage/v1/b/"
+                          "bucket/o?"
+                          "uploadType=resumable&name=path%2Fappendable\n"
+                          "Auth Token: fake_token\n"
+                          "Header X-Upload-Content-Length: 27\n"
+                          "Post: yes\n"
+                          "Timeouts: 5 1 10\n",
+                          "",
+                          {{"Location",
+                            "https://custom/upload/"
+                            "location"}}),
+      // Uploads entire file again.
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-26/27\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[0], contents[1], contents[2], "\n"),
+          ""),
+      new FakeHttpRequest(
+          "Uri: https://www.googleapis.com/upload/storage/v1/b/bucket/o?"
+          "uploadType=resumable&name=path%2Fappendable\n"
+          "Auth Token: fake_token\n"
+          "Header X-Upload-Content-Length: 36\n"
+          "Post: yes\n"
+          "Timeouts: 5 1 10\n",
+          "", {{"Location", "https://custom/upload/location"}}),
+      // Uploads entire file again.
+      new FakeHttpRequest(
+          strings::StrCat("Uri: https://custom/upload/location\n"
+                          "Auth Token: fake_token\n"
+                          "Header Content-Range: bytes 0-35/36\n"
+                          "Timeouts: 5 1 30\n"
+                          "Put body: ",
+                          contents[0], contents[1], contents[2], contents[3],
+                          "\n"),
+          ""),
+  });
+  GcsFileSystem fs(
+      std::unique_ptr(new FakeAuthProvider),
+      std::unique_ptr(
+          new FakeHttpRequestFactory(&requests)),
+      std::unique_ptr(new FakeZoneProvider), 32 /* block size */,
+      32 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */,
+      0 /* stat cache max entries */, 0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, kTestRetryConfig,
+      kTestTimeoutConfig, *kAllowedLocationsDefault,
+      nullptr /* gcs additional header */, false /* compose append */);
+
+  // Create an appendable file. This should read the file from GCS, and pull its
+  // contents into the block cache.
+  std::unique_ptr wfile;
+  TF_EXPECT_OK(fs.NewAppendableFile("gs://bucket/path/appendable", &wfile));
+  TF_EXPECT_OK(wfile->Append(contents[1]));
+  TF_EXPECT_OK(wfile->Flush());
+  TF_EXPECT_OK(wfile->Append(contents[2]));
+  TF_EXPECT_OK(wfile->Flush());
+  TF_EXPECT_OK(wfile->Append(contents[3]));
+  TF_EXPECT_OK(wfile->Close());
+}
+
+TEST(GcsFileSystemTest, AppendModeCompose) {
+  unsetenv("GCS_APPEND_MODE");
+  setenv("GCS_APPEND_MODE", "compose", 1);
+  GcsFileSystem fs1;
+  EXPECT_EQ(true, fs1.compose_append());
+}
+
+TEST(GcsFileSystemTest, AppendModeDefault) {
+  unsetenv("GCS_APPEND_MODE");
+  GcsFileSystem fs1;
+  EXPECT_EQ(false, fs1.compose_append());
+}
+
 }  // namespace
 }  // namespace tensorflow

From 030e65acd56b55084a5c8a9f7e5ed3db3dc63093 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Fri, 19 Jun 2020 14:17:36 -0700
Subject: [PATCH 0657/1390] Move GetTensorData calls directly to LstmStep call
 in the int8x8_8 version.

PiperOrigin-RevId: 317378494
Change-Id: I880de60ee5c963468b2eb47d8cca7a3b9765f57c
---
 tensorflow/lite/kernels/lstm_eval.cc | 102 ++++++++-------------------
 1 file changed, 31 insertions(+), 71 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index f38fdc95f3e..b2f3d77912b 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -2091,50 +2091,6 @@ TfLiteStatus EvalInteger8x8_8(
   const int n_cell = input_to_output_weights->dims->data[0];
   const int n_output = recurrent_to_output_weights->dims->data[1];
 
-  // Weights and states.
-  const int8_t* input_to_input_weight_ptr =
-      GetTensorData(input_to_input_weights);
-  const int8_t* recurrent_to_input_weight_ptr =
-      GetTensorData(recurrent_to_input_weights);
-  const int8_t* cell_to_input_weight_ptr =
-      GetTensorData(cell_to_input_weights);
-  const int8_t* input_to_forget_weight_ptr =
-      GetTensorData(input_to_forget_weights);
-  const int8_t* recurrent_to_forget_weight_ptr =
-      GetTensorData(recurrent_to_forget_weights);
-  const int8_t* cell_to_forget_weight_ptr =
-      GetTensorData(cell_to_forget_weights);
-  const int8_t* input_to_cell_weight_ptr =
-      GetTensorData(input_to_cell_weights);
-  const int8_t* recurrent_to_cell_weight_ptr =
-      GetTensorData(recurrent_to_cell_weights);
-  const int8_t* input_to_output_weight_ptr =
-      GetTensorData(input_to_output_weights);
-  const int8_t* recurrent_to_output_weight_ptr =
-      GetTensorData(recurrent_to_output_weights);
-  const int8_t* cell_to_output_weight_ptr =
-      GetTensorData(cell_to_output_weights);
-  const int8_t* projection_weight_ptr =
-      GetTensorData(projection_weights);
-  const int16_t* layer_norm_input_weight_ptr =
-      GetTensorData(input_layer_norm_coefficients);
-  const int16_t* layer_norm_forget_weight_ptr =
-      GetTensorData(forget_layer_norm_coefficients);
-  const int16_t* layer_norm_cell_weight_ptr =
-      GetTensorData(cell_layer_norm_coefficients);
-  const int16_t* layer_norm_output_weight_ptr =
-      GetTensorData(output_layer_norm_coefficients);
-  const int32_t* input_gate_bias_ptr = GetTensorData(input_gate_bias);
-  const int32_t* forget_gate_bias_ptr =
-      GetTensorData(forget_gate_bias);
-  const int32_t* cell_gate_bias_ptr = GetTensorData(cell_gate_bias);
-  const int32_t* output_gate_bias_ptr =
-      GetTensorData(output_gate_bias);
-  const int32_t* projection_bias_ptr = GetTensorData(projection_bias);
-  int16_t* cell_ptr = GetTensorData(cell_state);
-  int8_t* output_state_ptr = GetTensorData(output_state);
-  int8_t* output_ptr = nullptr;
-
   const int32_t input_zp = input->params.zero_point;
   const int32_t output_state_zp = output_state->params.zero_point;
 
@@ -2146,89 +2102,93 @@ TfLiteStatus EvalInteger8x8_8(
 
   for (int t = 0; t < max_time; t++) {
     const int t_rel = t;
-    output_ptr = output->data.int8 + t_rel * output_step;
-
+    int8_t* output_ptr = GetTensorData(output) + t_rel * output_step;
     // Input can be int8 asymmetric or int16 symmetric.
-    const int8_t* input_ptr = input->data.int8 + t_rel * input_step;
+    const int8_t* input_ptr = GetTensorData(input) + t_rel * input_step;
     lstm_eval::LstmStepInteger(
         input_ptr, input_zp,
 
-        input_to_input_weight_ptr,
+        GetTensorData(input_to_input_weights),
         integer_lstm_param->effective_input_to_input_scale_a,
         integer_lstm_param->effective_input_to_input_scale_b,
 
-        input_to_forget_weight_ptr,
+        GetTensorData(input_to_forget_weights),
         integer_lstm_param->effective_input_to_forget_scale_a,
         integer_lstm_param->effective_input_to_forget_scale_b,
 
-        input_to_cell_weight_ptr,
+        GetTensorData(input_to_cell_weights),
         integer_lstm_param->effective_input_to_cell_scale_a,
         integer_lstm_param->effective_input_to_cell_scale_b,
 
-        input_to_output_weight_ptr,
+        GetTensorData(input_to_output_weights),
         integer_lstm_param->effective_input_to_output_scale_a,
         integer_lstm_param->effective_input_to_output_scale_b,
 
-        recurrent_to_input_weight_ptr,
+        GetTensorData(recurrent_to_input_weights),
         integer_lstm_param->effective_recurrent_to_input_scale_a,
         integer_lstm_param->effective_recurrent_to_input_scale_b,
 
-        recurrent_to_forget_weight_ptr,
+        GetTensorData(recurrent_to_forget_weights),
         integer_lstm_param->effective_recurrent_to_forget_scale_a,
         integer_lstm_param->effective_recurrent_to_forget_scale_b,
 
-        recurrent_to_cell_weight_ptr,
+        GetTensorData(recurrent_to_cell_weights),
         integer_lstm_param->effective_recurrent_to_cell_scale_a,
         integer_lstm_param->effective_recurrent_to_cell_scale_b,
 
-        recurrent_to_output_weight_ptr,
+        GetTensorData(recurrent_to_output_weights),
         integer_lstm_param->effective_recurrent_to_output_scale_a,
         integer_lstm_param->effective_recurrent_to_output_scale_b,
 
-        cell_to_input_weight_ptr,
+        GetTensorData(cell_to_input_weights),
         integer_lstm_param->effective_cell_to_input_scale_a,
         integer_lstm_param->effective_cell_to_input_scale_b,
 
-        cell_to_forget_weight_ptr,
+        GetTensorData(cell_to_forget_weights),
         integer_lstm_param->effective_cell_to_forget_scale_a,
         integer_lstm_param->effective_cell_to_forget_scale_b,
 
-        cell_to_output_weight_ptr,
+        GetTensorData(cell_to_output_weights),
         integer_lstm_param->effective_cell_to_output_scale_a,
         integer_lstm_param->effective_cell_to_output_scale_b,
 
-        projection_weight_ptr, integer_lstm_param->effective_proj_scale_a,
+        GetTensorData(projection_weights),
+        integer_lstm_param->effective_proj_scale_a,
         integer_lstm_param->effective_proj_scale_b,
 
-        layer_norm_input_weight_ptr,
+        GetTensorData(input_layer_norm_coefficients),
         integer_lstm_param->layer_norm_input_scale_a,
         integer_lstm_param->layer_norm_input_scale_b,
 
-        layer_norm_forget_weight_ptr,
+        GetTensorData(forget_layer_norm_coefficients),
         integer_lstm_param->layer_norm_forget_scale_a,
         integer_lstm_param->layer_norm_forget_scale_b,
 
-        layer_norm_cell_weight_ptr, integer_lstm_param->layer_norm_cell_scale_a,
+        GetTensorData(cell_layer_norm_coefficients),
+        integer_lstm_param->layer_norm_cell_scale_a,
         integer_lstm_param->layer_norm_cell_scale_b,
 
-        layer_norm_output_weight_ptr,
+        GetTensorData(output_layer_norm_coefficients),
         integer_lstm_param->layer_norm_output_scale_a,
         integer_lstm_param->layer_norm_output_scale_b,
 
-        input_gate_bias_ptr, forget_gate_bias_ptr, cell_gate_bias_ptr,
-        output_gate_bias_ptr, projection_bias_ptr,
+        GetTensorData(input_gate_bias),
+        GetTensorData(forget_gate_bias),
+        GetTensorData(cell_gate_bias),
+        GetTensorData(output_gate_bias),
+        GetTensorData(projection_bias),
 
         params, integer_lstm_param->intermediate_scale_a,
         integer_lstm_param->intermediate_scale_b,
         integer_lstm_param->intermediate_zp,
         integer_lstm_param->quantized_cell_clip,
         integer_lstm_param->quantized_proj_clip, n_batch, n_cell, n_input,
-        n_output, output_batch_leading_dim, output_state_ptr, output_state_zp,
-        cell_ptr, output_ptr, GetTensorData(scratch0),
-        GetTensorData(scratch1), GetTensorData(scratch2),
-        GetTensorData(scratch3), GetTensorData(scratch4),
-        GetTensorData(scratch5), GetTensorData(scratch6),
-        GetTensorData(scratch7));
+        n_output, output_batch_leading_dim, GetTensorData(output_state),
+        output_state_zp, GetTensorData(cell_state), output_ptr,
+        GetTensorData(scratch0), GetTensorData(scratch1),
+        GetTensorData(scratch2), GetTensorData(scratch3),
+        GetTensorData(scratch4), GetTensorData(scratch5),
+        GetTensorData(scratch6), GetTensorData(scratch7));
   }
 
   return kTfLiteOk;

From 9e7df609bca2e363aa3217b70c6f95578ba07f10 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 14:21:31 -0700
Subject: [PATCH 0658/1390] Fix tokenization tests and update testing_utils to
 transfer state between layer creation.

PiperOrigin-RevId: 317379253
Change-Id: I786c2eb0506239de0e7f1a5f314a8f1b0bda10d4
---
 tensorflow/python/keras/testing_utils.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/keras/testing_utils.py b/tensorflow/python/keras/testing_utils.py
index 1928588fea1..cceaabe37a5 100644
--- a/tensorflow/python/keras/testing_utils.py
+++ b/tensorflow/python/keras/testing_utils.py
@@ -94,7 +94,8 @@ def layer_test(layer_cls,
                expected_output_shape=None,
                validate_training=True,
                adapt_data=None,
-               custom_objects=None):
+               custom_objects=None,
+               test_harness=None):
   """Test routine for a layer with a single input and single output.
 
   Arguments:
@@ -114,6 +115,8 @@ def layer_test(layer_cls,
       be tested for this layer. This is only relevant for PreprocessingLayers.
     custom_objects: Optional dictionary mapping name strings to custom objects
       in the layer class. This is helpful for testing custom layers.
+    test_harness: The Tensorflow test, if any, that this function is being
+      called in.
 
   Returns:
     The output data (Numpy array) returned by the layer, for additional
@@ -143,9 +146,15 @@ def layer_test(layer_cls,
     expected_output_dtype = input_dtype
 
   if dtypes.as_dtype(expected_output_dtype) == dtypes.string:
-    assert_equal = string_test
+    if test_harness:
+      assert_equal = test_harness.assertAllEqual
+    else:
+      assert_equal = string_test
   else:
-    assert_equal = numeric_test
+    if test_harness:
+      assert_equal = test_harness.assertAllClose
+    else:
+      assert_equal = numeric_test
 
   # instantiation
   kwargs = kwargs or {}
@@ -228,6 +237,7 @@ def layer_test(layer_cls,
   # test training mode (e.g. useful for dropout tests)
   # Rebuild the model to avoid the graph being reused between predict() and
   # See b/120160788 for more details. This should be mitigated after 2.0.
+  layer_weights = layer.get_weights()  # Get the layer weights BEFORE training.
   if validate_training:
     model = models.Model(x, layer(x))
     if _thread_local_data.run_eagerly is not None:
@@ -252,6 +262,8 @@ def layer_test(layer_cls,
   model = models.Sequential()
   model.add(layers.Input(shape=input_shape[1:], dtype=input_dtype))
   model.add(layer)
+
+  layer.set_weights(layer_weights)
   actual_output = model.predict(input_data)
   actual_output_shape = actual_output.shape
   for expected_dim, actual_dim in zip(computed_output_shape,

From e51d4027b323bef5d6042ee0b18d617193211ac8 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Fri, 19 Jun 2020 21:49:53 +0000
Subject: [PATCH 0659/1390] Add annotations to fallback ops

---
 tensorflow/python/framework/python_op_gen.cc | 56 ++++++++++----------
 1 file changed, 27 insertions(+), 29 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 062a9aa01e4..cdef72a155d 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -168,7 +168,8 @@ class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp {
   bool AddEagerFallbackCode(const string& parameters,
                             const std::vector& output_sizes,
                             const string& num_outputs_expr,
-                            const string& eager_not_allowed_error);
+                            const string& eager_not_allowed_error,
+                            std::unordered_map& type_annotations);
   void AddEagerFastPathExecute();
 
   void AddEagerInferredAttrs(const string& indentation);
@@ -355,18 +356,19 @@ string GenEagerPythonOp::Code() {
   }
 
   string parameters;
+  // Param can be an input or an attr
   for (const auto& param : params_no_default_) {
     if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
     strings::StrAppend(¶meters, param.GetRenameTo());
 
     // Add type annotations to param
     if (type_annotations.find(param.GetName()) != type_annotations.end()) {
-      if (!type_annotations[param.GetName()].empty()) {
-        strings::StrAppend(¶meters, ": ", type_annotations[param.GetName()]);
-      }
+      strings::StrAppend(¶meters, ": ", type_annotations[param.GetName()]);
     }
   }
 
+  // Append to parameters and parameters_with_defaults because multiple functions
+  // are generated (op and fallback op)
   string parameters_with_defaults = parameters;
   for (const auto& param_and_default : params_with_default_) {
     if (!parameters.empty()) strings::StrAppend(¶meters, ", ");
@@ -375,14 +377,12 @@ string GenEagerPythonOp::Code() {
 
     // Add type annotations to param_and_default
     if (type_annotations.find(param_and_default.first.GetName()) != type_annotations.end()) {
-      if (!type_annotations[param_and_default.first.GetName()].empty()) {
-        strings::StrAppend(¶meters, ": ", type_annotations[param_and_default.first.GetName()]);
-        strings::StrAppend(¶meters_with_defaults,
-                           param_and_default.first.GetRenameTo(), ": ",
-                           type_annotations[param_and_default.first.GetName()], " ",
-                           "= ", param_and_default.second);
-        continue;
-      }
+      const string param_type = type_annotations[param_and_default.first.GetName()];
+      strings::StrAppend(¶meters, param_and_default.first.GetRenameTo(), ": ", param_type);
+      strings::StrAppend(¶meters_with_defaults,
+                         param_and_default.first.GetRenameTo(), ": ",
+                         param_type, " = ", param_and_default.second);
+      continue;
     }
 
     strings::StrAppend(¶meters, param_and_default.first.GetRenameTo());
@@ -425,7 +425,7 @@ string GenEagerPythonOp::Code() {
   }
 
   if (!AddEagerFallbackCode(parameters, output_sizes, num_outputs_expr,
-                            eager_not_allowed_error)) {
+                            eager_not_allowed_error, type_annotations)) {
     return result_;
   }
 
@@ -449,35 +449,29 @@ std::unordered_map GenEagerPythonOp::GetTypeAnnotationMap() {
   for (const auto& arg : op_def_.input_arg()) {
     // Do not add type annotations to args that accept a sequence of Tensors
     if (!arg.number_attr().empty()) continue;
-    string type_annotation;
     if (type_annotations.find(arg.type_attr()) != type_annotations.end()) {
       // Get the correct TypeVar if input maps to an attr
-      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_annotations[arg.type_attr()], "]");
+      type_annotations[arg.name()] = "_ops.Tensor[" + type_annotations[arg.type_attr()] + "]";
     } else {
       // Get the dtype of the Tensor
       const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
       if (dtype_type.find(py_dtype) != dtype_type.end()) {
-        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtype_type[py_dtype], "]");
+        type_annotations[arg.name()] = "_ops.Tensor[" + dtype_type[py_dtype] + "]";
       }
     }
-
-    type_annotations[arg.name()] = type_annotation;
   }
 
   // Mapping output Tensor to its type
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
-    string type_annotation;
     if (type_annotations.find(arg.type_attr()) != type_annotations.end()) {
-      strings::StrAppend(&type_annotation, "_ops.Tensor[", type_annotations[arg.type_attr()], "]");
+      type_annotations[arg.name()] = "_ops.Tensor[" + type_annotations[arg.type_attr()] + "]";
     } else {
       const string py_dtype = python_op_gen_internal::DataTypeToPython(arg.type(), "_dtypes.");
       if (dtype_type.find(py_dtype) != dtype_type.end()) {
-        strings::StrAppend(&type_annotation, "_ops.Tensor[", dtype_type[py_dtype], "]");
+        type_annotations[arg.name()] = "_ops.Tensor[" + dtype_type[py_dtype] + "]";
       }
     }
-
-    type_annotations[arg.name()] = type_annotation;
   }
 
   return type_annotations;
@@ -521,19 +515,20 @@ void GenEagerPythonOp::GenerateTypeVars(std::unordered_map& type
   if (added_typevar) strings::StrAppend(&result_, "\n");
 }
 
+// TODO(rahulkamat): Modify AddDefLine() to add return type annotation
 void GenEagerPythonOp::AddReturnTypeAnnotation(std::unordered_map& type_annotations) {
   if (op_def_.output_arg_size() == 1) {
     const auto& arg = op_def_.output_arg(0);
     // Add type annotations to param
     if (type_annotations.find(arg.name()) != type_annotations.end()) {
-      if (!type_annotations[arg.name()].empty()) {
-        result_.erase(result_.length() - 2);
-        strings::StrAppend(&result_, " -> ", type_annotations[arg.name()], ":\n");
-      }
+      result_.erase(result_.length() - 2);
+      strings::StrAppend(&result_, " -> ", type_annotations[arg.name()], ":\n");
     }
   }
 }
 
+
+
 void GenEagerPythonOp::HandleGraphMode(
     const string& function_setup, const std::vector& output_sizes) {
   strings::StrAppend(&result_, "  # Add nodes to the TensorFlow graph.\n");
@@ -903,11 +898,14 @@ bool GenEagerPythonOp::AddEagerFastPathAndGraphCode(
 
 bool GenEagerPythonOp::AddEagerFallbackCode(
     const string& parameters, const std::vector& output_sizes,
-    const string& num_outputs_expr, const string& eager_not_allowed_error) {
+    const string& num_outputs_expr, const string& eager_not_allowed_error,
+    std::unordered_map& type_annotations) {
   AddDefLine(
       strings::StrCat(function_name_, kEagerFallbackSuffix),
       strings::StrCat(parameters, parameters.empty() ? "" : ", ", "ctx"));
-
+  if (type_annotate_ops.find(op_def_.name()) != type_annotate_ops.end()) {
+    AddReturnTypeAnnotation(type_annotations);
+  }
   if (!eager_not_allowed_error.empty()) {
     strings::StrAppend(&result_, "  ", eager_not_allowed_error);
     return true;

From f129485019052480c99442f774fb7e4a59ae5227 Mon Sep 17 00:00:00 2001
From: Bixia Zheng 
Date: Fri, 19 Jun 2020 14:59:47 -0700
Subject: [PATCH 0660/1390] [TF:TRT] Cosmetic fix.

Rewrite two lines of #if into #if GOOGLE_CUDA && GOOGLE_TENSORRT.

PiperOrigin-RevId: 317386436
Change-Id: Icc8ae27a17900b6f0a198d32c6d73345084eab50
---
 tensorflow/compiler/tf2tensorrt/common/utils.h       |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_graph.cc    |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_graph.h     |  6 ++----
 .../tf2tensorrt/convert/convert_graph_test.cc        |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_nodes.cc    |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_nodes.h     |  6 ++----
 .../tf2tensorrt/convert/convert_nodes_test.cc        |  6 ++----
 .../compiler/tf2tensorrt/convert/logger_registry.cc  |  6 ++----
 .../compiler/tf2tensorrt/convert/logger_registry.h   |  5 +++--
 .../tf2tensorrt/convert/trt_optimization_pass.cc     |  6 ++----
 .../tf2tensorrt/convert/trt_optimization_pass.h      |  6 ++----
 .../tf2tensorrt/kernels/get_calibration_data_op.cc   |  6 ++----
 .../compiler/tf2tensorrt/kernels/trt_engine_op.cc    |  6 ++----
 .../tf2tensorrt/kernels/trt_engine_op_test.cc        |  6 ++----
 .../tf2tensorrt/kernels/trt_engine_resource_ops.cc   |  6 ++----
 .../kernels/trt_engine_resource_ops_test.cc          |  6 ++----
 .../tf2tensorrt/ops/get_calibration_data_op.cc       |  6 ++----
 tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc |  6 ++----
 .../tf2tensorrt/ops/trt_engine_resource_ops.cc       |  6 ++----
 .../compiler/tf2tensorrt/plugin/plugin_cast.cu.cc    |  6 ++----
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc |  6 ++----
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h  |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/segment.cc   |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/segment.h    |  6 ++----
 .../compiler/tf2tensorrt/segment/segment_test.cc     |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/union_find.h |  6 ++----
 tensorflow/compiler/tf2tensorrt/tensorrt_test.cc     |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_allocator.cc      | 12 ++++--------
 .../compiler/tf2tensorrt/utils/trt_allocator.h       | 12 ++++--------
 .../compiler/tf2tensorrt/utils/trt_engine_utils.cc   |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_engine_utils.h    |  6 ++----
 .../tf2tensorrt/utils/trt_int8_calibrator.cc         |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_int8_calibrator.h |  6 ++----
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc  |  6 ++----
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.h   |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_lru_cache.cc      |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_lru_cache.h       |  6 ++----
 .../utils/trt_shape_optimization_profiles.h          |  6 ++----
 .../utils/trt_shape_optimization_profiles_test.cc    |  6 ++----
 39 files changed, 83 insertions(+), 162 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/common/utils.h b/tensorflow/compiler/tf2tensorrt/common/utils.h
index 9ab0145e1ec..b428733ecd4 100644
--- a/tensorflow/compiler/tf2tensorrt/common/utils.h
+++ b/tensorflow/compiler/tf2tensorrt/common/utils.h
@@ -16,8 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 #define TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/platform/logging.h"
 
@@ -29,7 +28,6 @@ namespace tensorrt {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
index 1c51d51f1c9..5429aaf3362 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
@@ -53,8 +53,7 @@ limitations under the License.
 #include "tensorflow/core/util/device_name_utils.h"
 #include "tensorflow/tools/graph_transforms/transform_utils.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 namespace tensorflow {
@@ -884,5 +883,4 @@ Status ConvertAfterShapes(const ConversionParams& params) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
index 53ab84a6fa9..d3897e864fa 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
@@ -24,8 +24,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -66,7 +65,6 @@ Status RegisterGraphToFunctionLibrary(const GraphDef& segment_graph_def,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
index a1f523d6bfa..54fb1d56441 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
@@ -34,8 +34,7 @@ limitations under the License.
 #include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -231,5 +230,4 @@ TEST_F(ConvertAfterShapesTest, DirectlyConnectedEngines) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
index 96cec556942..2ec616ba621 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
@@ -59,8 +59,7 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/strided_slice_op.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 #include "third_party/tensorrt/NvInferPlugin.h"
 
@@ -6258,5 +6257,4 @@ bool OutputEdgeValidator::operator()(const Edge* out_edge) const {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
index 7a1276c645c..a621735fad1 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
@@ -33,8 +33,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -694,7 +693,6 @@ BinaryOperationMap();
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
index c24b169f651..53ec9ee7ada 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
@@ -21,8 +21,7 @@ limitations under the License.
 #include 
 #include 
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include 
 #include 
@@ -6636,5 +6635,4 @@ TEST_F(OpConverterTest, ConvertPad) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
index 82e68cbb28d..07c9c2f1ea0 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
@@ -12,8 +12,7 @@ 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.
 ==============================================================================*/
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 
@@ -58,5 +57,4 @@ LoggerRegistry* GetLoggerRegistry() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
index 45b302742d0..2a265cf7caa 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
@@ -19,7 +19,8 @@ limitations under the License.
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -53,5 +54,5 @@ class RegisterLogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_LOGGER_REGISTRY_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
index 72f4fe5ef9b..1cf98d135cb 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
@@ -28,8 +28,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stacktrace.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
@@ -302,5 +301,4 @@ static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
index f79048bb5f6..e0aaa5500ab 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
@@ -23,8 +23,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -83,6 +82,5 @@ class TRTOptimizationPass : public grappler::CustomGraphOptimizer {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
index 3143b06817e..76fb40b9520 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
@@ -22,8 +22,7 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/lib/core/refcount.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -67,5 +66,4 @@ REGISTER_KERNEL_BUILDER(Name("GetCalibrationDataOp").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index 98d199ca9ab..1094555a622 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -48,8 +48,7 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -1009,5 +1008,4 @@ REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
index a06010de1c7..71193dc24cf 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
@@ -50,8 +50,7 @@ limitations under the License.
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/public/version.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -306,5 +305,4 @@ TYPED_TEST(TRTEngineOpTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
index 2c5821df6ac..3b6e7e91d3b 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
@@ -33,8 +33,7 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -251,5 +250,4 @@ REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
index 4a24160569d..6a073ee24d0 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
@@ -48,8 +48,7 @@ limitations under the License.
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -246,5 +245,4 @@ TEST_F(TRTEngineResourceOpsTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
index 573172b92e6..2af3164c3e2 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -34,5 +33,4 @@ Returns calibration data for the given resource name
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
index bd3c2b299a9..2527fe9b910 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -59,5 +58,4 @@ REGISTER_OP("TRTEngineOp")
     .Attr("static_engine: bool = true");
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
index 01911de66ec..3141092de03 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -46,5 +45,4 @@ REGISTER_OP("SerializeTRTResource")
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
index 4c0d8b0392a..141a7d1f462 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
@@ -17,8 +17,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #define EIGEN_USE_GPU  // For definition of Eigen::GpuDevice.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
@@ -234,5 +233,4 @@ REGISTER_TFTRT_PLUGIN(CastPluginCreator);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
index 563ce724f43..83d5f9b5965 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
@@ -17,8 +17,7 @@ limitations under the License.
 
 #include 
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -30,5 +29,4 @@ const char* kTfTrtPluginNamespace = "TF";
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
index bdb046e6c71..600ac6683da 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
@@ -20,8 +20,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -90,7 +89,6 @@ class TrtPluginRegistrar {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
index 32e30006f58..d9080b6f69a 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
@@ -35,8 +35,7 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/env_var.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -1062,5 +1061,4 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h
index 7295c8f0d9d..3f79983cfd2 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.h
@@ -25,8 +25,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -67,7 +66,6 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
index 2437481a9c4..f3bc5bfbee6 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
@@ -26,8 +26,7 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -522,5 +521,4 @@ TEST_F(SegmentTest, IncompatibleBatchSizes) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/segment/union_find.h b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
index 70e83c12fca..b53615ec019 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/union_find.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
@@ -19,8 +19,7 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/types/optional.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -217,7 +216,6 @@ UnionFind* UnionFind::FindRoot() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_
diff --git a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
index 510591bfe00..e994d20df33 100644
--- a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
@@ -18,8 +18,7 @@ limitations under the License.
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/test.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -164,5 +163,4 @@ TEST(TensorrtTest, BasicFunctions) {
 }  // namespace
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
index 617ea7fad5c..d4f3a524577 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
@@ -17,11 +17,9 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -52,8 +50,7 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -113,5 +110,4 @@ void TRTDeviceAllocator::free(void* memory) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
index 4ab8b52f523..d219a8a14e8 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
@@ -20,11 +20,9 @@ limitations under the License.
 
 #include "tensorflow/core/framework/allocator.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -33,8 +31,7 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -69,6 +66,5 @@ class TRTDeviceAllocator : public TRTBaseAllocator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
index ed997b267b1..8ccfb8b06f0 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
@@ -25,8 +25,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/errors.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -257,5 +256,4 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
index a471749877a..1ea4fe28cb4 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
@@ -24,8 +24,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/status.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -91,7 +90,6 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
index 554c127fa37..24271e352a7 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
@@ -20,8 +20,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 
 namespace tensorflow {
@@ -147,5 +146,4 @@ TRTInt8Calibrator::~TRTInt8Calibrator() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
index 06b39716490..4c670e85f52 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
@@ -22,8 +22,7 @@ limitations under the License.
 #include 
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -101,6 +100,5 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
index 193687ebc8c..e34bf5e7397 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
@@ -15,8 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 #include "tensorflow/core/platform/logging.h"
@@ -68,5 +67,4 @@ REGISTER_TENSORRT_LOGGER("DefaultLogger", Logger::GetLogger());
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
index 2ade1b48f47..ce6552e8fe9 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
@@ -18,8 +18,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -40,7 +39,6 @@ class Logger : public nvinfer1::ILogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
index fbcdaad52c0..ee7e6272372 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
@@ -23,8 +23,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -141,5 +140,4 @@ EngineContext* TRTEngineCacheResource::GetEngineContext(const int profile_id) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
index 8e345254f75..991b9a949e4 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
@@ -115,8 +115,7 @@ class LRUCache {
   }
 };
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 struct EngineContext {
   EngineContext() {}  // Creates an empty context.
@@ -223,8 +222,7 @@ class TRTEngineCacheResource : public ResourceBase {
   TrtShapeOptimizationProfile profiles_;
 };
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
index 40c7f5dcf31..fc688b14139 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
@@ -29,8 +29,7 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -173,6 +172,5 @@ class TrtShapeOptimizationProfile {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
index 501810587e0..32c2200fb71 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include 
 
@@ -214,5 +213,4 @@ TEST_F(TrtShapeOptimizationProfileTest, Dynamic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

From 32d63d0a3efb5e0b65dc6f590e54248054cf9f14 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Fri, 15 May 2020 11:40:22 -0500
Subject: [PATCH 0661/1390] Removed TENSOR_OP disable env vars.

* TF_DISABLE_CUBLAS_TENSOR_OP_MATH
* TF_DISABLE_CUDNN_TENSOR_OP_MATH
* TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH
---
 tensorflow/stream_executor/cuda/cuda_blas.cc | 21 ++------
 tensorflow/stream_executor/cuda/cuda_dnn.cc  | 55 +++++---------------
 2 files changed, 16 insertions(+), 60 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index c9f0fc462c9..65c07e72154 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -101,18 +101,6 @@ static std::string ToString(cublasStatus_t status) {
   }
 }
 
-// Decide whether to enable TENSOR_OP_MATH
-static bool TensorOpMathEnabled() {
-  static bool is_enabled = [] {
-    bool is_disabled;
-    TF_CHECK_OK(
-        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUBLAS_TENSOR_OP_MATH",
-                                       /*default_val=*/false, &is_disabled));
-    return !is_disabled;
-  }();
-  return is_enabled;
-}
-
 // cuBLAS has interfaces that permit pointers to be passed from either the host
 // memory space or the device memory space; however, you must instruct it as to
 // which address space those pointers are in with cublasSetPointerMode.
@@ -1640,7 +1628,7 @@ bool CUDABlas::DoBlasGemm(
                                                                    &cc_minor);
 
   // GPUs < sm_70 don't support tensor ops.
-  if (cc_major >= 7 && TensorOpMathEnabled()) {
+  if (cc_major >= 7) {
     use_tensor_ops = true;
   }
 #endif
@@ -1921,8 +1909,7 @@ static bool TensorOpsAvailable(int cc_major) {
   // strictly correct.  We can't simply enable it, though, as that would change
   // clients' behavior significantly: Using tensor ops on fp32 inputs cause them
   // to be rounded to fp16.
-  if (cc_major >= 7 && TensorOpMathEnabled() &&
-      std::is_same::value) {
+  if (cc_major >= 7 && std::is_same::value) {
     return true;
   }
 #endif
@@ -2270,7 +2257,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal(
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
           &cc_major, &cc_minor) &&
       cc_major >= 5) {
-    bool use_tensor_ops = TensorOpMathEnabled() && data_type == CUDA_R_16F;
+    bool use_tensor_ops = data_type == CUDA_R_16F;
     cublasGemmAlgo_t algo =
         (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
     cudaDataType_t compute_type =
@@ -2425,7 +2412,7 @@ bool CUDABlas::DoBlasGemmStridedBatched(
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
           &cc_major, &cc_minor)) {
     // GPUs < sm_70 don't support tensor ops.
-    if (cc_major >= 7 && TensorOpMathEnabled()) {
+    if (cc_major >= 7) {
       use_tensor_ops = true;
     }
 #if CUDA_VERSION >= 9010
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index be18c989861..e46c271443b 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -601,31 +601,6 @@ class CudnnFilterDescriptor {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnFilterDescriptor);
 };
 
-// A helper function to decide whether to enable the TENSOR_OP_MATH math type
-bool TensorOpMathEnabled() {
-  static bool is_enabled = [] {
-    bool is_disabled = false;
-    TF_CHECK_OK(
-        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_TENSOR_OP_MATH",
-                                       /*default_val=*/false, &is_disabled));
-    return !is_disabled;
-  }();
-  return is_enabled;
-}
-
-// A helper function to decide whether to enable the TENSOR_OP_MATH math type
-// for RNNs.
-bool RnnTensorOpMathEnabled() {
-  static bool is_enabled = [] {
-    bool is_disabled = false;
-    TF_CHECK_OK(
-        tensorflow::ReadBoolFromEnvVar("TF_DISABLE_CUDNN_RNN_TENSOR_OP_MATH",
-                                       /*default_val=*/false, &is_disabled));
-    return !is_disabled;
-  }();
-  return is_enabled;
-}
-
 // A helper function to decide whether to use
 // CUDNN_BATCHNORM_SPATIAL_PERSISTENT in batchnorm. This mode can be faster in
 // some tasks because an optimized path may be selected for CUDNN_DATA_FLOAT
@@ -749,9 +724,7 @@ class CudnnConvolutionDescriptor {
 #if CUDNN_VERSION >= 7000
     cudnnMathType_t math_type =
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
-    if (TensorOpMathEnabled()) {
-      CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
-    }
+    CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
 #endif
   }
 
@@ -1155,21 +1128,19 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
     // in profile mode, which is run with algorithms returned from
     // GetRnnAlgorithms() (which are non-default and explicitly set whether to
     // use tensor ops). CuDNN 7.2.1 fixed this issue
-    if (RnnTensorOpMathEnabled()) {
-      cudnnMathType_t math_type;
-      if (algorithm_config.algorithm().has_value()) {
-        math_type = algorithm_config.algorithm()->tensor_ops_enabled()
-                        ? CUDNN_TENSOR_OP_MATH
-                        : CUDNN_DEFAULT_MATH;
-      } else {
+    cudnnMathType_t math_type;
+    if (algorithm_config.algorithm().has_value()) {
+      math_type = algorithm_config.algorithm()->tensor_ops_enabled()
+                      ? CUDNN_TENSOR_OP_MATH
+                      : CUDNN_DEFAULT_MATH;
+    } else {
 #if CUDNN_VERSION >= 7201
-        math_type = CUDNN_TENSOR_OP_MATH;
+      math_type = CUDNN_TENSOR_OP_MATH;
 #else
-        math_type = CUDNN_DEFAULT_MATH;
+      math_type = CUDNN_DEFAULT_MATH;
 #endif  // CUDNN_VERSION >= 7201
-      }
-      CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
     }
+    CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
 #endif  // CUDNN_VERSION >= 7000
 
     return CudnnRnnDescriptor(cudnn, std::move(rnn_desc), std::move(rnn_plan),
@@ -2686,7 +2657,7 @@ AllocateCudnnConvolutionBackwardFilterWorkspace(
 }
 
 static bool TensorOpMathAvailable(int cc_major) {
-  return cc_major >= 7 && CUDNN_VERSION >= 7000 && TensorOpMathEnabled();
+  return cc_major >= 7 && CUDNN_VERSION >= 7000;
 }
 
 port::StatusOr GetCudnnConvolutionForwardAlgorithm(
@@ -3480,9 +3451,7 @@ bool CudnnSupport::GetRnnAlgorithms(
   for (auto i : algo_types) {
     out_algorithms->push_back({i, /*use_tensor_ops=*/false});
 #if CUDNN_VERSION >= 7100
-    if (RnnTensorOpMathEnabled()) {
-      out_algorithms->push_back({i, /*use_tensor_ops=*/true});
-    }
+    out_algorithms->push_back({i, /*use_tensor_ops=*/true});
 #endif
   }
   return true;

From d2afc9ce83b48170f0821b6b3b5debddd857d320 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Fri, 15 May 2020 11:46:41 -0500
Subject: [PATCH 0662/1390] Add global setting control TF32 execution

---
 tensorflow/core/platform/BUILD         |  7 +++++++
 tensorflow/core/platform/tf32_utils.cc | 27 ++++++++++++++++++++++++++
 tensorflow/core/platform/tf32_utils.h  | 27 ++++++++++++++++++++++++++
 3 files changed, 61 insertions(+)
 create mode 100644 tensorflow/core/platform/tf32_utils.cc
 create mode 100644 tensorflow/core/platform/tf32_utils.h

diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD
index 70bb8a89417..33a1e7cfe0a 100644
--- a/tensorflow/core/platform/BUILD
+++ b/tensorflow/core/platform/BUILD
@@ -938,6 +938,13 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "tf32_utils",
+    srcs = ["tf32_utils.cc"],
+    hdrs = ["tf32_utils.h"],
+    copts = tf_copts(),
+)
+
 tf_cc_tests(
     name = "low_level_library_tests",
     size = "small",
diff --git a/tensorflow/core/platform/tf32_utils.cc b/tensorflow/core/platform/tf32_utils.cc
new file mode 100644
index 00000000000..715b5996dc3
--- /dev/null
+++ b/tensorflow/core/platform/tf32_utils.cc
@@ -0,0 +1,27 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/platform/tf32_utils.h"
+
+namespace tensorflow {
+
+// TODO(nluehr): enable tf32 execution by default after TF32 Ampere testing.
+static bool tf32_enabled = false;
+
+void allow_tf32_execution(bool allow) { tf32_enabled = allow; }
+
+bool tf32_execution_allowed() { return tf32_enabled; }
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/platform/tf32_utils.h b/tensorflow/core/platform/tf32_utils.h
new file mode 100644
index 00000000000..a0ce58f9bbd
--- /dev/null
+++ b/tensorflow/core/platform/tf32_utils.h
@@ -0,0 +1,27 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_
+#define TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_
+
+namespace tensorflow {
+
+void allow_tf32_execution(bool allow);
+
+bool tf32_execution_allowed();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_PLATFORM_TF32_UTILS_H_

From 16033c0b3484409a965acc0dd3054695145311a8 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Fri, 15 May 2020 13:33:02 -0500
Subject: [PATCH 0663/1390] Python tf.config tf32 interface

---
 tensorflow/python/BUILD               | 11 +++++++++++
 tensorflow/python/framework/config.py | 26 ++++++++++++++++++++++++++
 tensorflow/python/util/tf32.cc        | 22 ++++++++++++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 tensorflow/python/util/tf32.cc

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index de9cf9a24c7..5f9e2dfb1ff 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -788,6 +788,16 @@ tf_python_pybind_extension(
     ],
 )
 
+tf_python_pybind_extension(
+    name = "_pywrap_tf32_execution",
+    srcs = ["util/tf32.cc"],
+    module_name = "_pywrap_tf32_execution",
+    deps = [
+        "//tensorflow/core/platform:tf32_utils",
+        "@pybind11",
+    ],
+)
+
 tf_python_pybind_extension(
     name = "_pywrap_util_port",
     srcs = ["util/port_wrapper.cc"],
@@ -5678,6 +5688,7 @@ py_library(
         "//tensorflow:composite_tensor_whitelist",
     ],
     deps = [
+        ":_pywrap_tf32_execution",
         ":tf_decorator",
         ":tf_export",
         ":tf_stack",
diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py
index 9ff16f2a327..cb95965dfb2 100644
--- a/tensorflow/python/framework/config.py
+++ b/tensorflow/python/framework/config.py
@@ -18,10 +18,36 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python import _pywrap_tf32_execution
 from tensorflow.python.eager import context
 from tensorflow.python.util import deprecation
 from tensorflow.python.util.tf_export import tf_export
 
+def tensor_float32_execution_allowed():
+  """Get if TensorFloat-32 operations are enabled on supported hardware.
+
+  Returns:
+    True if TensorFloat-32 execution is enabled and False otherwise.
+  """
+  return _pywrap_tf32_execution.is_allowed()
+
+def allow_tensor_float_32_execution(allow):
+  """Allow use of TensorFloat-32 with float32 ops on supported hardware.
+
+  TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture.
+  TensorFloat-32 kernels take float32 inputs and produce float32 outputs.
+  Internally, the inputs are cast to a custom representation with 10-bit
+  mantissa (similar to float16) and 8-bit exponent (similar to float32) and are
+  executed using TensorCores with float32 accumulation. For more information,
+  see https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/.
+
+  TensorFloat-32 execution is disabled by default, but this may change in a
+  future version.
+  
+  Args:
+    allow: whether to allow TensorFloat-32 execution
+  """
+  _pywrap_tf32_execution.allow(allow)
 
 @tf_export('config.threading.get_intra_op_parallelism_threads')
 def get_intra_op_parallelism_threads():
diff --git a/tensorflow/python/util/tf32.cc b/tensorflow/python/util/tf32.cc
new file mode 100644
index 00000000000..7dece6ccdae
--- /dev/null
+++ b/tensorflow/python/util/tf32.cc
@@ -0,0 +1,22 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "pybind11/pybind11.h"
+#include "tensorflow/core/platform/tf32_utils.h"
+
+PYBIND11_MODULE(_pywrap_tf32_execution, m) {
+  m.def("allow", &tensorflow::allow_tf32_execution);
+  m.def("is_allowed", &tensorflow::tf32_execution_allowed);
+}

From 376efd71b10eb7d7b900c3f7e7aff99bf15b0196 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Tue, 19 May 2020 14:58:30 -0500
Subject: [PATCH 0664/1390] Convolution TF32 Plumbing

---
 tensorflow/stream_executor/cuda/BUILD       |   1 +
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 200 +++++++++++++-------
 2 files changed, 135 insertions(+), 66 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/BUILD b/tensorflow/stream_executor/cuda/BUILD
index c3cf9f5db15..cdc0de7c72a 100644
--- a/tensorflow/stream_executor/cuda/BUILD
+++ b/tensorflow/stream_executor/cuda/BUILD
@@ -356,6 +356,7 @@ cc_library(
         "@local_config_cuda//cuda:cudnn_header",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "//tensorflow/core/platform:tf32_utils",
         "//tensorflow/stream_executor:dnn",
         "//tensorflow/stream_executor:event",
         "//tensorflow/stream_executor:plugin_registry",
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index e46c271443b..53ba31d7d0d 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -20,8 +20,8 @@ limitations under the License.
 #include 
 
 #include "absl/strings/str_cat.h"
-#include "third_party/eigen3/Eigen/Core"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/tf32_utils.h"
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/cuda/cuda_activation.h"
 #include "tensorflow/stream_executor/cuda/cuda_diagnostics.h"
@@ -42,6 +42,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/scratch_allocator.h"
 #include "tensorflow/stream_executor/stream.h"
 #include "tensorflow/stream_executor/stream_executor_pimpl.h"
+#include "third_party/eigen3/Eigen/Core"
 // clang-format off
 #include "third_party/gpus/cudnn/cudnn.h"
 #include "absl/strings/string_view.h"
@@ -705,10 +706,6 @@ class CudnnConvolutionDescriptor {
             : CUDNN_CROSS_CORRELATION,
         data_type));
 
-    // NOTE(benbarsdell): This only applies if tensor op math is enabled
-    //                      and algo selection is set to Default.
-    this->set_use_tensor_op_math(true);
-
 #if CUDNN_MAJOR >= 7
     VLOG(2) << "Requesting grouped convolution: "
             << convolution_descriptor.group_count();
@@ -720,10 +717,14 @@ class CudnnConvolutionDescriptor {
 #endif
   }
 
-  void set_use_tensor_op_math(bool use_tensor_op_math) const {
+  void set_use_tensor_op_math(bool use_tensor_op_math) {
 #if CUDNN_VERSION >= 7000
     cudnnMathType_t math_type =
+#if CUDNN_VERSION >= 8000
+        (use_tensor_op_math ? CUDNN_DEFAULT_MATH : CUDNN_FMA_MATH);
+#else
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
+#endif
     CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
 #endif
   }
@@ -736,6 +737,38 @@ class CudnnConvolutionDescriptor {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnConvolutionDescriptor);
 };
 
+// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math
+// set
+static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) {
+  cudnnMathType_t math_type;
+  CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type));
+#if CUDNN_VERSION >= 8000
+  return math_type != CUDNN_FMA_MATH;
+#else
+  return math_type == CUDNN_TENSOR_OP_MATH;
+#endif
+}
+
+static bool TensorOpMathAvailable(int cc_major) {
+  return cc_major >= 7 && CUDNN_VERSION >= 7000;
+}
+
+static bool IsTensorMathAllowed(Stream* stream, dnn::DataType input_type) {
+  int cc_major, cc_minor;
+  std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
+  if (!TensorOpMathAvailable(cc_major)) {
+    return false;
+  }
+  if (input_type == dnn::DataType::kFloat) {
+    if (CUDNN_VERSION < 8000) {
+      return false;
+    } else if (!tensorflow::tf32_execution_allowed()) {
+      return false;
+    }
+  }
+  return true;
+}
+
 // Turns a PoolingDescriptor structure into a cudnn pooling descriptor handle
 // within a scope.
 class CudnnPoolingDescriptor {
@@ -2531,10 +2564,11 @@ port::StatusOr> AllocateCudnnConvolutionForwardWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  // TODO(csigg): This has side effects on the convolution descriptor. It is
-  // functionally correct because the convolution is run with the algorithm of
-  // the last call to this function, but should be fixed anyway.
-  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
+  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
+    return port::Status(
+        port::error::INTERNAL,
+        "Mismatch between cudnn conv and algorithm descriptors.");
+  }
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2574,10 +2608,11 @@ AllocateCudnnConvolutionBackwardDataWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  // TODO(csigg): This has side effects on the convolution descriptor. It is
-  // functionally correct because the convolution is run with the algorithm of
-  // the last call to this function, but should be fixed anyway.
-  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
+  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
+    return port::Status(
+        port::error::INTERNAL,
+        "Mismatch between cudnn conv and algorithm descriptors.");
+  }
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2619,10 +2654,11 @@ AllocateCudnnConvolutionBackwardFilterWorkspace(
     const CudnnTensorDescriptor& output_nd,
     const dnn::AlgorithmDesc& algorithm_desc,
     ScratchAllocator* scratch_allocator) {
-  // TODO(csigg): This has side effects on the convolution descriptor. It is
-  // functionally correct because the convolution is run with the algorithm of
-  // the last call to this function, but should be fixed anyway.
-  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
+  if (IsTensorMathOpSet(conv) != algorithm_desc.tensor_ops_enabled()) {
+    return port::Status(
+        port::error::INTERNAL,
+        "Mismatch between cudnn conv and algorithm descriptors.");
+  }
 
   // Query the size of the workspace and allocate it.
   size_t size_in_bytes;
@@ -2656,18 +2692,39 @@ AllocateCudnnConvolutionBackwardFilterWorkspace(
   return scratch_allocator->AllocateBytes(size_in_bytes);
 }
 
-static bool TensorOpMathAvailable(int cc_major) {
-  return cc_major >= 7 && CUDNN_VERSION >= 7000;
+port::StatusOr UseTensorOps(Stream* stream, dnn::DataType type,
+                                  absl::optional desc) {
+  bool use_tensor_ops;
+  if (desc.has_value()) {
+    use_tensor_ops = desc->tensor_ops_enabled();
+    if (use_tensor_ops && !IsTensorMathAllowed(stream, type)) {
+      return port::Status(port::error::INVALID_ARGUMENT,
+                          "Algo requests disallowed tensor op evaluation.");
+    }
+  } else {
+    use_tensor_ops = IsTensorMathAllowed(stream, type);
+  }
+  return use_tensor_ops;
 }
 
 port::StatusOr GetCudnnConvolutionForwardAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    const CudnnConvolutionDescriptor& conv,
+    dnn::DataType element_type,
+    const dnn::ConvolutionDescriptor& convolution_descriptor,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
+
+  CudnnConvolutionDescriptor conv(
+      convolution_descriptor,
+      ToCudnnDataType(GetConvAccumulatorType(element_type)));
+  bool use_tensor_ops;
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
+
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2680,10 +2737,7 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm(
                         GetCudnnConvolutionForwardAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    int cc_major, cc_minor;
-    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
-    algo_desc = dnn::AlgorithmDesc(
-        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
+    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
   }
 
   const auto scratch_or = AllocateCudnnConvolutionForwardWorkspace(
@@ -2707,6 +2761,9 @@ port::StatusOr GetCudnnConvolutionForwardAlgorithm(
                      "Returned status: ", scratch_or.status().ToString()));
   }
 
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionForwardWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -2717,10 +2774,19 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    const CudnnConvolutionDescriptor& conv,
+    dnn::DataType element_type,
+    const dnn::ConvolutionDescriptor& convolution_descriptor,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
+  CudnnConvolutionDescriptor conv(
+      convolution_descriptor,
+      ToCudnnDataType(GetConvAccumulatorType(element_type)));
+  bool use_tensor_ops;
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
+
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2733,10 +2799,7 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
                         GetCudnnConvolutionBackwardDataAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    int cc_major, cc_minor;
-    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
-    algo_desc = dnn::AlgorithmDesc(
-        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
+    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
   }
 
   const auto scratch_or = AllocateCudnnConvolutionBackwardDataWorkspace(
@@ -2759,6 +2822,9 @@ port::StatusOr GetCudnnConvolutionBackwardDataAlgorithm(
         "while a secondary algorithm is not provided.");
   }
 
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardDataWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -2769,10 +2835,19 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,
     const CudnnTensorDescriptor& input_nd, const CudnnFilterDescriptor& filter,
-    const CudnnConvolutionDescriptor& conv,
+    dnn::DataType element_type,
+    const dnn::ConvolutionDescriptor& convolution_descriptor,
     const CudnnTensorDescriptor& output_nd, ScratchAllocator* scratch_allocator,
     DeviceMemory* scratch) {
   absl::optional algo_desc = algorithm_config.algorithm();
+  CudnnConvolutionDescriptor conv(
+      convolution_descriptor,
+      ToCudnnDataType(GetConvAccumulatorType(element_type)));
+  bool use_tensor_ops;
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
+
   if (!algo_desc.has_value()) {
     // Pick fastest algorithm within memory limit according to cuDNN's
     // heuristics.
@@ -2785,10 +2860,7 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
                         GetCudnnConvolutionBackwardFilterAlgo(
                             cudnn, input_nd, filter, conv, output_nd,
                             specify_workspace_limit, memory_limit_bytes));
-    int cc_major, cc_minor;
-    std::tie(cc_major, cc_minor) = GetCcMajorMinor(stream);
-    algo_desc = dnn::AlgorithmDesc(
-        algo, /*use_tensor_ops=*/TensorOpMathAvailable(cc_major));
+    algo_desc = dnn::AlgorithmDesc(algo, use_tensor_ops);
   }
 
   auto scratch_or = AllocateCudnnConvolutionBackwardFilterWorkspace(
@@ -2811,6 +2883,9 @@ port::StatusOr GetCudnnConvolutionBackwardFilterAlgorithm(
         "while a secondary algorithm is not provided.");
   }
 
+  SE_ASSIGN_OR_RETURN(use_tensor_ops,
+                      UseTensorOps(stream, element_type, algo_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
   SE_ASSIGN_OR_RETURN(*scratch, AllocateCudnnConvolutionBackwardFilterWorkspace(
                                     stream, cudnn, input_nd, filter, conv,
                                     output_nd, *algo_desc, scratch_allocator));
@@ -2975,35 +3050,32 @@ port::Status CudnnSupport::DoPrepareForConvolution(
   CudnnTensorDescriptor output_nd(
       output_descriptor,
       ToCudnnDataType(element_type, output_descriptor.layout()));
-  CudnnConvolutionDescriptor conv(
-      convolution_descriptor,
-      ToCudnnDataType(GetConvAccumulatorType(element_type)));
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
 
   switch (kind) {
     case dnn::ConvolutionKind::FORWARD: {
-      SE_ASSIGN_OR_RETURN(
-          *algorithm_desc,
-          GetCudnnConvolutionForwardAlgorithm(
-              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
-              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(*algorithm_desc,
+                          GetCudnnConvolutionForwardAlgorithm(
+                              stream, cudnn, algorithm_config, input_nd,
+                              filter_nd, element_type, convolution_descriptor,
+                              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     case dnn::ConvolutionKind::BACKWARD_DATA: {
-      SE_ASSIGN_OR_RETURN(
-          *algorithm_desc,
-          GetCudnnConvolutionBackwardDataAlgorithm(
-              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
-              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(*algorithm_desc,
+                          GetCudnnConvolutionBackwardDataAlgorithm(
+                              stream, cudnn, algorithm_config, input_nd,
+                              filter_nd, element_type, convolution_descriptor,
+                              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     case dnn::ConvolutionKind::BACKWARD_FILTER: {
-      SE_ASSIGN_OR_RETURN(
-          *algorithm_desc,
-          GetCudnnConvolutionBackwardFilterAlgorithm(
-              stream, cudnn, algorithm_config, input_nd, filter_nd, conv,
-              output_nd, scratch_allocator, scratch_memory));
+      SE_ASSIGN_OR_RETURN(*algorithm_desc,
+                          GetCudnnConvolutionBackwardFilterAlgorithm(
+                              stream, cudnn, algorithm_config, input_nd,
+                              filter_nd, element_type, convolution_descriptor,
+                              output_nd, scratch_allocator, scratch_memory));
       break;
     }
     default:
@@ -3032,8 +3104,9 @@ port::Status CudnnSupport::DoConvolve(
   auto accumulator_type = GetConvAccumulatorType(element_type);
   CudnnConvolutionDescriptor conv(convolution_descriptor,
                                   ToCudnnDataType(accumulator_type));
-  // Set use_tensor_math param to correct value
-  conv.set_use_tensor_op_math(algorithm_desc.tensor_ops_enabled());
+  SE_ASSIGN_OR_RETURN(bool use_tensor_ops,
+                      UseTensorOps(stream, element_type, algorithm_desc));
+  conv.set_use_tensor_op_math(use_tensor_ops);
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
   // Alpha is the scaling factor for input.
@@ -3266,14 +3339,6 @@ port::Status CudnnSupport::DoConvolve(
   return port::Status::OK();
 }
 
-// A helper function to query if a CudnnConvolutionDescriptor has tensor_op_math
-// set
-static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) {
-  cudnnMathType_t math_type;
-  CHECK_CUDNN_OK(cudnnGetConvolutionMathType(conv.handle(), &math_type));
-  return math_type == CUDNN_TENSOR_OP_MATH;
-}
-
 template 
 port::Status CudnnSupport::DoFusedConvolveImpl(
@@ -3307,8 +3372,6 @@ port::Status CudnnSupport::DoFusedConvolveImpl(
       filter_descriptor,
       GetCudnnDataType(conv_input_descriptor.layout()));
   CudnnTensorDescriptor bias_nd(bias_descriptor, GetCudnnDataType());
-  CudnnConvolutionDescriptor conv(convolution_descriptor,
-                                  ToCudnnDataType(accumulator_type));
 
   auto cudnn = cudnn_->GetHandle(parent_, stream);
 
@@ -3318,9 +3381,14 @@ port::Status CudnnSupport::DoFusedConvolveImpl(
   SE_ASSIGN_OR_RETURN(
       dnn::AlgorithmDesc algo_desc,
       GetCudnnConvolutionForwardAlgorithm(
-          stream, cudnn, algorithm_config, conv_input_nd, filter, conv,
+          stream, cudnn, algorithm_config, conv_input_nd, filter,
+          dnn::ToDataType::value, convolution_descriptor,
           output_nd, scratch_allocator, &scratch));
 
+  CudnnConvolutionDescriptor conv(convolution_descriptor,
+                                  ToCudnnDataType(accumulator_type));
+  conv.set_use_tensor_op_math(algo_desc.tensor_ops_enabled());
+
   std::unique_ptr timer;
   if (is_profiling) {
     timer.reset(new GpuTimer(parent_));  // NOLINT

From 107e6348236d35a9fcdb4a1375ff07bf4975131c Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Tue, 19 May 2020 15:54:10 -0500
Subject: [PATCH 0665/1390] Plumb TF32 for RNN

---
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 30 ++++++++++++++-------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 53ba31d7d0d..820d39c7201 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -1161,17 +1161,26 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
     // in profile mode, which is run with algorithms returned from
     // GetRnnAlgorithms() (which are non-default and explicitly set whether to
     // use tensor ops). CuDNN 7.2.1 fixed this issue
-    cudnnMathType_t math_type;
+    bool allow_tensor_ops =
+        data_type != CUDNN_DATA_FLOAT || tensorflow::tf32_execution_allowed();
+    bool use_tensor_ops;
     if (algorithm_config.algorithm().has_value()) {
-      math_type = algorithm_config.algorithm()->tensor_ops_enabled()
-                      ? CUDNN_TENSOR_OP_MATH
-                      : CUDNN_DEFAULT_MATH;
+      use_tensor_ops = algorithm_config.algorithm()->tensor_ops_enabled();
     } else {
-#if CUDNN_VERSION >= 7201
-      math_type = CUDNN_TENSOR_OP_MATH;
-#else
-      math_type = CUDNN_DEFAULT_MATH;
-#endif  // CUDNN_VERSION >= 7201
+      use_tensor_ops = CUDNN_VERSION >= 7201 && allow_tensor_ops;
+    }
+
+    if (use_tensor_ops && !allow_tensor_ops) {
+      return port::Status(port::error::INVALID_ARGUMENT,
+                          "Algo requests disallowed tensor op evaluation.");
+    }
+
+    cudnnMathType_t math_type;
+    if (use_tensor_ops) {
+      math_type =
+          CUDNN_VERSION >= 8000 ? CUDNN_DEFAULT_MATH : CUDNN_TENSOR_OP_MATH;
+    } else {
+      math_type = CUDNN_VERSION >= 8000 ? CUDNN_FMA_MATH : CUDNN_DEFAULT_MATH;
     }
     CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
 #endif  // CUDNN_VERSION >= 7000
@@ -2707,6 +2716,9 @@ port::StatusOr UseTensorOps(Stream* stream, dnn::DataType type,
   return use_tensor_ops;
 }
 
+cudnnDataType_t GetRnnComputeType(dnn::DataType data_type);
+dnn::DataType GetConvAccumulatorType(dnn::DataType data_type);
+
 port::StatusOr GetCudnnConvolutionForwardAlgorithm(
     Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config,

From 03d836e2616574c657f7702e24e4fc79c661115b Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Wed, 20 May 2020 10:06:35 -0500
Subject: [PATCH 0666/1390] Plumb TF32 for cublas gemm

---
 tensorflow/stream_executor/cuda/BUILD        |  1 +
 tensorflow/stream_executor/cuda/cuda_blas.cc | 84 +++++++++-----------
 tensorflow/stream_executor/cuda/cuda_blas.h  |  8 +-
 3 files changed, 43 insertions(+), 50 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/BUILD b/tensorflow/stream_executor/cuda/BUILD
index cdc0de7c72a..3a14be9ad50 100644
--- a/tensorflow/stream_executor/cuda/BUILD
+++ b/tensorflow/stream_executor/cuda/BUILD
@@ -251,6 +251,7 @@ cc_library(
         "@local_config_cuda//cuda:cuda_headers",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "//tensorflow/core/platform:tf32_utils",
         "//tensorflow/stream_executor",
         "//tensorflow/stream_executor:event",
         "//tensorflow/stream_executor:host_or_device_scalar",
diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index 65c07e72154..e2cbb0b75df 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -48,7 +48,7 @@ limitations under the License.
 
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
-#include "third_party/eigen3/Eigen/Core"
+#include "tensorflow/core/platform/tf32_utils.h"
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/cuda/cuda_activation.h"
 #include "tensorflow/stream_executor/cuda/cuda_gpu_executor.h"
@@ -66,6 +66,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/plugin_registry.h"
 #include "tensorflow/stream_executor/scratch_allocator.h"
 #include "tensorflow/stream_executor/stream_executor.h"
+#include "third_party/eigen3/Eigen/Core"
 
 namespace stream_executor {
 namespace gpu {
@@ -225,6 +226,18 @@ bool CUDABlas::Init() {
     return false;
   }
 
+#if CUDA_VERSION >= 9000
+#if CUBLAS_VER_MAJOR >= 11
+  ret = cublasSetMathMode(blas_, CUBLAS_TF32_TENSOR_OP_MATH);
+#else
+  ret = cublasSetMathMode(blas_, CUBLAS_TENSOR_OP_MATH);
+#endif
+  if (ret != CUBLAS_STATUS_SUCCESS) {
+    LOG(ERROR) << "failed to set cublas default math mode: " << ToString(ret);
+    return false;
+  }
+#endif
+
   return true;
 }
 
@@ -387,7 +400,7 @@ cudaDataType_t CUDAComputationType(blas::ComputationType ty) {
 template 
 bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                                   bool pointer_mode_host, bool err_on_failure,
-                                  bool use_tensor_op_math, Args... args) {
+                                  Args... args) {
   absl::MutexLock lock(&mu_);
 
   CHECK(blas_ != nullptr);
@@ -401,10 +414,10 @@ bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                                            : CUBLAS_POINTER_MODE_DEVICE)) {
     return false;
   }
-#if CUDA_VERSION >= 9000
+#if CUBLAS_VER_MAJOR >= 11
   ScopedCublasMathMode math_mode{blas_};
-  if (use_tensor_op_math) {
-    if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) {
+  if (!tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_DEFAULT_MATH)) {
       return false;
     }
   }
@@ -1621,21 +1634,9 @@ bool CUDABlas::DoBlasGemm(
     }
   }
 
-  bool use_tensor_ops = false;
-#if CUDA_VERSION >= 9000
-  int cc_major, cc_minor;
-  stream->parent()->GetDeviceDescription().cuda_compute_capability(&cc_major,
-                                                                   &cc_minor);
-
-  // GPUs < sm_70 don't support tensor ops.
-  if (cc_major >= 7) {
-    use_tensor_ops = true;
-  }
-#endif
-
   return DoBlasInternalImpl(
       cublasSgemmEx, stream, true /* = pointer_mode_host */,
-      true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa),
+      true /* = err_on_failure= */, CUDABlasTranspose(transa),
       CUDABlasTranspose(transb), m, n, k, &alpha, GpuMemory(a),
       SE_CUDA_DATA_HALF, lda, GpuMemory(b), SE_CUDA_DATA_HALF, ldb, &beta,
       GpuMemoryMutable(c), SE_CUDA_DATA_HALF, ldc);
@@ -2257,7 +2258,8 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal(
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
           &cc_major, &cc_minor) &&
       cc_major >= 5) {
-    bool use_tensor_ops = data_type == CUDA_R_16F;
+    bool use_tensor_ops =
+        data_type == CUDA_R_16F || tensorflow::tf32_execution_allowed();
     cublasGemmAlgo_t algo =
         (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
     cudaDataType_t compute_type =
@@ -2271,7 +2273,7 @@ port::Status CUDABlas::DoBlasGemmBatchedInternal(
     bool ok;
     ok = DoBlasInternalImpl(
         AS_LAMBDA(cublasGemmBatchedEx), stream, true /* = pointer_mode_host */,
-        true /* = err_on_failure */, use_tensor_ops, CUDABlasTranspose(transa),
+        true /* = err_on_failure */, CUDABlasTranspose(transa),
         CUDABlasTranspose(transb), m, n, k, &alpha, a_void_ptrs, data_type, lda,
         b_void_ptrs, data_type, ldb, &beta, c_void_ptrs, data_type, ldc,
         batch_count, compute_type, algo);
@@ -2406,33 +2408,25 @@ bool CUDABlas::DoBlasGemmStridedBatched(
     int lda, int64 stride_a, const DeviceMemory &b, int ldb,
     int64 stride_b, float beta, DeviceMemory *c, int ldc,
     int64 stride_c, int batch_count) {
-  bool use_tensor_ops = false;
-#if CUDA_VERSION >= 9000
+#if CUDA_VERSION >= 9010
   int cc_major, cc_minor;
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
-          &cc_major, &cc_minor)) {
-    // GPUs < sm_70 don't support tensor ops.
-    if (cc_major >= 7) {
-      use_tensor_ops = true;
+          &cc_major, &cc_minor) &&
+      cc_major >= 5) {
+    cublasGemmAlgo_t algo =
+        (cc_major >= 7 ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
+    bool ok = DoBlasInternalImpl(
+        AS_LAMBDA(cublasGemmStridedBatchedEx), stream,
+        true /* = pointer_mode_host */, true /* = err_on_failure */,
+        CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha,
+        GpuMemory(a), CUDA_R_16F, lda, stride_a, GpuMemory(b), CUDA_R_16F, ldb,
+        stride_b, &beta, GpuMemoryMutable(c), CUDA_R_16F, ldc, stride_c,
+        batch_count, CUDA_R_32F, algo);
+    if (ok) {
+      return true;
     }
-#if CUDA_VERSION >= 9010
-    if (cc_major >= 5) {
-      cublasGemmAlgo_t algo =
-          (use_tensor_ops ? CUBLAS_GEMM_DFALT_TENSOR_OP : CUBLAS_GEMM_DFALT);
-      bool ok = DoBlasInternalImpl(
-          AS_LAMBDA(cublasGemmStridedBatchedEx), stream,
-          true /* = pointer_mode_host */, true /* = err_on_failure */,
-          use_tensor_ops, CUDABlasTranspose(transa), CUDABlasTranspose(transb),
-          m, n, k, &alpha, GpuMemory(a), CUDA_R_16F, lda, stride_a,
-          GpuMemory(b), CUDA_R_16F, ldb, stride_b, &beta, GpuMemoryMutable(c),
-          CUDA_R_16F, ldc, stride_c, batch_count, CUDA_R_32F, algo);
-      if (ok) {
-        return true;
-      }
-      LOG(ERROR) << "failed BLAS call, see log for details";
-      return false;
-    }
-#endif
+    LOG(ERROR) << "failed BLAS call, see log for details";
+    return false;
   }
 #endif
   // Either CUDA_VERSION < 9.1 or SM < 5.0. Fall back to a loop.
@@ -2445,7 +2439,7 @@ bool CUDABlas::DoBlasGemmStridedBatched(
         reinterpret_cast<__half *>(GpuMemoryMutable(c) + batch * stride_c);
     bool ok = DoBlasInternalImpl(
         cublasSgemmEx, stream, true /* = pointer_mode_host */,
-        true /* = err_on_failure= */, use_tensor_ops, CUDABlasTranspose(transa),
+        true /* = err_on_failure= */, CUDABlasTranspose(transa),
         CUDABlasTranspose(transb), m, n, k, &alpha, a_matrix, SE_CUDA_DATA_HALF,
         lda, b_matrix, SE_CUDA_DATA_HALF, ldb, &beta, c_matrix,
         SE_CUDA_DATA_HALF, ldc);
diff --git a/tensorflow/stream_executor/cuda/cuda_blas.h b/tensorflow/stream_executor/cuda/cuda_blas.h
index 817bdb72777..556456c83db 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.h
+++ b/tensorflow/stream_executor/cuda/cuda_blas.h
@@ -83,7 +83,7 @@ class CUDABlas : public blas::BlasSupport {
   template 
   bool DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                           bool pointer_mode_host, bool err_on_failure,
-                          bool use_tensor_op_math, Args... args);
+                          Args... args);
 
   // Convenience functions that call DoBlasInternalImpl with different values
   // for err_on_failure.
@@ -91,8 +91,7 @@ class CUDABlas : public blas::BlasSupport {
   bool DoBlasInternal(FuncT cublas_func, Stream *stream, bool pointer_mode_host,
                       Args... args) {
     return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host,
-                              /*err_on_failure=*/true, /*use_tensor_ops=*/false,
-                              args...);
+                              /*err_on_failure=*/true, args...);
   }
   template 
   bool DoBlasInternalFailureOK(FuncT cublas_func, Stream *stream,
@@ -100,8 +99,7 @@ class CUDABlas : public blas::BlasSupport {
     // Tensor ops are hard-coded off in this path, but can still be enabled with
     // a specific algorithm choice as in DoBlasGemmWithAlgorithmImpl().
     return DoBlasInternalImpl(cublas_func, stream, pointer_mode_host,
-                              /*err_on_failure=*/false,
-                              /*use_tensor_ops=*/false, args...);
+                              /*err_on_failure=*/false, args...);
   }
 
   // A helper function to implement DoBlasGemmBatched interfaces for generic

From 7791e36a57b3d67e625e298a7e8a1063a26571fd Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Mon, 8 Jun 2020 11:21:49 -0500
Subject: [PATCH 0667/1390] Address review comments

---
 tensorflow/core/platform/tf32_utils.cc | 10 ++++++----
 tensorflow/core/platform/tf32_utils.h  |  2 +-
 tensorflow/python/framework/config.py  |  7 +++++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/platform/tf32_utils.cc b/tensorflow/core/platform/tf32_utils.cc
index 715b5996dc3..4456e768c0a 100644
--- a/tensorflow/core/platform/tf32_utils.cc
+++ b/tensorflow/core/platform/tf32_utils.cc
@@ -14,14 +14,16 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/platform/tf32_utils.h"
+#include 
 
 namespace tensorflow {
 
-// TODO(nluehr): enable tf32 execution by default after TF32 Ampere testing.
-static bool tf32_enabled = false;
+// Whether TensorFloat-32 should be used where supported.
+// TODO(nluehr): Maybe enable by default after TF32 Ampere testing.
+static std::atomic tf32_allowed{false};
 
-void allow_tf32_execution(bool allow) { tf32_enabled = allow; }
+void allow_tf32_execution(bool allowed) { tf32_allowed = allowed; }
 
-bool tf32_execution_allowed() { return tf32_enabled; }
+bool tf32_execution_allowed() { return tf32_allowed; }
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/platform/tf32_utils.h b/tensorflow/core/platform/tf32_utils.h
index a0ce58f9bbd..7a158d00ad3 100644
--- a/tensorflow/core/platform/tf32_utils.h
+++ b/tensorflow/core/platform/tf32_utils.h
@@ -18,7 +18,7 @@ limitations under the License.
 
 namespace tensorflow {
 
-void allow_tf32_execution(bool allow);
+void allow_tf32_execution(bool allowed);
 
 bool tf32_execution_allowed();
 
diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py
index cb95965dfb2..e80ad1d72c4 100644
--- a/tensorflow/python/framework/config.py
+++ b/tensorflow/python/framework/config.py
@@ -23,6 +23,8 @@ from tensorflow.python.eager import context
 from tensorflow.python.util import deprecation
 from tensorflow.python.util.tf_export import tf_export
 
+
+# No tf_export until TF is built against CUDA11 which is required for TF32.
 def tensor_float32_execution_allowed():
   """Get if TensorFloat-32 operations are enabled on supported hardware.
 
@@ -31,7 +33,8 @@ def tensor_float32_execution_allowed():
   """
   return _pywrap_tf32_execution.is_allowed()
 
-def allow_tensor_float_32_execution(allow):
+# No tf_export until TF is built against CUDA11 which is required for TF32.
+def allow_tensor_float_32_execution(allowed):
   """Allow use of TensorFloat-32 with float32 ops on supported hardware.
 
   TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture.
@@ -47,7 +50,7 @@ def allow_tensor_float_32_execution(allow):
   Args:
     allow: whether to allow TensorFloat-32 execution
   """
-  _pywrap_tf32_execution.allow(allow)
+  _pywrap_tf32_execution.allow(allowed)
 
 @tf_export('config.threading.get_intra_op_parallelism_threads')
 def get_intra_op_parallelism_threads():

From 6ff37474dd16d77f65d510a9f3bf0f56d08122b9 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Tue, 9 Jun 2020 15:35:47 -0500
Subject: [PATCH 0668/1390] Use CUDNN_TENSOR_OP_MATH to enable tensor cores.

---
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 820d39c7201..166fa0e32d0 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -721,7 +721,7 @@ class CudnnConvolutionDescriptor {
 #if CUDNN_VERSION >= 7000
     cudnnMathType_t math_type =
 #if CUDNN_VERSION >= 8000
-        (use_tensor_op_math ? CUDNN_DEFAULT_MATH : CUDNN_FMA_MATH);
+        (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH);
 #else
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
 #endif
@@ -1177,8 +1177,7 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
 
     cudnnMathType_t math_type;
     if (use_tensor_ops) {
-      math_type =
-          CUDNN_VERSION >= 8000 ? CUDNN_DEFAULT_MATH : CUDNN_TENSOR_OP_MATH;
+      math_type = CUDNN_TENSOR_OP_MATH;
     } else {
       math_type = CUDNN_VERSION >= 8000 ? CUDNN_FMA_MATH : CUDNN_DEFAULT_MATH;
     }

From 8384080492a80b486fd000e3fbd6626cae0fb2cc Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Thu, 11 Jun 2020 15:35:04 -0500
Subject: [PATCH 0669/1390] Make python names consistent

---
 tensorflow/python/framework/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py
index e80ad1d72c4..c19e514a932 100644
--- a/tensorflow/python/framework/config.py
+++ b/tensorflow/python/framework/config.py
@@ -34,7 +34,7 @@ def tensor_float32_execution_allowed():
   return _pywrap_tf32_execution.is_allowed()
 
 # No tf_export until TF is built against CUDA11 which is required for TF32.
-def allow_tensor_float_32_execution(allowed):
+def allow_tensor_float32_execution(allowed):
   """Allow use of TensorFloat-32 with float32 ops on supported hardware.
 
   TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture.

From 5fe636fb4a74b58bd7da6de9ae1f286ad395c272 Mon Sep 17 00:00:00 2001
From: Reed 
Date: Thu, 11 Jun 2020 17:16:28 -0700
Subject: [PATCH 0670/1390] Use float_32 instead of float32 in function names

---
 tensorflow/python/framework/config.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/framework/config.py b/tensorflow/python/framework/config.py
index c19e514a932..767a1a99f4f 100644
--- a/tensorflow/python/framework/config.py
+++ b/tensorflow/python/framework/config.py
@@ -25,7 +25,7 @@ from tensorflow.python.util.tf_export import tf_export
 
 
 # No tf_export until TF is built against CUDA11 which is required for TF32.
-def tensor_float32_execution_allowed():
+def tensor_float_32_execution_allowed():
   """Get if TensorFloat-32 operations are enabled on supported hardware.
 
   Returns:
@@ -34,7 +34,7 @@ def tensor_float32_execution_allowed():
   return _pywrap_tf32_execution.is_allowed()
 
 # No tf_export until TF is built against CUDA11 which is required for TF32.
-def allow_tensor_float32_execution(allowed):
+def allow_tensor_float_32_execution(allowed):
   """Allow use of TensorFloat-32 with float32 ops on supported hardware.
 
   TensorFloat-32 is a math mode introduced with the NVIDIA Ampere architecture.

From 434f49d54fc8165c2f019bd114272ce2200202da Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Fri, 19 Jun 2020 15:44:04 -0500
Subject: [PATCH 0671/1390] Rework gemm TF32

---
 tensorflow/stream_executor/cuda/cuda_blas.cc | 142 +++++++++++++++----
 1 file changed, 113 insertions(+), 29 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index e2cbb0b75df..e9e3635d8c1 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -226,18 +226,6 @@ bool CUDABlas::Init() {
     return false;
   }
 
-#if CUDA_VERSION >= 9000
-#if CUBLAS_VER_MAJOR >= 11
-  ret = cublasSetMathMode(blas_, CUBLAS_TF32_TENSOR_OP_MATH);
-#else
-  ret = cublasSetMathMode(blas_, CUBLAS_TENSOR_OP_MATH);
-#endif
-  if (ret != CUBLAS_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set cublas default math mode: " << ToString(ret);
-    return false;
-  }
-#endif
-
   return true;
 }
 
@@ -1634,6 +1622,13 @@ bool CUDABlas::DoBlasGemm(
     }
   }
 
+#if CUDA_VERSION < 11000
+  ScopedCublasMathMode math_mode{blas_};
+  if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) {
+    return false;
+  }
+#endif
+
   return DoBlasInternalImpl(
       cublasSgemmEx, stream, true /* = pointer_mode_host */,
       true /* = err_on_failure= */, CUDABlasTranspose(transa),
@@ -1681,6 +1676,16 @@ bool CUDABlas::DoBlasGemm(Stream *stream, blas::Transpose transa,
                       "precondition violation";
     }
   }
+
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
+
   return DoBlasInternal(cublasSgemm, stream, true /* = pointer_mode_host */,
                         CUDABlasTranspose(transa), CUDABlasTranspose(transb), m,
                         n, k, &alpha, GpuMemory(a), lda, GpuMemory(b), ldb,
@@ -1707,6 +1712,16 @@ bool CUDABlas::DoBlasGemm(Stream *stream, blas::Transpose transa,
                           DeviceMemory> *c, int ldc) {
   auto cb_alpha = GpuComplexValue(alpha);
   auto cb_beta = GpuComplexValue(beta);
+
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
+
   return DoBlasInternal(cublasCgemm, stream, true /* = pointer_mode_host */,
                         CUDABlasTranspose(transa), CUDABlasTranspose(transb), m,
                         n, k, GpuComplex(&cb_alpha), GpuComplex(GpuMemory(a)),
@@ -1903,20 +1918,6 @@ static bool UsesTensorOps(blas::AlgorithmType algo) {
 #endif
 }
 
-template 
-static bool TensorOpsAvailable(int cc_major) {
-#if CUDA_VERSION >= 9000
-  // cublas *does* allow tensor ops on inputs that are not fp16, so this is not
-  // strictly correct.  We can't simply enable it, though, as that would change
-  // clients' behavior significantly: Using tensor ops on fp32 inputs cause them
-  // to be rounded to fp16.
-  if (cc_major >= 7 && std::is_same::value) {
-    return true;
-  }
-#endif
-  return false;
-}
-
 template 
 bool CUDABlas::DoBlasGemmWithAlgorithmImpl(
     Stream *stream, blas::Transpose transa, blas::Transpose transb, uint64 m,
@@ -1935,17 +1936,52 @@ bool CUDABlas::DoBlasGemmWithAlgorithmImpl(
     return false;
   }
 
-  if (UsesTensorOps(algorithm) && !TensorOpsAvailable(cc_major)) {
-    if (std::is_same::value) {
+  bool algo_uses_tensor_ops = UsesTensorOps(algorithm);
+  cublasMath_t math_type = CUBLAS_DEFAULT_MATH;
+  if (algo_uses_tensor_ops) {
+    if (cc_major < 7) {
       VLOG(2) << "DoBlasGemmWithAlgorithm returning false because algorithm "
               << algorithm
               << " uses tensor ops, but tensor ops are not available in sm"
               << cc_major << "X devices.";
+      return false;
+    } else if (std::is_same::value) {
+#if CUDA_VERSION < 11000
+      VLOG(2) << "DoBlasGemmWithAlgorithm returning false because algorithm "
+              << algorithm
+              << " uses tensor ops, but tensor ops are not available for fp32"
+              << " inputs.";
+      return false;
+#else
+      if (cc_major < 8) {
+        VLOG(2) << "DoBlasGemmWithAlgorithm returning false because algorithm "
+                << algorithm
+                << " uses tensor ops, but tensor ops are not available in sm"
+                << cc_major << "X devices for float input types.";
+        return false;
+      } else if (!tensorflow::tf32_execution_allowed()) {
+        VLOG(2) << "DoBlasGemmWithAlgorithm returning false because algorithm "
+                << algorithm
+                << " uses tensor ops, but tensor ops are disabled for fp32"
+                << " inputs.";
+        return false;
+      }
+      math_type = CUBLAS_TF32_TENSOR_OP_MATH;
+#endif
+    } else if (std::is_same::value) {
+#if CUDA_VERSION < 11000
+      math_type = CUBLAS_TENSOR_OP_MATH;
+#endif
     } else {
       VLOG(2) << "DoBlasGemmWithAlgorithm returning false because algorithm "
               << algorithm
-              << " uses tensor ops, but the input data type is not fp16.";
+              << " uses tensor ops, which are not supported for InT.";
+      return false;
     }
+  }
+
+  ScopedCublasMathMode math_mode{blas_};
+  if (!math_mode.Init(math_type)) {
     return false;
   }
 
@@ -2325,6 +2361,13 @@ bool CUDABlas::DoBlasGemmBatched(
     int ldc, int batch_count, ScratchAllocator *scratch_allocator) {
   // Note: The func passed here (cublasSgemmBatched) is not actually called,
   // due to special handling of fp16 inside DoBlasGemmBatchedInternal.
+#if CUDA_VERSION < 11000
+  ScopedCublasMathMode math_mode{blas_};
+  if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) {
+    return false;
+  }
+#endif
+
   port::Status status = DoBlasGemmBatchedInternal(
       cublasSgemmBatched, stream, transa, transb, m, n, k, alpha, a_array, lda,
       b_array, ldb, beta, c_array, ldc, batch_count, scratch_allocator);
@@ -2341,6 +2384,15 @@ bool CUDABlas::DoBlasGemmBatched(
     const port::ArraySlice *> &b_array, int ldb, float beta,
     const port::ArraySlice *> &c_array, int ldc,
     int batch_count, ScratchAllocator *scratch_allocator) {
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
+
   port::Status status = DoBlasGemmBatchedInternal(
       cublasSgemmBatched, stream, transa, transb, m, n, k, alpha, a_array, lda,
       b_array, ldb, beta, c_array, ldc, batch_count, scratch_allocator);
@@ -2375,6 +2427,15 @@ bool CUDABlas::DoBlasGemmBatched(
     int ldb, std::complex beta,
     const port::ArraySlice> *> &c_array,
     int ldc, int batch_count, ScratchAllocator *scratch_allocator) {
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
+
   port::Status status = DoBlasGemmBatchedInternal(
       cublasCgemmBatched, stream, transa, transb, m, n, k, alpha, a_array, lda,
       b_array, ldb, beta, c_array, ldc, batch_count, scratch_allocator);
@@ -2408,6 +2469,13 @@ bool CUDABlas::DoBlasGemmStridedBatched(
     int lda, int64 stride_a, const DeviceMemory &b, int ldb,
     int64 stride_b, float beta, DeviceMemory *c, int ldc,
     int64 stride_c, int batch_count) {
+#if CUDA_VERSION < 11000
+  ScopedCublasMathMode math_mode{blas_};
+  if (!math_mode.Init(CUBLAS_TENSOR_OP_MATH)) {
+    return false;
+  }
+#endif
+
 #if CUDA_VERSION >= 9010
   int cc_major, cc_minor;
   if (stream->parent()->GetDeviceDescription().cuda_compute_capability(
@@ -2457,6 +2525,14 @@ bool CUDABlas::DoBlasGemmStridedBatched(
     int64 stride_a, const DeviceMemory &b, int ldb, int64 stride_b,
     float beta, DeviceMemory *c, int ldc, int64 stride_c,
     int batch_count) {
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
   return DoBlasInternal(
       cublasSgemmStridedBatched, stream, true /* = pointer_mode_host */,
       CUDABlasTranspose(transa), CUDABlasTranspose(transb), m, n, k, &alpha,
@@ -2484,6 +2560,14 @@ bool CUDABlas::DoBlasGemmStridedBatched(
     const DeviceMemory> &b, int ldb, int64 stride_b,
     std::complex beta, DeviceMemory> *c, int ldc,
     int64 stride_c, int batch_count) {
+#if CUBLAS_VER_MAJOR >= 11
+  ScopedCublasMathMode math_mode{blas_};
+  if (tensorflow::tf32_execution_allowed()) {
+    if (!math_mode.Init(CUBLAS_TF32_TENSOR_OP_MATH)) {
+      return false;
+    }
+  }
+#endif
   auto cb_alpha = GpuComplexValue(alpha);
   auto cb_beta = GpuComplexValue(beta);
   return DoBlasInternal(

From 90443c2b17a321f3589039dd560773b8d47a5cb7 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar 
Date: Fri, 19 Jun 2020 15:15:51 -0700
Subject: [PATCH 0672/1390] Enable test for generating tf_ops

PiperOrigin-RevId: 317389351
Change-Id: I54f6ac53a974cf603ed5fe1a30b6fbb464c80d28
---
 tensorflow/compiler/mlir/tensorflow/BUILD | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index bcabb13d301..11a455eec72 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -57,6 +57,7 @@ gentbl(
     td_srcs = [
         ":tensorflow_ops_td_files",
     ],
+    test = True,
 )
 
 gentbl(
@@ -88,6 +89,7 @@ gentbl(
     td_srcs = [
         ":tensorflow_ops_td_files",
     ],
+    test = True,
 )
 
 gentbl(
@@ -112,6 +114,7 @@ gentbl(
         "@llvm-project//mlir:include/mlir/IR/OpBase.td",
         "@llvm-project//mlir:include/mlir/Dialect/StandardOps/IR/Ops.td",
     ],
+    test = True,
 )
 
 gentbl(
@@ -137,6 +140,7 @@ gentbl(
         "@llvm-project//mlir:include/mlir/Dialect/StandardOps/IR/Ops.td",
         "@llvm-project//mlir:include/mlir/IR/OpBase.td",
     ],
+    test = True,
 )
 
 gentbl(
@@ -161,6 +165,7 @@ gentbl(
         "@llvm-project//mlir:include/mlir/IR/OpBase.td",
         "@llvm-project//mlir:include/mlir/Dialect/StandardOps/IR/Ops.td",
     ],
+    test = True,
 )
 
 gentbl(

From fb5a3c3c6a4894a15bcc88ba9655c6afff8ff3f5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 15:30:49 -0700
Subject: [PATCH 0673/1390] Fix issues caught by
 -Wtautological-constant-compare

This fixes a check that is always evaluates to true.  The check expression is a
conditional operator that returns 1 or 2, both of which convert to true.  This
fix puts in the enum value that causes the check to pass.

PiperOrigin-RevId: 317391929
Change-Id: I12f3f9b06f494b4bcf0e4ff194d4dc4edafd52e2
---
 tensorflow/core/kernels/hexagon/graph_transferer.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/hexagon/graph_transferer.cc b/tensorflow/core/kernels/hexagon/graph_transferer.cc
index 7f15f3ab20d..9d6d0563b5f 100644
--- a/tensorflow/core/kernels/hexagon/graph_transferer.cc
+++ b/tensorflow/core/kernels/hexagon/graph_transferer.cc
@@ -667,7 +667,7 @@ void GraphTransferer::RegisterNodeWithPaddingAndStrides(
       << "Op " << node.type_string() << " not found in map(id = " << op_type_id
       << ")";
   // Safety check of padding id
-  CHECK(padding == Padding::VALID ? 1 : 2);
+  CHECK(padding == Padding::SAME);
   AppendNodeParamsWithIoParams(
       shape_refiner, node, node.name(), id, node.type_string(), op_type_id,
       static_cast(padding), node.num_inputs(), extra_inputs,

From ca7c0657c2e22343810957951cfb27aa5bf4c165 Mon Sep 17 00:00:00 2001
From: Yasir Modak <42785357+ymodak@users.noreply.github.com>
Date: Fri, 19 Jun 2020 15:43:04 -0700
Subject: [PATCH 0674/1390] update arg doc for sigmoid_cross_entropy

fixes #40593
---
 tensorflow/python/ops/losses/losses_impl.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/losses/losses_impl.py b/tensorflow/python/ops/losses/losses_impl.py
index 6a7b4b68420..2887b3d78f9 100644
--- a/tensorflow/python/ops/losses/losses_impl.py
+++ b/tensorflow/python/ops/losses/losses_impl.py
@@ -676,8 +676,9 @@ def sigmoid_cross_entropy(
       `{0, 1}`.
     logits: Float `[batch_size, num_classes]` logits outputs of the network.
     weights: Optional `Tensor` whose rank is either 0, or the same rank as
-      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
-      be either `1`, or the same as the corresponding `losses` dimension).
+    `multi_class_labels`, and must be broadcastable to `multi_class_labels` 
+    (i.e., all dimensions must be either `1`, or the same as the 
+    corresponding `losses` dimension).
     label_smoothing: If greater than `0` then smooth the labels.
     scope: The scope for the operations performed in computing the loss.
     loss_collection: collection to which the loss will be added.

From 18c43b7bd488f0216ce11127158df7084e26cdb4 Mon Sep 17 00:00:00 2001
From: Haoyu Zhang 
Date: Fri, 19 Jun 2020 15:43:22 -0700
Subject: [PATCH 0675/1390] Clear cancel callback when gRPC eager call returns
 with state.

PiperOrigin-RevId: 317393892
Change-Id: Ife800821494dd4cc2992eec9a5470d989596a6d7
---
 .../core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
index 924112e0d96..1d65f945f27 100644
--- a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
+++ b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
@@ -81,6 +81,7 @@ class GrpcEagerServiceImpl : public AsyncServiceInterface {
       local_impl_.RunComponentFunction(call_opts.get(), &call->request,
                                        &call->response,
                                        [call, call_opts](const Status& s) {
+                                         call->ClearCancelCallback();
                                          call->SendResponse(ToGrpcStatus(s));
                                        });
     });

From e3809400b295e8a66fd8e2b20a348356c78a0435 Mon Sep 17 00:00:00 2001
From: Kuangyuan Chen 
Date: Fri, 19 Jun 2020 15:51:19 -0700
Subject: [PATCH 0676/1390] Import initialization graph in SignatureDef
 SavedModels as an MLIR function in TF saved model dialect. -Mark the init
 function referenced by SessionInitializerOp as an exported function with the
 reserved name "__tf_saved_model_session_initializer". -Remove variable init
 logic in init function if it is already imported as a global tensor. -Add a
 canonicalizer to SessionInitializerOp to remove empty init functions. -Return
 error on SessionInitializerOp in IREE compiler pipeline.

PiperOrigin-RevId: 317395165
Change-Id: Idb7e54cac08add9fc8f2ccfbaf341135fdf59e3b
---
 tensorflow/compiler/mlir/tensorflow/BUILD     |   4 +-
 .../mlir/tensorflow/ir/tf_saved_model.cc      |  75 +++++-
 .../mlir/tensorflow/ir/tf_saved_model.h       |   4 +
 .../mlir/tensorflow/ir/tf_saved_model_ops.td  |  26 ++
 .../tests/tf_saved_model/common_v1.py         |  12 +-
 .../tests/tf_saved_model/hash_table_v1.py     |  92 +++++++
 .../tf_saved_model/remove_init_variable_v1.py |  74 ++++++
 .../tensorflow/tests/tf_saved_model_ops.mlir  |  16 ++
 .../tests/tf_saved_model_ops_invalid.mlir     |  68 +++++
 .../mlir/tensorflow/translate/import_model.cc | 235 ++++++++++++++----
 10 files changed, 549 insertions(+), 57 deletions(-)
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py

diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 11a455eec72..b159815d5eb 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -666,7 +666,9 @@ cc_library(
         ":tensorflow_types",
         ":translate_utils",
         "//tensorflow/cc/saved_model:bundle_v2",
+        "//tensorflow/cc/saved_model:constants",
         "//tensorflow/cc/saved_model:loader_lite",
+        "//tensorflow/cc/saved_model:loader_util",
         "//tensorflow/compiler/jit:shape_inference_helpers",
         "//tensorflow/compiler/mlir:op_or_arg_name_mapper",
         "//tensorflow/compiler/tf2xla:functionalize_control_flow",
@@ -678,6 +680,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/utils:transitive_fanin",
+        "//tensorflow/core/platform:protobuf_internal",
         "//tensorflow/core/platform:types",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
@@ -687,7 +690,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:StandardOps",
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index 7db0eed7713..ef55761686e 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "mlir/IR/Identifier.h"  // from @llvm-project
 #include "mlir/IR/Module.h"  // from @llvm-project
 #include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/IR/StandardTypes.h"  // from @llvm-project
 #include "mlir/IR/SymbolTable.h"  // from @llvm-project
 #include "mlir/IR/TypeUtilities.h"  // from @llvm-project
@@ -76,6 +77,23 @@ static LogicalResult Verify(GlobalTensorOp global_tensor) {
   return success();
 }
 
+static LogicalResult Verify(SessionInitializerOp session_initializer) {
+  mlir::SymbolTable symbol_table(
+      session_initializer.getParentOfType());
+
+  auto init_func_op =
+      symbol_table.lookup(session_initializer.initializer());
+  if (!init_func_op)
+    return session_initializer.emitOpError()
+           << "the initializer function does not exist";
+
+  if (!init_func_op.getType().getResults().empty())
+    return session_initializer.emitOpError()
+           << "the initializer function should have no output";
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc.inc"
 
@@ -214,12 +232,12 @@ static LogicalResult VerifySavedModelModule(
   for (auto func : module.getOps()) {
     const bool is_exported = IsExported(func);
 
-    if (is_exported && !func.isPublic()) {
+    if (is_exported && func.getVisibility() != FuncOp::Visibility::Public) {
       return func.emitError()
              << "exported function @" << func.getName() << " should be public";
     }
 
-    if (!is_exported && func.isPublic()) {
+    if (!is_exported && func.getVisibility() == FuncOp::Visibility::Public) {
       return func.emitError() << "non-exported function @" << func.getName()
                               << " should be private";
     }
@@ -232,6 +250,19 @@ static LogicalResult VerifySavedModelModule(
       }
     }
   }
+
+  auto session_initializers = module.getOps();
+  if (!session_initializers.empty() &&
+      !llvm::hasSingleElement(session_initializers)) {
+    return (*++session_initializers.begin()).emitError()
+           << "there must be no more than one session_initializer op";
+  }
+
+  auto is_init = [&session_initializers](mlir::FuncOp func) {
+    if (session_initializers.empty()) return false;
+    return (*session_initializers.begin()).initializer() == func.getName();
+  };
+
   SymbolTable symbol_table(module);
   auto symbol_uses = SymbolTable::getSymbolUses(&module.getBodyRegion());
   if (!symbol_uses.hasValue()) {
@@ -242,6 +273,12 @@ static LogicalResult VerifySavedModelModule(
     auto func = symbol_table.lookup(
         symbol_use.getSymbolRef().cast().getValue());
     if (func && IsExported(func)) {
+      // If it is an init function, then it can be used by the unique
+      // session_initializer op.
+      if (is_init(func) &&
+          llvm::isa(symbol_use.getUser()))
+        continue;
+
       return symbol_use.getUser()
           ->emitError("exported function cannot be internally referenced")
           .attachNote(func.getLoc())
@@ -361,5 +398,39 @@ GlobalTensorOp LookupBoundInput(FuncOp func, int arg_index,
   return symbol_table.lookup(attr.getValue());
 }
 
+SessionInitializerOp GetSessionInitializerOp(mlir::ModuleOp op) {
+  auto initializers = op.getOps();
+  if (initializers.empty()) return {};
+  return *initializers.begin();
+}
+
+class OptimizeSessionInitializerPattern
+    : public OpRewritePattern {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SessionInitializerOp op,
+                                PatternRewriter &rewriter) const override {
+    SymbolTable symbol_table(op.getParentOfType());
+    auto init_func_op = symbol_table.lookup(op.initializer());
+
+    // The init function can only be referenced from the SessionInitializerOp.
+    // And there is at most one SessionInitializerOp in the module. So both ops
+    // have no other uses and can be simply erased.
+    if (init_func_op.front().begin()->isKnownTerminator()) {
+      rewriter.eraseOp(init_func_op);
+      rewriter.eraseOp(op);
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+void SessionInitializerOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert(context);
+}
+
 }  // namespace tf_saved_model
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h
index 47ebb1a1be5..b6f8753cc51 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h
@@ -61,6 +61,10 @@ GlobalTensorOp LookupBoundInput(FuncOp func, int arg_index,
 // should have.
 Type GetBoundInputArgTypeFor(GlobalTensorOp global_tensor);
 
+// Returns the session initializer of this module if it exists. Returns null
+// otherwise.
+SessionInitializerOp GetSessionInitializerOp(mlir::ModuleOp op);
+
 }  // namespace tf_saved_model
 }  // namespace mlir
 
diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
index 4431a160edf..dc1210a4d2a 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td
@@ -128,4 +128,30 @@ def TfSavedModel_GlobalTensorOp : TfSavedModel_Op<"global_tensor"> {
   let verifier = [{ return Verify(*this); }];
 }
 
+def TfSavedModel_SessionInitializerOp: TfSavedModel_Op<"session_initializer"> {
+  let summary = "Initializes TensorFlow session state.";
+  let description = [{
+    The session initializer op marks a function that must be called by an
+    external agent exactly once to initialize TensorFlow session state, and this
+    must happen before any other exported functions are called. There must be no
+    more than one session initializer in a saved model.
+
+    The `initializer` represents the initialization function. The function have
+    no output and this function should be only called once.
+
+    This is used, for example, to initialize hash tables stored in resources and
+    accessed by resource name (rather than as resource handles or bound inputs
+    which is how `global_tensor`s are referenced)
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$initializer
+  );
+
+
+  let verifier = [{ return Verify(*this); }];
+
+  let hasCanonicalizer = 1;
+}
+
 #endif // SAVED_MODEL_DIALECT
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
index 7171f63bb05..5bfcfa5378a 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/common_v1.py
@@ -46,7 +46,10 @@ def set_tf_options():
 # This function needs to take a "create_module_fn", as opposed to just the
 # module itself, because the creation of the module has to be delayed until
 # after absl and tensorflow have run various initialization steps.
-def do_test(signature_def_map, show_debug_info=False):
+def do_test(signature_def_map,
+            init_op=None,
+            canonicalize=False,
+            show_debug_info=False):
   """Runs test.
 
   1. Performs absl and tf "main"-like initialization that must run before almost
@@ -61,6 +64,9 @@ def do_test(signature_def_map, show_debug_info=False):
   Args:
     signature_def_map: A map from string key to signature_def. The key will be
       used as function name in the resulting MLIR.
+    init_op: The initializer op for the saved model. If set, it will generate a
+      initializer graph in the resulting MLIR.
+    canonicalize: If true, canonicalizer will be run on the resulting MLIR.
     show_debug_info: If true, shows debug locations in the resulting MLIR.
   """
 
@@ -84,6 +90,7 @@ def do_test(signature_def_map, show_debug_info=False):
     builder.add_meta_graph_and_variables(
         sess, [tf.saved_model.tag_constants.SERVING],
         signature_def_map,
+        main_op=init_op,
         strip_default_attrs=True)
     builder.save()
 
@@ -97,6 +104,9 @@ def do_test(signature_def_map, show_debug_info=False):
     mlir = pywrap_mlir.experimental_run_pass_pipeline(mlir,
                                                       'tf-standard-pipeline',
                                                       show_debug_info)
+    if canonicalize:
+      mlir = pywrap_mlir.experimental_run_pass_pipeline(mlir, 'canonicalize',
+                                                        show_debug_info)
     print(mlir)
 
   app.run(app_main)
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
new file mode 100644
index 00000000000..16290455608
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py
@@ -0,0 +1,92 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# RUN: %p/hash_table_v1 | FileCheck %s
+
+# pylint: disable=missing-docstring,line-too-long
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v1 as tf
+from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
+
+# Verify that the tf.versions attribute exists. It is difficult to enforce
+# contents, since the version numbers change over time. The conversion logic
+# itself is verified in the common graphdef converter, so here just assert
+# it is being invoked.
+# CHECK: module
+# CHECK-SAME: tf.versions
+# CHECK-SAME: bad_consumers
+# CHECK-SAME: min_consumer
+# CHECK-SAME: producer
+
+# CHECK: "tf_saved_model.session_initializer"() {initializer = [[init:@.*]]} : () -> ()
+# CHECK: "tf_saved_model.global_tensor"()
+
+# CHECK:      func [[init]]
+# CHECK-NEXT: [[R5:%.*]] = "tf.Const"()
+# CHECK-NEXT: [[R6:%.*]] = "tf.Const"()
+# CHECK-NEXT: [[R7:%.*]] = "tf.HashTableV2"()
+# CHECK-SAME: shared_name = "[[hash_table:.*]]"
+# CHECK-NEXT: "tf.LookupTableImportV2"([[R7]], [[R5]], [[R6]])
+
+# CHECK:      func {{@[a-zA-Z_0-9]+}}(
+# CHECK-SAME: [[ARG0:%.*]]: tensor
+# CHECK-SAME: [[ARG1:%.*]]: tensor, value = {{.*}} : tensor<1x3xf32>} : () -> ()
+# CHECK-NOT: session_initializer
+
+# CHECK:      func {{@[a-zA-Z_0-9]+}}(
+# CHECK-SAME:   [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},
+# CHECK-SAME:   [[ARG1:%.*]]: tensor>> {tf_saved_model.bound_input = @[[VAR]]})
+# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
+# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]
+
+# CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor>>) -> tensor<1x3xf32>
+# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
+# CHECK-NEXT: return [[R1]] : tensor<3x3xf32>
+
+
+def Test():
+
+  x = tf.constant([[1.0], [1.0], [1.0]])
+  y = tf.compat.v1.get_variable(
+      name='y',
+      shape=(1, 3),
+      initializer=tf.random_normal_initializer(),
+      trainable=True)
+  r = tf.matmul(x, y)
+
+  tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
+  tensor_info_r = tf.compat.v1.saved_model.utils.build_tensor_info(r)
+
+  return {
+      'key': (tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
+          inputs={'x': tensor_info_x},
+          outputs={'r': tensor_info_r},
+          method_name='some_function'))
+  }
+
+
+if __name__ == '__main__':
+  common_v1.set_tf_options()
+  common_v1.do_test(
+      Test(), tf.initializers.global_variables(), canonicalize=True)
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
index e2dc5785cf4..26cdf025a10 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
@@ -2,6 +2,11 @@
 
 module attributes {tf_saved_model.semantics} {
 
+  // CHECK: tf_saved_model.session_initializer
+  "tf_saved_model.session_initializer"() {
+    initializer = @init
+  } : () -> ()
+
   // Representation for constants: (immutable) global tensor.
   // CHECK: tf_saved_model.global_tensor
   "tf_saved_model.global_tensor"() {
@@ -39,6 +44,17 @@ module attributes {tf_saved_model.semantics} {
     return
   }
 
+  // Representation for init functions
+  // CHECK: func @init
+  // CHECK-SAME: exported_names = ["__tf_saved_model_session_initializer"]
+  func @init(
+    %arg1: tensor>> {tf_saved_model.bound_input = @some_constant}
+  ) attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"]}
+  {
+    "tf.some_call"(%arg1) : (tensor>>) -> ()
+    return
+  }
+
 }
 
 // -----
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
index 7287fcf66c8..260174b184f 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
@@ -261,6 +261,39 @@ module attributes {tf_saved_model.semantics} {
 
 // -----
 
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function does not exist}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function should have no output}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  // expected-error@+1 {{there must be no more than one session_initializer op}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> tensor<1xf32> attributes {sym_visibility = "private"} {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// -----
+
 module attributes {tf_saved_model.semantics} {
 
   // expected-error@+1 {{exported function @f should be public}}
@@ -284,3 +317,38 @@ module attributes {tf_saved_model.semantics} {
   }
 
 }
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function does not exist}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  // expected-error@+1 {{the initializer function should have no output}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> (tensor<1xf32> {tf_saved_model.index_path = ["output"]})
+    attributes { tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"] } {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  // expected-error@+1 {{there must be no more than one session_initializer op}}
+  "tf_saved_model.session_initializer"() { initializer = @init } : () -> ()
+  func @init() -> (tensor<1xf32> {tf_saved_model.index_path = ["output"]})
+    attributes { tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"] } {
+    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
+    return %0 : tensor<1xf32>
+  }
+}
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
index 820d0ce31fb..fea809c0798 100644
--- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc
@@ -60,6 +60,8 @@ limitations under the License.
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Verifier.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
@@ -116,6 +118,7 @@ using mlir::NamedAttrList;
 using mlir::TensorType;
 using mlir::TF::VarHandleOp;
 using mlir::tf_saved_model::GlobalTensorOp;
+using mlir::tf_saved_model::SessionInitializerOp;
 using stream_executor::port::StatusOr;
 
 namespace {
@@ -2955,6 +2958,13 @@ void SortSavedModelModule(mlir::ModuleOp module) {
     named_global_tensor.global_tensor.getOperation()->moveBefore(
         &module.getBody()->front());
   }
+
+  auto initializers = module.getOps();
+  if (!initializers.empty()) {
+    (*initializers.begin())
+        .getOperation()
+        ->moveBefore(&module.getBody()->front());
+  }
 }
 
 Status CreateSavedModelIR(
@@ -3241,17 +3251,32 @@ class SavedModelSignatureDefImporter {
                                  absl::Span exported_names,
                                  mlir::MLIRContext* context)
       : bundle_(bundle),
+        flib_def_(OpRegistry::Global(), graph_def().library()),
+        debug_info_(),
         exported_names_(exported_names),
-        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {}
+        module_(mlir::ModuleOp::create(mlir::UnknownLoc::get(context))) {
+    // debug_info might not be loaded with loader_lite.
+    if (bundle_.debug_info != nullptr) debug_info_ = *bundle_.debug_info;
+  }
 
   // Converts the SavedModel to the SavedModel dialect. Creates an MLIR function
   // for each signature.
   StatusOr ConvertSignatures();
-  Status ConvertSignature(const GraphDef& graphdef,
-                          const std::string& sig_def_key,
-                          const SignatureDef& signature_def,
-                          const GraphDebugInfo& debug_info,
-                          const FunctionLibraryDefinition& flib_def);
+  Status ConvertSignature(const std::string& sig_def_key,
+                          const SignatureDef& signature_def);
+
+  // Converts the initialization graph in the SavedModel to an MLIR function.
+  Status ConvertInitializer();
+
+  // Converts a graph with feeds and fetches to an MLIR function.
+  StatusOr ConvertGraph(
+      const std::string& name,
+      const std::vector>& inputs,
+      const std::vector>& outputs,
+      const std::vector control_outputs);
+
+  // Coarsens the islands in `module_`.
+  Status CoarsenIslands();
 
   // Creates GlobalTensorOp for each variable and moves each VarHandle op to
   // the enclosing function's arguments.
@@ -3262,6 +3287,10 @@ class SavedModelSignatureDefImporter {
   // tensor's shape is used to provide the most accurate nested shape.
   void LiftVariable(VarHandleOp op, GlobalTensorOp global_tensor);
 
+  // Removes the variable and related ops in the init function if it is already
+  // imported as a global tensor.
+  void RemoveVariable(VarHandleOp op);
+
   using VarGlobalMap = llvm::MapVector<
       llvm::StringRef,
       std::pair>>;
@@ -3273,18 +3302,68 @@ class SavedModelSignatureDefImporter {
   GraphImportConfig::InputArrays ParseInputArrays(
       const std::vector>& inputs);
 
+  const GraphDef& graph_def() const {
+    return bundle_.meta_graph_def.graph_def();
+  }
+  const FunctionLibraryDefinition& flib_def() const { return flib_def_; }
+  const GraphDebugInfo& debug_info() const { return debug_info_; }
+
   const SavedModelBundle& bundle_;
+  FunctionLibraryDefinition flib_def_;
+  GraphDebugInfo debug_info_;
   absl::Span exported_names_;
   mlir::OwningModuleRef module_;
 };
 
+Status SavedModelSignatureDefImporter::ConvertInitializer() {
+  std::vector asset_file_defs;
+  TF_RETURN_IF_ERROR(
+      internal::GetAssetFileDefs(bundle_.meta_graph_def, &asset_file_defs));
+
+  if (!asset_file_defs.empty())
+    return errors::Unimplemented(
+        absl::StrCat("Assets are not supported in signaturedef importer"));
+
+  std::string init_node_name;
+  TF_RETURN_IF_ERROR(
+      internal::GetInitOp("", bundle_.meta_graph_def, &init_node_name));
+
+  if (init_node_name.empty()) return Status::OK();
+
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(init_node_name, {}, {}, {init_node_name}));
+
+  mlir::SymbolTable symbol_table(*sub_module);
+
+  auto init_func_op = symbol_table.lookup(init_node_name);
+
+  init_func_op.removeAttr("tf.entry_function");
+
+  mlir::OpBuilder builder(module_->getBodyRegion());
+
+  // Set the exported name of init function to an reserved name for
+  // tf_saved_model.
+  init_func_op.setAttr(
+      "tf_saved_model.exported_names",
+      builder.getStrArrayAttr({"__tf_saved_model_session_initializer"}));
+
+  builder.create(
+      module_->getLoc(), builder.getSymbolRefAttr(init_func_op.getName()));
+
+  // Move the converted functions to top level MLIR module.
+  auto* block = module_->getBody();
+  auto* sub_block = sub_module->getBody();
+  block->getOperations().splice(
+      mlir::Block::iterator(block->getTerminator()), sub_block->getOperations(),
+      sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
+
+  return Status::OK();
+}
+
 StatusOr
 SavedModelSignatureDefImporter::ConvertSignatures() {
   const auto& signatures = bundle_.GetSignatures();
-  const auto& graphdef = bundle_.meta_graph_def.graph_def();
-  PopulateTfVersions(module_.get(), graphdef.versions());
-
-  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graphdef.library());
+  PopulateTfVersions(module_.get(), graph_def().versions());
 
   // debug_info might not be loaded with loader_lite.
   GraphDebugInfo debug_info;
@@ -3307,23 +3386,49 @@ SavedModelSignatureDefImporter::ConvertSignatures() {
       continue;
     }
 
-    TF_RETURN_IF_ERROR(ConvertSignature(graphdef, sig_def_key, signature_def,
-                                        debug_info, flib_def));
+    TF_RETURN_IF_ERROR(ConvertSignature(sig_def_key, signature_def));
   }
-  TF_RETURN_IF_ERROR(LiftVariables());
+
+  TF_RETURN_IF_ERROR(ConvertInitializer());
 
   mlir::OpBuilder builder(module_->getBodyRegion());
   module_->setAttr("tf_saved_model.semantics", builder.getUnitAttr());
+
+  TF_RETURN_IF_ERROR(CoarsenIslands());
+  TF_RETURN_IF_ERROR(LiftVariables());
+
   SortSavedModelModule(*module_);
   MarkSavedModelFunctionVisibility(*module_);
 
   return std::move(module_);
 }
 
+StatusOr SavedModelSignatureDefImporter::ConvertGraph(
+    const std::string& name,
+    const std::vector>& inputs,
+    const std::vector>& outputs,
+    const std::vector control_outputs) {
+  GraphImportConfig specs;
+  specs.prune_unused_nodes = true;
+  specs.inputs = ParseInputArrays(inputs);
+  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
+  specs.control_outputs = control_outputs;
+
+  // Convert sub-graphdef to sub-graph.
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.add_default_attributes = true;
+  Graph graph(OpRegistry::Global());
+
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, graph_def(), &graph));
+
+  // Convert sub-graph to MLIR module.true
+  return GraphDefImporter::Convert(module_->getContext(), graph, debug_info(),
+                                   flib_def(), specs, name);
+}
+
 Status SavedModelSignatureDefImporter::ConvertSignature(
-    const GraphDef& graphdef, const std::string& sig_def_key,
-    const SignatureDef& signature_def, const GraphDebugInfo& debug_info,
-    const FunctionLibraryDefinition& flib_def) {
+    const std::string& sig_def_key, const SignatureDef& signature_def) {
   // Create local vectors for the input and output and sort them to be
   // deterministic. We don't want anyone to really depend on the order, client
   // should lookup argument/result mapping by attribute name.
@@ -3339,34 +3444,9 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
     return lhs.first.size() < rhs.first.size() || lhs.first > rhs.first;
   });
 
-  GraphImportConfig specs;
-  specs.prune_unused_nodes = true;
-  specs.inputs = ParseInputArrays(inputs);
-  for (auto& output : outputs) specs.outputs.push_back(output.second.name());
-
-  // Remove unused nodes and create sub-graphdef.
-  GraphDef sub_graph_def;
-  TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(
-      graphdef, &sub_graph_def,
-      /*terminal_nodes=*/{specs.outputs.begin(), specs.outputs.end()}));
-
-  // Set the function library definitions in the pruned graphdef.
-  *sub_graph_def.mutable_library() = flib_def.ToProto();
-
-  // Convert sub-graphdef to sub-graph.
-  GraphConstructorOptions options;
-  options.allow_internal_ops = true;
-  options.add_default_attributes = true;
-  Graph sub_graph(OpRegistry::Global());
-
-  TF_RETURN_IF_ERROR(
-      ConvertGraphDefToGraph(options, sub_graph_def, &sub_graph));
-
   // Convert sub-graph to MLIR module.
-  TF_ASSIGN_OR_RETURN(
-      auto sub_module,
-      GraphDefImporter::Convert(module_->getContext(), sub_graph, debug_info,
-                                flib_def, specs, sig_def_key));
+  TF_ASSIGN_OR_RETURN(auto sub_module,
+                      ConvertGraph(sig_def_key, inputs, outputs, {}));
   mlir::OpBuilder builder(sub_module->getBodyRegion());
 
   // Find the FuncOp which corresponds to current SignatureDef.
@@ -3399,16 +3479,28 @@ Status SavedModelSignatureDefImporter::ConvertSignature(
       sub_block->begin(), mlir::Block::iterator(sub_block->getTerminator()));
 
   return Status::OK();
-}
+}  // namespace
 
 Status SavedModelSignatureDefImporter::LiftVariables() {
   VarGlobalMap var_globals;
+  llvm::SmallVector init_vars;
 
-  auto walker = [&var_globals](mlir::Operation* op) {
-    if (auto var_handle_op = llvm::dyn_cast(op))
-      var_globals[var_handle_op.shared_name()].second.push_back(var_handle_op);
-    else if (op->getName().getStringRef() == "tf.VariableV2")
+  auto session_initializer =
+      mlir::tf_saved_model::GetSessionInitializerOp(*module_);
+
+  auto walker = [&var_globals, &init_vars,
+                 &session_initializer](mlir::Operation* op) {
+    if (auto var_handle_op = llvm::dyn_cast(op)) {
+      if (session_initializer &&
+          session_initializer.initializer() ==
+              var_handle_op.getParentOfType().getName())
+        init_vars.push_back(var_handle_op);
+      else
+        var_globals[var_handle_op.shared_name()].second.push_back(
+            var_handle_op);
+    } else if (op->getName().getStringRef() == "tf.VariableV2") {
       return mlir::WalkResult::interrupt();
+    }
     return mlir::WalkResult::advance();
   };
   bool contains_ref_variable = module_->walk(walker).wasInterrupted();
@@ -3425,9 +3517,51 @@ Status SavedModelSignatureDefImporter::LiftVariables() {
     for (VarHandleOp var_handle : it.second.second)
       LiftVariable(var_handle, it.second.first);
 
+  for (auto op : init_vars) RemoveVariable(op);
+
   return Status::OK();
 }
 
+Status SavedModelSignatureDefImporter::CoarsenIslands() {
+  mlir::StatusScopedDiagnosticHandler diag_handler(module_->getContext());
+
+  mlir::PassManager pm(module_->getContext());
+  pm.addNestedPass(
+      mlir::tf_executor::CreateTFExecutorIslandCoarseningPass());
+  if (mlir::failed(pm.run(*module_)))
+    return diag_handler.Combine(
+        errors::Internal("failed to coarsening islands."));
+
+  return Status::OK();
+}
+
+void SavedModelSignatureDefImporter::RemoveVariable(VarHandleOp op) {
+  llvm::SmallVector work_list;
+  work_list.push_back(op);
+  while (!work_list.empty()) {
+    auto* op = work_list.back();
+    work_list.pop_back();
+
+    for (mlir::Value res : op->getResults()) {
+      for (mlir::Operation* user : res.getUsers()) {
+        work_list.push_back(user);
+      }
+    }
+
+    for (auto& use : op->getOpOperands()) {
+      if (mlir::Value value = use.get()) {
+        mlir::Operation* def = value.getDefiningOp();
+        work_list.push_back(def);
+      }
+    }
+
+    op->dropAllReferences();
+    op->dropAllDefinedValueUses();
+
+    op->erase();
+  }
+}
+
 void SavedModelSignatureDefImporter::LiftVariable(
     VarHandleOp op, GlobalTensorOp global_tensor) {
   mlir::OpBuilder builder(&module_->getBodyRegion());
@@ -3460,12 +3594,7 @@ void SavedModelSignatureDefImporter::LiftVariable(
   // Add the newly added function param to entry block's arguments.
   auto new_value = func_op.front().addArgument(resource_type);
 
-  // Remove the VarHandleOp also updating the containing island's return type.
-  DCHECK(llvm::isa(op.getParentOp()));
-  DCHECK(llvm::cast(op.getParentOp())
-             .WrapsSingleOp());
   op.getOperation()->replaceAllUsesWith(llvm::ArrayRef(new_value));
-  op.getParentOp()->getResult(0).setType(resource_type);
   op.getOperation()->erase();
 }
 

From 1823f877359bb138c57a005c30aba8832dfa79fb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 15:55:58 -0700
Subject: [PATCH 0677/1390] Fix issues with `TypeIndex` on MacOS, i.e. hash on
 the type name where available since this otherwise causes problems when
 loading different shared libraries with `RTLD_LOCAL`.

PiperOrigin-RevId: 317395983
Change-Id: I14b3add5fa19725b2150b68813364d16b8320130
---
 tensorflow/core/framework/type_index.h | 38 ++++++++++++++++++++------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/framework/type_index.h b/tensorflow/core/framework/type_index.h
index fd27d8bcb35..fcf68677a12 100644
--- a/tensorflow/core/framework/type_index.h
+++ b/tensorflow/core/framework/type_index.h
@@ -24,6 +24,10 @@ limitations under the License.
 
 #include "tensorflow/core/platform/types.h"
 
+#if defined(MACOS) || defined(TARGET_OS_MAC)
+#include "tensorflow/core/platform/hash.h"
+#endif  // defined(MACOS) || defined(TARGET_OS_MAC)
+
 namespace tensorflow {
 
 // On some platforms, we would like to avoid using RTTI in order to have smaller
@@ -53,10 +57,33 @@ class TypeIndex {
 
   // Returns a TypeIndex object that corresponds to a typename.
   template 
-  static TypeIndex Make(const char* name) {
+  static TypeIndex Make() {
     static bool hash_bit[1];
+
+#if defined(__GXX_RTTI) || defined(_CPPRTTI)
+
+#if defined(MACOS) || defined(TARGET_OS_MAC)
+    // Use a hash based on the type name to avoid issues due to RTLD_LOCAL on
+    // MacOS (b/156979412).
+    return TypeIndex(Hash64(typeid(T).name()), typeid(T).name());
+#else
+    // Use the real type name if we have RTTI.
     return TypeIndex(static_cast(reinterpret_cast(hash_bit)),
-                     name);
+                     typeid(T).name());
+#endif  // defined(MACOS) || defined(TARGET_OS_MAC)
+
+#else
+#if defined(MACOS) || defined(TARGET_OS_MAC)
+    // Warn MacOS users that not using RTTI can cause problems (b/156979412).
+#warning \
+    "Compiling with RTTI disabled on MacOS can cause problems when comparing " \
+    "types across shared libraries."
+#endif  // defined(MACOS) || defined(TARGET_OS_MAC)
+
+    // No type names available.
+    return TypeIndex(static_cast(reinterpret_cast(hash_bit)),
+                     "[RTTI disabled]");
+#endif  // __GXX_RTTI
   }
 
  private:
@@ -70,12 +97,7 @@ class TypeIndex {
 
 template 
 inline TypeIndex MakeTypeIndex() {
-#if defined(__GXX_RTTI) || defined(_CPPRTTI)
-  // Use the real type name if we have RTTI.
-  return TypeIndex::Make(typeid(T).name());
-#else
-  return TypeIndex::Make("[RTTI disabled]");
-#endif  // __GXX_RTTI
+  return TypeIndex::Make();
 }
 
 }  // namespace tensorflow

From 54a7bd1d738e0f2b8b9b729f48f8e90f2f39b3c6 Mon Sep 17 00:00:00 2001
From: Nathan Luehr 
Date: Fri, 19 Jun 2020 18:02:40 -0500
Subject: [PATCH 0678/1390] Remove unnecessary tf32_execution_allowed check

---
 tensorflow/stream_executor/cuda/cuda_blas.cc | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_blas.cc b/tensorflow/stream_executor/cuda/cuda_blas.cc
index e9e3635d8c1..e387690da26 100644
--- a/tensorflow/stream_executor/cuda/cuda_blas.cc
+++ b/tensorflow/stream_executor/cuda/cuda_blas.cc
@@ -402,14 +402,6 @@ bool CUDABlas::DoBlasInternalImpl(FuncT cublas_func, Stream *stream,
                                            : CUBLAS_POINTER_MODE_DEVICE)) {
     return false;
   }
-#if CUBLAS_VER_MAJOR >= 11
-  ScopedCublasMathMode math_mode{blas_};
-  if (!tensorflow::tf32_execution_allowed()) {
-    if (!math_mode.Init(CUBLAS_DEFAULT_MATH)) {
-      return false;
-    }
-  }
-#endif
   cublasStatus_t ret = cublas_func(blas_, args...);
   if ((err_on_failure || VLOG_IS_ON(3)) && ret != CUBLAS_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to run cuBLAS routine: " << ToString(ret);

From 866eb4828af9aa63e75b51fa05e84cc6b3176e5b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 15:58:05 -0700
Subject: [PATCH 0679/1390] tf.numpy avoid using typing module given Kokoro
 breakage.

PiperOrigin-RevId: 317396321
Change-Id: Ifa05dcbbdd8998708e52ba52cdb03e33e4952f21
---
 tensorflow/python/ops/numpy_ops/np_array_ops.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 47236d45561..74f1bc8cbef 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -21,7 +21,6 @@ from __future__ import print_function
 
 import math
 import numbers
-from typing import Sequence
 import numpy as np
 import six
 
@@ -1684,10 +1683,10 @@ def _slice_helper(tensor, slice_spec):
 def _as_spec_tuple(slice_spec):
   """Convert slice_spec to tuple."""
   if isinstance(slice_spec,
-                Sequence) and not isinstance(slice_spec, np.ndarray):
+                (list, tuple)) and not isinstance(slice_spec, np.ndarray):
     is_index = True
     for s in slice_spec:
-      if s is None or s is Ellipsis or isinstance(s, (Sequence, slice)):
+      if s is None or s is Ellipsis or isinstance(s, (list, tuple, slice)):
         is_index = False
         break
       elif isinstance(s, (np_arrays.ndarray, np.ndarray)) and s.ndim != 0:

From fac99746cb7382af83fc9922c345a1bd68caa516 Mon Sep 17 00:00:00 2001
From: Chenkai Kuang 
Date: Fri, 19 Jun 2020 16:08:23 -0700
Subject: [PATCH 0680/1390] Override "map_resources" in AggregatingVariable.

PiperOrigin-RevId: 317398115
Change-Id: Ic57325aacf0fb45a66a469428f544fb94f3cc031
---
 tensorflow/python/distribute/ps_values.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/python/distribute/ps_values.py b/tensorflow/python/distribute/ps_values.py
index 37cd6e12d90..5fb2d42b626 100644
--- a/tensorflow/python/distribute/ps_values.py
+++ b/tensorflow/python/distribute/ps_values.py
@@ -166,6 +166,14 @@ class AggregatingVariable(variables_lib.Variable, core.Tensor):
   def _gather_saveables_for_checkpoint(self):
     return {trackable.VARIABLE_VALUE_KEY: self._v}
 
+  def _map_resources(self):
+    """For implementing `Trackable`."""
+    # By delegating this method to the wrapped variable, SavedModel with
+    # AggregatingVariable are identical to SavedModel with normal variables.
+    obj_map, resource_map = self._v._map_resources()  # pylint:disable=protected-access
+    obj_map[self] = obj_map[self._v]
+    return obj_map, resource_map
+
   # pylint: disable=multiple-statements
   def __add__(self, o):
     return self._v + o

From a70ad66828f198f4f495c3f29445b396045b6c3f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Fri, 19 Jun 2020 16:23:33 -0700
Subject: [PATCH 0681/1390] GitHub Issue #40462

PiperOrigin-RevId: 317400375
Change-Id: I13891b5c2f41ac97674ddfda679c4273d53b25ef
---
 tensorflow/python/keras/preprocessing/image_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/preprocessing/image_dataset.py b/tensorflow/python/keras/preprocessing/image_dataset.py
index d287c4ef372..6e78ac3c03a 100644
--- a/tensorflow/python/keras/preprocessing/image_dataset.py
+++ b/tensorflow/python/keras/preprocessing/image_dataset.py
@@ -147,7 +147,7 @@ def image_dataset_from_directory(directory,
           'directory. If you wish to infer the labels from the subdirectory '
           'names in the target directory, pass `labels="inferred"`. '
           'If you wish to get a dataset that only contains images '
-          '(no labels), pass `labels=None`.')
+          '(no labels), pass `labels_mode=None`.')
     if class_names:
       raise ValueError('You can only pass `class_names` if the labels are '
                        'inferred from the subdirectory names in the target '

From 62683d061cf31d05588a94cc333b53542cea9568 Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Fri, 19 Jun 2020 16:25:08 -0700
Subject: [PATCH 0682/1390] [XLA] Rollback of rollback of "Implement
 LocalClient::Run which supports buffer donation"

PiperOrigin-RevId: 317400695
Change-Id: I56f1f8df347d5a3b2bad9526c7315c63ad6ddadb
---
 .../compiler/xla/client/local_client.cc       | 26 ++++++++++++-------
 tensorflow/compiler/xla/client/local_client.h | 21 +++++++++++++++
 .../tests/multiple_devices_on_host_test.cc    |  3 ++-
 tensorflow/compiler/xla/tests/while_test.cc   |  6 +++--
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index afe115deda8..aa252067e19 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -176,15 +176,23 @@ StatusOr LocalExecutable::Run(
   for (const ShapedBuffer* const arg : arguments) {
     argument_shapes.push_back(&arg->on_host_shape());
   }
-  TF_ASSIGN_OR_RETURN(auto options_and_stream,
-                      RunHelper(argument_shapes, run_options));
-  ExecutableRunOptions options = options_and_stream.first.run_options();
-  options.set_device_ordinal(-1);
-  auto result = RunAsync(arguments, options);
-  Status block_status = options.stream()->BlockHostUntilDone();
-  TF_RETURN_IF_ERROR(result.status());
-  TF_RETURN_IF_ERROR(block_status);
-  return result;
+  return AsyncCallAndBlockHostUntilDone(
+      argument_shapes, run_options, [&](const ExecutableRunOptions& options) {
+        return RunAsync(arguments, options);
+      });
+}
+
+StatusOr LocalExecutable::Run(
+    std::vector arguments, ExecutableRunOptions run_options) {
+  std::vector argument_shapes;
+  argument_shapes.reserve(arguments.size());
+  for (const ExecutionInput& arg : arguments) {
+    argument_shapes.push_back(&arg.shape());
+  }
+  return AsyncCallAndBlockHostUntilDone(
+      argument_shapes, run_options, [&](const ExecutableRunOptions& options) {
+        return RunAsync(argument_shapes, std::move(arguments), options);
+      });
 }
 
 static std::shared_ptr DumpArguments(
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index 7cdeb9dcbf6..3241ac73d54 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -51,6 +51,11 @@ class LocalExecutable {
       const absl::Span arguments,
       ExecutableRunOptions run_options);
 
+  // Similar to Run(), but allows for donating argument buffers to the
+  // executable.
+  StatusOr Run(std::vector arguments,
+                                ExecutableRunOptions run_options);
+
   // Similar to Run(), but need not block the host waiting for the computation
   // to complete before returning.
   StatusOr RunAsync(
@@ -90,6 +95,22 @@ class LocalExecutable {
   // Backend::devices_equivalent).
   int build_device_ordinal() const { return build_options_.device_ordinal(); }
 
+  template 
+  StatusOr AsyncCallAndBlockHostUntilDone(
+      absl::Span argument_shapes,
+      const ExecutableRunOptions& run_options,
+      std::function(const ExecutableRunOptions&)> async_callback) {
+    TF_ASSIGN_OR_RETURN(auto options_and_stream,
+                        RunHelper(argument_shapes, run_options));
+    ExecutableRunOptions options = options_and_stream.first.run_options();
+    options.set_device_ordinal(-1);
+    StatusOr result = async_callback(options);
+    Status block_status = options.stream()->BlockHostUntilDone();
+    TF_RETURN_IF_ERROR(result.status());
+    TF_RETURN_IF_ERROR(block_status);
+    return result;
+  }
+
   // Compiled computation.
   std::unique_ptr executable_;
 
diff --git a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
index 2b19aaded9c..2231fc6feab 100644
--- a/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
+++ b/tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
@@ -45,7 +45,8 @@ void CompileAndExecute(
       xla::ClientLibrary::GetXlaService(client->platform())
           ->backend()
           .memory_allocator());
-  StatusOr result = executable->Run({}, execute_options);
+  StatusOr result =
+      executable->Run(absl::Span(), execute_options);
   {
     absl::MutexLock lock(results_mutex);
     results->emplace_back(device_ordinal, std::move(result));
diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc
index d575bbb1f3e..8e8c3605cc7 100644
--- a/tensorflow/compiler/xla/tests/while_test.cc
+++ b/tensorflow/compiler/xla/tests/while_test.cc
@@ -1324,14 +1324,16 @@ void BM_WhileLoop(int num_iters) {
   options.set_allocator(&allocator);
   const int kWarmups = 2;
   for (int i = 0; i < kWarmups; ++i) {
-    auto result = executable->Run({}, options);
+    auto result =
+        executable->Run(absl::Span(), options);
     ASSERT_TRUE(result.ok());
   }
 
   // Run benchmark.
   tensorflow::testing::StartTiming();
   for (int i = 0; i < num_iters; ++i) {
-    auto result = executable->Run({}, options);
+    auto result =
+        executable->Run(absl::Span(), options);
     ASSERT_TRUE(result.ok());
   }
 }

From 3c9a2f200e7c439ce5e3345bdd9055b88b9777ca Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Fri, 19 Jun 2020 16:34:05 -0700
Subject: [PATCH 0683/1390] [TF-numpy] Changed all callsites of
 np_doc/np_doc_only to use string names, to avoid blocking imports when some
 numpy symbols are missing (e.g. because of an older version of numpy).

Also moved `np_fun_name` to be the first positional argument of np_doc/np_doc_only.

PiperOrigin-RevId: 317402093
Change-Id: I8b3ba54909e9507c7ab062f5bf5d7f13ad549317
---
 .../python/ops/numpy_ops/np_array_ops.py      | 156 ++++++------
 .../python/ops/numpy_ops/np_math_ops.py       | 228 +++++++++---------
 tensorflow/python/ops/numpy_ops/np_utils.py   |  91 ++++---
 .../python/ops/numpy_ops/np_utils_test.py     |  10 +-
 4 files changed, 254 insertions(+), 231 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 74f1bc8cbef..a87c72ed763 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -41,17 +41,17 @@ from tensorflow.python.ops.numpy_ops import np_utils
 from tensorflow.python.util import nest
 
 
-@np_utils.np_doc(np.empty)
+@np_utils.np_doc('empty')
 def empty(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return zeros(shape, dtype)
 
 
-@np_utils.np_doc(np.empty_like)
+@np_utils.np_doc('empty_like')
 def empty_like(a, dtype=None):
   return zeros_like(a, dtype)
 
 
-@np_utils.np_doc(np.zeros)
+@np_utils.np_doc('zeros')
 def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
   dtype = (
       np_utils.result_type(dtype) if dtype else np_dtypes.default_float_type())
@@ -60,7 +60,7 @@ def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.zeros(shape, dtype=dtype))
 
 
-@np_utils.np_doc(np.zeros_like)
+@np_utils.np_doc('zeros_like')
 def zeros_like(a, dtype=None):  # pylint: disable=missing-docstring
   if isinstance(a, np_arrays.ndarray):
     a = a.data
@@ -75,7 +75,7 @@ def zeros_like(a, dtype=None):  # pylint: disable=missing-docstring
   return np_arrays.tensor_to_ndarray(array_ops.zeros_like(a, dtype))
 
 
-@np_utils.np_doc(np.ones)
+@np_utils.np_doc('ones')
 def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
   if dtype:
     dtype = np_utils.result_type(dtype)
@@ -84,7 +84,7 @@ def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
   return np_arrays.tensor_to_ndarray(array_ops.ones(shape, dtype=dtype))
 
 
-@np_utils.np_doc(np.ones_like)
+@np_utils.np_doc('ones_like')
 def ones_like(a, dtype=None):
   if isinstance(a, np_arrays.ndarray):
     a = a.data
@@ -95,7 +95,7 @@ def ones_like(a, dtype=None):
   return np_arrays.tensor_to_ndarray(array_ops.ones_like(a, dtype))
 
 
-@np_utils.np_doc(np.eye)
+@np_utils.np_doc('eye')
 def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-docstring
   if dtype:
     dtype = np_utils.result_type(dtype)
@@ -127,12 +127,12 @@ def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-do
       array_ops.matrix_diag(diagonal=diagonal_, num_rows=N, num_cols=M, k=k))
 
 
-@np_utils.np_doc(np.identity)
+@np_utils.np_doc('identity')
 def identity(n, dtype=float):
   return eye(N=n, M=n, dtype=dtype)
 
 
-@np_utils.np_doc(np.full)
+@np_utils.np_doc('full')
 def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
   if not isinstance(shape, np_arrays.ndarray):
     shape = asarray(np_arrays.convert_to_tensor(shape, dtype_hint=np.int32))
@@ -144,7 +144,7 @@ def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
 
 # Using doc only here since np full_like signature doesn't seem to have the
 # shape argument (even though it exists in the documentation online).
-@np_utils.np_doc_only(np.full_like)
+@np_utils.np_doc_only('full_like')
 def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  # pylint: disable=missing-docstring,redefined-outer-name
   """order, subok and shape arguments mustn't be changed."""
   if order != 'K':
@@ -163,7 +163,7 @@ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  #
 
 # TODO(wangpeng): investigate whether we can make `copy` default to False.
 # pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
-@np_utils.np_doc_only(np.array)
+@np_utils.np_doc_only('array')
 def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
   """Since Tensors are immutable, a copy is made only if val is placed on a
 
@@ -224,7 +224,7 @@ def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-out
 # pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
 
 
-@np_utils.np_doc(np.asarray)
+@np_utils.np_doc('asarray')
 def asarray(a, dtype=None):
   if dtype:
     dtype = np_utils.result_type(dtype)
@@ -233,18 +233,18 @@ def asarray(a, dtype=None):
   return array(a, dtype, copy=False)
 
 
-@np_utils.np_doc(np.asanyarray)
+@np_utils.np_doc('asanyarray')
 def asanyarray(a, dtype=None):
   return asarray(a, dtype)
 
 
-@np_utils.np_doc(np.ascontiguousarray)
+@np_utils.np_doc('ascontiguousarray')
 def ascontiguousarray(a, dtype=None):
   return array(a, dtype, ndmin=1)
 
 
 # Numerical ranges.
-@np_utils.np_doc(np.arange)
+@np_utils.np_doc('arange')
 def arange(start, stop=None, step=1, dtype=None):
   """Returns `step`-separated values in the range [start, stop).
 
@@ -286,7 +286,7 @@ def arange(start, stop=None, step=1, dtype=None):
 
 
 # Building matrices.
-@np_utils.np_doc(np.diag)
+@np_utils.np_doc('diag')
 def diag(v, k=0):  # pylint: disable=missing-docstring
   """Raises an error if input is not 1- or 2-d."""
   v = asarray(v).data
@@ -321,7 +321,7 @@ def diag(v, k=0):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(result)
 
 
-@np_utils.np_doc(np.diagonal)
+@np_utils.np_doc('diagonal')
 def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstring
   a = asarray(a).data
 
@@ -352,7 +352,7 @@ def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstrin
   return a
 
 
-@np_utils.np_doc(np.diagflat)
+@np_utils.np_doc('diagflat')
 def diagflat(v, k=0):
   v = asarray(v)
   return diag(array_ops.reshape(v.data, [-1]), k)
@@ -363,21 +363,21 @@ def _promote_dtype(*arrays):
   return [asarray(a, dtype=dtype) for a in arrays]
 
 
-@np_utils.np_doc(np.all)
+@np_utils.np_doc('all')
 def all(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_all(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
-@np_utils.np_doc(np.any)
+@np_utils.np_doc('any')
 def any(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
   a = asarray(a, dtype=bool)
   return np_utils.tensor_to_ndarray(
       math_ops.reduce_any(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
-@np_utils.np_doc(np.compress)
+@np_utils.np_doc('compress')
 def compress(condition, a, axis=None):  # pylint: disable=redefined-outer-name,missing-function-docstring
   condition = asarray(condition, dtype=bool)
   a = asarray(a)
@@ -408,7 +408,7 @@ def compress(condition, a, axis=None):  # pylint: disable=redefined-outer-name,m
       array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis))
 
 
-@np_utils.np_doc(np.copy)
+@np_utils.np_doc('copy')
 def copy(a):
   return array(a, copy=True)
 
@@ -424,7 +424,7 @@ def _maybe_promote_to_int(a):
   return a
 
 
-@np_utils.np_doc(np.cumprod)
+@np_utils.np_doc('cumprod')
 def cumprod(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   a = asarray(a, dtype=dtype)
 
@@ -440,7 +440,7 @@ def cumprod(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(math_ops.cumprod(a.data, axis))
 
 
-@np_utils.np_doc(np.cumsum)
+@np_utils.np_doc('cumsum')
 def cumsum(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   a = asarray(a, dtype=dtype)
 
@@ -456,7 +456,7 @@ def cumsum(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(math_ops.cumsum(a.data, axis))
 
 
-@np_utils.np_doc(np.imag)
+@np_utils.np_doc('imag')
 def imag(a):
   a = asarray(a)
   # TODO(srbs): np.imag returns a scalar if a is a scalar, whereas we always
@@ -536,7 +536,7 @@ def _reduce(tf_fn,
       tf_fn(input_tensor=a.data, axis=axis, keepdims=keepdims))
 
 
-@np_utils.np_doc(np.sum)
+@np_utils.np_doc('sum')
 def sum(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=redefined-builtin
   return _reduce(
       math_ops.reduce_sum,
@@ -547,7 +547,7 @@ def sum(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=redefined-b
       tf_bool_fn=math_ops.reduce_any)
 
 
-@np_utils.np_doc(np.prod)
+@np_utils.np_doc('prod')
 def prod(a, axis=None, dtype=None, keepdims=None):
   return _reduce(
       math_ops.reduce_prod,
@@ -558,7 +558,7 @@ def prod(a, axis=None, dtype=None, keepdims=None):
       tf_bool_fn=math_ops.reduce_all)
 
 
-@np_utils.np_doc(np.mean)
+@np_utils.np_doc('mean')
 def mean(a, axis=None, dtype=None, keepdims=None):
   return _reduce(
       math_ops.reduce_mean,
@@ -569,7 +569,7 @@ def mean(a, axis=None, dtype=None, keepdims=None):
       promote_int=_TO_FLOAT)
 
 
-@np_utils.np_doc(np.amax)
+@np_utils.np_doc('amax')
 def amax(a, axis=None, keepdims=None):
   return _reduce(
       math_ops.reduce_max,
@@ -582,7 +582,7 @@ def amax(a, axis=None, keepdims=None):
       preserve_bool=True)
 
 
-@np_utils.np_doc(np.amin)
+@np_utils.np_doc('amin')
 def amin(a, axis=None, keepdims=None):
   return _reduce(
       math_ops.reduce_min,
@@ -595,7 +595,7 @@ def amin(a, axis=None, keepdims=None):
       preserve_bool=True)
 
 
-@np_utils.np_doc(np.var)
+@np_utils.np_doc('var')
 def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: disable=missing-docstring
   if dtype:
     working_dtype = np_utils.result_type(a, dtype)
@@ -642,7 +642,7 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: d
   return np_utils.tensor_to_ndarray(result)
 
 
-@np_utils.np_doc(np.std)
+@np_utils.np_doc('std')
 def std(a, axis=None, keepdims=None):  # pylint: disable=missing-function-docstring
   return _reduce(
       math_ops.reduce_std,
@@ -653,7 +653,7 @@ def std(a, axis=None, keepdims=None):  # pylint: disable=missing-function-docstr
       promote_int=_TO_FLOAT)
 
 
-@np_utils.np_doc(np.ravel)
+@np_utils.np_doc('ravel')
 def ravel(a):  # pylint: disable=missing-docstring
   a = asarray(a)
   out = np_utils.cond(
@@ -665,7 +665,7 @@ def ravel(a):  # pylint: disable=missing-docstring
 setattr(np_arrays.ndarray, 'ravel', ravel)
 
 
-@np_utils.np_doc(np.real)
+@np_utils.np_doc('real')
 def real(val):
   val = asarray(val)
   # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always
@@ -673,7 +673,7 @@ def real(val):
   return np_utils.tensor_to_ndarray(math_ops.real(val.data))
 
 
-@np_utils.np_doc(np.repeat)
+@np_utils.np_doc('repeat')
 def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
   a = asarray(a).data
   original_shape = a._shape_as_list()  # pylint: disable=protected-access
@@ -704,7 +704,7 @@ def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(result)
 
 
-@np_utils.np_doc(np.around)
+@np_utils.np_doc('around')
 def around(a, decimals=0):  # pylint: disable=missing-docstring
   a = asarray(a)
   dtype = a.dtype
@@ -726,7 +726,7 @@ def around(a, decimals=0):  # pylint: disable=missing-docstring
 setattr(np_arrays.ndarray, '__round__', around)
 
 
-@np_utils.np_doc(np.reshape)
+@np_utils.np_doc('reshape')
 def reshape(a, newshape, order='C'):
   """order argument can only b 'C' or 'F'."""
   if order not in {'C', 'F'}:
@@ -758,19 +758,19 @@ def _reshape_method_wrapper(a, *newshape, **kwargs):
   return reshape(a, newshape, order=order)
 
 
-@np_utils.np_doc(np.expand_dims)
+@np_utils.np_doc('expand_dims')
 def expand_dims(a, axis):
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.expand_dims(a.data, axis=axis))
 
 
-@np_utils.np_doc(np.squeeze)
+@np_utils.np_doc('squeeze')
 def squeeze(a, axis=None):
   a = asarray(a)
   return np_utils.tensor_to_ndarray(array_ops.squeeze(a, axis))
 
 
-@np_utils.np_doc(np.transpose)
+@np_utils.np_doc('transpose')
 def transpose(a, axes=None):
   a = asarray(a)
   if axes is not None:
@@ -778,7 +778,7 @@ def transpose(a, axes=None):
   return np_utils.tensor_to_ndarray(array_ops.transpose(a=a.data, perm=axes))
 
 
-@np_utils.np_doc(np.swapaxes)
+@np_utils.np_doc('swapaxes')
 def swapaxes(a, axis1, axis2):  # pylint: disable=missing-docstring
   a = asarray(a).data
 
@@ -794,7 +794,7 @@ def swapaxes(a, axis1, axis2):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(a)
 
 
-@np_utils.np_doc(np.moveaxis)
+@np_utils.np_doc('moveaxis')
 def moveaxis(a, source, destination):  # pylint: disable=missing-docstring
   """Raises ValueError if source, destination not in (-ndim(a), ndim(a))."""
   if not source and not destination:
@@ -914,7 +914,7 @@ setattr(np_arrays.ndarray, 'reshape', _reshape_method_wrapper)
 setattr(np_arrays.ndarray, '__setitem__', _setitem)
 
 
-@np_utils.np_doc(np.pad)
+@np_utils.np_doc('pad')
 def pad(ary, pad_width, mode, constant_values=0):
   """Only supports modes 'constant', 'reflect' and 'symmetric' currently."""
   if not (mode == 'constant' or mode == 'reflect' or mode == 'symmetric'):
@@ -930,7 +930,7 @@ def pad(ary, pad_width, mode, constant_values=0):
           constant_values=constant_values))
 
 
-@np_utils.np_doc(np.take)
+@np_utils.np_doc('take')
 def take(a, indices, axis=None, out=None, mode='clip'):
   """out argument is not supported, and default mode is clip."""
   if out is not None:
@@ -957,7 +957,7 @@ def take(a, indices, axis=None, out=None, mode='clip'):
   return np_utils.tensor_to_ndarray(array_ops.gather(a, indices, axis=axis))
 
 
-@np_utils.np_doc_only(np.where)
+@np_utils.np_doc_only('where')
 def where(condition, x=None, y=None):
   """Raises ValueError if exactly one of x or y is not None."""
   condition = asarray(condition, dtype=np.bool_)
@@ -970,7 +970,7 @@ def where(condition, x=None, y=None):
   raise ValueError('Both x and y must be ndarrays, or both must be None.')
 
 
-@np_utils.np_doc(np.select)
+@np_utils.np_doc('select')
 def select(condlist, choicelist, default=0):  # pylint: disable=missing-docstring
   if len(condlist) != len(choicelist):
     msg = 'condlist must have length equal to choicelist ({} vs {})'
@@ -987,19 +987,19 @@ def select(condlist, choicelist, default=0):  # pylint: disable=missing-docstrin
   return output
 
 
-@np_utils.np_doc(np.shape)
+@np_utils.np_doc('shape')
 def shape(a):
   a = asarray(a)
   return a.shape
 
 
-@np_utils.np_doc(np.ndim)
+@np_utils.np_doc('ndim')
 def ndim(a):
   a = asarray(a)
   return a.ndim
 
 
-@np_utils.np_doc(np.isscalar)
+@np_utils.np_doc('isscalar')
 def isscalar(a):
   return ndim(a) == 0
 
@@ -1034,7 +1034,7 @@ def _boundaries_to_sizes(a, boundaries, axis):
   return sizes
 
 
-@np_utils.np_doc(np.split)
+@np_utils.np_doc('split')
 def split(ary, indices_or_sections, axis=0):
   ary = asarray(ary)
   if not isinstance(indices_or_sections, six.integer_types):
@@ -1043,26 +1043,26 @@ def split(ary, indices_or_sections, axis=0):
   return [np_utils.tensor_to_ndarray(a) for a in result]
 
 
-def _split_on_axis(np_fun, axis):
+def _split_on_axis(np_fun_name, axis):
 
-  @np_utils.np_doc(np_fun)
+  @np_utils.np_doc(np_fun_name)
   def f(ary, indices_or_sections):
     return split(ary, indices_or_sections, axis=axis)
 
   return f
 
 
-vsplit = _split_on_axis(np.vsplit, axis=0)
-hsplit = _split_on_axis(np.hsplit, axis=1)
-dsplit = _split_on_axis(np.dsplit, axis=2)
+vsplit = _split_on_axis('vsplit', axis=0)
+hsplit = _split_on_axis('hsplit', axis=1)
+dsplit = _split_on_axis('dsplit', axis=2)
 
 
-@np_utils.np_doc(np.broadcast_to)
+@np_utils.np_doc('broadcast_to')
 def broadcast_to(array, shape):  # pylint: disable=redefined-outer-name
   return full(shape, array)
 
 
-@np_utils.np_doc(np.stack)
+@np_utils.np_doc('stack')
 def stack(arrays, axis=0):  # pylint: disable=missing-function-docstring
   if isinstance(arrays, (np_arrays.ndarray, ops.Tensor)):
     arrays = asarray(arrays)
@@ -1077,7 +1077,7 @@ def stack(arrays, axis=0):  # pylint: disable=missing-function-docstring
   return asarray(array_ops.stack(unwrapped_arrays, axis))
 
 
-@np_utils.np_doc(np.hstack)
+@np_utils.np_doc('hstack')
 def hstack(tup):
   arrays = [atleast_1d(a) for a in tup]
   arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
@@ -1091,7 +1091,7 @@ def hstack(tup):
       lambda: array_ops.concat(unwrapped_arrays, axis=1))
 
 
-@np_utils.np_doc(np.vstack)
+@np_utils.np_doc('vstack')
 def vstack(tup):
   arrays = [atleast_2d(a) for a in tup]
   arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
@@ -1101,7 +1101,7 @@ def vstack(tup):
   return array_ops.concat(unwrapped_arrays, axis=0)
 
 
-@np_utils.np_doc(np.dstack)
+@np_utils.np_doc('dstack')
 def dstack(tup):
   arrays = [atleast_3d(a) for a in tup]
   arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
@@ -1148,17 +1148,17 @@ def _atleast_nd(n, new_shape, *arys):
     return arys
 
 
-@np_utils.np_doc(np.atleast_1d)
+@np_utils.np_doc('atleast_1d')
 def atleast_1d(*arys):
   return _atleast_nd(1, _pad_left_to, *arys)
 
 
-@np_utils.np_doc(np.atleast_2d)
+@np_utils.np_doc('atleast_2d')
 def atleast_2d(*arys):
   return _atleast_nd(2, _pad_left_to, *arys)
 
 
-@np_utils.np_doc(np.atleast_3d)
+@np_utils.np_doc('atleast_3d')
 def atleast_3d(*arys):  # pylint: disable=missing-docstring
 
   def new_shape(_, old_shape):
@@ -1175,7 +1175,7 @@ def atleast_3d(*arys):  # pylint: disable=missing-docstring
   return _atleast_nd(3, new_shape, *arys)
 
 
-@np_utils.np_doc(np.nonzero)
+@np_utils.np_doc('nonzero')
 def nonzero(a):
   a = atleast_1d(a).data
   if a.shape.rank is None:
@@ -1189,7 +1189,7 @@ def nonzero(a):
           axis=1))
 
 
-@np_utils.np_doc(np.diag_indices)
+@np_utils.np_doc('diag_indices')
 def diag_indices(n, ndim=2):  # pylint: disable=missing-docstring,redefined-outer-name
   if n < 0:
     raise ValueError(
@@ -1202,7 +1202,7 @@ def diag_indices(n, ndim=2):  # pylint: disable=missing-docstring,redefined-oute
   return (math_ops.range(n),) * ndim
 
 
-@np_utils.np_doc(np.tri)
+@np_utils.np_doc('tri')
 def tri(N, M=None, k=0, dtype=None):  # pylint: disable=invalid-name,missing-docstring
   M = M if M is not None else N
   if dtype is not None:
@@ -1229,7 +1229,7 @@ def tri(N, M=None, k=0, dtype=None):  # pylint: disable=invalid-name,missing-doc
   return np_utils.tensor_to_ndarray(r)
 
 
-@np_utils.np_doc(np.tril)
+@np_utils.np_doc('tril')
 def tril(m, k=0):  # pylint: disable=missing-docstring
   m = asarray(m).data
   if m.shape.ndims is None:
@@ -1251,7 +1251,7 @@ def tril(m, k=0):  # pylint: disable=missing-docstring
           array_ops.broadcast_to(mask, array_ops.shape(m)), m, z))
 
 
-@np_utils.np_doc(np.triu)
+@np_utils.np_doc('triu')
 def triu(m, k=0):  # pylint: disable=missing-docstring
   m = asarray(m).data
   if m.shape.ndims is None:
@@ -1273,7 +1273,7 @@ def triu(m, k=0):  # pylint: disable=missing-docstring
           array_ops.broadcast_to(mask, array_ops.shape(m)), z, m))
 
 
-@np_utils.np_doc(np.flip)
+@np_utils.np_doc('flip')
 def flip(m, axis=None):  # pylint: disable=missing-docstring
   m = asarray(m).data
 
@@ -1286,17 +1286,17 @@ def flip(m, axis=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(array_ops.reverse(m, [axis]))
 
 
-@np_utils.np_doc(np.flipud)
+@np_utils.np_doc('flipud')
 def flipud(m):  # pylint: disable=missing-docstring
   return flip(m, 0)
 
 
-@np_utils.np_doc(np.fliplr)
+@np_utils.np_doc('fliplr')
 def fliplr(m):  # pylint: disable=missing-docstring
   return flip(m, 1)
 
 
-@np_utils.np_doc(np.roll)
+@np_utils.np_doc('roll')
 def roll(a, shift, axis=None):  # pylint: disable=missing-docstring
   a = asarray(a).data
 
@@ -1309,7 +1309,7 @@ def roll(a, shift, axis=None):  # pylint: disable=missing-docstring
   return np_utils.tensor_to_ndarray(array_ops.reshape(a, original_shape))
 
 
-@np_utils.np_doc(np.rot90)
+@np_utils.np_doc('rot90')
 def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
   m_rank = array_ops.rank(m)
   ax1, ax2 = np_utils._canonicalize_axes(axes, m_rank)  # pylint: disable=protected-access
@@ -1329,7 +1329,7 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
       return flip(transpose(m, perm), ax2)
 
 
-@np_utils.np_doc(np.vander)
+@np_utils.np_doc('vander')
 def vander(x, N=None, increasing=False):  # pylint: disable=missing-docstring,invalid-name
   x = asarray(x).data
 
@@ -1368,7 +1368,7 @@ def vander(x, N=None, increasing=False):  # pylint: disable=missing-docstring,in
           x, math_ops.cast(math_ops.range(start, limit, delta), dtype=x.dtype)))
 
 
-@np_utils.np_doc(np.ix_)
+@np_utils.np_doc('ix_')
 def ix_(*args):  # pylint: disable=missing-docstring
   n = len(args)
   output = []
@@ -1400,7 +1400,7 @@ def ix_(*args):  # pylint: disable=missing-docstring
   return output
 
 
-@np_utils.np_doc(np.broadcast_arrays)
+@np_utils.np_doc('broadcast_arrays')
 def broadcast_arrays(*args, **kwargs):  # pylint: disable=missing-docstring
   subok = kwargs.pop('subok', False)
   if subok:
@@ -1413,7 +1413,7 @@ def broadcast_arrays(*args, **kwargs):  # pylint: disable=missing-docstring
   return [np_utils.tensor_to_ndarray(arg) for arg in args]
 
 
-@np_utils.np_doc_only(np.sign)
+@np_utils.np_doc_only('sign')
 def sign(x, out=None, where=None, **kwargs):  # pylint: disable=missing-docstring,redefined-outer-name
   if out:
     raise ValueError('tf.numpy doesnt support setting out.')
@@ -1434,7 +1434,7 @@ def sign(x, out=None, where=None, **kwargs):  # pylint: disable=missing-docstrin
 
 # Note that np.take_along_axis may not be present in some supported versions of
 # numpy.
-@np_utils.np_doc(None, np_fun_name='take_along_axis')
+@np_utils.np_doc('take_along_axis')
 def take_along_axis(arr, indices, axis):  # pylint: disable=missing-docstring
   arr = asarray(arr)
   indices = asarray(indices)
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 361bfb50dec..3cf26095dd8 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -40,7 +40,7 @@ from tensorflow.python.ops.numpy_ops import np_dtypes
 from tensorflow.python.ops.numpy_ops import np_utils
 
 
-@np_utils.np_doc_only(np.dot)
+@np_utils.np_doc_only('dot')
 def dot(a, b):  # pylint: disable=missing-docstring
 
   def f(a, b):  # pylint: disable=missing-docstring
@@ -67,7 +67,7 @@ def _bin_op(tf_fun, a, b, promote=True):
   return np_utils.tensor_to_ndarray(tf_fun(a.data, b.data))
 
 
-@np_utils.np_doc(np.add)
+@np_utils.np_doc('add')
 def add(x1, x2):
 
   def add_or_or(x1, x2):
@@ -79,12 +79,12 @@ def add(x1, x2):
   return _bin_op(add_or_or, x1, x2)
 
 
-@np_utils.np_doc(np.subtract)
+@np_utils.np_doc('subtract')
 def subtract(x1, x2):
   return _bin_op(math_ops.subtract, x1, x2)
 
 
-@np_utils.np_doc(np.multiply)
+@np_utils.np_doc('multiply')
 def multiply(x1, x2):
 
   def mul_or_and(x1, x2):
@@ -96,7 +96,7 @@ def multiply(x1, x2):
   return _bin_op(mul_or_and, x1, x2)
 
 
-@np_utils.np_doc(np.true_divide)
+@np_utils.np_doc('true_divide')
 def true_divide(x1, x2):  # pylint: disable=missing-function-docstring
 
   def _avoid_float64(x1, x2):
@@ -123,7 +123,7 @@ def true_divide(x1, x2):  # pylint: disable=missing-function-docstring
 divide = true_divide
 
 
-@np_utils.np_doc(np.floor_divide)
+@np_utils.np_doc('floor_divide')
 def floor_divide(x1, x2):  # pylint: disable=missing-function-docstring
 
   def f(x1, x2):
@@ -136,7 +136,7 @@ def floor_divide(x1, x2):  # pylint: disable=missing-function-docstring
   return _bin_op(f, x1, x2)
 
 
-@np_utils.np_doc(np.mod)
+@np_utils.np_doc('mod')
 def mod(x1, x2):  # pylint: disable=missing-function-docstring
 
   def f(x1, x2):
@@ -152,12 +152,12 @@ def mod(x1, x2):  # pylint: disable=missing-function-docstring
 remainder = mod
 
 
-@np_utils.np_doc(np.divmod)
+@np_utils.np_doc('divmod')
 def divmod(x1, x2):  # pylint: disable=redefined-builtin
   return floor_divide(x1, x2), mod(x1, x2)
 
 
-@np_utils.np_doc(np.maximum)
+@np_utils.np_doc('maximum')
 def maximum(x1, x2):
 
   def max_or_or(x1, x2):
@@ -169,7 +169,7 @@ def maximum(x1, x2):
   return _bin_op(max_or_or, x1, x2)
 
 
-@np_utils.np_doc(np.minimum)
+@np_utils.np_doc('minimum')
 def minimum(x1, x2):
 
   def min_or_and(x1, x2):
@@ -181,7 +181,7 @@ def minimum(x1, x2):
   return _bin_op(min_or_and, x1, x2)
 
 
-@np_utils.np_doc(np.clip)
+@np_utils.np_doc('clip')
 def clip(a, a_min, a_max):  # pylint: disable=missing-docstring
   if a_min is None and a_max is None:
     raise ValueError('Not more than one of `a_min` and `a_max` may be `None`.')
@@ -196,7 +196,7 @@ def clip(a, a_min, a_max):  # pylint: disable=missing-docstring
             *np_utils.tf_broadcast(a.data, a_min.data, a_max.data)))
 
 
-@np_utils.np_doc(np.matmul)
+@np_utils.np_doc('matmul')
 def matmul(x1, x2):  # pylint: disable=missing-docstring
 
   def f(x1, x2):
@@ -215,12 +215,12 @@ def matmul(x1, x2):  # pylint: disable=missing-docstring
   return _bin_op(f, x1, x2)
 
 
-@np_utils.np_doc(np.tensordot)
+@np_utils.np_doc('tensordot')
 def tensordot(a, b, axes=2):
   return _bin_op(lambda a, b: math_ops.tensordot(a, b, axes=axes), a, b)
 
 
-@np_utils.np_doc_only(np.inner)
+@np_utils.np_doc_only('inner')
 def inner(a, b):  # pylint: disable=missing-function-docstring
 
   def f(a, b):
@@ -233,7 +233,7 @@ def inner(a, b):  # pylint: disable=missing-function-docstring
   return _bin_op(f, a, b)
 
 
-@np_utils.np_doc(np.cross)
+@np_utils.np_doc('cross')
 def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):  # pylint: disable=missing-docstring
 
   def f(a, b):  # pylint: disable=missing-docstring
@@ -309,7 +309,7 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):  # pylint: disable=mis
   return _bin_op(f, a, b)
 
 
-@np_utils.np_doc_only(np.vdot)
+@np_utils.np_doc_only('vdot')
 def vdot(a, b):  # pylint: disable=missing-docstring
   a, b = np_array_ops._promote_dtype(a, b)
   a = np_array_ops.reshape(a, [-1])
@@ -319,27 +319,27 @@ def vdot(a, b):  # pylint: disable=missing-docstring
   return dot(a, b)
 
 
-@np_utils.np_doc(np.power)
+@np_utils.np_doc('power')
 def power(x1, x2):
   return _bin_op(math_ops.pow, x1, x2)
 
 
-@np_utils.np_doc(np.float_power)
+@np_utils.np_doc('float_power')
 def float_power(x1, x2):
   return power(x1, x2)
 
 
-@np_utils.np_doc(np.arctan2)
+@np_utils.np_doc('arctan2')
 def arctan2(x1, x2):
   return _bin_op(math_ops.atan2, x1, x2)
 
 
-@np_utils.np_doc(np.nextafter)
+@np_utils.np_doc('nextafter')
 def nextafter(x1, x2):
   return _bin_op(math_ops.nextafter, x1, x2)
 
 
-@np_utils.np_doc(np.heaviside)
+@np_utils.np_doc('heaviside')
 def heaviside(x1, x2):  # pylint: disable=missing-function-docstring
 
   def f(x1, x2):
@@ -353,12 +353,12 @@ def heaviside(x1, x2):  # pylint: disable=missing-function-docstring
   return y
 
 
-@np_utils.np_doc(np.hypot)
+@np_utils.np_doc('hypot')
 def hypot(x1, x2):
   return sqrt(square(x1) + square(x2))
 
 
-@np_utils.np_doc(np.kron)
+@np_utils.np_doc('kron')
 def kron(a, b):  # pylint: disable=missing-function-docstring
   # pylint: disable=protected-access,g-complex-comprehension
   a, b = np_array_ops._promote_dtype(a, b)
@@ -389,7 +389,7 @@ def kron(a, b):  # pylint: disable=missing-function-docstring
   return np_array_ops.reshape(a_reshaped * b_reshaped, out_shape)
 
 
-@np_utils.np_doc(np.outer)
+@np_utils.np_doc('outer')
 def outer(a, b):
 
   def f(a, b):
@@ -399,7 +399,7 @@ def outer(a, b):
 
 
 # This can also be implemented via tf.reduce_logsumexp
-@np_utils.np_doc(np.logaddexp)
+@np_utils.np_doc('logaddexp')
 def logaddexp(x1, x2):
   amax = maximum(x1, x2)
   delta = x1 - x2
@@ -409,7 +409,7 @@ def logaddexp(x1, x2):
       amax + log1p(exp(-abs(delta))))
 
 
-@np_utils.np_doc(np.logaddexp2)
+@np_utils.np_doc('logaddexp2')
 def logaddexp2(x1, x2):
   amax = maximum(x1, x2)
   delta = x1 - x2
@@ -419,7 +419,7 @@ def logaddexp2(x1, x2):
       amax + log1p(exp2(-abs(delta))) / np.log(2))
 
 
-@np_utils.np_doc(np.polyval)
+@np_utils.np_doc('polyval')
 def polyval(p, x):  # pylint: disable=missing-function-docstring
 
   def f(p, x):
@@ -437,7 +437,7 @@ def polyval(p, x):  # pylint: disable=missing-function-docstring
   return _bin_op(f, p, x)
 
 
-@np_utils.np_doc(np.isclose)
+@np_utils.np_doc('isclose')
 def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):  # pylint: disable=missing-docstring
 
   def f(a, b):  # pylint: disable=missing-docstring
@@ -455,7 +455,7 @@ def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):  # pylint: disable=m
   return _bin_op(f, a, b)
 
 
-@np_utils.np_doc(np.allclose)
+@np_utils.np_doc('allclose')
 def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
   return np_array_ops.all(
       isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
@@ -490,13 +490,13 @@ def _tf_gcd(x1, x2):  # pylint: disable=missing-function-docstring
 
 
 # Note that np.gcd may not be present in some supported versions of numpy.
-@np_utils.np_doc(None, np_fun_name='gcd')
+@np_utils.np_doc('gcd')
 def gcd(x1, x2):
   return _bin_op(_tf_gcd, x1, x2)
 
 
 # Note that np.lcm may not be present in some supported versions of numpy.
-@np_utils.np_doc(None, np_fun_name='lcm')
+@np_utils.np_doc('lcm')
 def lcm(x1, x2):  # pylint: disable=missing-function-docstring
 
   def f(x1, x2):
@@ -527,22 +527,22 @@ def _bitwise_binary_op(tf_fn, x1, x2):  # pylint: disable=missing-function-docst
   return _bin_op(f, x1, x2)
 
 
-@np_utils.np_doc(np.bitwise_and)
+@np_utils.np_doc('bitwise_and')
 def bitwise_and(x1, x2):
   return _bitwise_binary_op(bitwise_ops.bitwise_and, x1, x2)
 
 
-@np_utils.np_doc(np.bitwise_or)
+@np_utils.np_doc('bitwise_or')
 def bitwise_or(x1, x2):
   return _bitwise_binary_op(bitwise_ops.bitwise_or, x1, x2)
 
 
-@np_utils.np_doc(np.bitwise_xor)
+@np_utils.np_doc('bitwise_xor')
 def bitwise_xor(x1, x2):
   return _bitwise_binary_op(bitwise_ops.bitwise_xor, x1, x2)
 
 
-@np_utils.np_doc(np.bitwise_not)
+@np_utils.np_doc('bitwise_not')
 def bitwise_not(x):
 
   def f(x):
@@ -574,62 +574,62 @@ def _scalar(tf_fn, x, promote_to_float=False):
   return np_utils.tensor_to_ndarray(tf_fn(x.data))
 
 
-@np_utils.np_doc(np.log)
+@np_utils.np_doc('log')
 def log(x):
   return _scalar(math_ops.log, x, True)
 
 
-@np_utils.np_doc(np.exp)
+@np_utils.np_doc('exp')
 def exp(x):
   return _scalar(math_ops.exp, x, True)
 
 
-@np_utils.np_doc(np.sqrt)
+@np_utils.np_doc('sqrt')
 def sqrt(x):
   return _scalar(math_ops.sqrt, x, True)
 
 
-@np_utils.np_doc(np.abs)
+@np_utils.np_doc('abs')
 def abs(x):  # pylint: disable=redefined-builtin
   return _scalar(math_ops.abs, x)
 
 
-@np_utils.np_doc(np.absolute)
+@np_utils.np_doc('absolute')
 def absolute(x):
   return abs(x)
 
 
-@np_utils.np_doc(np.fabs)
+@np_utils.np_doc('fabs')
 def fabs(x):
   return abs(x)
 
 
-@np_utils.np_doc(np.ceil)
+@np_utils.np_doc('ceil')
 def ceil(x):
   return _scalar(math_ops.ceil, x, True)
 
 
-@np_utils.np_doc(np.floor)
+@np_utils.np_doc('floor')
 def floor(x):
   return _scalar(math_ops.floor, x, True)
 
 
-@np_utils.np_doc(np.conj)
+@np_utils.np_doc('conj')
 def conj(x):
   return _scalar(math_ops.conj, x)
 
 
-@np_utils.np_doc(np.negative)
+@np_utils.np_doc('negative')
 def negative(x):
   return _scalar(math_ops.negative, x)
 
 
-@np_utils.np_doc(np.reciprocal)
+@np_utils.np_doc('reciprocal')
 def reciprocal(x):
   return _scalar(math_ops.reciprocal, x)
 
 
-@np_utils.np_doc(np.signbit)
+@np_utils.np_doc('signbit')
 def signbit(x):
 
   def f(x):
@@ -640,67 +640,67 @@ def signbit(x):
   return _scalar(f, x)
 
 
-@np_utils.np_doc(np.sin)
+@np_utils.np_doc('sin')
 def sin(x):
   return _scalar(math_ops.sin, x, True)
 
 
-@np_utils.np_doc(np.cos)
+@np_utils.np_doc('cos')
 def cos(x):
   return _scalar(math_ops.cos, x, True)
 
 
-@np_utils.np_doc(np.tan)
+@np_utils.np_doc('tan')
 def tan(x):
   return _scalar(math_ops.tan, x, True)
 
 
-@np_utils.np_doc(np.sinh)
+@np_utils.np_doc('sinh')
 def sinh(x):
   return _scalar(math_ops.sinh, x, True)
 
 
-@np_utils.np_doc(np.cosh)
+@np_utils.np_doc('cosh')
 def cosh(x):
   return _scalar(math_ops.cosh, x, True)
 
 
-@np_utils.np_doc(np.tanh)
+@np_utils.np_doc('tanh')
 def tanh(x):
   return _scalar(math_ops.tanh, x, True)
 
 
-@np_utils.np_doc(np.arcsin)
+@np_utils.np_doc('arcsin')
 def arcsin(x):
   return _scalar(math_ops.asin, x, True)
 
 
-@np_utils.np_doc(np.arccos)
+@np_utils.np_doc('arccos')
 def arccos(x):
   return _scalar(math_ops.acos, x, True)
 
 
-@np_utils.np_doc(np.arctan)
+@np_utils.np_doc('arctan')
 def arctan(x):
   return _scalar(math_ops.atan, x, True)
 
 
-@np_utils.np_doc(np.arcsinh)
+@np_utils.np_doc('arcsinh')
 def arcsinh(x):
   return _scalar(math_ops.asinh, x, True)
 
 
-@np_utils.np_doc(np.arccosh)
+@np_utils.np_doc('arccosh')
 def arccosh(x):
   return _scalar(math_ops.acosh, x, True)
 
 
-@np_utils.np_doc(np.arctanh)
+@np_utils.np_doc('arctanh')
 def arctanh(x):
   return _scalar(math_ops.atanh, x, True)
 
 
-@np_utils.np_doc(np.deg2rad)
+@np_utils.np_doc('deg2rad')
 def deg2rad(x):
 
   def f(x):
@@ -709,7 +709,7 @@ def deg2rad(x):
   return _scalar(f, x, True)
 
 
-@np_utils.np_doc(np.rad2deg)
+@np_utils.np_doc('rad2deg')
 def rad2deg(x):
   return x * (180.0 / np.pi)
 
@@ -719,7 +719,7 @@ _tf_float_types = [
 ]
 
 
-@np_utils.np_doc(np.angle)
+@np_utils.np_doc('angle')
 def angle(z, deg=False):  # pylint: disable=missing-function-docstring
 
   def f(x):
@@ -735,7 +735,7 @@ def angle(z, deg=False):  # pylint: disable=missing-function-docstring
   return y
 
 
-@np_utils.np_doc(np.cbrt)
+@np_utils.np_doc('cbrt')
 def cbrt(x):
 
   def f(x):
@@ -746,12 +746,12 @@ def cbrt(x):
   return _scalar(f, x, True)
 
 
-@np_utils.np_doc(np.conjugate)
+@np_utils.np_doc('conjugate')
 def conjugate(x):
   return _scalar(math_ops.conj, x)
 
 
-@np_utils.np_doc(np.exp2)
+@np_utils.np_doc('exp2')
 def exp2(x):
 
   def f(x):
@@ -760,12 +760,12 @@ def exp2(x):
   return _scalar(f, x, True)
 
 
-@np_utils.np_doc(np.expm1)
+@np_utils.np_doc('expm1')
 def expm1(x):
   return _scalar(math_ops.expm1, x, True)
 
 
-@np_utils.np_doc(np.fix)
+@np_utils.np_doc('fix')
 def fix(x):
 
   def f(x):
@@ -774,36 +774,36 @@ def fix(x):
   return _scalar(f, x, True)
 
 
-@np_utils.np_doc(np.iscomplex)
+@np_utils.np_doc('iscomplex')
 def iscomplex(x):
   return np_array_ops.imag(x) != 0
 
 
-@np_utils.np_doc(np.isreal)
+@np_utils.np_doc('isreal')
 def isreal(x):
   return np_array_ops.imag(x) == 0
 
 
-@np_utils.np_doc(np.iscomplexobj)
+@np_utils.np_doc('iscomplexobj')
 def iscomplexobj(x):
   x = np_array_ops.array(x)
   return np.issubdtype(x.dtype, np.complexfloating)
 
 
-@np_utils.np_doc(np.isrealobj)
+@np_utils.np_doc('isrealobj')
 def isrealobj(x):
   return not iscomplexobj(x)
 
 
-@np_utils.np_doc(np.isnan)
+@np_utils.np_doc('isnan')
 def isnan(x):
   return _scalar(math_ops.is_nan, x, True)
 
 
-def _make_nan_reduction(onp_reduction, reduction, init_val):
+def _make_nan_reduction(np_fun_name, reduction, init_val):
   """Helper to generate nan* functions."""
 
-  @np_utils.np_doc(onp_reduction)
+  @np_utils.np_doc(np_fun_name)
   def nan_reduction(a, axis=None, dtype=None, keepdims=False):
     a = np_array_ops.array(a)
     v = np_array_ops.array(init_val, dtype=a.dtype)
@@ -816,11 +816,11 @@ def _make_nan_reduction(onp_reduction, reduction, init_val):
   return nan_reduction
 
 
-nansum = _make_nan_reduction(np.nansum, np_array_ops.sum, 0)
-nanprod = _make_nan_reduction(np.nanprod, np_array_ops.prod, 1)
+nansum = _make_nan_reduction('nansum', np_array_ops.sum, 0)
+nanprod = _make_nan_reduction('nanprod', np_array_ops.prod, 1)
 
 
-@np_utils.np_doc(np.nanmean)
+@np_utils.np_doc('nanmean')
 def nanmean(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=missing-docstring
   a = np_array_ops.array(a)
   if np.issubdtype(a.dtype, np.bool_) or np.issubdtype(a.dtype, np.integer):
@@ -833,47 +833,47 @@ def nanmean(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=missing
   return nansum(a, axis=axis, dtype=dtype, keepdims=keepdims) / normalizer
 
 
-@np_utils.np_doc(np.isfinite)
+@np_utils.np_doc('isfinite')
 def isfinite(x):
   return _scalar(math_ops.is_finite, x, True)
 
 
-@np_utils.np_doc(np.isinf)
+@np_utils.np_doc('isinf')
 def isinf(x):
   return _scalar(math_ops.is_inf, x, True)
 
 
-@np_utils.np_doc(np.isneginf)
+@np_utils.np_doc('isneginf')
 def isneginf(x):
   return x == np_array_ops.full_like(x, -np.inf)
 
 
-@np_utils.np_doc(np.isposinf)
+@np_utils.np_doc('isposinf')
 def isposinf(x):
   return x == np_array_ops.full_like(x, np.inf)
 
 
-@np_utils.np_doc(np.log2)
+@np_utils.np_doc('log2')
 def log2(x):
   return log(x) / np.log(2)
 
 
-@np_utils.np_doc(np.log10)
+@np_utils.np_doc('log10')
 def log10(x):
   return log(x) / np.log(10)
 
 
-@np_utils.np_doc(np.log1p)
+@np_utils.np_doc('log1p')
 def log1p(x):
   return _scalar(math_ops.log1p, x, True)
 
 
-@np_utils.np_doc(np.positive)
+@np_utils.np_doc('positive')
 def positive(x):
   return _scalar(lambda x: x, x)
 
 
-@np_utils.np_doc(np.sinc)
+@np_utils.np_doc('sinc')
 def sinc(x):
 
   def f(x):
@@ -884,12 +884,12 @@ def sinc(x):
   return _scalar(f, x, True)
 
 
-@np_utils.np_doc(np.square)
+@np_utils.np_doc('square')
 def square(x):
   return _scalar(math_ops.square, x)
 
 
-@np_utils.np_doc(np.diff)
+@np_utils.np_doc('diff')
 def diff(a, n=1, axis=-1):  # pylint: disable=missing-function-docstring
 
   def f(a):
@@ -964,37 +964,37 @@ def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):
   return np_utils.tensor_to_ndarray(tf_fun(x1, x2))
 
 
-@np_utils.np_doc(np.equal)
+@np_utils.np_doc('equal')
 def equal(x1, x2):
   return _comparison(math_ops.equal, x1, x2)
 
 
-@np_utils.np_doc(np.not_equal)
+@np_utils.np_doc('not_equal')
 def not_equal(x1, x2):
   return _comparison(math_ops.not_equal, x1, x2)
 
 
-@np_utils.np_doc(np.greater)
+@np_utils.np_doc('greater')
 def greater(x1, x2):
   return _comparison(math_ops.greater, x1, x2, True)
 
 
-@np_utils.np_doc(np.greater_equal)
+@np_utils.np_doc('greater_equal')
 def greater_equal(x1, x2):
   return _comparison(math_ops.greater_equal, x1, x2, True)
 
 
-@np_utils.np_doc(np.less)
+@np_utils.np_doc('less')
 def less(x1, x2):
   return _comparison(math_ops.less, x1, x2, True)
 
 
-@np_utils.np_doc(np.less_equal)
+@np_utils.np_doc('less_equal')
 def less_equal(x1, x2):
   return _comparison(math_ops.less_equal, x1, x2, True)
 
 
-@np_utils.np_doc(np.array_equal)
+@np_utils.np_doc('array_equal')
 def array_equal(a1, a2):  # pylint: disable=missing-function-docstring
 
   def f(x1, x2):
@@ -1017,22 +1017,22 @@ def _logical_binary_op(tf_fun, x1, x2):
   return np_utils.tensor_to_ndarray(tf_fun(x1.data, x2.data))
 
 
-@np_utils.np_doc(np.logical_and)
+@np_utils.np_doc('logical_and')
 def logical_and(x1, x2):
   return _logical_binary_op(math_ops.logical_and, x1, x2)
 
 
-@np_utils.np_doc(np.logical_or)
+@np_utils.np_doc('logical_or')
 def logical_or(x1, x2):
   return _logical_binary_op(math_ops.logical_or, x1, x2)
 
 
-@np_utils.np_doc(np.logical_xor)
+@np_utils.np_doc('logical_xor')
 def logical_xor(x1, x2):
   return _logical_binary_op(math_ops.logical_xor, x1, x2)
 
 
-@np_utils.np_doc(np.logical_not)
+@np_utils.np_doc('logical_not')
 def logical_not(x):
   x = np_array_ops.array(x, dtype=np.bool_)
   return np_utils.tensor_to_ndarray(math_ops.logical_not(x.data))
@@ -1047,7 +1047,7 @@ setattr(np_arrays.ndarray, '__eq__', _wrap(equal))
 setattr(np_arrays.ndarray, '__ne__', _wrap(not_equal))
 
 
-@np_utils.np_doc(np.linspace)
+@np_utils.np_doc('linspace')
 def linspace(  # pylint: disable=missing-docstring
     start,
     stop,
@@ -1086,7 +1086,7 @@ def linspace(  # pylint: disable=missing-docstring
     return np_arrays.tensor_to_ndarray(result)
 
 
-@np_utils.np_doc(np.logspace)
+@np_utils.np_doc('logspace')
 def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
   dtype = np_utils.result_type(start, stop, dtype)
   result = linspace(
@@ -1097,7 +1097,7 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
   return np_arrays.tensor_to_ndarray(result)
 
 
-@np_utils.np_doc(np.geomspace)
+@np_utils.np_doc('geomspace')
 def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):  # pylint: disable=missing-docstring
   dtype = dtype or np_utils.result_type(start, stop, float(num),
                                         np_array_ops.zeros((), dtype))
@@ -1121,13 +1121,13 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):  # pylint
   return np_utils.tensor_to_ndarray(math_ops.cast(res, dtype))
 
 
-@np_utils.np_doc(np.ptp)
+@np_utils.np_doc('ptp')
 def ptp(a, axis=None, keepdims=None):
   return (np_array_ops.amax(a, axis=axis, keepdims=keepdims) -
           np_array_ops.amin(a, axis=axis, keepdims=keepdims))
 
 
-@np_utils.np_doc_only(np.concatenate)
+@np_utils.np_doc_only('concatenate')
 def concatenate(arys, axis=0):
   if not isinstance(arys, (list, tuple)):
     arys = [arys]
@@ -1138,7 +1138,7 @@ def concatenate(arys, axis=0):
   return np_arrays.tensor_to_ndarray(array_ops.concat(arys, axis))
 
 
-@np_utils.np_doc_only(np.tile)
+@np_utils.np_doc_only('tile')
 def tile(a, reps):  # pylint: disable=missing-function-docstring
   a = np_array_ops.array(a).data
   reps = np_array_ops.array(reps, dtype=dtypes.int32).reshape([-1]).data
@@ -1155,13 +1155,13 @@ def tile(a, reps):  # pylint: disable=missing-function-docstring
   return np_arrays.tensor_to_ndarray(array_ops.tile(a, reps))
 
 
-@np_utils.np_doc(np.count_nonzero)
+@np_utils.np_doc('count_nonzero')
 def count_nonzero(a, axis=None):
   return np_arrays.tensor_to_ndarray(
       math_ops.count_nonzero(np_array_ops.array(a).data, axis))
 
 
-@np_utils.np_doc(np.argsort)
+@np_utils.np_doc('argsort')
 def argsort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
   # TODO(nareshmodi): make string tensors also work.
   if kind not in ('quicksort', 'stable'):
@@ -1186,7 +1186,7 @@ def argsort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missin
   return np_array_ops.array(tf_ans, dtype=np.intp)
 
 
-@np_utils.np_doc(np.sort)
+@np_utils.np_doc('sort')
 def sort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
   if kind != 'quicksort':
     raise ValueError("Only 'quicksort' is supported.")
@@ -1212,17 +1212,17 @@ def _argminmax(fn, a, axis=None):
   return np_utils.tensor_to_ndarray(fn(input=a_t, axis=axis))
 
 
-@np_utils.np_doc(np.argmax)
+@np_utils.np_doc('argmax')
 def argmax(a, axis=None):
   return _argminmax(math_ops.argmax, a, axis)
 
 
-@np_utils.np_doc(np.argmin)
+@np_utils.np_doc('argmin')
 def argmin(a, axis=None):
   return _argminmax(math_ops.argmin, a, axis)
 
 
-@np_utils.np_doc(np.append)
+@np_utils.np_doc('append')
 def append(arr, values, axis=None):
   if axis is None:
     return concatenate([np_array_ops.ravel(arr), np_array_ops.ravel(values)], 0)
@@ -1230,7 +1230,7 @@ def append(arr, values, axis=None):
     return concatenate([arr, values], axis=axis)
 
 
-@np_utils.np_doc(np.average)
+@np_utils.np_doc('average')
 def average(a, axis=None, weights=None, returned=False):  # pylint: disable=missing-docstring
   if axis is not None and not isinstance(axis, six.integer_types):
     # TODO(wangpeng): Support tuple of ints as `axis`
@@ -1293,7 +1293,7 @@ def average(a, axis=None, weights=None, returned=False):  # pylint: disable=miss
   return avg
 
 
-@np_utils.np_doc(np.trace)
+@np_utils.np_doc('trace')
 def trace(a, offset=0, axis1=0, axis2=1, dtype=None):  # pylint: disable=missing-docstring
   if dtype:
     dtype = np_utils.result_type(dtype)
@@ -1311,7 +1311,7 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None):  # pylint: disable=missing
   return np_array_ops.sum(a, -1, dtype)
 
 
-@np_utils.np_doc(np.meshgrid)
+@np_utils.np_doc('meshgrid')
 def meshgrid(*xi, **kwargs):
   """This currently requires copy=True and sparse=False."""
   sparse = kwargs.get('sparse', False)
diff --git a/tensorflow/python/ops/numpy_ops/np_utils.py b/tensorflow/python/ops/numpy_ops/np_utils.py
index 186e56816fe..04ec38d611c 100644
--- a/tensorflow/python/ops/numpy_ops/np_utils.py
+++ b/tensorflow/python/ops/numpy_ops/np_utils.py
@@ -233,28 +233,71 @@ def _is_compatible_param_kind(a, b):
   return relax(a) == relax(b)
 
 
-def np_doc(np_fun, np_fun_name=None):
-  """Attachs numpy docstring to a function.
+def _prepare_np_fun_name_and_fun(np_fun_name, np_fun):
+  """Mutually propagates information between `np_fun_name` and `np_fun`.
+
+  If one is None and the other is not, we'll try to make the former not None in
+  a best effort.
 
   Args:
-    np_fun: the numpy function whose docstring will be used.
-    np_fun_name: optional name for the np_fun symbol. At least one of np_fun or
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
       np_fun_name shoud be set.
+    np_fun: the numpy function whose docstring will be used.
 
   Returns:
-    A function decorator that attaches the docstring from `np_fun` to the
-    decorated function.
+    Processed `np_fun_name` and `np_fun`.
   """
+  if np_fun_name is not None:
+    assert isinstance(np_fun_name, str)
+  if np_fun is not None:
+    assert not isinstance(np_fun, str)
   if np_fun is None:
     assert np_fun_name is not None
     try:
       np_fun = getattr(np, str(np_fun_name))
     except AttributeError:
       np_fun = None
-  np_sig = _np_signature(np_fun)
   if np_fun_name is None:
     assert np_fun is not None
     np_fun_name = np_fun.__name__
+  return np_fun_name, np_fun
+
+
+def _np_doc_helper(f, np_f, np_fun_name=None, unsupported_params=None):
+  """Helper to get docs."""
+  assert np_f or np_fun_name
+  if not np_fun_name:
+    np_fun_name = np_f.__name__
+  doc = 'TensorFlow variant of `numpy.%s`.\n\n' % np_fun_name
+  if unsupported_params:
+    doc += 'Unsupported arguments: ' + ', '.join(
+        '`' + name + '`' for name in unsupported_params) + '.\n\n'
+  if _has_docstring(f):
+    doc += f.__doc__
+    doc = _add_blank_line(doc)
+  if _has_docstring(np_f):
+    doc += 'Documentation for `numpy.%s`:\n\n' % np_f.__name__
+    # TODO(wangpeng): It looks like code snippets in numpy doc don't work
+    # correctly with doctest. Fix that and remove the reformatting of the np_f
+    # comment.
+    doc += np_f.__doc__.replace('>>>', '>')
+  return doc
+
+
+def np_doc(np_fun_name, np_fun=None):
+  """Attachs numpy docstring to a function.
+
+  Args:
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
+      np_fun_name shoud be set.
+    np_fun: (optional) the numpy function whose docstring will be used.
+
+  Returns:
+    A function decorator that attaches the docstring from `np_fun` to the
+    decorated function.
+  """
+  np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
+  np_sig = _np_signature(np_fun)
 
   def decorator(f):
     """The decorator."""
@@ -294,44 +337,24 @@ def np_doc(np_fun, np_fun_name=None):
   return decorator
 
 
-def _np_doc_helper(f, np_f, np_fun_name=None, unsupported_params=None):
-  """Helper to get docs."""
-  if not unsupported_params and not _has_docstring(f) and _has_docstring(np_f):
-    # TODO(wangpeng): It looks like code snippets in numpy doc don't work
-    # correctly with doctest. Fix that and remove the reformatting of the np_f
-    # comment, here and below.
-    return np_f.__doc__.replace('>>>', '>')
-  assert np_f or np_fun_name
-  if not np_fun_name:
-    np_fun_name = np_f.__name__
-  doc = 'TensorFlow variant of `numpy.%s`.\n\n' % np_fun_name
-  if unsupported_params:
-    doc += 'Unsupported arguments: ' + ', '.join(
-        '`' + name + '`' for name in unsupported_params) + '.\n\n'
-  if _has_docstring(f):
-    doc += f.__doc__
-    doc = _add_blank_line(doc)
-  if _has_docstring(np_f):
-    doc += 'Documentation for `numpy.%s`:\n\n' % np_f.__name__
-    doc += np_f.__doc__.replace('>>>', '>')
-  return doc
-
-
-def np_doc_only(np_f):
+def np_doc_only(np_fun_name, np_fun=None):
   """Attachs numpy docstring to a function.
 
   This differs from np_doc in that it doesn't check for a match in signature.
 
   Args:
-    np_f: the numpy function whose docstring will be used.
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
+      np_fun_name shoud be set.
+    np_fun: (optional) the numpy function whose docstring will be used.
 
   Returns:
-    A function decorator that attaches the docstring from `np_f` to the
+    A function decorator that attaches the docstring from `np_fun` to the
     decorated function.
   """
+  np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
 
   def decorator(f):
-    f.__doc__ = _np_doc_helper(f, np_f)
+    f.__doc__ = _np_doc_helper(f, np_fun, np_fun_name=np_fun_name)
     return f
 
   return decorator
diff --git a/tensorflow/python/ops/numpy_ops/np_utils_test.py b/tensorflow/python/ops/numpy_ops/np_utils_test.py
index 6d0dfa51185..38b51f05e6e 100644
--- a/tensorflow/python/ops/numpy_ops/np_utils_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_utils_test.py
@@ -31,7 +31,7 @@ class UtilsTest(test.TestCase):
       """np_fun docstring."""
       return
 
-    @np_utils.np_doc(np_fun)
+    @np_utils.np_doc(None, np_fun=np_fun)
     def f():
       """f docstring."""
       return
@@ -47,7 +47,7 @@ np_fun docstring."""
 
   def testNpDocName(self):
 
-    @np_utils.np_doc(None, np_fun_name='foo')
+    @np_utils.np_doc('foo')
     def f():
       """f docstring."""
       return
@@ -70,20 +70,20 @@ f docstring.
     # pylint: disable=unused-variable
     with self.assertRaisesRegexp(TypeError, 'Cannot find parameter'):
 
-      @np_utils.np_doc(np_fun)
+      @np_utils.np_doc(None, np_fun=np_fun)
       def f1(a):
         return
 
     with self.assertRaisesRegexp(TypeError, 'is of kind'):
 
-      @np_utils.np_doc(np_fun)
+      @np_utils.np_doc(None, np_fun=np_fun)
       def f2(x, kwargs):
         return
 
     with self.assertRaisesRegexp(TypeError,
                                  'Parameter "y" should have a default value'):
 
-      @np_utils.np_doc(np_fun)
+      @np_utils.np_doc(None, np_fun=np_fun)
       def f3(x, y):
         return
 

From f5a0fdaa0aeff548623811b887c6da34303ab25f Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Sat, 20 Jun 2020 01:41:44 +0200
Subject: [PATCH 0684/1390] Use executing_eagerly_outside_functions

---
 .../keras/mixed_precision/experimental/autocast_variable.py   | 4 ++--
 .../mixed_precision/experimental/autocast_variable_test.py    | 3 +++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
index b60100c7b48..a717fbb41e2 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py
@@ -190,7 +190,7 @@ class AutoCastVariable(variables.Variable, core.Tensor):
 
   def _apply_assign_update(
       self, update_fn, value, use_locking=None, name=None, read_value=True):
-    if context.executing_eagerly() or ops.inside_function():
+    if ops.executing_eagerly_outside_functions():
       assign_op = update_fn(value, use_locking, name, False)
       return self if read_value else assign_op
 
@@ -202,7 +202,7 @@ class AutoCastVariable(variables.Variable, core.Tensor):
 
   def _apply_update(self, update_fn, *args, **kwargs):
     update_var = update_fn(*args, **kwargs)
-    if context.executing_eagerly() or ops.inside_function():
+    if ops.executing_eagerly_outside_functions():
       return self
 
     # Fallback to wrapping the returned variable in graph mode if possible
diff --git a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
index 9036109af96..cb5a5d7cb3f 100644
--- a/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
+++ b/tensorflow/python/keras/mixed_precision/experimental/autocast_variable_test.py
@@ -345,6 +345,9 @@ class AutoCastVariableTest(test.TestCase, parameterized.TestCase):
 
   @combinations.generate(maybe_distribute)
   def test_assign_tf_function(self, distribution):
+    if not context.executing_eagerly():
+      self.skipTest("Test is not compatible with graph mode")
+
     with distribution.scope():
       x = get_var(0., dtypes.float32)
       x = autocast_variable.create_autocast_variable(x)

From 2c5e5a643cc6ba68da46e4e2058f178434848dd8 Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Fri, 19 Jun 2020 16:37:27 -0700
Subject: [PATCH 0685/1390] [TF-numpy] Adds an accessor class for numpy_ops, in
 order to be tf_exported'ed.

PiperOrigin-RevId: 317402589
Change-Id: I6bb5f4f9d3b42cf8c2653d60d80c20b37d0bd59f
---
 tensorflow/python/ops/numpy_ops/BUILD         |  1 +
 .../python/ops/numpy_ops/np_accessor.py       | 32 +++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 tensorflow/python/ops/numpy_ops/np_accessor.py

diff --git a/tensorflow/python/ops/numpy_ops/BUILD b/tensorflow/python/ops/numpy_ops/BUILD
index 3f18a7b3e01..3479a622bc0 100644
--- a/tensorflow/python/ops/numpy_ops/BUILD
+++ b/tensorflow/python/ops/numpy_ops/BUILD
@@ -11,6 +11,7 @@ py_library(
     name = "numpy",
     srcs = [
         "__init__.py",
+        "np_accessor.py",
         "np_array_ops.py",
         "np_arrays.py",
         "np_dtypes.py",
diff --git a/tensorflow/python/ops/numpy_ops/np_accessor.py b/tensorflow/python/ops/numpy_ops/np_accessor.py
new file mode 100644
index 00000000000..64786d2c50a
--- /dev/null
+++ b/tensorflow/python/ops/numpy_ops/np_accessor.py
@@ -0,0 +1,32 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""An accessor class for numpy_ops contents."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.ops import numpy_ops
+
+
+class Numpy:
+  """An accessor class that forwards attribute accesses to module `numpy_ops`.
+  """
+
+  def __getattr__(self, attr):
+    return getattr(numpy_ops, attr)
+
+
+numpy = Numpy()

From f51b649394b8b0c2cbf8179ebd4b64b5a915110b Mon Sep 17 00:00:00 2001
From: Austin Anderson 
Date: Fri, 19 Jun 2020 16:39:28 -0700
Subject: [PATCH 0686/1390] Experimental internal CI changes

PiperOrigin-RevId: 317402932
Change-Id: Ibaef72f01e06f4518b85b21f82000599eed7e4bd
---
 .../per_release/scripts/nonpip_gpu.sh         | 75 +++++++++++++++++++
 1 file changed, 75 insertions(+)
 create mode 100644 tensorflow/tools/ci_build/per_release/scripts/nonpip_gpu.sh

diff --git a/tensorflow/tools/ci_build/per_release/scripts/nonpip_gpu.sh b/tensorflow/tools/ci_build/per_release/scripts/nonpip_gpu.sh
new file mode 100644
index 00000000000..6fd7c3d5854
--- /dev/null
+++ b/tensorflow/tools/ci_build/per_release/scripts/nonpip_gpu.sh
@@ -0,0 +1,75 @@
+#!/bin/bash
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+set -e
+set -x
+
+if [[ -n "${KOKORO_ARTIFACTS_DIR}" ]]; then
+  cd "${KOKORO_ARTIFACTS_DIR}"
+  ls
+  source "$(find "${KOKORO_ARTIFACTS_DIR}" -name "common_google.sh")"
+  cd git/gob-tensorflow
+
+fi
+
+if [[ -z "${TF_KOKORO_PY_VERSION}" ]]; then
+  echo "You must set TF_KOKORO_PY_VERSION, e.g. '3.7', indicating the "
+  echo "Python version to be used for this build."
+  exit 2
+fi
+
+source tensorflow/tools/ci_build/release/common.sh
+
+install_ubuntu_16_pip_deps "pip${TF_KOKORO_PY_VERSION}"
+# Update bazel
+install_bazelisk
+
+# Run configure.
+export TF_NEED_GCP=1
+export TF_NEED_HDFS=1
+export TF_NEED_S3=1
+export TF_NEED_CUDA=1
+export TF_CUDA_VERSION=10
+export TF_CUDNN_VERSION=7
+export TF_NEED_TENSORRT=1
+export TENSORRT_INSTALL_PATH=/usr/local/tensorrt
+export CC_OPT_FLAGS='-mavx'
+export PYTHON_BIN_PATH=$(which "python${TF_KOKORO_PY_VERSION}")
+export TF2_BEHAVIOR=1
+export PROJECT_NAME="tensorflow_gpu"
+export LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$TENSORRT_INSTALL_PATH/lib"
+export TF_CUDA_COMPUTE_CAPABILITIES=sm_35,sm_37,sm_52,sm_60,sm_61,compute_70
+
+yes "" | "$PYTHON_BIN_PATH" configure.py
+
+# Get the default test targets for bazel.
+source tensorflow/tools/ci_build/build_scripts/PRESUBMIT_BUILD_TARGETS.sh
+
+# Exclude -no_oss_py36, for example
+tag_filters="gpu,requires-gpu,-no_gpu,-no_oss,-oss_serial,-no_oss_py${TF_KOKORO_PY_VERSION//.}"
+
+set +e
+bazel test --config=cuda --config=opt \
+  --crosstool_top=//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda10.1:toolchain \
+  --linkopt=-lrt \
+  --action_env=TF2_BEHAVIOR="${TF2_BEHAVIOR}" \
+  --test_lang_filters=py \
+  --build_tag_filters=${tag_filters} \
+  --test_tag_filters=${tag_filters} \
+  --test_timeout="300,450,1200,3600" --local_test_jobs=4 \
+  --test_output=errors --verbose_failures=true --keep_going \
+  --run_under=//tensorflow/tools/ci_build/gpu_build:parallel_gpu_execute \
+  -- ${DEFAULT_BAZEL_TARGETS} -//tensorflow/lite/...
+test_xml_summary_exit

From 6116b7f9114f28dcffd685222285a8c5f7db3daa Mon Sep 17 00:00:00 2001
From: George Karpenkov 
Date: Fri, 19 Jun 2020 17:06:02 -0700
Subject: [PATCH 0687/1390] [XLA] [client] Implement a RunAsync overload which
 does not need a vector of shapes

PiperOrigin-RevId: 317406952
Change-Id: I69d8cc8a68ffdfbf70e2969f5df5e6adba7d2e1d
---
 tensorflow/compiler/xla/client/local_client.cc | 10 ++++++++++
 tensorflow/compiler/xla/client/local_client.h  |  3 +++
 2 files changed, 13 insertions(+)

diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index aa252067e19..5fc9909fa2a 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -320,6 +320,16 @@ StatusOr LocalExecutable::RunAsync(
   return std::move(outputs);
 }
 
+StatusOr LocalExecutable::RunAsync(
+    std::vector arguments, ExecutableRunOptions run_options) {
+  std::vector argument_shapes;
+  argument_shapes.reserve(arguments.size());
+  for (const ExecutionInput& arg : arguments) {
+    argument_shapes.push_back(&arg.shape());
+  }
+  return RunAsync(argument_shapes, std::move(arguments), run_options);
+}
+
 se::Platform* LocalClient::platform() const {
   return local_service_->backend().platform();
 }
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index 3241ac73d54..8b91f4a1739 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -68,6 +68,9 @@ class LocalExecutable {
       absl::Span argument_host_shapes,
       std::vector arguments, ExecutableRunOptions run_options);
 
+  StatusOr RunAsync(std::vector arguments,
+                                     ExecutableRunOptions run_options);
+
   // Return the options used to build the executable.
   const ExecutableBuildOptions& build_options() const { return build_options_; }
 

From f840a6226841eadee32ada80898279b5d9a7ca3b Mon Sep 17 00:00:00 2001
From: Sanjoy Das 
Date: Fri, 19 Jun 2020 17:08:17 -0700
Subject: [PATCH 0688/1390] Rollback "[TF:TRT] Cosmetic fix."

PiperOrigin-RevId: 317407274
Change-Id: I73cd486acf9091e6678e553ab9b0545288f73324
---
 tensorflow/compiler/tf2tensorrt/common/utils.h       |  6 ++++--
 .../compiler/tf2tensorrt/convert/convert_graph.cc    |  6 ++++--
 .../compiler/tf2tensorrt/convert/convert_graph.h     |  6 ++++--
 .../tf2tensorrt/convert/convert_graph_test.cc        |  6 ++++--
 .../compiler/tf2tensorrt/convert/convert_nodes.cc    |  6 ++++--
 .../compiler/tf2tensorrt/convert/convert_nodes.h     |  6 ++++--
 .../tf2tensorrt/convert/convert_nodes_test.cc        |  6 ++++--
 .../compiler/tf2tensorrt/convert/logger_registry.cc  |  6 ++++--
 .../compiler/tf2tensorrt/convert/logger_registry.h   |  5 ++---
 .../tf2tensorrt/convert/trt_optimization_pass.cc     |  6 ++++--
 .../tf2tensorrt/convert/trt_optimization_pass.h      |  6 ++++--
 .../tf2tensorrt/kernels/get_calibration_data_op.cc   |  6 ++++--
 .../compiler/tf2tensorrt/kernels/trt_engine_op.cc    |  6 ++++--
 .../tf2tensorrt/kernels/trt_engine_op_test.cc        |  6 ++++--
 .../tf2tensorrt/kernels/trt_engine_resource_ops.cc   |  6 ++++--
 .../kernels/trt_engine_resource_ops_test.cc          |  6 ++++--
 .../tf2tensorrt/ops/get_calibration_data_op.cc       |  6 ++++--
 tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc |  6 ++++--
 .../tf2tensorrt/ops/trt_engine_resource_ops.cc       |  6 ++++--
 .../compiler/tf2tensorrt/plugin/plugin_cast.cu.cc    |  6 ++++--
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc |  6 ++++--
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h  |  6 ++++--
 tensorflow/compiler/tf2tensorrt/segment/segment.cc   |  6 ++++--
 tensorflow/compiler/tf2tensorrt/segment/segment.h    |  6 ++++--
 .../compiler/tf2tensorrt/segment/segment_test.cc     |  6 ++++--
 tensorflow/compiler/tf2tensorrt/segment/union_find.h |  6 ++++--
 tensorflow/compiler/tf2tensorrt/tensorrt_test.cc     |  6 ++++--
 .../compiler/tf2tensorrt/utils/trt_allocator.cc      | 12 ++++++++----
 .../compiler/tf2tensorrt/utils/trt_allocator.h       | 12 ++++++++----
 .../compiler/tf2tensorrt/utils/trt_engine_utils.cc   |  6 ++++--
 .../compiler/tf2tensorrt/utils/trt_engine_utils.h    |  6 ++++--
 .../tf2tensorrt/utils/trt_int8_calibrator.cc         |  6 ++++--
 .../compiler/tf2tensorrt/utils/trt_int8_calibrator.h |  6 ++++--
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc  |  6 ++++--
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.h   |  6 ++++--
 .../compiler/tf2tensorrt/utils/trt_lru_cache.cc      |  6 ++++--
 .../compiler/tf2tensorrt/utils/trt_lru_cache.h       |  6 ++++--
 .../utils/trt_shape_optimization_profiles.h          |  6 ++++--
 .../utils/trt_shape_optimization_profiles_test.cc    |  6 ++++--
 39 files changed, 162 insertions(+), 83 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/common/utils.h b/tensorflow/compiler/tf2tensorrt/common/utils.h
index b428733ecd4..9ab0145e1ec 100644
--- a/tensorflow/compiler/tf2tensorrt/common/utils.h
+++ b/tensorflow/compiler/tf2tensorrt/common/utils.h
@@ -16,7 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 #define TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "tensorflow/core/platform/logging.h"
 
@@ -28,6 +29,7 @@ namespace tensorrt {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif
+#endif
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
index 5429aaf3362..1c51d51f1c9 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
@@ -53,7 +53,8 @@ limitations under the License.
 #include "tensorflow/core/util/device_name_utils.h"
 #include "tensorflow/tools/graph_transforms/transform_utils.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 namespace tensorflow {
@@ -883,4 +884,5 @@ Status ConvertAfterShapes(const ConversionParams& params) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
index d3897e864fa..53ab84a6fa9 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
@@ -24,7 +24,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -65,6 +66,7 @@ Status RegisterGraphToFunctionLibrary(const GraphDef& segment_graph_def,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
index 54fb1d56441..a1f523d6bfa 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
@@ -34,7 +34,8 @@ limitations under the License.
 #include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -230,4 +231,5 @@ TEST_F(ConvertAfterShapesTest, DirectlyConnectedEngines) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
index 2ec616ba621..96cec556942 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
@@ -59,7 +59,8 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/strided_slice_op.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 #include "third_party/tensorrt/NvInferPlugin.h"
 
@@ -6257,4 +6258,5 @@ bool OutputEdgeValidator::operator()(const Edge* out_edge) const {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
index a621735fad1..7a1276c645c 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
@@ -33,7 +33,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -693,6 +694,7 @@ BinaryOperationMap();
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
index 53ec9ee7ada..c24b169f651 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
@@ -21,7 +21,8 @@ limitations under the License.
 #include 
 #include 
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include 
 #include 
@@ -6635,4 +6636,5 @@ TEST_F(OpConverterTest, ConvertPad) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
index 07c9c2f1ea0..82e68cbb28d 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
@@ -12,7 +12,8 @@ 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.
 ==============================================================================*/
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 
@@ -57,4 +58,5 @@ LoggerRegistry* GetLoggerRegistry() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
index 2a265cf7caa..45b302742d0 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
@@ -19,8 +19,7 @@ limitations under the License.
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
-
+#if GOOGLE_CUDA
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -54,5 +53,5 @@ class RegisterLogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_LOGGER_REGISTRY_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
index 1cf98d135cb..72f4fe5ef9b 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
@@ -28,7 +28,8 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stacktrace.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
@@ -301,4 +302,5 @@ static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif
+#endif
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
index e0aaa5500ab..f79048bb5f6 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
@@ -23,7 +23,8 @@ limitations under the License.
 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -82,5 +83,6 @@ class TRTOptimizationPass : public grappler::CustomGraphOptimizer {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
index 76fb40b9520..3143b06817e 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
@@ -22,7 +22,8 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/lib/core/refcount.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -66,4 +67,5 @@ REGISTER_KERNEL_BUILDER(Name("GetCalibrationDataOp").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index 1094555a622..98d199ca9ab 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -48,7 +48,8 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -1008,4 +1009,5 @@ REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
index 71193dc24cf..a06010de1c7 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
@@ -50,7 +50,8 @@ limitations under the License.
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/public/version.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -305,4 +306,5 @@ TYPED_TEST(TRTEngineOpTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
index 3b6e7e91d3b..2c5821df6ac 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
@@ -33,7 +33,8 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -250,4 +251,5 @@ REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
index 6a073ee24d0..4a24160569d 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
@@ -48,7 +48,8 @@ limitations under the License.
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -245,4 +246,5 @@ TEST_F(TRTEngineResourceOpsTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
index 2af3164c3e2..573172b92e6 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
@@ -13,7 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -33,4 +34,5 @@ Returns calibration data for the given resource name
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
index 2527fe9b910..bd3c2b299a9 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
@@ -13,7 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -58,4 +59,5 @@ REGISTER_OP("TRTEngineOp")
     .Attr("static_engine: bool = true");
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
index 3141092de03..01911de66ec 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
@@ -13,7 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -45,4 +46,5 @@ REGISTER_OP("SerializeTRTResource")
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
index 141a7d1f462..4c0d8b0392a 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
@@ -17,7 +17,8 @@ limitations under the License.
 #include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #define EIGEN_USE_GPU  // For definition of Eigen::GpuDevice.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
@@ -233,4 +234,5 @@ REGISTER_TFTRT_PLUGIN(CastPluginCreator);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
index 83d5f9b5965..563ce724f43 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
@@ -17,7 +17,8 @@ limitations under the License.
 
 #include 
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -29,4 +30,5 @@ const char* kTfTrtPluginNamespace = "TF";
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
index 600ac6683da..bdb046e6c71 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
@@ -20,7 +20,8 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -89,6 +90,7 @@ class TrtPluginRegistrar {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
index d9080b6f69a..32e30006f58 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
@@ -35,7 +35,8 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/env_var.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -1061,4 +1062,5 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h
index 3f79983cfd2..7295c8f0d9d 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.h
@@ -25,7 +25,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -66,6 +67,7 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
index f3bc5bfbee6..2437481a9c4 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
@@ -26,7 +26,8 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -521,4 +522,5 @@ TEST_F(SegmentTest, IncompatibleBatchSizes) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/segment/union_find.h b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
index b53615ec019..70e83c12fca 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/union_find.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
@@ -19,7 +19,8 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/types/optional.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -216,6 +217,7 @@ UnionFind* UnionFind::FindRoot() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_
diff --git a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
index e994d20df33..510591bfe00 100644
--- a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
@@ -18,7 +18,8 @@ limitations under the License.
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/test.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -163,4 +164,5 @@ TEST(TensorrtTest, BasicFunctions) {
 }  // namespace
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
index d4f3a524577..617ea7fad5c 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
@@ -17,9 +17,11 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 namespace tensorflow {
 namespace tensorrt {
@@ -50,7 +52,8 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -110,4 +113,5 @@ void TRTDeviceAllocator::free(void* memory) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
index d219a8a14e8..4ab8b52f523 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
@@ -20,9 +20,11 @@ limitations under the License.
 
 #include "tensorflow/core/framework/allocator.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 namespace tensorflow {
 namespace tensorrt {
@@ -31,7 +33,8 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -66,5 +69,6 @@ class TRTDeviceAllocator : public TRTBaseAllocator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
index 8ccfb8b06f0..ed997b267b1 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
@@ -25,7 +25,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/errors.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -256,4 +257,5 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
index 1ea4fe28cb4..a471749877a 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
@@ -24,7 +24,8 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/status.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -90,6 +91,7 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
index 24271e352a7..554c127fa37 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
@@ -20,7 +20,8 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 
 namespace tensorflow {
@@ -146,4 +147,5 @@ TRTInt8Calibrator::~TRTInt8Calibrator() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif
+#endif
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
index 4c670e85f52..06b39716490 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
@@ -22,7 +22,8 @@ limitations under the License.
 #include 
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -100,5 +101,6 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif
+#endif
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
index e34bf5e7397..193687ebc8c 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
@@ -15,7 +15,8 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 #include "tensorflow/core/platform/logging.h"
@@ -67,4 +68,5 @@ REGISTER_TENSORRT_LOGGER("DefaultLogger", Logger::GetLogger());
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
index ce6552e8fe9..2ade1b48f47 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
@@ -18,7 +18,8 @@ limitations under the License.
 
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -39,6 +40,7 @@ class Logger : public nvinfer1::ILogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
index ee7e6272372..fbcdaad52c0 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
@@ -23,7 +23,8 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -140,4 +141,5 @@ EngineContext* TRTEngineCacheResource::GetEngineContext(const int profile_id) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
index 991b9a949e4..8e345254f75 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
@@ -115,7 +115,8 @@ class LRUCache {
   }
 };
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 struct EngineContext {
   EngineContext() {}  // Creates an empty context.
@@ -222,7 +223,8 @@ class TRTEngineCacheResource : public ResourceBase {
   TrtShapeOptimizationProfile profiles_;
 };
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
index fc688b14139..40c7f5dcf31 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
@@ -29,7 +29,8 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -172,5 +173,6 @@ class TrtShapeOptimizationProfile {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
index 32c2200fb71..501810587e0 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
@@ -13,7 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA && GOOGLE_TENSORRT
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
 
 #include 
 
@@ -213,4 +214,5 @@ TEST_F(TrtShapeOptimizationProfileTest, Dynamic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA

From 715b02167d188d15f0205273b93ea6bbd606f4c5 Mon Sep 17 00:00:00 2001
From: Wenhao Jia 
Date: Fri, 19 Jun 2020 17:27:50 -0700
Subject: [PATCH 0689/1390] Restore TpuPlatform auto registration code.

PiperOrigin-RevId: 317409587
Change-Id: If44d7a39a45c4c7026f70a4d79d965a54c4db295
---
 tensorflow/stream_executor/tpu/tpu_platform.cc | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/stream_executor/tpu/tpu_platform.cc b/tensorflow/stream_executor/tpu/tpu_platform.cc
index 13a845829c1..4bccd822e91 100644
--- a/tensorflow/stream_executor/tpu/tpu_platform.cc
+++ b/tensorflow/stream_executor/tpu/tpu_platform.cc
@@ -134,4 +134,12 @@ void RegisterTpuPlatform() {
   }
 }
 
+REGISTER_MODULE_INITIALIZER(tpu_platform, RegisterTpuPlatform());
+
+// Note that module initialization sequencing is not supported in the
+// open-source project, so this will be a no-op there.
+REGISTER_MODULE_INITIALIZER_SEQUENCE(tpu_platform, multi_platform_manager);
+REGISTER_MODULE_INITIALIZER_SEQUENCE(multi_platform_manager_listener,
+                                     tpu_platform);
+
 }  // namespace tensorflow

From 3427843d707cf166cbd07755bd11a1c7dea76730 Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Fri, 19 Jun 2020 17:33:29 -0700
Subject: [PATCH 0690/1390] Dist strat interop

PiperOrigin-RevId: 317410192
Change-Id: Ibfd1e3ac143422ccffa5f240075f5ae93a90ad07
---
 .../python/ops/numpy_ops/np_interop_test.py   | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_interop_test.py b/tensorflow/python/ops/numpy_ops/np_interop_test.py
index f52d3dae78b..9580b787202 100644
--- a/tensorflow/python/ops/numpy_ops/np_interop_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_interop_test.py
@@ -22,8 +22,13 @@ from __future__ import print_function
 import numpy as onp
 
 
+from tensorflow.python.distribute import distribution_strategy_context
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.distribute import reduce_util
 from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import control_flow_ops
@@ -36,6 +41,17 @@ from tensorflow.python.platform import test
 
 class InteropTest(test.TestCase):
 
+  def setUp(self):
+    super(InteropTest, self).setUp()
+    physical_devices = config.list_physical_devices('CPU')
+    configs = config.get_logical_device_configuration(physical_devices[0])
+    if configs is None:
+      logical_devices = [
+          context.LogicalDeviceConfiguration() for _ in range(3)
+      ]
+      config.set_logical_device_configuration(physical_devices[0],
+                                              logical_devices)
+
   def testGradientTapeInterop(self):
     with backprop.GradientTape() as t:
       x = np_array_ops.asarray(3.0)
@@ -139,6 +155,39 @@ class InteropTest(test.TestCase):
 
 #     self.assertEqual(t.numpy(), [1., 2., 3.])
 
+  def testDistStratInterop(self):
+    strategy = mirrored_strategy.MirroredStrategy(
+        devices=['CPU:0', 'CPU:1', 'CPU:2'])
+
+    multiplier = np_array_ops.asarray(5.)
+
+    with strategy.scope():
+      @def_function.function
+      def run():
+        ctx = distribution_strategy_context.get_replica_context()
+        val = np_array_ops.asarray(ctx.replica_id_in_sync_group)
+        return val * multiplier
+
+      distributed_values = strategy.run(run)
+      reduced = strategy.reduce(reduce_util.ReduceOp.SUM,
+                                distributed_values, axis=None)
+
+    values = distributed_values.values
+
+    # Note that this should match the number of virtual CPUs.
+    self.assertLen(values, 3)
+    self.assertIsInstance(values[0], np_arrays.ndarray)
+    self.assertIsInstance(values[1], np_arrays.ndarray)
+    self.assertIsInstance(values[2], np_arrays.ndarray)
+    self.assertAllClose(values[0], 0)
+    self.assertAllClose(values[1], 5)
+    self.assertAllClose(values[2], 10)
+
+    # "strategy.reduce" doesn't rewrap in ndarray.
+    # self.assertIsInstance(reduced, np_arrays.ndarray)
+    self.assertAllClose(reduced, 15)
+
+
 if __name__ == '__main__':
   ops.enable_eager_execution()
   test.main()

From d737ef92f9c9e251ee3e9ad21090c489281a9dc9 Mon Sep 17 00:00:00 2001
From: XingyuLong 
Date: Fri, 19 Jun 2020 20:45:26 -0400
Subject: [PATCH 0691/1390] Update

---
 .../filesystem/plugins/gcs/gcs_helper.cc      |  34 ++
 .../filesystem/plugins/gcs/gcs_helper.h       |  33 ++
 .../tf_saved_model/remove_init_variable_v1.py |  74 +++
 .../transforms/device_index_selector.cc       |  85 +++
 .../tests/tf_device_index_selector.mlir       |  25 +
 tensorflow/compiler/tests/case_test.py        |  87 +++
 .../compiler/tf2tensorrt/common/utils.h       |  35 ++
 .../service/cpu/test_target_triple_helper.h   |  28 +
 .../xla/service/gpu/reduction_splitter.cc     | 117 ++++
 .../xla/service/gpu/reduction_splitter.h      |  49 ++
 .../service/gpu/reduction_splitter_test.cc    | 140 +++++
 .../xla/tests/manifest_checking_test.cc       | 129 +++++
 .../xla/tests/manifest_checking_test.h        |  35 ++
 .../api_def_BandedTriangularSolve.pbtxt       |   4 +
 .../api_def/base_api/api_def_BesselI0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselI1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselJ0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselJ1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK0e.pbtxt  |   4 +
 .../api_def/base_api/api_def_BesselK1.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselK1e.pbtxt  |   4 +
 .../api_def/base_api/api_def_BesselY0.pbtxt   |   4 +
 .../api_def/base_api/api_def_BesselY1.pbtxt   |   4 +
 ...tatelessParameterizedTruncatedNormal.pbtxt |  54 ++
 .../api_def_BandedTriangularSolve.pbtxt       |   4 +
 .../kernels/banded_triangular_solve_op.cc     | 293 ++++++++++
 .../banded_triangular_solve_op_test.cc        | 180 ++++++
 tensorflow/core/kernels/cwise_op_neg_1.cc     |  44 ++
 tensorflow/core/kernels/cwise_op_neg_2.cc     |  26 +
 .../mlir_generated_op_gpu_tanh_test.cc        |  85 +++
 .../special_math/special_math_op_bessel.cc    |  78 +++
 .../special_math_op_gpu_bessel.cu.cc          |  41 ++
 .../core/kernels/topk_op_gpu_uint32.cu.cc     |  28 +
 .../core/kernels/topk_op_gpu_uint64.cu.cc     |  28 +
 .../BandedTriangularSolve.pbtxt               |  42 ++
 .../ops/compat/ops_history_v2/BesselI0.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselI1.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselJ0.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselJ1.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselK0.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselK0e.pbtxt |  23 +
 .../ops/compat/ops_history_v2/BesselK1.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselK1e.pbtxt |  23 +
 .../ops/compat/ops_history_v2/BesselY0.pbtxt  |  23 +
 .../ops/compat/ops_history_v2/BesselY1.pbtxt  |  23 +
 ...tatelessParameterizedTruncatedNormal.pbtxt |  65 +++
 .../utils/op_metrics_db_utils_test.cc         |  46 ++
 .../tpu_compilation_cache_entry_impl.h        | 108 ++++
 .../tpu_compilation_cache_entry_unloader.h    |  69 +++
 .../kernels/tpu_compilation_cache_interface.h | 355 ++++++++++++
 tensorflow/core/tpu/kernels/tpu_op_consts.cc  |  24 +
 tensorflow/core/tpu/kernels/tpu_op_consts.h   |  39 ++
 tensorflow/core/tpu/kernels/tpu_op_util.cc    | 151 +++++
 tensorflow/core/tpu/kernels/tpu_op_util.h     |  40 ++
 tensorflow/core/tpu/tpu_library_init_fns.inc  | 166 ++++++
 .../support/metadata/build_defs.bzl           |  43 ++
 .../metadata/cc/metadata_parser.h.template    |  28 +
 .../metadata/cc/test/metadata_parser_test.cc  |  33 ++
 .../lite/support/metadata/MetadataParser.java |  27 +
 .../metadata/metadata_parser.py.template      |  26 +
 .../support/metadata/metadata_parser_test.py  |  38 ++
 .../himax_we1_evb/detection_responder.cc      |  33 ++
 .../himax_we1_evb/image_provider.cc           |  41 ++
 .../lite/micro/himax_we1_evb/debug_log.cc     |  32 ++
 .../make/targets/himax_we1_evb_makefile.inc   |  91 +++
 .../tools/optimize/testdata/mixed16x8.bin     | Bin 0 -> 1184 bytes
 .../tools/optimize/testdata/transpose.bin     | Bin 0 -> 544 bytes
 .../python/framework/python_op_gen_test.cc    |  42 ++
 .../distribute/mirrored_strategy_test.py      |  89 +++
 .../distribute/mirrored_variable_test.py      | 106 ++++
 .../integration_test/tpu_strategy_test.py     |  69 +++
 .../banded_triangular_solve_op_test.py        | 232 ++++++++
 .../parameterized_truncated_normal_op_test.py | 520 ++++++++++++++++++
 tensorflow/python/ops/numpy_ops/README.md     |  94 ++++
 .../python/ops/numpy_ops/np_accessor.py       |  32 ++
 .../python/ops/numpy_ops/np_interop_test.py   | 144 +++++
 tensorflow/security/fuzzing/op_fuzzing/BUILD  |  39 ++
 .../fuzzing/op_fuzzing/fuzz_session.h         | 156 ++++++
 .../fuzzing/op_fuzzing/identity_fuzz.cc       |  45 ++
 .../security/fuzzing/status_group_fuzz.cc     |  66 +++
 ...flow.distribute.-distributed-dataset.pbtxt |  16 +
 ...low.distribute.-distributed-iterator.pbtxt |  20 +
 .../linux/mkl/install_openmpi_horovod.sh      |  80 +++
 .../per_release/scripts/nonpip_gpu.sh         |  75 +++
 85 files changed, 5328 insertions(+)
 create mode 100644 tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
 create mode 100644 tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
 create mode 100644 tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py
 create mode 100644 tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
 create mode 100644 tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
 create mode 100644 tensorflow/compiler/tests/case_test.py
 create mode 100644 tensorflow/compiler/tf2tensorrt/common/utils.h
 create mode 100644 tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter.h
 create mode 100644 tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc
 create mode 100644 tensorflow/compiler/xla/tests/manifest_checking_test.cc
 create mode 100644 tensorflow/compiler/xla/tests/manifest_checking_test.h
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselI1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselJ0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselJ1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK0e.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselK1e.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselY0.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_BesselY1.pbtxt
 create mode 100644 tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
 create mode 100644 tensorflow/core/api_def/python_api/api_def_BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/kernels/banded_triangular_solve_op.cc
 create mode 100644 tensorflow/core/kernels/banded_triangular_solve_op_test.cc
 create mode 100644 tensorflow/core/kernels/cwise_op_neg_1.cc
 create mode 100644 tensorflow/core/kernels/cwise_op_neg_2.cc
 create mode 100644 tensorflow/core/kernels/mlir_generated_op_gpu_tanh_test.cc
 create mode 100644 tensorflow/core/kernels/special_math/special_math_op_bessel.cc
 create mode 100644 tensorflow/core/kernels/special_math/special_math_op_gpu_bessel.cu.cc
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BandedTriangularSolve.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselI0.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselI1.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselJ0.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselJ1.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselK0.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselK0e.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselK1.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselK1e.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselY0.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/BesselY1.pbtxt
 create mode 100644 tensorflow/core/ops/compat/ops_history_v2/StatelessParameterizedTruncatedNormal.pbtxt
 create mode 100644 tensorflow/core/profiler/utils/op_metrics_db_utils_test.cc
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_impl.h
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compilation_cache_entry_unloader.h
 create mode 100644 tensorflow/core/tpu/kernels/tpu_compilation_cache_interface.h
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_consts.cc
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_consts.h
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_util.cc
 create mode 100644 tensorflow/core/tpu/kernels/tpu_op_util.h
 create mode 100644 tensorflow/core/tpu/tpu_library_init_fns.inc
 create mode 100644 tensorflow/lite/experimental/support/metadata/build_defs.bzl
 create mode 100644 tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
 create mode 100644 tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
 create mode 100644 tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
 create mode 100644 tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
 create mode 100644 tensorflow/lite/experimental/support/metadata/metadata_parser_test.py
 create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/detection_responder.cc
 create mode 100644 tensorflow/lite/micro/examples/person_detection_experimental/himax_we1_evb/image_provider.cc
 create mode 100644 tensorflow/lite/micro/himax_we1_evb/debug_log.cc
 create mode 100644 tensorflow/lite/micro/tools/make/targets/himax_we1_evb_makefile.inc
 create mode 100644 tensorflow/lite/tools/optimize/testdata/mixed16x8.bin
 create mode 100644 tensorflow/lite/tools/optimize/testdata/transpose.bin
 create mode 100644 tensorflow/python/framework/python_op_gen_test.cc
 create mode 100644 tensorflow/python/keras/distribute/mirrored_strategy_test.py
 create mode 100644 tensorflow/python/keras/distribute/mirrored_variable_test.py
 create mode 100644 tensorflow/python/keras/integration_test/tpu_strategy_test.py
 create mode 100644 tensorflow/python/kernel_tests/banded_triangular_solve_op_test.py
 create mode 100644 tensorflow/python/kernel_tests/random/parameterized_truncated_normal_op_test.py
 create mode 100644 tensorflow/python/ops/numpy_ops/README.md
 create mode 100644 tensorflow/python/ops/numpy_ops/np_accessor.py
 create mode 100644 tensorflow/python/ops/numpy_ops/np_interop_test.py
 create mode 100644 tensorflow/security/fuzzing/op_fuzzing/BUILD
 create mode 100644 tensorflow/security/fuzzing/op_fuzzing/fuzz_session.h
 create mode 100644 tensorflow/security/fuzzing/op_fuzzing/identity_fuzz.cc
 create mode 100644 tensorflow/security/fuzzing/status_group_fuzz.cc
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-dataset.pbtxt
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.distribute.-distributed-iterator.pbtxt
 create mode 100755 tensorflow/tools/ci_build/linux/mkl/install_openmpi_horovod.sh
 create mode 100644 tensorflow/tools/ci_build/per_release/scripts/nonpip_gpu.sh

diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
new file mode 100644
index 00000000000..4504a9f3b35
--- /dev/null
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
@@ -0,0 +1,34 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
+
+#include 
+
+#include 
+#include 
+#include 
+
+TempFile::TempFile(const char* temp_file_name, std::ios::openmode mode)
+    : std::fstream(temp_file_name, mode), name_(temp_file_name) {}
+
+TempFile::TempFile(TempFile&& rhs)
+    : std::fstream(std::move(rhs)), name_(std::move(rhs.name_)) {}
+
+TempFile::~TempFile() {
+  std::fstream::close();
+  std::remove(name_.c_str());
+}
+
+const std::string TempFile::getName() const { return name_; }
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
new file mode 100644
index 00000000000..1a521ca4f1e
--- /dev/null
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
+
+#include 
+#include 
+
+class TempFile : public std::fstream {
+ public:
+  // We should specify openmode each time we call TempFile.
+  TempFile(const char* temp_file_name, std::ios::openmode mode);
+  TempFile(TempFile&& rhs);
+  ~TempFile() override;
+  const std::string getName() const;
+
+ private:
+  const std::string name_;
+};
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py
new file mode 100644
index 00000000000..117132649d7
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py
@@ -0,0 +1,74 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+# RUN: %p/remove_init_variable_v1 | FileCheck %s
+
+# pylint: disable=missing-docstring,line-too-long
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v1 as tf
+from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1
+
+# Verify that the tf.versions attribute exists. It is difficult to enforce
+# contents, since the version numbers change over time. The conversion logic
+# itself is verified in the common graphdef converter, so here just assert
+# it is being invoked.
+# CHECK: module
+# CHECK-SAME: tf.versions
+# CHECK-SAME: bad_consumers
+# CHECK-SAME: min_consumer
+# CHECK-SAME: producer
+
+# CHECK: "tf_saved_model.global_tensor"() {is_mutable, sym_name = "[[VAR:[a-zA-Z_0-9]+]]", type = tensor<1x3xf32>, value = {{.*}} : tensor<1x3xf32>} : () -> ()
+# CHECK-NOT: session_initializer
+
+# CHECK:      func {{@[a-zA-Z_0-9]+}}(
+# CHECK-SAME:   [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},
+# CHECK-SAME:   [[ARG1:%.*]]: tensor>> {tf_saved_model.bound_input = @[[VAR]]})
+# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
+# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]
+
+# CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor>>) -> tensor<1x3xf32>
+# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
+# CHECK-NEXT: return [[R1]] : tensor<3x3xf32>
+
+
+def Test():
+
+  x = tf.constant([[1.0], [1.0], [1.0]])
+  y = tf.compat.v1.get_variable(
+      name='y',
+      shape=(1, 3),
+      initializer=tf.random_normal_initializer(),
+      trainable=True)
+  r = tf.matmul(x, y)
+
+  tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
+  tensor_info_r = tf.compat.v1.saved_model.utils.build_tensor_info(r)
+
+  return {
+      'key': (tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
+          inputs={'x': tensor_info_x},
+          outputs={'r': tensor_info_r},
+          method_name='some_function'))
+  }
+
+
+if __name__ == '__main__':
+  common_v1.set_tf_options()
+  common_v1.do_test(
+      Test(), tf.initializers.global_variables(), canonicalize=True)
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
new file mode 100644
index 00000000000..550647a915a
--- /dev/null
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
@@ -0,0 +1,85 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+// Converts DeviceIndex to constant device.
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+
+namespace mlir {
+namespace TF {
+namespace {
+
+// Folds the DeviceIndex op to a constant value. The DeviceIndex return the
+// index of the device the op should run on. The user can use this to provide
+// different op specializations. E.g.,
+//
+// ```mlir
+//  %1 = "tf.DeviceIndex"()
+//          {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+//  %4 = "tf.Case"(%1, %arg0, %arg1)
+//          {branches = [@foo, @baz], output_shapes = [#tf.shape<>]} :
+//            (tensor, tensor, tensor) -> tensor
+// ```
+//
+// Shows an example where there are 2 different functions which could be
+// executed to produce the same values but with different functions optimized
+// for CPU or GPU.
+struct DeviceIndexSelector
+    : public PassWrapper> {
+  void runOnOperation() override;
+};
+
+}  // namespace
+
+void DeviceIndexSelector::runOnOperation() {
+  FuncOp func = getOperation();
+  // Convert all the DeviceIndex ops to constant values.
+  func.getBody().walk([](TF::DeviceIndexOp op) {
+    // This just selects the default in all cases where DeviceIndex feeds into
+    // tf.Case. This could be enhanced to have some sort of policy in the
+    // future.
+    OpBuilder b(op);
+    RankedTensorType type = RankedTensorType::get({}, b.getIntegerType(32));
+    int index = op.device_names().size();
+    for (auto use : op.getOperation()->getUsers()) {
+      // Skip if it doesn't feed into case. Alternatively this could always
+      // return the CPU device index if it exists.
+      if (!isa(use)) return;
+    }
+    DenseElementsAttr attr =
+        DenseElementsAttr::get(type, b.getI32IntegerAttr(index));
+    auto constant = b.create(op.getLoc(), type, attr);
+    op.replaceAllUsesWith(constant.getOperation());
+    op.erase();
+  });
+}
+
+// Creates an instance of the TensorFlow DeviceIndex selector pass.
+std::unique_ptr> CreateDeviceIndexSelectorPass() {
+  return std::make_unique();
+}
+
+static PassRegistration pass(
+    "tf-device-index-selector", "Fold tf.DeviceIndex to constant");
+
+}  // namespace TF
+}  // namespace mlir
diff --git a/tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir b/tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
new file mode 100644
index 00000000000..7fc2b210f91
--- /dev/null
+++ b/tensorflow/compiler/tensorflow/tests/tf_device_index_selector.mlir
@@ -0,0 +1,25 @@
+// Test DeviceIndex selector.
+
+// RUN: tf-opt --tf-device-index-selector %s | FileCheck %s
+
+// CHECK-LABEL: func @select
+func @select(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) {
+  // CHECK:  %[[first:.*]] = "tf.DeviceIndex"
+  // CHECK: constant dense<2>
+  // CHECK:  return %[[first]],
+  %0 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+  %1 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor
+  %4 = "tf.Case"(%1, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf.shape<>]} : (tensor, tensor, tensor) -> tensor
+
+  return %0, %4 : tensor, tensor
+}
+
+func @add(%i: tensor, %arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
+  %0 = "tf.Add"(%arg0, %arg1): (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
+
+func @sub(%i: tensor, %arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
+  %0 = "tf.Sub"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  return %0 : tensor<*xf32>
+}
diff --git a/tensorflow/compiler/tests/case_test.py b/tensorflow/compiler/tests/case_test.py
new file mode 100644
index 00000000000..3b2dff537da
--- /dev/null
+++ b/tensorflow/compiler/tests/case_test.py
@@ -0,0 +1,87 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""Tests for while loops in XLA."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import test
+
+
+class CaseTest(xla_test.XLATestCase):
+
+  def testCaseBasic(self):
+
+    @def_function.function(experimental_compile=True)
+    def switch_case_test(branch_index):
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        return array_ops.constant(31)
+
+      def f3():
+        return array_ops.constant(-1)
+
+      return control_flow_ops.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f3)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test(array_ops.constant(0)).numpy(), 17)
+      self.assertEqual(switch_case_test(array_ops.constant(1)).numpy(), 31)
+      self.assertEqual(switch_case_test(array_ops.constant(2)).numpy(), -1)
+      self.assertEqual(switch_case_test(array_ops.constant(3)).numpy(), -1)
+
+  def testBranchIsPruned(self):
+
+    @def_function.function(experimental_compile=True)
+    def switch_case_test():
+      branch_index = array_ops.constant(0)
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        # Some operations that XLA cannot compile.
+        image_ops.decode_image(io_ops.read_file('/tmp/bmp'))
+        return array_ops.constant(31)
+
+      # This tests that we do not try to compile all branches if the branch
+      # index in trivially constant.
+      return control_flow_ops.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f2)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test().numpy(), 17)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()
diff --git a/tensorflow/compiler/tf2tensorrt/common/utils.h b/tensorflow/compiler/tf2tensorrt/common/utils.h
new file mode 100644
index 00000000000..9ab0145e1ec
--- /dev/null
+++ b/tensorflow/compiler/tf2tensorrt/common/utils.h
@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
+#define TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+#define LOG_WARNING_WITH_PREFIX LOG(WARNING) << "TF-TRT Warning: "
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif
+#endif
+
+#endif  // TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
diff --git a/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h
new file mode 100644
index 00000000000..857de4a8143
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/test_target_triple_helper.h
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
+#define TENSORFLOW_TEST_TARGET_TRIPLE_HELPER_H_
+
+#if (defined(__powerpc__) || \
+     defined(__ppc__) && (__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__))
+static const char kTargetCpuForHost[] = "ppc";
+static const char kTargetTripleForHost[] = "ppc64le-ibm-linux-gnu";
+#else
+static const char kTargetCpuForHost[] = "";
+static const char kTargetTripleForHost[] = "x86_64-pc-linux";
+#endif
+
+#endif
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc b/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
new file mode 100644
index 00000000000..b68213ec35f
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter.cc
@@ -0,0 +1,117 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/gpu/reduction_splitter.h"
+
+#include 
+
+#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h"
+#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
+#include "tensorflow/compiler/xla/service/hlo_instructions.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+
+namespace xla {
+namespace gpu {
+
+class ReductionSplitterVisitor : public DfsHloRewriteVisitor {
+ public:
+  Status HandleReduce(HloInstruction *reduce) override {
+    VLOG(4) << "Input: " << reduce->ToString();
+
+    // Reductions with contiguous dimensions are lowered to efficient code. No
+    // need to split such ops.
+    if (IsReductionFromOrToContiguousDimensions(*reduce)) {
+      return Status::OK();
+    }
+    if (reduce->dimensions().size() < 2) {
+      return Status::OK();
+    }
+    if (!reduce->shape().IsArray()) {
+      // TODO(cheshire): Handle variadic reduction.
+      return Status::OK();
+    }
+
+    HloInstruction *operand = reduce->mutable_operand(0);
+    const Shape &shape = operand->shape();
+    CHECK(shape == LayoutUtil::GetWithDefaultLayout(shape))
+        << "Default layout should be enforced on reduction operand";
+    // Verify that contiguous dimensions have been grouped by the
+    // ReductionDimensionGrouper pass.
+    for (int64 i = 0; i < reduce->dimensions().size(); ++i) {
+      for (int64 j = i + 1; j < reduce->dimensions().size(); ++j) {
+        CHECK(abs(reduce->dimensions(i) - reduce->dimensions(j)) > 1)
+            << "Reduction dimensions must not be consecutive";
+      }
+    }
+
+    // The reduce op has non-contiguous dimensions. Look for the dimension with
+    // the largest shape dimension. Reducing along this dimension first will
+    // reduce the output size most effectively.
+    int64 max_shape_dim = 0;
+    int64 max_reduce_dim = 0;
+    const auto &input_shape = reduce->operand(0)->shape();
+    for (int64 i = 0; i < reduce->dimensions().size(); ++i) {
+      if (input_shape.dimensions(reduce->dimensions(i)) > max_shape_dim) {
+        max_reduce_dim = reduce->dimensions(i);
+        max_shape_dim = input_shape.dimensions(max_reduce_dim);
+      }
+    }
+    // TODO(tjoerg): Run microbenchmarks to tune this threshold.
+    if (max_shape_dim < 128) {
+      return Status::OK();
+    }
+
+    // Split the reduction into a pre-reduction and a final reduction.
+    VLOG(3) << "Splitting reduction " << reduce->name() << " at dimension "
+            << max_reduce_dim;
+    std::vector pre_reduce_dims;
+    pre_reduce_dims.push_back(max_reduce_dim);
+    std::vector pre_reduce_shape_dims(input_shape.dimensions().begin(),
+                                             input_shape.dimensions().end());
+    pre_reduce_shape_dims.erase(pre_reduce_shape_dims.begin() + max_reduce_dim);
+    Shape pre_reduce_shape = ShapeUtil::MakeShape(
+        reduce->shape().element_type(), pre_reduce_shape_dims);
+    std::unique_ptr pre_reduce = HloInstruction::CreateReduce(
+        pre_reduce_shape, reduce->mutable_operand(0),
+        reduce->mutable_operand(1), pre_reduce_dims, reduce->to_apply());
+    pre_reduce->set_metadata(reduce->metadata());
+
+    std::vector final_reduce_dims(reduce->dimensions().begin(),
+                                         reduce->dimensions().end());
+    final_reduce_dims.erase(
+        std::remove(final_reduce_dims.begin(), final_reduce_dims.end(),
+                    max_reduce_dim),
+        final_reduce_dims.end());
+    for (int64 i = 0; i < final_reduce_dims.size(); ++i) {
+      if (final_reduce_dims[i] > max_reduce_dim) {
+        final_reduce_dims[i]--;
+      }
+    }
+    std::unique_ptr final_reduce = HloInstruction::CreateReduce(
+        reduce->shape(),
+        reduce->parent()->AddInstruction(std::move(pre_reduce)),
+        reduce->mutable_operand(1), final_reduce_dims, reduce->to_apply());
+    return ReplaceWithNewInstruction(reduce, std::move(final_reduce));
+  }
+};
+
+StatusOr ReductionSplitter::Run(HloModule *module) {
+  TF_ASSIGN_OR_RETURN(bool changed,
+                      ReductionSplitterVisitor().RunOnModule(module));
+  return changed;
+}
+
+}  // namespace gpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter.h b/tensorflow/compiler/xla/service/gpu/reduction_splitter.h
new file mode 100644
index 00000000000..f161b579eb8
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter.h
@@ -0,0 +1,49 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
+
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
+
+namespace xla {
+namespace gpu {
+
+// Splits a reduce op into two consecutive reduce ops if
+// * the reduce dimensions are not contiguous and
+// * at least one reduce dimension is large (i.e. corresponds to a large input
+//   shape dimension).
+//
+// Reductions with non-contiguous dimensions are emitted as simple element-wise
+// loops. This is inefficient when reducing large input shape dimensions.
+// Splitting such reductions allows using more efficient reduction emitters.
+//
+// This pass splits reduce ops into two consecutive reduce ops. Run it to a
+// fixpoint to split reduce ops along multiple large dimensions.
+//
+// Precondition: ReductionDimensionGrouper has been run and adjacent reduce
+// dimentsions have been grouped. Reduction layouts have been normalized.
+
+class ReductionSplitter : public HloModulePass {
+ public:
+  absl::string_view name() const override { return "reduction-splitter"; }
+
+  StatusOr Run(HloModule* module) override;
+};
+
+}  // namespace gpu
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_REDUCTION_SPLITTER_H_
diff --git a/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc b/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc
new file mode 100644
index 00000000000..1be55b84204
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/reduction_splitter_test.cc
@@ -0,0 +1,140 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/gpu/reduction_splitter.h"
+
+#include "tensorflow/compiler/xla/service/hlo_matchers.h"
+#include "tensorflow/compiler/xla/service/hlo_parser.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/test_helpers.h"
+#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+
+namespace xla {
+namespace gpu {
+namespace {
+
+namespace op = xla::testing::opcode_matchers;
+
+class ReductionSplitterTest : public HloTestBase {};
+
+TEST_F(ReductionSplitterTest, SplitReductionAtDimensionTwo) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f16[6,16,512,64]{3,2,1,0} parameter(0)
+    transpose.1781 = f16[6,512,16,64]{3,1,2,0} transpose(param_0), dimensions={0,2,1,3}
+    convert.6986 = f32[6,512,16,64]{3,1,2,0} convert(transpose.1781)
+    bitcast.2136 = f32[6,16,512,64]{3,2,1,0} bitcast(convert.6986)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[16,64]{1,0} reduce(bitcast.2136, constant_11111), dimensions={0,2}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  ASSERT_TRUE(ReductionSplitter().Run(module.get()).ValueOrDie());
+  SCOPED_TRACE(module->ToString());
+  const HloInstruction* root_reduction =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(root_reduction, op::Reduce(op::Reduce(), op::Constant()));
+
+  auto* pre_reduction = root_reduction->operand(0);
+  EXPECT_THAT(pre_reduction->dimensions(), std::vector({2}));
+  EXPECT_THAT(pre_reduction->shape(), ShapeUtil::MakeShape(F32, {6, 16, 64}));
+  EXPECT_THAT(root_reduction->dimensions(), std::vector({0}));
+  EXPECT_THAT(root_reduction->shape(), ShapeUtil::MakeShape(F32, {16, 64}));
+}
+
+TEST_F(ReductionSplitterTest, SplitReductionAtDimensionZero) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[1024,16,512,64,128]{4,3,2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[16,64]{1,0} reduce(param_0, constant_11111), dimensions={2,0,4}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  ASSERT_TRUE(ReductionSplitter().Run(module.get()).ValueOrDie());
+  SCOPED_TRACE(module->ToString());
+  const HloInstruction* root_reduction =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(root_reduction, op::Reduce(op::Reduce(), op::Constant()));
+
+  auto* pre_reduction = root_reduction->operand(0);
+  EXPECT_THAT(pre_reduction->dimensions(), std::vector({0}));
+  EXPECT_THAT(pre_reduction->shape(),
+              ShapeUtil::MakeShape(F32, {16, 512, 64, 128}));
+  EXPECT_THAT(root_reduction->dimensions(), std::vector({1, 3}));
+  EXPECT_THAT(root_reduction->shape(), ShapeUtil::MakeShape(F32, {16, 64}));
+}
+
+TEST_F(ReductionSplitterTest, DontSplitReductionWithSmallDimensions) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[8,1024,8]{2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    ROOT reduce.982 = f32[1024]{0} reduce(param_0, constant_11111), dimensions={2,0}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  EXPECT_FALSE(ReductionSplitter().Run(module.get()).ValueOrDie());
+}
+
+TEST_F(ReductionSplitterTest, DontSplitReductionsWithContiguousDimensions) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+  HloModule test
+
+  add_computation {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT add = f32[] add(x, y)
+  }
+
+  ENTRY entry_computation {
+    param_0 = f32[128,128,64,128]{3,2,1,0} parameter(0)
+    constant_11111 = f32[] constant(0)
+    // The dimenstions to keep (1 and 2) are contiguous.
+    ROOT reduce.982 = f32[128,64]{1,0} reduce(param_0, constant_11111), dimensions={3,0}, to_apply=add_computation
+  }
+  )")
+                    .ValueOrDie();
+  EXPECT_FALSE(ReductionSplitter().Run(module.get()).ValueOrDie());
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/manifest_checking_test.cc b/tensorflow/compiler/xla/tests/manifest_checking_test.cc
new file mode 100644
index 00000000000..8806290472d
--- /dev/null
+++ b/tensorflow/compiler/xla/tests/manifest_checking_test.cc
@@ -0,0 +1,129 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/tests/manifest_checking_test.h"
+
+#include 
+#include 
+#include 
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_split.h"
+#include "tensorflow/compiler/xla/tests/test_macros.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/regexp.h"
+
+namespace xla {
+
+namespace {
+
+// Mapping from test name; i.e. MyTest.MyTestCase to platforms on which it is
+// disabled - a sequence of regexps.
+using ManifestT = absl::flat_hash_map>;
+
+ManifestT ReadManifest() {
+  ManifestT manifest;
+
+  absl::string_view path = absl::NullSafeStringView(kDisabledManifestPath);
+  if (path.empty()) {
+    return manifest;
+  }
+
+  // Note: parens are required to disambiguate vs function decl.
+  std::ifstream file_stream((std::string(path)));
+  std::string contents((std::istreambuf_iterator(file_stream)),
+                       std::istreambuf_iterator());
+
+  std::vector lines = absl::StrSplit(contents, '\n');
+  for (std::string& line : lines) {
+    auto comment = line.find("//");
+    if (comment != std::string::npos) {
+      line = line.substr(0, comment);
+    }
+    if (line.empty()) {
+      continue;
+    }
+    absl::StripTrailingAsciiWhitespace(&line);
+    std::vector pieces = absl::StrSplit(line, ' ');
+    CHECK_GE(pieces.size(), 1);
+    auto& platforms = manifest[pieces[0]];
+    for (size_t i = 1; i < pieces.size(); ++i) {
+      platforms.push_back(pieces[i]);
+    }
+  }
+  return manifest;
+}
+
+}  // namespace
+
+void ManifestCheckingTest::SetUp() {
+  const testing::TestInfo* test_info =
+      testing::UnitTest::GetInstance()->current_test_info();
+  absl::string_view test_case_name = test_info->test_suite_name();
+  absl::string_view test_name = test_info->name();
+  VLOG(1) << "test_case_name: " << test_case_name;
+  VLOG(1) << "test_name: " << test_name;
+
+  // Remove the type suffix from the test case name.
+  if (const char* type_param = test_info->type_param()) {
+    VLOG(1) << "type_param: " << type_param;
+    size_t last_slash = test_case_name.rfind('/');
+    test_case_name = test_case_name.substr(0, last_slash);
+    VLOG(1) << "test_case_name: " << test_case_name;
+  }
+
+  // Remove the test instantiation name if it is present.
+  auto first_slash = test_case_name.find('/');
+  if (first_slash != test_case_name.npos) {
+    test_case_name.remove_prefix(first_slash + 1);
+    VLOG(1) << "test_case_name: " << test_case_name;
+  }
+
+  ManifestT manifest = ReadManifest();
+
+  // If the test name ends with a slash followed by one or more characters,
+  // strip that off.
+  auto last_slash = test_name.rfind('/');
+  if (last_slash != test_name.npos) {
+    test_name = test_name.substr(0, last_slash);
+    VLOG(1) << "test_name: " << test_name;
+  }
+
+  // First try full match: test_case_name.test_name
+  // If that fails, try to find just the test_case_name; this would disable all
+  // tests in the test case.
+  auto it = manifest.find(absl::StrCat(test_case_name, ".", test_name));
+  if (it == manifest.end()) {
+    it = manifest.find(test_case_name);
+    if (it == manifest.end()) {
+      return;
+    }
+  }
+
+  // Expect a full match vs. one of the platform regexps to disable the test.
+  const std::vector& disabled_platforms = it->second;
+  auto platform_string = kTestPlatform;
+  for (const auto& s : disabled_platforms) {
+    if (RE2::FullMatch(/*text=*/platform_string, /*re=*/s)) {
+      GTEST_SKIP();
+      return;
+    }
+  }
+
+  // We didn't hit in the disabled manifest entries, so don't disable it.
+}
+
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/manifest_checking_test.h b/tensorflow/compiler/xla/tests/manifest_checking_test.h
new file mode 100644
index 00000000000..4f44ed76a3e
--- /dev/null
+++ b/tensorflow/compiler/xla/tests/manifest_checking_test.h
@@ -0,0 +1,35 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
+#define TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
+
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+
+// This class allows us to intercept the test name and use an arbitrary
+// heuristic to decide whether the test case should be disabled. We
+// determine whether the test case should be disabled by resolving the (test
+// case name, test name) in a manifest file.
+class ManifestCheckingTest : public ::testing::Test {
+ protected:
+  // This method runs before each test runs.
+  void SetUp() override;
+};
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_TESTS_MANIFEST_CHECKING_TEST_H_
diff --git a/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt b/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
new file mode 100644
index 00000000000..ba5e1bdcaf2
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BandedTriangularSolve.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BandedTriangularSolve"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
new file mode 100644
index 00000000000..2c47960429c
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselI0.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselI0"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselI1.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselI1.pbtxt
new file mode 100644
index 00000000000..e0007b44162
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselI1.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselI1"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselJ0.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselJ0.pbtxt
new file mode 100644
index 00000000000..4010afadcb8
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselJ0.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselJ0"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselJ1.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselJ1.pbtxt
new file mode 100644
index 00000000000..12d16910227
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselJ1.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselJ1"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselK0.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselK0.pbtxt
new file mode 100644
index 00000000000..31d701c821b
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselK0.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselK0"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselK0e.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselK0e.pbtxt
new file mode 100644
index 00000000000..fac0c1b3459
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselK0e.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselK0e"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselK1.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselK1.pbtxt
new file mode 100644
index 00000000000..de80f304540
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselK1.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselK1"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselK1e.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselK1e.pbtxt
new file mode 100644
index 00000000000..c565a85def2
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselK1e.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselK1e"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselY0.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselY0.pbtxt
new file mode 100644
index 00000000000..af57e504d65
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselY0.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselY0"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_BesselY1.pbtxt b/tensorflow/core/api_def/base_api/api_def_BesselY1.pbtxt
new file mode 100644
index 00000000000..b2cd9827f6f
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_BesselY1.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "BesselY1"
+  visibility: HIDDEN
+}
diff --git a/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt b/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
new file mode 100644
index 00000000000..15bd4670cef
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_StatelessParameterizedTruncatedNormal.pbtxt
@@ -0,0 +1,54 @@
+op {
+  graph_op_name: "StatelessParameterizedTruncatedNormal"
+  visibility: HIDDEN
+  in_arg {
+    name: "shape"
+    description: <
\n","  \n","  \n","
\n","    Run in Google Colab\n","  \n","    View source on GitHub\n","  
\n"]},{"cell_type":"markdown","metadata":{"id":"XaVtYN4nlCft","colab_type":"text"},"source":["**Training is much faster using GPU acceleration.** Before you proceed, ensure you are using a GPU runtime by going to **Runtime -> Change runtime type** and set **Hardware accelerator: GPU**. Training 15,000 iterations will take 1.5 - 2 hours on a GPU runtime.\n","\n","## Configure Defaults\n","\n","**MODIFY** the following constants for your specific use case."]},{"cell_type":"code","metadata":{"id":"ludfxbNIaegy","colab_type":"code","colab":{}},"source":["# A comma-delimited list of the words you want to train for.\n","# The options are: yes,no,up,down,left,right,on,off,stop,go\n","# All the other words will be used to train an \"unknown\" label and silent\n","# audio data with no spoken words will be used to train a \"silence\" label.\n","WANTED_WORDS = \"yes,no\"\n","\n","# The number of steps and learning rates can be specified as comma-separated\n","# lists to define the rate at each stage. For example,\n","# TRAINING_STEPS=12000,3000 and LEARNING_RATE=0.001,0.0001\n","# will run 12,000 training loops in total, with a rate of 0.001 for the first\n","# 8,000, and 0.0001 for the final 3,000.\n","TRAINING_STEPS = \"12000,3000\"\n","LEARNING_RATE = \"0.001,0.0001\"\n","\n","# Calculate the total number of steps, which is used to identify the checkpoint\n","# file name.\n","TOTAL_STEPS = str(sum(map(lambda string: int(string), TRAINING_STEPS.split(\",\"))))\n","\n","# Print the configuration to confirm it\n","print(\"Training these words: %s\" % WANTED_WORDS)\n","print(\"Training steps in each stage: %s\" % TRAINING_STEPS)\n","print(\"Learning rate in each stage: %s\" % LEARNING_RATE)\n","print(\"Total number of training steps: %s\" % TOTAL_STEPS)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"gCgeOpvY9pAi","colab_type":"text"},"source":["**DO NOT MODIFY** the following constants as they include filepaths used in this notebook and data that is shared during training and inference."]},{"cell_type":"code","metadata":{"id":"Nd1iM1o2ymvA","colab_type":"code","colab":{}},"source":["# Calculate the percentage of 'silence' and 'unknown' training samples required\n","# to ensure that we have equal number of samples for each label.\n","number_of_labels = WANTED_WORDS.count(',') + 1\n","number_of_total_labels = number_of_labels + 2 # for 'silence' and 'unknown' label\n","equal_percentage_of_training_samples = int(100.0/(number_of_total_labels))\n","SILENT_PERCENTAGE = equal_percentage_of_training_samples\n","UNKNOWN_PERCENTAGE = equal_percentage_of_training_samples\n","\n","# Constants which are shared during training and inference\n","PREPROCESS = 'micro'\n","WINDOW_STRIDE = 20\n","MODEL_ARCHITECTURE = 'tiny_conv' # Other options include: single_fc, conv,\n","                      # low_latency_conv, low_latency_svdf, tiny_embedding_conv\n","\n","# Constants used during training only\n","VERBOSITY = 'WARN'\n","EVAL_STEP_INTERVAL = '1000'\n","SAVE_STEP_INTERVAL = '1000'\n","\n","# Constants for training directories and filepaths\n","DATASET_DIR =  'dataset/'\n","LOGS_DIR = 'logs/'\n","TRAIN_DIR = 'train/' # for training checkpoints and other files.\n","\n","# Constants for inference directories and filepaths\n","import os\n","MODELS_DIR = 'models'\n","if not os.path.exists(MODELS_DIR):\n","  os.mkdir(MODELS_DIR)\n","MODEL_TF = os.path.join(MODELS_DIR, 'model.pb')\n","MODEL_TFLITE = os.path.join(MODELS_DIR, 'model.tflite')\n","FLOAT_MODEL_TFLITE = os.path.join(MODELS_DIR, 'float_model.tflite')\n","MODEL_TFLITE_MICRO = os.path.join(MODELS_DIR, 'model.cc')\n","SAVED_MODEL = os.path.join(MODELS_DIR, 'saved_model')\n","\n","QUANT_INPUT_MIN = 0.0\n","QUANT_INPUT_MAX = 26.0\n","QUANT_INPUT_RANGE = QUANT_INPUT_MAX - QUANT_INPUT_MIN"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"6rLYpvtg9P4o","colab_type":"text"},"source":["## Setup Environment\n","\n","Install Dependencies"]},{"cell_type":"code","metadata":{"id":"ed_XpUrU5DvY","colab_type":"code","colab":{}},"source":["%tensorflow_version 1.x\n","import tensorflow as tf"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"T9Ty5mR58E4i","colab_type":"text"},"source":["**DELETE** any old data from previous runs\n"]},{"cell_type":"code","metadata":{"id":"APGx0fEh7hFF","colab_type":"code","colab":{}},"source":["!rm -rf {DATASET_DIR} {LOGS_DIR} {TRAIN_DIR} {MODELS_DIR}"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"GfEUlfFBizio","colab_type":"text"},"source":["Clone the TensorFlow Github Repository, which contains the relevant code required to run this tutorial."]},{"cell_type":"code","metadata":{"id":"yZArmzT85SLq","colab_type":"code","colab":{}},"source":["!git clone -q --depth 1 https://github.com/tensorflow/tensorflow"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"nS9swHLSi7Bi","colab_type":"text"},"source":["Load TensorBoard to visualize the accuracy and loss as training proceeds.\n"]},{"cell_type":"code","metadata":{"id":"q4qF1VxP3UE4","colab_type":"code","colab":{}},"source":["%load_ext tensorboard\n","%tensorboard --logdir {LOGS_DIR}"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"x1J96Ron-O4R","colab_type":"text"},"source":["## Training\n","\n","The following script downloads the dataset and begin training."]},{"cell_type":"code","metadata":{"id":"VJsEZx6lynbY","colab_type":"code","colab":{}},"source":["!python tensorflow/tensorflow/examples/speech_commands/train.py \\\n","--data_dir={DATASET_DIR} \\\n","--wanted_words={WANTED_WORDS} \\\n","--silence_percentage={SILENT_PERCENTAGE} \\\n","--unknown_percentage={UNKNOWN_PERCENTAGE} \\\n","--preprocess={PREPROCESS} \\\n","--window_stride={WINDOW_STRIDE} \\\n","--model_architecture={MODEL_ARCHITECTURE} \\\n","--how_many_training_steps={TRAINING_STEPS} \\\n","--learning_rate={LEARNING_RATE} \\\n","--train_dir={TRAIN_DIR} \\\n","--summaries_dir={LOGS_DIR} \\\n","--verbosity={VERBOSITY} \\\n","--eval_step_interval={EVAL_STEP_INTERVAL} \\\n","--save_step_interval={SAVE_STEP_INTERVAL}"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"UczQKtqLi7OJ","colab_type":"text"},"source":["# Skipping the training\n","\n","If you don't want to spend an hour or two training the model from scratch, you can download pretrained checkpoints by uncommenting the lines below (removing the '#'s at the start of each line) and running them."]},{"cell_type":"code","metadata":{"id":"RZw3VNlnla-J","colab_type":"code","colab":{}},"source":["#!curl -O \"https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_micro_train_2020_05_10.tgz\"\n","#!tar xzf speech_micro_train_2020_05_10.tgz"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"XQUJLrdS-ftl","colab_type":"text"},"source":["## Generate a TensorFlow Model for Inference\n","\n","Combine relevant training results (graph, weights, etc) into a single file for inference. This process is known as freezing a model and the resulting model is known as a frozen model/graph, as it cannot be further re-trained after this process."]},{"cell_type":"code","metadata":{"id":"xyc3_eLh9sAg","colab_type":"code","colab":{}},"source":["!rm -rf {SAVED_MODEL}\n","!python tensorflow/tensorflow/examples/speech_commands/freeze.py \\\n","--wanted_words=$WANTED_WORDS \\\n","--window_stride_ms=$WINDOW_STRIDE \\\n","--preprocess=$PREPROCESS \\\n","--model_architecture=$MODEL_ARCHITECTURE \\\n","--start_checkpoint=$TRAIN_DIR$MODEL_ARCHITECTURE'.ckpt-'{TOTAL_STEPS} \\\n","--save_format=saved_model \\\n","--output_file={SAVED_MODEL}"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"_DBGDxVI-nKG","colab_type":"text"},"source":["## Generate a TensorFlow Lite Model\n","\n","Convert the frozen graph into a TensorFlow Lite model, which is fully quantized for use with embedded devices.\n","\n","The following cell will also print the model size, which will be under 20 kilobytes."]},{"cell_type":"code","metadata":{"id":"RIitkqvGWmre","colab_type":"code","colab":{}},"source":["import sys\n","# We add this path so we can import the speech processing modules.\n","sys.path.append(\"/content/tensorflow/tensorflow/examples/speech_commands/\")\n","import input_data\n","import models\n","import numpy as np"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"kzqECqMxgBh4","colab_type":"code","colab":{}},"source":["SAMPLE_RATE = 16000\n","CLIP_DURATION_MS = 1000\n","WINDOW_SIZE_MS = 30.0\n","FEATURE_BIN_COUNT = 40\n","BACKGROUND_FREQUENCY = 0.8\n","BACKGROUND_VOLUME_RANGE = 0.1\n","TIME_SHIFT_MS = 100.0\n","\n","DATA_URL = 'https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz'\n","VALIDATION_PERCENTAGE = 10\n","TESTING_PERCENTAGE = 10"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"rNQdAplJV1fz","colab_type":"code","colab":{}},"source":["model_settings = models.prepare_model_settings(\n","    len(input_data.prepare_words_list(WANTED_WORDS.split(','))),\n","    SAMPLE_RATE, CLIP_DURATION_MS, WINDOW_SIZE_MS,\n","    WINDOW_STRIDE, FEATURE_BIN_COUNT, PREPROCESS)\n","audio_processor = input_data.AudioProcessor(\n","    DATA_URL, DATASET_DIR,\n","    SILENT_PERCENTAGE, UNKNOWN_PERCENTAGE,\n","    WANTED_WORDS.split(','), VALIDATION_PERCENTAGE,\n","    TESTING_PERCENTAGE, model_settings, LOGS_DIR)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"lBj_AyCh1cC0","colab_type":"code","colab":{}},"source":["with tf.Session() as sess:\n","  float_converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n","  float_tflite_model = float_converter.convert()\n","  float_tflite_model_size = open(FLOAT_MODEL_TFLITE, \"wb\").write(float_tflite_model)\n","  print(\"Float model is %d bytes\" % float_tflite_model_size)\n","\n","  converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n","  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n","  converter.inference_input_type = tf.lite.constants.INT8\n","  converter.inference_output_type = tf.lite.constants.INT8\n","  def representative_dataset_gen():\n","    for i in range(100):\n","      data, _ = audio_processor.get_data(1, i*1, model_settings,\n","                                         BACKGROUND_FREQUENCY, \n","                                         BACKGROUND_VOLUME_RANGE,\n","                                         TIME_SHIFT_MS,\n","                                         'testing',\n","                                         sess)\n","      flattened_data = np.array(data.flatten(), dtype=np.float32).reshape(1, 1960)\n","      yield [flattened_data]\n","  converter.representative_dataset = representative_dataset_gen\n","  tflite_model = converter.convert()\n","  tflite_model_size = open(MODEL_TFLITE, \"wb\").write(tflite_model)\n","  print(\"Quantized model is %d bytes\" % tflite_model_size)\n"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"EeLiDZTbLkzv","colab_type":"text"},"source":["# Testing the TensorFlow Lite model's accuracy\n","\n","Verify that the model we've exported is still accurate, using the TF Lite Python API and our test set."]},{"cell_type":"code","metadata":{"id":"wQsEteKRLryJ","colab_type":"code","colab":{}},"source":["with tf.Session() as sess:\n","  test_data, test_labels = audio_processor.get_data(\n","      -1, 0, model_settings, BACKGROUND_FREQUENCY, BACKGROUND_VOLUME_RANGE,\n","      TIME_SHIFT_MS, 'testing', sess)\n","\n","float_interpreter = tf.lite.Interpreter(FLOAT_MODEL_TFLITE)\n","float_interpreter.allocate_tensors()\n","\n","float_input_index = float_interpreter.get_input_details()[0][\"index\"]\n","\n","float_output_index = float_interpreter.get_output_details()[0][\"index\"]\n","float_model_output = float_interpreter.tensor(float_output_index)\n","\n","float_correct_predictions = 0\n","for i in range(len(test_data)):\n","  current_input = test_data[i]\n","  current_label = test_labels[i]\n","  flattened_input = np.array(current_input.flatten(), dtype=np.float32).reshape(1, 1960)\n","  float_interpreter.set_tensor(float_input_index, flattened_input)\n","  float_interpreter.invoke()\n","  top_prediction = float_model_output()[0].argmax()\n","  if top_prediction == current_label:\n","    float_correct_predictions += 1\n","\n","print('Float accuracy is %f%% (N=%d)' % ((float_correct_predictions * 100) / len(test_data), len(test_data)))\n","\n","interpreter = tf.lite.Interpreter(MODEL_TFLITE)\n","interpreter.allocate_tensors()\n","\n","input_index = interpreter.get_input_details()[0][\"index\"]\n","\n","output_index = interpreter.get_output_details()[0][\"index\"]\n","model_output = interpreter.tensor(output_index)\n","\n","with tf.Session() as sess:\n","  test_data, test_labels = audio_processor.get_data(\n","      -1, 0, model_settings, BACKGROUND_FREQUENCY, BACKGROUND_VOLUME_RANGE,\n","      TIME_SHIFT_MS, 'testing', sess)\n","\n","correct_predictions = 0\n","for i in range(len(test_data)):\n","  current_input = test_data[i]\n","  current_label = test_labels[i]\n","  quantized_input = np.zeros((1960), np.int8)\n","  for index, input_value in enumerate(current_input.flatten()):\n","    # These scaling values are derived from those used in input_data.py in the\n","    # training pipeline.\n","    value = ((input_value - QUANT_INPUT_MIN) * 256) / QUANT_INPUT_RANGE\n","    value -= 128\n","    if value < -128:\n","      value = -128\n","    if value > 127:\n","      value = 127\n","    quantized_input[index] = value\n","  flattened_input = np.array(quantized_input.flatten(), dtype=np.int8).reshape(1, 1960)\n","  interpreter.set_tensor(input_index, flattened_input)\n","  interpreter.invoke()\n","  top_prediction = model_output()[0].argmax()\n","  if top_prediction == current_label:\n","    correct_predictions += 1\n","\n","print('Quantized accuracy is %f%% (N=%d)' % ((correct_predictions * 100) / len(test_data), len(test_data)))\n"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"dt6Zqbxu-wIi","colab_type":"text"},"source":["## Generate a TensorFlow Lite for MicroControllers Model\n","Convert the TensorFlow Lite model into a C source file that can be loaded by TensorFlow Lite for Microcontrollers."]},{"cell_type":"code","metadata":{"id":"XohZOTjR8ZyE","colab_type":"code","colab":{}},"source":["# Install xxd if it is not available\n","!apt-get update && apt-get -qq install xxd\n","# Convert to a C source file\n","!xxd -i {MODEL_TFLITE} > {MODEL_TFLITE_MICRO}\n","# Update variable names\n","REPLACE_TEXT = MODEL_TFLITE.replace('/', '_').replace('.', '_')\n","!sed -i 's/'{REPLACE_TEXT}'/g_model/g' {MODEL_TFLITE_MICRO}"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"2pQnN0i_-0L2","colab_type":"text"},"source":["## Deploy to a Microcontroller\n","\n","Follow the instructions in the [micro_speech](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/micro_speech) README.md for [TensorFlow Lite for MicroControllers](https://www.tensorflow.org/lite/microcontrollers/overview) to deploy this model on a specific microcontroller.\n","\n","**Reference Model:** If you have not modified this notebook, you can follow the instructions as is, to deploy the model. Refer to the [`micro_speech/train/models`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train/models) directory to access the models generated in this notebook. \n","\n","**New Model:** If you have generated a new model to identify different words: (i) Update `kCategoryCount` and `kCategoryLabels` in [`micro_speech/micro_features/micro_model_settings.h`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/micro_features/micro_model_settings.h) and (ii) Update the values assigned to the variables defined in [`micro_speech/micro_features/model.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/micro_features/model.cc) with values displayed after running the following cell."]},{"cell_type":"code","metadata":{"id":"eoYyh0VU8pca","colab_type":"code","colab":{}},"source":["# Print the C source file\n","!cat {MODEL_TFLITE_MICRO}"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"iYlIKpO2mkhv","colab_type":"code","colab":{}},"source":[""],"execution_count":0,"outputs":[]}]}
\ No newline at end of file
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "train_micro_speech_model.ipynb",
+      "provenance": [],
+      "collapsed_sections": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "pO4-CY_TCZZS",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Train a Simple Audio Recognition Model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BaFfr7DHRmGF",
+        "colab_type": "text"
+      },
+      "source": [
+        "This notebook demonstrates how to train a 20 kB [Simple Audio Recognition](https://www.tensorflow.org/tutorials/sequences/audio_recognition) model to recognize keywords in speech.\n",
+        "\n",
+        "The model created in this notebook is used in the [micro_speech](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/micro_speech) example for [TensorFlow Lite for MicroControllers](https://www.tensorflow.org/lite/microcontrollers/overview).\n",
+        "\n",
+        "\n",
+        "  \n",
+        "  \n",
+        "
\n",
+        "    Run in Google Colab\n",
+        "  \n",
+        "    View source on GitHub\n",
+        "  
\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XaVtYN4nlCft",
+        "colab_type": "text"
+      },
+      "source": [
+        "**Training is much faster using GPU acceleration.** Before you proceed, ensure you are using a GPU runtime by going to **Runtime -> Change runtime type** and set **Hardware accelerator: GPU**. Training 15,000 iterations will take 1.5 - 2 hours on a GPU runtime.\n",
+        "\n",
+        "## Configure Defaults\n",
+        "\n",
+        "**MODIFY** the following constants for your specific use case."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ludfxbNIaegy",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# A comma-delimited list of the words you want to train for.\n",
+        "# The options are: yes,no,up,down,left,right,on,off,stop,go\n",
+        "# All the other words will be used to train an \"unknown\" label and silent\n",
+        "# audio data with no spoken words will be used to train a \"silence\" label.\n",
+        "WANTED_WORDS = \"yes,no\"\n",
+        "\n",
+        "# The number of steps and learning rates can be specified as comma-separated\n",
+        "# lists to define the rate at each stage. For example,\n",
+        "# TRAINING_STEPS=12000,3000 and LEARNING_RATE=0.001,0.0001\n",
+        "# will run 12,000 training loops in total, with a rate of 0.001 for the first\n",
+        "# 8,000, and 0.0001 for the final 3,000.\n",
+        "TRAINING_STEPS = \"12000,3000\"\n",
+        "LEARNING_RATE = \"0.001,0.0001\"\n",
+        "\n",
+        "# Calculate the total number of steps, which is used to identify the checkpoint\n",
+        "# file name.\n",
+        "TOTAL_STEPS = str(sum(map(lambda string: int(string), TRAINING_STEPS.split(\",\"))))\n",
+        "\n",
+        "# Print the configuration to confirm it\n",
+        "print(\"Training these words: %s\" % WANTED_WORDS)\n",
+        "print(\"Training steps in each stage: %s\" % TRAINING_STEPS)\n",
+        "print(\"Learning rate in each stage: %s\" % LEARNING_RATE)\n",
+        "print(\"Total number of training steps: %s\" % TOTAL_STEPS)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gCgeOpvY9pAi",
+        "colab_type": "text"
+      },
+      "source": [
+        "**DO NOT MODIFY** the following constants as they include filepaths used in this notebook and data that is shared during training and inference."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "Nd1iM1o2ymvA",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Calculate the percentage of 'silence' and 'unknown' training samples required\n",
+        "# to ensure that we have equal number of samples for each label.\n",
+        "number_of_labels = WANTED_WORDS.count(',') + 1\n",
+        "number_of_total_labels = number_of_labels + 2 # for 'silence' and 'unknown' label\n",
+        "equal_percentage_of_training_samples = int(100.0/(number_of_total_labels))\n",
+        "SILENT_PERCENTAGE = equal_percentage_of_training_samples\n",
+        "UNKNOWN_PERCENTAGE = equal_percentage_of_training_samples\n",
+        "\n",
+        "# Constants which are shared during training and inference\n",
+        "PREPROCESS = 'micro'\n",
+        "WINDOW_STRIDE = 20\n",
+        "MODEL_ARCHITECTURE = 'tiny_conv' # Other options include: single_fc, conv,\n",
+        "                      # low_latency_conv, low_latency_svdf, tiny_embedding_conv\n",
+        "\n",
+        "# Constants used during training only\n",
+        "VERBOSITY = 'WARN'\n",
+        "EVAL_STEP_INTERVAL = '1000'\n",
+        "SAVE_STEP_INTERVAL = '1000'\n",
+        "\n",
+        "# Constants for training directories and filepaths\n",
+        "DATASET_DIR =  'dataset/'\n",
+        "LOGS_DIR = 'logs/'\n",
+        "TRAIN_DIR = 'train/' # for training checkpoints and other files.\n",
+        "\n",
+        "# Constants for inference directories and filepaths\n",
+        "import os\n",
+        "MODELS_DIR = 'models'\n",
+        "if not os.path.exists(MODELS_DIR):\n",
+        "  os.mkdir(MODELS_DIR)\n",
+        "MODEL_TF = os.path.join(MODELS_DIR, 'model.pb')\n",
+        "MODEL_TFLITE = os.path.join(MODELS_DIR, 'model.tflite')\n",
+        "FLOAT_MODEL_TFLITE = os.path.join(MODELS_DIR, 'float_model.tflite')\n",
+        "MODEL_TFLITE_MICRO = os.path.join(MODELS_DIR, 'model.cc')\n",
+        "SAVED_MODEL = os.path.join(MODELS_DIR, 'saved_model')\n",
+        "\n",
+        "QUANT_INPUT_MIN = 0.0\n",
+        "QUANT_INPUT_MAX = 26.0\n",
+        "QUANT_INPUT_RANGE = QUANT_INPUT_MAX - QUANT_INPUT_MIN"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6rLYpvtg9P4o",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Setup Environment\n",
+        "\n",
+        "Install Dependencies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "ed_XpUrU5DvY",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "%tensorflow_version 1.x\n",
+        "import tensorflow as tf"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "T9Ty5mR58E4i",
+        "colab_type": "text"
+      },
+      "source": [
+        "**DELETE** any old data from previous runs\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "APGx0fEh7hFF",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!rm -rf {DATASET_DIR} {LOGS_DIR} {TRAIN_DIR} {MODELS_DIR}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GfEUlfFBizio",
+        "colab_type": "text"
+      },
+      "source": [
+        "Clone the TensorFlow Github Repository, which contains the relevant code required to run this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "yZArmzT85SLq",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!git clone -q --depth 1 https://github.com/tensorflow/tensorflow"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nS9swHLSi7Bi",
+        "colab_type": "text"
+      },
+      "source": [
+        "Load TensorBoard to visualize the accuracy and loss as training proceeds.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "q4qF1VxP3UE4",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "%load_ext tensorboard\n",
+        "%tensorboard --logdir {LOGS_DIR}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "x1J96Ron-O4R",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Training\n",
+        "\n",
+        "The following script downloads the dataset and begin training."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "VJsEZx6lynbY",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!python tensorflow/tensorflow/examples/speech_commands/train.py \\\n",
+        "--data_dir={DATASET_DIR} \\\n",
+        "--wanted_words={WANTED_WORDS} \\\n",
+        "--silence_percentage={SILENT_PERCENTAGE} \\\n",
+        "--unknown_percentage={UNKNOWN_PERCENTAGE} \\\n",
+        "--preprocess={PREPROCESS} \\\n",
+        "--window_stride={WINDOW_STRIDE} \\\n",
+        "--model_architecture={MODEL_ARCHITECTURE} \\\n",
+        "--how_many_training_steps={TRAINING_STEPS} \\\n",
+        "--learning_rate={LEARNING_RATE} \\\n",
+        "--train_dir={TRAIN_DIR} \\\n",
+        "--summaries_dir={LOGS_DIR} \\\n",
+        "--verbosity={VERBOSITY} \\\n",
+        "--eval_step_interval={EVAL_STEP_INTERVAL} \\\n",
+        "--save_step_interval={SAVE_STEP_INTERVAL}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UczQKtqLi7OJ",
+        "colab_type": "text"
+      },
+      "source": [
+        "# Skipping the training\n",
+        "\n",
+        "If you don't want to spend an hour or two training the model from scratch, you can download pretrained checkpoints by uncommenting the lines below (removing the '#'s at the start of each line) and running them."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "RZw3VNlnla-J",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "#!curl -O \"https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_micro_train_2020_05_10.tgz\"\n",
+        "#!tar xzf speech_micro_train_2020_05_10.tgz"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XQUJLrdS-ftl",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Generate a TensorFlow Model for Inference\n",
+        "\n",
+        "Combine relevant training results (graph, weights, etc) into a single file for inference. This process is known as freezing a model and the resulting model is known as a frozen model/graph, as it cannot be further re-trained after this process."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "xyc3_eLh9sAg",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "!rm -rf {SAVED_MODEL}\n",
+        "!python tensorflow/tensorflow/examples/speech_commands/freeze.py \\\n",
+        "--wanted_words=$WANTED_WORDS \\\n",
+        "--window_stride_ms=$WINDOW_STRIDE \\\n",
+        "--preprocess=$PREPROCESS \\\n",
+        "--model_architecture=$MODEL_ARCHITECTURE \\\n",
+        "--start_checkpoint=$TRAIN_DIR$MODEL_ARCHITECTURE'.ckpt-'{TOTAL_STEPS} \\\n",
+        "--save_format=saved_model \\\n",
+        "--output_file={SAVED_MODEL}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_DBGDxVI-nKG",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Generate a TensorFlow Lite Model\n",
+        "\n",
+        "Convert the frozen graph into a TensorFlow Lite model, which is fully quantized for use with embedded devices.\n",
+        "\n",
+        "The following cell will also print the model size, which will be under 20 kilobytes."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "RIitkqvGWmre",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "import sys\n",
+        "# We add this path so we can import the speech processing modules.\n",
+        "sys.path.append(\"/content/tensorflow/tensorflow/examples/speech_commands/\")\n",
+        "import input_data\n",
+        "import models\n",
+        "import numpy as np"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "kzqECqMxgBh4",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "SAMPLE_RATE = 16000\n",
+        "CLIP_DURATION_MS = 1000\n",
+        "WINDOW_SIZE_MS = 30.0\n",
+        "FEATURE_BIN_COUNT = 40\n",
+        "BACKGROUND_FREQUENCY = 0.8\n",
+        "BACKGROUND_VOLUME_RANGE = 0.1\n",
+        "TIME_SHIFT_MS = 100.0\n",
+        "\n",
+        "DATA_URL = 'https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz'\n",
+        "VALIDATION_PERCENTAGE = 10\n",
+        "TESTING_PERCENTAGE = 10"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "rNQdAplJV1fz",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "model_settings = models.prepare_model_settings(\n",
+        "    len(input_data.prepare_words_list(WANTED_WORDS.split(','))),\n",
+        "    SAMPLE_RATE, CLIP_DURATION_MS, WINDOW_SIZE_MS,\n",
+        "    WINDOW_STRIDE, FEATURE_BIN_COUNT, PREPROCESS)\n",
+        "audio_processor = input_data.AudioProcessor(\n",
+        "    DATA_URL, DATASET_DIR,\n",
+        "    SILENT_PERCENTAGE, UNKNOWN_PERCENTAGE,\n",
+        "    WANTED_WORDS.split(','), VALIDATION_PERCENTAGE,\n",
+        "    TESTING_PERCENTAGE, model_settings, LOGS_DIR)"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "lBj_AyCh1cC0",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "with tf.Session() as sess:\n",
+        "  float_converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n",
+        "  float_tflite_model = float_converter.convert()\n",
+        "  float_tflite_model_size = open(FLOAT_MODEL_TFLITE, \"wb\").write(float_tflite_model)\n",
+        "  print(\"Float model is %d bytes\" % float_tflite_model_size)\n",
+        "\n",
+        "  converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL)\n",
+        "  converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
+        "  converter.inference_input_type = tf.lite.constants.INT8\n",
+        "  converter.inference_output_type = tf.lite.constants.INT8\n",
+        "  def representative_dataset_gen():\n",
+        "    for i in range(100):\n",
+        "      data, _ = audio_processor.get_data(1, i*1, model_settings,\n",
+        "                                         BACKGROUND_FREQUENCY, \n",
+        "                                         BACKGROUND_VOLUME_RANGE,\n",
+        "                                         TIME_SHIFT_MS,\n",
+        "                                         'testing',\n",
+        "                                         sess)\n",
+        "      flattened_data = np.array(data.flatten(), dtype=np.float32).reshape(1, 1960)\n",
+        "      yield [flattened_data]\n",
+        "  converter.representative_dataset = representative_dataset_gen\n",
+        "  tflite_model = converter.convert()\n",
+        "  tflite_model_size = open(MODEL_TFLITE, \"wb\").write(tflite_model)\n",
+        "  print(\"Quantized model is %d bytes\" % tflite_model_size)\n"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EeLiDZTbLkzv",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Testing the TensorFlow Lite model's accuracy\n",
+        "\n",
+        "Verify that the model we've exported is still accurate, using the TF Lite Python API and our test set."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "wQsEteKRLryJ",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Helper function to run inference\n",
+        "def run_tflite_inference(tflite_model_path, model_type=\"Float\"):\n",
+        "  # Load test data\n",
+        "  np.random.seed(0) # set random seed for reproducible test results.\n",
+        "  with tf.Session() as sess:\n",
+        "    test_data, test_labels = audio_processor.get_data(\n",
+        "        -1, 0, model_settings, BACKGROUND_FREQUENCY, BACKGROUND_VOLUME_RANGE,\n",
+        "        TIME_SHIFT_MS, 'testing', sess)\n",
+        "  test_data = np.expand_dims(test_data, axis=1).astype(np.float32)\n",
+        "\n",
+        "  # Initialize the interpreter\n",
+        "  interpreter = tf.lite.Interpreter(tflite_model_path)\n",
+        "  interpreter.allocate_tensors()\n",
+        "\n",
+        "  input_details = interpreter.get_input_details()[0]\n",
+        "  output_details = interpreter.get_output_details()[0]\n",
+        "\n",
+        "  # For quantized models, manually quantize the input data from float to integer\n",
+        "  if model_type == \"Quantized\":\n",
+        "    input_scale, input_zero_point = input_details[\"quantization\"]\n",
+        "    test_data = test_data / input_scale + input_zero_point\n",
+        "    test_data = test_data.astype(input_details[\"dtype\"])\n",
+        "\n",
+        "  correct_predictions = 0\n",
+        "  for i in range(len(test_data)):\n",
+        "    interpreter.set_tensor(input_details[\"index\"], test_data[i])\n",
+        "    interpreter.invoke()\n",
+        "    output = interpreter.get_tensor(output_details[\"index\"])[0]\n",
+        "    top_prediction = output.argmax()\n",
+        "    correct_predictions += (top_prediction == test_labels[i])\n",
+        "\n",
+        "  print('%s model accuracy is %f%% (Number of test samples=%d)' % (\n",
+        "      model_type, (correct_predictions * 100) / len(test_data), len(test_data)))"
+      ],
+      "execution_count": 110,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "l-pD52Na6jRa",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Compute float model accuracy\n",
+        "run_tflite_inference(FLOAT_MODEL_TFLITE)\n",
+        "\n",
+        "# Compute quantized model accuracy\n",
+        "run_tflite_inference(MODEL_TFLITE, model_type='Quantized')"
+      ],
+      "execution_count": 111,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dt6Zqbxu-wIi",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Generate a TensorFlow Lite for MicroControllers Model\n",
+        "Convert the TensorFlow Lite model into a C source file that can be loaded by TensorFlow Lite for Microcontrollers."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "XohZOTjR8ZyE",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Install xxd if it is not available\n",
+        "!apt-get update && apt-get -qq install xxd\n",
+        "# Convert to a C source file\n",
+        "!xxd -i {MODEL_TFLITE} > {MODEL_TFLITE_MICRO}\n",
+        "# Update variable names\n",
+        "REPLACE_TEXT = MODEL_TFLITE.replace('/', '_').replace('.', '_')\n",
+        "!sed -i 's/'{REPLACE_TEXT}'/g_model/g' {MODEL_TFLITE_MICRO}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2pQnN0i_-0L2",
+        "colab_type": "text"
+      },
+      "source": [
+        "## Deploy to a Microcontroller\n",
+        "\n",
+        "Follow the instructions in the [micro_speech](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/micro/examples/micro_speech) README.md for [TensorFlow Lite for MicroControllers](https://www.tensorflow.org/lite/microcontrollers/overview) to deploy this model on a specific microcontroller.\n",
+        "\n",
+        "**Reference Model:** If you have not modified this notebook, you can follow the instructions as is, to deploy the model. Refer to the [`micro_speech/train/models`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/train/models) directory to access the models generated in this notebook.\n",
+        "\n",
+        "**New Model:** If you have generated a new model to identify different words: (i) Update `kCategoryCount` and `kCategoryLabels` in [`micro_speech/micro_features/micro_model_settings.h`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/micro_features/micro_model_settings.h) and (ii) Update the values assigned to the variables defined in [`micro_speech/micro_features/model.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/micro/examples/micro_speech/micro_features/model.cc) with values displayed after running the following cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "eoYyh0VU8pca",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# Print the C source file\n",
+        "!cat {MODEL_TFLITE_MICRO}"
+      ],
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file

From 5d4a29eaf590b4a3068ef4d0b7bea9d4f7bd9369 Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Mon, 22 Jun 2020 16:50:08 -0700
Subject: [PATCH 0836/1390] Special case wrapping of ndarrays in the gradient
 tape code.

PiperOrigin-RevId: 317762474
Change-Id: Ie848ad90a88aff5b2faef4069c3f05887038c367
---
 tensorflow/python/eager/BUILD                |  1 +
 tensorflow/python/eager/backprop.py          | 24 ++++++++++++++++++--
 tensorflow/python/ops/numpy_ops/np_arrays.py |  8 ++-----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index f51bd97e488..408d784ae82 100644
--- a/tensorflow/python/eager/BUILD
+++ b/tensorflow/python/eager/BUILD
@@ -588,6 +588,7 @@ py_library(
         "//tensorflow/python:tensor_shape",
         "//tensorflow/python:unconnected_gradients",
         "//tensorflow/python:util",
+        "//tensorflow/python/ops/numpy_ops:numpy",
         "//tensorflow/python/ops/parallel_for:control_flow_ops",
         "@six_archive//:six",
     ],
diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index 8da3f71360a..5800a51f89a 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -62,6 +62,9 @@ from tensorflow.python.util.tf_export import tf_export
 pfor_ops = LazyLoader(
     "pfor_ops", globals(),
     "tensorflow.python.ops.parallel_for.control_flow_ops")
+np_arrays = LazyLoader(
+    "np_arrays", globals(),
+    "tensorflow.python.ops.numpy_ops.np_arrays")
 
 function = LazyLoader("function", globals(),
                       "tensorflow.python.eager.function")
@@ -721,9 +724,11 @@ pywrap_tfe.TFE_Py_RegisterVSpace(_default_vspace)
 
 
 def _handle_or_self(x):
-  """If x is ResourceVariable, return its handle, else x."""
+  """Unwrap resource variable/ndarray to return tensors."""
   if resource_variable_ops.is_resource_variable(x):
-    x = x.handle
+    return x.handle
+  if isinstance(x, np_arrays.ndarray):
+    return x.data
   return x
 
 
@@ -1023,6 +1028,7 @@ class GradientTape(object):
             "gradient in order to compute higher order "
             "derivatives.", 1)
 
+    num_ndarrays = 0
     flat_targets = []
     for t in nest.flatten(target):
       if not backprop_util.IsTrainable(t):
@@ -1033,7 +1039,12 @@ class GradientTape(object):
       if resource_variable_ops.is_resource_variable(t):
         with self:
           t = ops.convert_to_tensor(t)
+      elif isinstance(t, np_arrays.ndarray):
+        t = t.data
+        num_ndarrays += 1
       flat_targets.append(t)
+    # Only rewrap if all targets are ndarray. If not, prefer tensors.
+    rewrap_as_ndarray = num_ndarrays == len(flat_targets)
 
     flat_sources = nest.flatten(sources)
     flat_sources_raw = flat_sources
@@ -1066,6 +1077,9 @@ class GradientTape(object):
       self._watched_variables = self._tape.watched_variables()
       self._tape = None
 
+    if rewrap_as_ndarray:
+      flat_grad = nest.map_structure(np_arrays.tensor_to_ndarray, flat_grad)
+
     grad = nest.pack_sequence_as(sources, flat_grad)
     return grad
 
@@ -1120,6 +1134,10 @@ class GradientTape(object):
       ValueError: If vectorization of jacobian computation fails.
     """
     flat_sources = nest.flatten(sources)
+    rewrap_as_ndarray = False
+    if isinstance(target, np_arrays.ndarray):
+      target = target.data
+      rewrap_as_ndarray = True
     target_static_shape = target.shape
     target_shape = array_ops.shape(target)
     # Note that we push and pop the tape here and below. This is needed since we
@@ -1169,6 +1187,8 @@ class GradientTape(object):
         out = array_ops.reshape(out, new_shape)
         if context.executing_eagerly():
           out.set_shape(target_static_shape.concatenate(flat_sources[i].shape))
+      if rewrap_as_ndarray:
+        out = np_arrays.tensor_to_ndarray(out)
       output[i] = out
 
     return nest.pack_sequence_as(sources, output)
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index fd26318bea9..eca84421d1b 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -82,10 +82,10 @@ class NdarraySpec(type_spec.BatchableTypeSpec):
     return (self._data_spec,)
 
   def _batch(self, batch_size):
-    return NdarraySpec(self._data_spec.batch(batch_size))
+    return NdarraySpec(self._data_spec._batch(batch_size))  # pylint: disable=protected-access
 
   def _unbatch(self):
-    return NdarraySpec(self._data_spec.unbatch())
+    return NdarraySpec(self._data_spec._unbatch())  # pylint: disable=protected-access
 
 
 class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
@@ -306,10 +306,6 @@ class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
   def __repr__(self):
     return 'ndarray<{}>'.format(self.data.__repr__())
 
-  @property
-  def _id(self):
-    return self.data._id  # pylint: disable=protected-access
-
 
 def tensor_to_ndarray(tensor):
   return ndarray.from_tensor(tensor)

From 6b0b4876f8fe024485d02080d0be2d541ece6b17 Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Mon, 22 Jun 2020 16:54:12 -0700
Subject: [PATCH 0837/1390] [TF-numpy] Adds __matmul__ method to ndarray.

PiperOrigin-RevId: 317763209
Change-Id: I3a2ab8e07b040144239ecee58d4ba9267f5b8977
---
 tensorflow/python/ops/numpy_ops/np_math_ops.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index b242e51a2e5..427aa96e5a4 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -954,6 +954,7 @@ setattr(np_arrays.ndarray, '__pow__', _wrap(power))
 setattr(np_arrays.ndarray, '__rpow__', _wrap(power, True))
 setattr(np_arrays.ndarray, '__truediv__', _wrap(true_divide))
 setattr(np_arrays.ndarray, '__rtruediv__', _wrap(true_divide, True))
+setattr(np_arrays.ndarray, '__matmul__', _wrap(matmul))
 
 
 def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):

From 8bec96c5a8a14239a4b607e5e4c4ba1c5a9971f8 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Tue, 23 Jun 2020 00:44:39 +0000
Subject: [PATCH 0838/1390] Add tests to verify type annotations are correctly
 added

---
 .../python/framework/python_op_gen_test.cc    | 251 +++++++++++++++++-
 1 file changed, 249 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen_test.cc b/tensorflow/python/framework/python_op_gen_test.cc
index 5185086fdd3..2561195c407 100644
--- a/tensorflow/python/framework/python_op_gen_test.cc
+++ b/tensorflow/python/framework/python_op_gen_test.cc
@@ -20,22 +20,269 @@ limitations under the License.
 #include "tensorflow/core/framework/op_gen_lib.h"
 #include "tensorflow/core/platform/test.h"
 
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+
 namespace tensorflow {
 namespace {
 
+constexpr char kBaseOpDef[] = R"(
+op {
+  name: "Foo"
+  input_arg {
+    name: "x"
+    type_attr: "T"
+  }
+  input_arg {
+    name: "y"
+    type_attr: "T2"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "T"
+  }
+  attr {
+    name: "T"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_UINT8
+        type: DT_INT8
+      }
+    }
+  }
+  attr {
+    name: "T2"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_STRING
+        type: DT_FLOAT
+        type: DT_DOUBLE
+      }
+    }
+  }
+  summary: "Summary for op Foo."
+  description: "Description for op Foo."
+}
+op {
+  name: "Bar"
+  input_arg {
+    name: "x"
+    type: DT_STRING
+  }
+  input_arg {
+    name: "y"
+    type: DT_QINT8
+  }
+  output_arg {
+    name: "output"
+    type: DT_BOOL
+  }
+  summary: "Summary for op Bar."
+  description: "Description for op Bar."
+}
+op {
+  name: "FooBar"
+  input_arg {
+    name: "x"
+    type: DT_FLOAT
+  }
+  output_arg {
+    name: "output"
+    type: DT_BOOL
+  }
+  attr {
+    name: "t"
+    type: "type"
+    allowed_values {
+      list {
+        type: DT_HALF
+        type: DT_INT8
+      }
+    }
+  }
+  attr {
+    name: "var1"
+    type: "bool"
+    default_value {
+      b: false
+    }
+  }
+  attr {
+    name: "var2"
+    type: "int"
+    default_value {
+      i: 0
+    }
+  }
+  summary: "Summary for op FooBar."
+  description: "Description for op FooBar."
+}
+op {
+  name: "Baz"
+  input_arg {
+    name: "inputs"
+    number_attr: "N"
+    type_list_attr: "T"
+  }
+  output_arg {
+    name: "output1"
+    type: DT_BOOL
+  }
+  output_arg {
+    name: "output2"
+    type: DT_BOOL
+  }
+  attr {
+    name: "T"
+    type: "bool"
+  }
+  attr {
+    name: "N"
+    type: "int"
+  }
+  summary: "Summary for op Baz."
+  description: "Description for op Baz."
+}
+)";
+
+std::unordered_set type_annotate_ops {
+  "Foo",
+  "Bar",
+  "FooBar",
+  "Baz"
+};
+
+
+void ExpectHasSubstr(const string& s, const string& expected) {
+  EXPECT_TRUE(absl::StrContains(s, expected))
+      << "'Generated ops does not contain '" << expected << "'";
+}
+
+void ExpectDoesNotHaveSubstr(const string& s, const string& expected) {
+  EXPECT_FALSE(absl::StrContains(s, expected))
+      << "'Generated ops contains '" << expected << "'";
+}
+
+void ExpectSubstrOrder(const string& s, const string& before,
+                       const string& after) {
+  int before_pos = s.find(before);
+  int after_pos = s.find(after);
+  ASSERT_NE(std::string::npos, before_pos);
+  ASSERT_NE(std::string::npos, after_pos);
+  EXPECT_LT(before_pos, after_pos)
+      << before << "' is not before '" << after;
+}
+
 TEST(PythonOpGen, Basic) {
   OpList ops;
   OpRegistry::Global()->Export(false, &ops);
 
   ApiDefMap api_def_map(ops);
 
-  string code = GetPythonOps(ops, api_def_map, {}, "");
+  string code = GetPythonOps(ops, api_def_map, {}, "", {});
 
   EXPECT_TRUE(absl::StrContains(code, "def case"));
-
   // TODO(mdan): Add tests to verify type annotations are correctly added.
 }
 
+TEST(PythonOpGen, TypeAnnotateSingleTypeTensor) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string typed_bar = "def bar(x: _ops.Tensor[_dtypes.String], y: _ops.Tensor[_dtypes.QInt8], name=None) -> _ops.Tensor[_dtypes.Bool]:";
+  ExpectHasSubstr(code, typed_bar);
+
+  const string untyped_bar = "def bar(x, y, name=None):";
+  ExpectDoesNotHaveSubstr(code, untyped_bar);
+}
+
+TEST(PythonOpGen, TypeAnnotateMultiTypeTensor) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string typed_foo = "def foo(x: _ops.Tensor[TV_Foo_T], y: _ops.Tensor[TV_Foo_T2], name=None) -> _ops.Tensor[TV_Foo_T]:";
+  ExpectHasSubstr(code, typed_foo);
+}
+
+TEST(PythonOpGen, GenerateCorrectTypeVars) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string typevars_foo = R"(
+TV_Foo_T = TypeVar("TV_Foo_T", _dtypes.Int8, _dtypes.UInt8)
+TV_Foo_T2 = TypeVar("TV_Foo_T2", _dtypes.Float32, _dtypes.Float64, _dtypes.String)
+)";
+
+  ExpectHasSubstr(code, typevars_foo);
+}
+
+TEST(PythonOpGen, TypeAnnotateFallback) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string typed_foo_fallback = "def foo_eager_fallback(x: _ops.Tensor[TV_Foo_T], y: _ops.Tensor[TV_Foo_T2], name, ctx) -> _ops.Tensor[TV_Foo_T]:";
+  ExpectHasSubstr(code, typed_foo_fallback);
+}
+
+TEST(PythonOpGen, GenerateTypeVarAboveOp) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string typevar_foo = "TV_Foo_";
+  const string def_foo = "def foo";
+  ExpectSubstrOrder(code, typevar_foo, def_foo);
+}
+
+
+TEST(PythonOpGen, TypeAnnotateDefaultParams) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string params = "def foo_bar(x: _ops.Tensor[_dtypes.Float32], t: TV_FooBar_t, var1: bool = False, var2: int = 0, name=None)";
+  const string params_fallback = "def foo_bar_eager_fallback(x: _ops.Tensor[_dtypes.Float32], t: TV_FooBar_t, var1: bool, var2: int, name, ctx)";
+  ExpectHasSubstr(code, params);
+  ExpectHasSubstr(code, params_fallback);
+}
+
+TEST(PythonOpGen, NoTypingSequenceTensors) {
+  OpList op_defs;
+  OpRegistry::Global()->Export(false, &op_defs);
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  string code = GetPythonOps(op_defs, api_def_map, {}, "", type_annotate_ops);
+
+  const string baz_def_line = "def baz(inputs, name=None):";
+
+  ExpectHasSubstr(code, baz_def_line);
+}
+
 // TODO(mdan): Include more tests with synhtetic ops and api defs.
 
 }  // namespace

From 23e387975e683a971dd0cbe2198cb60c5cf71fc4 Mon Sep 17 00:00:00 2001
From: Frank Chen 
Date: Mon, 22 Jun 2020 16:58:50 -0700
Subject: [PATCH 0839/1390] Rename tpu_load_library to
 tpu_api_dlsym_initializer.

This better reflects what the load library module does.

PiperOrigin-RevId: 317763973
Change-Id: I66169394c2f1d77b260323cb81d76f0f4e167579
---
 tensorflow/core/BUILD                         |  2 +-
 tensorflow/core/framework/load_library.cc     |  2 +-
 tensorflow/core/tpu/BUILD                     | 24 ++++++--
 tensorflow/core/tpu/kernels/BUILD             |  4 +-
 .../core/tpu/kernels/tpu_configuration_ops.cc |  2 +-
 .../tpu/kernels/tpu_mesh_state_interface.h    |  2 +-
 tensorflow/core/tpu/tpu_api.cc                | 57 +++++++++++++++++++
 .../tpu/{tpu_library_loader.h => tpu_api.h}   | 12 +---
 ...loader.cc => tpu_api_dlsym_initializer.cc} | 40 +------------
 .../core/tpu/tpu_api_dlsym_initializer.h      | 38 +++++++++++++
 ...c => tpu_api_dlsym_initializer_windows.cc} | 20 +------
 tensorflow/stream_executor/tpu/BUILD          | 12 ++--
 .../stream_executor/tpu/tpu_executor.cc       |  2 +-
 .../stream_executor/tpu/tpu_node_context.cc   |  2 +-
 .../stream_executor/tpu/tpu_platform.cc       |  2 +-
 tensorflow/stream_executor/tpu/tpu_stream.h   |  2 +-
 tensorflow/stream_executor/tpu/tpu_timer.h    |  2 +-
 .../tpu/tpu_transfer_manager.cc               |  2 +-
 18 files changed, 141 insertions(+), 86 deletions(-)
 create mode 100644 tensorflow/core/tpu/tpu_api.cc
 rename tensorflow/core/tpu/{tpu_library_loader.h => tpu_api.h} (77%)
 rename tensorflow/core/tpu/{tpu_library_loader.cc => tpu_api_dlsym_initializer.cc} (71%)
 create mode 100644 tensorflow/core/tpu/tpu_api_dlsym_initializer.h
 rename tensorflow/core/tpu/{tpu_library_loader_windows.cc => tpu_api_dlsym_initializer_windows.cc} (64%)

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 695035c91e9..d0be6ee9597 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -2259,7 +2259,7 @@ tf_cuda_library(
         "//tensorflow/core/platform/default/build_config:platformlib",
         "//tensorflow/core/profiler/lib:annotated_traceme",
         "//tensorflow/core/profiler/lib:traceme",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api_dlsym_initializer",
         "//tensorflow/core/util:einsum_op_util",
         "//tensorflow/core/util:padding",
         "//tensorflow/core/util:port",
diff --git a/tensorflow/core/framework/load_library.cc b/tensorflow/core/framework/load_library.cc
index ab08d644074..22181e1c8be 100644
--- a/tensorflow/core/framework/load_library.cc
+++ b/tensorflow/core/framework/load_library.cc
@@ -22,7 +22,7 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/mem.h"
 #if !defined(IS_MOBILE_PLATFORM)
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
 #endif  // IS_MOBILE_PLATFORM
 
 namespace tensorflow {
diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD
index 9e89cd69235..aa811f23672 100644
--- a/tensorflow/core/tpu/BUILD
+++ b/tensorflow/core/tpu/BUILD
@@ -107,15 +107,31 @@ cc_library(
 )
 
 cc_library(
-    name = "tpu_library_loader",
+    name = "tpu_api",
+    srcs = ["tpu_api.cc"],
+    hdrs = ["tpu_api.h"],
+    deps = [
+        ":libtftpu_header",
+        ":tpu_config_c_api",
+        "//tensorflow/core/tpu/kernels:tpu_compile_c_api_hdrs",
+        "//tensorflow/core/tpu/kernels:tpu_mesh_state_c_api_hdrs",
+        "//tensorflow/core/tpu/kernels:tpu_util_c_api_hdrs",
+        "//tensorflow/stream_executor/tpu:tpu_executor_c_api_hdrs",
+        "//tensorflow/stream_executor/tpu:tpu_node_context_c_api_hdrs",
+    ],
+)
+
+cc_library(
+    name = "tpu_api_dlsym_initializer",
     srcs = if_windows(
-        ["tpu_library_loader_windows.cc"],
-        otherwise = ["tpu_library_loader.cc"],
+        ["tpu_api_dlsym_initializer_windows.cc"],
+        otherwise = ["tpu_api_dlsym_initializer.cc"],
     ),
-    hdrs = ["tpu_library_loader.h"],
+    hdrs = ["tpu_api_dlsym_initializer.h"],
     visibility = ["//visibility:public"],
     deps = [
         ":libtftpu_header",
+        ":tpu_api",
         ":tpu_config_c_api",
         ":tpu_library_init_fns",
         "//tensorflow/core/platform:errors",
diff --git a/tensorflow/core/tpu/kernels/BUILD b/tensorflow/core/tpu/kernels/BUILD
index a9f2202cd45..f69c97b81de 100644
--- a/tensorflow/core/tpu/kernels/BUILD
+++ b/tensorflow/core/tpu/kernels/BUILD
@@ -72,10 +72,10 @@ tf_kernel_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/core/tpu:tpu_config_c_api",
         "//tensorflow/core/tpu:tpu_configuration",
         "//tensorflow/core/tpu:tpu_defs",
-        "//tensorflow/core/tpu:tpu_library_loader",
         "//tensorflow/stream_executor/tpu:proto_helper",
     ],
     alwayslink = 1,
@@ -224,7 +224,7 @@ cc_library(
         "//tensorflow/compiler/xla/service",
         "//tensorflow/core:framework",
         "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
     ],
 )
 
diff --git a/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc b/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
index 12a3256a44f..583f1aec207 100644
--- a/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
+++ b/tensorflow/core/tpu/kernels/tpu_configuration_ops.cc
@@ -24,10 +24,10 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/core/tpu/tpu_config_c_api.h"
 #include "tensorflow/core/tpu/tpu_configuration.h"
 #include "tensorflow/core/tpu/tpu_defs.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h b/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
index 3eff3be4915..e2ac38b5f84 100644
--- a/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
+++ b/tensorflow/core/tpu/kernels/tpu_mesh_state_interface.h
@@ -21,7 +21,7 @@ limitations under the License.
 #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/tpu/tpu_api.cc b/tensorflow/core/tpu/tpu_api.cc
new file mode 100644
index 00000000000..8dad82b3029
--- /dev/null
+++ b/tensorflow/core/tpu/tpu_api.cc
@@ -0,0 +1,57 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/tpu/tpu_api.h"
+
+namespace tensorflow {
+namespace tpu {
+
+TfTpu_BaseFn* InitializeApiFn() {
+  static TfTpu_BaseFn base_fn;
+  return &base_fn;
+}
+
+TfTpu_ConfigApiFn* ConfigApiFn() {
+  static TfTpu_ConfigApiFn config_api_fn;
+  return &config_api_fn;
+}
+
+TfTpu_MeshStateApiFn* MeshStateApiFn() {
+  static TfTpu_MeshStateApiFn mesh_state_api_fn;
+  return &mesh_state_api_fn;
+}
+
+TfTpu_CompileApiFn* CompileApiFn() {
+  static TfTpu_CompileApiFn compile_api_fn;
+  return &compile_api_fn;
+}
+
+TfTpu_ExecutorApiFn* ExecutorApiFn() {
+  static TfTpu_ExecutorApiFn executor_api_fn;
+  return &executor_api_fn;
+}
+
+TfTpu_NodeContextApiFn* NodeContextApiFn() {
+  static TfTpu_NodeContextApiFn node_context_api_fn;
+  return &node_context_api_fn;
+}
+
+TfTpu_UtilApiFn* UtilApiFn() {
+  static TfTpu_UtilApiFn util_api_fn;
+  return &util_api_fn;
+}
+
+}  // namespace tpu
+}  // namespace tensorflow
diff --git a/tensorflow/core/tpu/tpu_library_loader.h b/tensorflow/core/tpu/tpu_api.h
similarity index 77%
rename from tensorflow/core/tpu/tpu_library_loader.h
rename to tensorflow/core/tpu/tpu_api.h
index ba6c324707d..c47ace6601d 100644
--- a/tensorflow/core/tpu/tpu_library_loader.h
+++ b/tensorflow/core/tpu/tpu_api.h
@@ -13,10 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
-#define TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
+#ifndef TENSORFLOW_CORE_TPU_TPU_API_H_
+#define TENSORFLOW_CORE_TPU_TPU_API_H_
 
-#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
 #include "tensorflow/core/tpu/kernels/tpu_util_c_api.h"
@@ -25,13 +24,9 @@ limitations under the License.
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 #include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
 
-// LINT.IfChange
 namespace tensorflow {
 namespace tpu {
 
-Status InitializeTpuLibrary(void* library_handle);
-
-// TODO(frankchn): Separate out API functions from the loader.
 TfTpu_BaseFn* InitializeApiFn();
 
 TfTpu_ConfigApiFn* ConfigApiFn();
@@ -48,6 +43,5 @@ TfTpu_UtilApiFn* UtilApiFn();
 
 }  // namespace tpu
 }  // namespace tensorflow
-// LINT.ThenChange(//tensorflow/core/tpu/tpu_library_loader_windows.cc)
 
-#endif  // TENSORFLOW_CORE_TPU_TPU_LIBRARY_LOADER_H_
+#endif  // TENSORFLOW_CORE_TPU_TPU_API_H_
diff --git a/tensorflow/core/tpu/tpu_library_loader.cc b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
similarity index 71%
rename from tensorflow/core/tpu/tpu_library_loader.cc
rename to tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
index 834b86e68a7..c6666421327 100644
--- a/tensorflow/core/tpu/tpu_library_loader.cc
+++ b/tensorflow/core/tpu/tpu_api_dlsym_initializer.cc
@@ -13,14 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// TODO(frankchn): Rename to `tpu_api_dlsym_initializer` or similar.
-
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
 
 #include 
 
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
 #include "tensorflow/stream_executor/tpu/tpu_platform.h"
 
@@ -39,41 +38,6 @@ namespace tpu {
 
 #include "tensorflow/core/tpu/tpu_library_init_fns.inc"
 
-TfTpu_BaseFn* InitializeApiFn() {
-  static TfTpu_BaseFn base_fn;
-  return &base_fn;
-}
-
-TfTpu_ConfigApiFn* ConfigApiFn() {
-  static TfTpu_ConfigApiFn config_api_fn;
-  return &config_api_fn;
-}
-
-TfTpu_MeshStateApiFn* MeshStateApiFn() {
-  static TfTpu_MeshStateApiFn mesh_state_api_fn;
-  return &mesh_state_api_fn;
-}
-
-TfTpu_CompileApiFn* CompileApiFn() {
-  static TfTpu_CompileApiFn compile_api_fn;
-  return &compile_api_fn;
-}
-
-TfTpu_ExecutorApiFn* ExecutorApiFn() {
-  static TfTpu_ExecutorApiFn executor_api_fn;
-  return &executor_api_fn;
-}
-
-TfTpu_NodeContextApiFn* NodeContextApiFn() {
-  static TfTpu_NodeContextApiFn node_context_api_fn;
-  return &node_context_api_fn;
-}
-
-TfTpu_UtilApiFn* UtilApiFn() {
-  static TfTpu_UtilApiFn util_api_fn;
-  return &util_api_fn;
-}
-
 Status InitializeTpuLibrary(void* library_handle) {
   bool shared_object_loaded = true;
   if (library_handle == nullptr) {
diff --git a/tensorflow/core/tpu/tpu_api_dlsym_initializer.h b/tensorflow/core/tpu/tpu_api_dlsym_initializer.h
new file mode 100644
index 00000000000..257fa25ad37
--- /dev/null
+++ b/tensorflow/core/tpu/tpu_api_dlsym_initializer.h
@@ -0,0 +1,38 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_
+#define TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_
+
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tpu/kernels/tpu_compile_c_api.h"
+#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
+#include "tensorflow/core/tpu/kernels/tpu_util_c_api.h"
+#include "tensorflow/core/tpu/libtftpu.h"
+#include "tensorflow/core/tpu/tpu_config_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
+#include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
+
+// LINT.IfChange
+namespace tensorflow {
+namespace tpu {
+
+Status InitializeTpuLibrary(void* library_handle);
+
+}  // namespace tpu
+}  // namespace tensorflow
+// LINT.ThenChange(//tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc)
+
+#endif  // TENSORFLOW_CORE_TPU_TPU_API_DLSYM_INITIALIZER_H_
diff --git a/tensorflow/core/tpu/tpu_library_loader_windows.cc b/tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc
similarity index 64%
rename from tensorflow/core/tpu/tpu_library_loader_windows.cc
rename to tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc
index 7cf1b5cdb1d..f453a98e558 100644
--- a/tensorflow/core/tpu/tpu_library_loader_windows.cc
+++ b/tensorflow/core/tpu/tpu_api_dlsym_initializer_windows.cc
@@ -15,28 +15,14 @@ limitations under the License.
 
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/status.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api_dlsym_initializer.h"
 
-// Reminder: Update tpu_library_loader.cc if you are adding new publicly
-// visible methods.
+// Reminder: Update tpu_api_dlsym_initializer_windows.cc if you are adding new
+// publicly visible methods.
 
 namespace tensorflow {
 namespace tpu {
 
-TfTpu_BaseFn* InitializeApiFn() { return nullptr; }
-
-TfTpu_ConfigApiFn* ConfigApiFn() { return nullptr; }
-
-TfTpu_MeshStateApiFn* MeshStateApiFn() { return nullptr; }
-
-TfTpu_CompileApiFn* CompileApiFn() { return nullptr; }
-
-TfTpu_ExecutorApiFn* ExecutorApiFn() { return nullptr; }
-
-TfTpu_NodeContextApiFn* NodeContextApiFn() { return nullptr; }
-
-TfTpu_UtilApiFn* UtilApiFn() { return nullptr; }
-
 Status InitializeTpuLibrary(void* library_handle) {
   return errors::Unimplemented(
       "Loading TPU library is not supported on Windows.");
diff --git a/tensorflow/stream_executor/tpu/BUILD b/tensorflow/stream_executor/tpu/BUILD
index bf88e9809d0..720ba6bc0c3 100644
--- a/tensorflow/stream_executor/tpu/BUILD
+++ b/tensorflow/stream_executor/tpu/BUILD
@@ -70,7 +70,7 @@ cc_library(
         ":status_helper",
         ":tpu_executor_c_api_hdrs",
         ":tpu_stream_interface",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:stream",
     ],
 )
@@ -81,7 +81,7 @@ cc_library(
     deps = [
         ":tpu_executor_c_api_hdrs",
         "//tensorflow/core/platform:types",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:stream",
     ],
 )
@@ -101,7 +101,7 @@ cc_library(
         ":tpu_timer",
         "//tensorflow/c:tf_status",
         "//tensorflow/core:lib",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:stream",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -151,7 +151,7 @@ cc_library(
         "//tensorflow/compiler/xla/service:stream_pool",
         "//tensorflow/compiler/xla/service:transfer_manager",
         "//tensorflow/core:framework",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:device_memory_allocator",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/memory",
@@ -169,7 +169,7 @@ cc_library(
         ":tpu_platform_interface",
         "//tensorflow/c:tf_status",
         "//tensorflow/core/platform:types",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:stream",
         "@com_google_absl//absl/container:flat_hash_map",
     ],
@@ -201,7 +201,7 @@ cc_library(
         "//tensorflow/compiler/xla:xla_data_proto_cc",
         "//tensorflow/compiler/xla/service:shaped_buffer",
         "//tensorflow/compiler/xla/service:transfer_manager",
-        "//tensorflow/core/tpu:tpu_library_loader",
+        "//tensorflow/core/tpu:tpu_api",
         "//tensorflow/stream_executor:stream",
     ],
 )
diff --git a/tensorflow/stream_executor/tpu/tpu_executor.cc b/tensorflow/stream_executor/tpu/tpu_executor.cc
index cb1410880eb..95c32714732 100644
--- a/tensorflow/stream_executor/tpu/tpu_executor.cc
+++ b/tensorflow/stream_executor/tpu/tpu_executor.cc
@@ -17,7 +17,7 @@ limitations under the License.
 
 #include "tensorflow/c/tf_status.h"
 #include "tensorflow/core/lib/gtl/cleanup.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/device_memory.h"
 #include "tensorflow/stream_executor/lib/status.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
diff --git a/tensorflow/stream_executor/tpu/tpu_node_context.cc b/tensorflow/stream_executor/tpu/tpu_node_context.cc
index 356ede40fb3..b502264cfc7 100644
--- a/tensorflow/stream_executor/tpu/tpu_node_context.cc
+++ b/tensorflow/stream_executor/tpu/tpu_node_context.cc
@@ -17,7 +17,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/backend.h"
 #include "tensorflow/compiler/xla/service/platform_util.h"
 #include "tensorflow/compiler/xla/service/transfer_manager.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/device_memory_allocator.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 #include "tensorflow/stream_executor/tpu/tpu_node_context_c_api.h"
diff --git a/tensorflow/stream_executor/tpu/tpu_platform.cc b/tensorflow/stream_executor/tpu/tpu_platform.cc
index 4bccd822e91..97a97a63351 100644
--- a/tensorflow/stream_executor/tpu/tpu_platform.cc
+++ b/tensorflow/stream_executor/tpu/tpu_platform.cc
@@ -16,7 +16,7 @@ limitations under the License.
 #include "tensorflow/stream_executor/tpu/tpu_platform.h"
 
 #include "tensorflow/c/tf_status.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/platform.h"
 #include "tensorflow/stream_executor/tpu/status_helper.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor.h"
diff --git a/tensorflow/stream_executor/tpu/tpu_stream.h b/tensorflow/stream_executor/tpu/tpu_stream.h
index e1aa1164248..209a624b462 100644
--- a/tensorflow/stream_executor/tpu/tpu_stream.h
+++ b/tensorflow/stream_executor/tpu/tpu_stream.h
@@ -16,7 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_STREAM_H_
 
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/stream_executor_internal.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
 #include "tensorflow/stream_executor/tpu/status_helper.h"
diff --git a/tensorflow/stream_executor/tpu/tpu_timer.h b/tensorflow/stream_executor/tpu/tpu_timer.h
index d7f8f660b37..0ad48ce8a80 100644
--- a/tensorflow/stream_executor/tpu/tpu_timer.h
+++ b/tensorflow/stream_executor/tpu/tpu_timer.h
@@ -17,7 +17,7 @@ limitations under the License.
 #define TENSORFLOW_STREAM_EXECUTOR_TPU_TPU_TIMER_H_
 
 #include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/stream_executor_internal.h"
 #include "tensorflow/stream_executor/tpu/tpu_executor_c_api.h"
 
diff --git a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
index 934fabbf54d..c55af7d58b9 100644
--- a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
+++ b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
@@ -17,7 +17,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
-#include "tensorflow/core/tpu/tpu_library_loader.h"
+#include "tensorflow/core/tpu/tpu_api.h"
 #include "tensorflow/stream_executor/device_memory.h"
 #include "tensorflow/stream_executor/tpu/c_api_conversions.h"
 #include "tensorflow/stream_executor/tpu/proto_helper.h"

From 3cc9264e307a1ec221680542c79e86e0f3d13807 Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Mon, 22 Jun 2020 17:00:43 -0700
Subject: [PATCH 0840/1390] [TF-numpy] Adds an accessor class for numpy_ops, in
 order to be tf_exported'ed.

PiperOrigin-RevId: 317764327
Change-Id: I7af7e421c89b0c754e836a2a84cb67a1a82687a1
---
 tensorflow/python/ops/numpy_ops/BUILD         |  1 -
 .../python/ops/numpy_ops/np_accessor.py       | 32 -------------------
 2 files changed, 33 deletions(-)
 delete mode 100644 tensorflow/python/ops/numpy_ops/np_accessor.py

diff --git a/tensorflow/python/ops/numpy_ops/BUILD b/tensorflow/python/ops/numpy_ops/BUILD
index a70e3f3918d..c4203840c61 100644
--- a/tensorflow/python/ops/numpy_ops/BUILD
+++ b/tensorflow/python/ops/numpy_ops/BUILD
@@ -11,7 +11,6 @@ py_library(
     name = "numpy",
     srcs = [
         "__init__.py",
-        "np_accessor.py",
         "np_array_ops.py",
         "np_arrays.py",
         "np_dtypes.py",
diff --git a/tensorflow/python/ops/numpy_ops/np_accessor.py b/tensorflow/python/ops/numpy_ops/np_accessor.py
deleted file mode 100644
index 64786d2c50a..00000000000
--- a/tensorflow/python/ops/numpy_ops/np_accessor.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"""An accessor class for numpy_ops contents."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.ops import numpy_ops
-
-
-class Numpy:
-  """An accessor class that forwards attribute accesses to module `numpy_ops`.
-  """
-
-  def __getattr__(self, attr):
-    return getattr(numpy_ops, attr)
-
-
-numpy = Numpy()

From 1a3b7af373ab365780729d927ff0117951b3ddfb Mon Sep 17 00:00:00 2001
From: Shanqing Cai 
Date: Mon, 22 Jun 2020 17:05:24 -0700
Subject: [PATCH 0841/1390] [tfdbg2] Support reading multiple DebugEvent file
 sets from the same tfdbg run

TensorFlow jobs that involve multiple hosts (e.g., parameter-server setups and
TPU coordinator-worker setups) can generate >1 DebugEvent file sets when
instrumented with tfdbg2's `tf.debugging.experimental.enable_dump_debug_info()`.

This CL adds capability to load these multiple file sets belonging to the
same tfdbg_run_id to DebugEventsReader and DebugDataReader.

PiperOrigin-RevId: 317765159
Change-Id: Ifcf593bd8b404e3e1c3a6f3f3be70bd6b8b73555
---
 .../python/debug/lib/debug_events_reader.py   | 290 +++++++++++-------
 .../debug/lib/debug_events_writer_test.py     |  87 +++++-
 .../python/debug/lib/debug_v2_ops_test.py     |  19 +-
 .../debug/lib/dumping_callback_test_lib.py    |   7 +-
 4 files changed, 282 insertions(+), 121 deletions(-)

diff --git a/tensorflow/python/debug/lib/debug_events_reader.py b/tensorflow/python/debug/lib/debug_events_reader.py
index 743cea7103a..915c65f7594 100644
--- a/tensorflow/python/debug/lib/debug_events_reader.py
+++ b/tensorflow/python/debug/lib/debug_events_reader.py
@@ -46,28 +46,37 @@ class DebugEventsReader(object):
   # penalty.
   _READER_RELEASE_PER = 100
 
+  _METADATA_SUFFIX = ".metadata"
+  _SOURCE_FILE_SUFFIX = ".source_files"
+  _STACK_FRAMES_SUFFIX = ".stack_frames"
+  _GRAPHS_SUFFIX = ".graphs"
+  _EXECUTION_SUFFIX = ".execution"
+  _GRAPH_EXECUTION_TRACES_SUFFIX = ".graph_execution_traces"
+
   def __init__(self, dump_root):
     if not file_io.is_directory(dump_root):
       raise ValueError("Specified dump_root is not a directory: %s" % dump_root)
-    metadata_paths = file_io.get_matching_files(
-        os.path.join(dump_root, "*.metadata"))
-    if not metadata_paths:
-      raise ValueError("Cannot find any metadata file in directory: %s" %
-                       dump_root)
-    elif len(metadata_paths) > 1:
-      raise ValueError(
-          "Unexpected: Found multiple (%d) metadata in directory: %s" %
-          (len(metadata_paths), dump_root))
-    self._metadata_path = compat.as_bytes(metadata_paths[0])
-    self._metadata_reader = None
+    self._dump_root = dump_root
+    self._metadata_paths = self._load_metadata_files()
 
-    prefix = metadata_paths[0][:-len(".metadata")]
-    self._source_files_path = compat.as_bytes("%s.source_files" % prefix)
-    self._stack_frames_path = compat.as_bytes("%s.stack_frames" % prefix)
-    self._graphs_path = compat.as_bytes("%s.graphs" % prefix)
-    self._execution_path = compat.as_bytes("%s.execution" % prefix)
-    self._graph_execution_traces_path = compat.as_bytes(
-        "%s.graph_execution_traces" % prefix)
+    prefixes = [
+        metadata_path[:-len(self._METADATA_SUFFIX)]
+        for metadata_path in self._metadata_paths
+    ]
+    prefix = prefixes[0]  # This is the prefix of the main file set.
+    self._source_files_path = compat.as_bytes(prefix + self._SOURCE_FILE_SUFFIX)
+    self._stack_frames_path = compat.as_bytes(prefix +
+                                              self._STACK_FRAMES_SUFFIX)
+    self._graphs_path = compat.as_bytes(prefix + self._GRAPHS_SUFFIX)
+    self._execution_path = compat.as_bytes(prefix + self._EXECUTION_SUFFIX)
+    # There can be multiple .graph_execution_trace files each belonging
+    # to a file set generated on an individual host, in the case of
+    # a distributed TensorFlow job.
+    # This is different from the other debug event files in the file set.
+    self._graph_execution_traces_paths = [
+        compat.as_bytes(prefix + self._GRAPH_EXECUTION_TRACES_SUFFIX)
+        for prefix in prefixes
+    ]
     self._readers = dict()  # A map from file path to reader.
     # A map from file path to current reading offset.
     self._reader_offsets = dict()
@@ -78,6 +87,91 @@ class DebugEventsReader(object):
 
     self._offsets = dict()
 
+  def _load_metadata_files(self):
+    """Load and parse metadata files in the dump root.
+
+    Check that all metadata files have a common tfdbg_run_id, and raise
+    a ValueError if their tfdbg_run_ids differ.
+
+    Returns:
+      A list of metadata file paths in ascending order of their starting
+        wall_time timestamp.
+    """
+
+    metadata_paths = file_io.get_matching_files(
+        os.path.join(self._dump_root, "*%s" % self._METADATA_SUFFIX))
+    if not metadata_paths:
+      raise ValueError("Cannot find any tfdbg metadata file in directory: %s" %
+                       self._dump_root)
+    wall_times = []
+    run_ids = []
+    tensorflow_versions = []
+    file_versions = []
+    for metadata_path in metadata_paths:
+      reader = tf_record.tf_record_random_reader(metadata_path)
+      try:
+        record = reader.read(0)[0]
+        debug_event = debug_event_pb2.DebugEvent.FromString(record)
+        wall_times.append(debug_event.wall_time)
+        run_ids.append(debug_event.debug_metadata.tfdbg_run_id)
+        tensorflow_versions.append(
+            debug_event.debug_metadata.tensorflow_version)
+        file_versions.append(debug_event.debug_metadata.file_version)
+      finally:
+        reader.close()
+    self._starting_wall_time = wall_times[0]
+    self._tfdbg_run_id = run_ids[0]
+    self._tensorflow_version = tensorflow_versions[0]
+    self._file_version = file_versions[0]
+    if len(metadata_paths) == 1:
+      # Fast path for a common case (only one DebugEvent file set.)
+      return metadata_paths
+
+    num_no_id = len([run_id for run_id in run_ids if not run_id])
+    if num_no_id:
+      paths_without_run_id = [
+          metadata_path
+          for metadata_path, run_id in zip(metadata_paths, run_ids)
+          if not run_id
+      ]
+      raise ValueError(
+          "Found %d tfdbg metadata files and %d of them do not "
+          "have tfdbg run ids. The metadata files without run ids are: %s" %
+          (len(run_ids), num_no_id, paths_without_run_id))
+    elif len(set(run_ids)) != 1:
+      raise ValueError(
+          "Unexpected: Found multiple (%d) tfdbg2 runs in directory %s" %
+          (len(set(run_ids)), self._dump_root))
+    # Return the metadata files in ascending order of their timestamps.
+    paths_and_timestamps = sorted(
+        zip(metadata_paths, wall_times), key=lambda t: t[1])
+    self._starting_wall_time = paths_and_timestamps[0][1]
+    return [path[0] for path in paths_and_timestamps]
+
+  def starting_wall_time(self):
+    """Get the starting timestamp of the instrumented TensorFlow program.
+
+    When there are multiple hosts (i.e., multiple tfdbg file sets), the earliest
+    timestamp among the file sets is returned. It is assumed to be the job that
+    starts first (e.g., the coordinator).
+
+    Returns:
+      Starting timestamp in seconds since the epoch, as a float.
+    """
+    return self._starting_wall_time
+
+  def tfdbg_run_id(self):
+    """Get the run ID of the instrumented TensorFlow program."""
+    return self._tfdbg_run_id
+
+  def tensorflow_version(self):
+    """Get the version string of TensorFlow that the debugged program ran on."""
+    return self._tensorflow_version
+
+  def tfdbg_file_version(self):
+    """Get the tfdbg file format version."""
+    return self._file_version
+
   def __enter__(self):
     return self
 
@@ -139,9 +233,6 @@ class DebugEventsReader(object):
           self._reader_offsets[file_path] = 0
     return self._readers[file_path]
 
-  def metadata_iterator(self):
-    return self._generic_iterator(self._metadata_path)
-
   def source_files_iterator(self):
     return self._generic_iterator(self._source_files_path)
 
@@ -193,14 +284,18 @@ class DebugEventsReader(object):
       proto_string = self._get_reader(self._execution_path).read(offset)[0]
     return debug_event_pb2.DebugEvent.FromString(proto_string)
 
-  def graph_execution_traces_iterator(self):
-    return self._generic_iterator(self._graph_execution_traces_path)
+  def graph_execution_traces_iterators(self):
+    return [
+        self._generic_iterator(path)
+        for path in self._graph_execution_traces_paths
+    ]
 
-  def read_graph_execution_traces_event(self, offset):
-    """Read DebugEvent at given offset from .graph_execution_traces file.
+  def read_graph_execution_traces_event(self, locator):
+    """Read DebugEvent at given offset from given .graph_execution_traces file.
 
     Args:
-      offset: Offset to read the DebugEvent proto from.
+      locator: A (file_index, offset) tuple that locates the DebugEvent
+        containing the graph execution trace.
 
     Returns:
       A DebugEventProto.
@@ -209,9 +304,11 @@ class DebugEventsReader(object):
       `errors.DataLossError` if offset is at a wrong location.
       `IndexError` if offset is out of range of the file.
     """
-    with self._reader_read_locks[self._graph_execution_traces_path]:
-      proto_string = self._get_reader(
-          self._graph_execution_traces_path).read(offset)[0]
+    file_index, offset = locator
+    graph_execution_traces_path = self._graph_execution_traces_paths[file_index]
+    with self._reader_read_locks[graph_execution_traces_path]:
+      proto_string = self._get_reader(graph_execution_traces_path).read(
+          offset)[0]
     return debug_event_pb2.DebugEvent.FromString(proto_string)
 
   def close(self):
@@ -227,21 +324,27 @@ class BaseDigest(object):
 
   Properties:
     wall_time: A timestamp for the digest as a `float` (unit: s).
-    offset: A offset number in the corresponding file that can be used for
-      fast random read access.
+    locator: A datum that allows tracng the digest to its original
+      location. It can be either of the two:
+       1. Bytes offset from the beginning of the file as a single integer,
+          for the case of all digests of the same kind coming from the same
+          file.
+       2. A tuple of a file index and a byte offset. This applies to case
+          in which the same type of debugger data may come from multple files,
+          e.g., graph execution traces.
   """
 
-  def __init__(self, wall_time, offset):
+  def __init__(self, wall_time, locator):
     self._wall_time = wall_time
-    self._offset = offset
+    self._locator = locator
 
   @property
   def wall_time(self):
     return self._wall_time
 
   @property
-  def offset(self):
-    return self._offset
+  def locator(self):
+    return self._locator
 
   def to_json(self):
     return {"wall_time": self.wall_time}
@@ -265,10 +368,10 @@ class ExecutionDigest(BaseDigest):
 
   def __init__(self,
                wall_time,
-               offset,
+               locator,
                op_type,
                output_tensor_device_ids=None):
-    super(ExecutionDigest, self).__init__(wall_time, offset)
+    super(ExecutionDigest, self).__init__(wall_time, locator)
     self._op_type = op_type
     self._output_tensor_device_ids = _tuple_or_none(output_tensor_device_ids)
 
@@ -332,7 +435,7 @@ class Execution(ExecutionDigest):
                debug_tensor_values=None):
     super(Execution, self).__init__(
         execution_digest.wall_time,
-        execution_digest.offset,
+        execution_digest.locator,
         execution_digest.op_type,
         output_tensor_device_ids=execution_digest.output_tensor_device_ids)
     self._host_name = host_name
@@ -556,7 +659,7 @@ class GraphOpCreationDigest(BaseDigest):
 
   def __init__(self,
                wall_time,
-               offset,
+               locator,
                graph_id,
                op_type,
                op_name,
@@ -565,7 +668,7 @@ class GraphOpCreationDigest(BaseDigest):
                stack_frame_ids,
                input_names=None,
                device_name=None):
-    super(GraphOpCreationDigest, self).__init__(wall_time, offset)
+    super(GraphOpCreationDigest, self).__init__(wall_time, locator)
     self._graph_id = graph_id
     self._op_type = op_type
     self._op_name = op_name
@@ -640,14 +743,9 @@ class GraphExecutionTraceDigest(BaseDigest):
       graph.
   """
 
-  def __init__(self,
-               wall_time,
-               offset,
-               op_type,
-               op_name,
-               output_slot,
+  def __init__(self, wall_time, locator, op_type, op_name, output_slot,
                graph_id):
-    super(GraphExecutionTraceDigest, self).__init__(wall_time, offset)
+    super(GraphExecutionTraceDigest, self).__init__(wall_time, locator)
     self._op_type = op_type
     self._op_name = op_name
     self._output_slot = output_slot
@@ -701,13 +799,13 @@ class GraphExecutionTrace(GraphExecutionTraceDigest):
                tensor_debug_mode,
                debug_tensor_value=None,
                device_name=None):
-    super(GraphExecutionTrace, self).__init__(
-        graph_execution_trace_digest.wall_time,
-        graph_execution_trace_digest.offset,
-        graph_execution_trace_digest.op_type,
-        graph_execution_trace_digest.op_name,
-        graph_execution_trace_digest.output_slot,
-        graph_execution_trace_digest.graph_id)
+    super(GraphExecutionTrace,
+          self).__init__(graph_execution_trace_digest.wall_time,
+                         graph_execution_trace_digest.locator,
+                         graph_execution_trace_digest.op_type,
+                         graph_execution_trace_digest.op_name,
+                         graph_execution_trace_digest.output_slot,
+                         graph_execution_trace_digest.graph_id)
     self._graph_ids = tuple(graph_ids)
     self._tensor_debug_mode = tensor_debug_mode
     self._debug_tensor_value = debug_tensor_value
@@ -780,17 +878,17 @@ def _parse_tensor_value(tensor_proto, return_list=False):
     return None
 
 
-def _execution_digest_from_debug_event_proto(debug_event, offset):
+def _execution_digest_from_debug_event_proto(debug_event, locator):
   """Convert a DebugEvent proto into an ExecutionDigest data object."""
   return ExecutionDigest(
       debug_event.wall_time,
-      offset,
+      locator,
       debug_event.execution.op_type,
-      output_tensor_device_ids=(
-          debug_event.execution.output_tensor_device_ids or None))
+      output_tensor_device_ids=(debug_event.execution.output_tensor_device_ids
+                                or None))
 
 
-def _execution_from_debug_event_proto(debug_event, offset):
+def _execution_from_debug_event_proto(debug_event, locator):
   """Convert a DebugEvent proto into an Execution data object."""
   execution_proto = debug_event.execution
 
@@ -806,7 +904,7 @@ def _execution_from_debug_event_proto(debug_event, offset):
       debug_tensor_values.append(
           _parse_tensor_value(tensor_proto, return_list=True))
   return Execution(
-      _execution_digest_from_debug_event_proto(debug_event, offset),
+      _execution_digest_from_debug_event_proto(debug_event, locator),
       execution_proto.code_location.host_name,
       tuple(execution_proto.code_location.stack_frame_ids),
       execution_proto.tensor_debug_mode,
@@ -832,7 +930,6 @@ class DebugDataReader(object):
 
   def __init__(self, dump_root):
     self._reader = DebugEventsReader(dump_root)
-    self._load_metadata()
 
     # TODO(cais): Implement pagination for memory constraints.
     self._execution_digests = []
@@ -858,13 +955,6 @@ class DebugDataReader(object):
   def _add_monitor(self, monitor):
     self._monitors.append(monitor)
 
-  def _load_metadata(self):
-    metadata_iter = self._reader.metadata_iterator()
-    debug_event = next(metadata_iter).debug_event
-    self._starting_wall_time = debug_event.wall_time
-    self._tensorflow_version = debug_event.debug_metadata.tensorflow_version
-    self._tfdbg_run_id = debug_event.debug_metadata.tfdbg_run_id
-
   def _load_source_files(self):
     """Incrementally read the .source_files DebugEvent file."""
     source_files_iter = self._reader.source_files_iterator()
@@ -944,37 +1034,32 @@ class DebugDataReader(object):
 
   def _load_graph_execution_traces(self):
     """Incrementally load the .graph_execution_traces file."""
-    traces_iter = self._reader.graph_execution_traces_iterator()
-    for debug_event, offset in traces_iter:
-      self._graph_execution_trace_digests.append(
-          self._graph_execution_trace_digest_from_debug_event_proto(
-              debug_event, offset))
-      if self._monitors:
-        graph_execution_trace = (
-            self._graph_execution_trace_from_debug_event_proto(
-                debug_event, offset))
-        for monitor in self._monitors:
-          monitor.on_graph_execution_trace(
-              len(self._graph_execution_trace_digests) - 1,
-              graph_execution_trace)
+    for i, traces_iter in enumerate(
+        self._reader.graph_execution_traces_iterators()):
+      for debug_event, offset in traces_iter:
+        self._graph_execution_trace_digests.append(
+            self._graph_execution_trace_digest_from_debug_event_proto(
+                debug_event, (i, offset)))
+        if self._monitors:
+          graph_execution_trace = (
+              self._graph_execution_trace_from_debug_event_proto(
+                  debug_event, (i, offset)))
+          for monitor in self._monitors:
+            monitor.on_graph_execution_trace(
+                len(self._graph_execution_trace_digests) - 1,
+                graph_execution_trace)
 
-  def _graph_execution_trace_digest_from_debug_event_proto(self,
-                                                           debug_event,
-                                                           offset):
+  def _graph_execution_trace_digest_from_debug_event_proto(
+      self, debug_event, locator):
     trace_proto = debug_event.graph_execution_trace
     op_name = trace_proto.op_name
     op_type = self._lookup_op_type(trace_proto.tfdbg_context_id, op_name)
     return GraphExecutionTraceDigest(
-        debug_event.wall_time,
-        offset,
-        op_type,
-        op_name,
+        debug_event.wall_time, locator, op_type, op_name,
         trace_proto.output_slot,
         debug_event.graph_execution_trace.tfdbg_context_id)
 
-  def _graph_execution_trace_from_debug_event_proto(self,
-                                                    debug_event,
-                                                    offset):
+  def _graph_execution_trace_from_debug_event_proto(self, debug_event, locator):
     """Convert a DebugEvent proto into a GraphExecutionTrace data object."""
     trace_proto = debug_event.graph_execution_trace
     graph_ids = [trace_proto.tfdbg_context_id]
@@ -995,7 +1080,7 @@ class DebugDataReader(object):
           trace_proto.tensor_proto, return_list=True)
     return GraphExecutionTrace(
         self._graph_execution_trace_digest_from_debug_event_proto(
-            debug_event, offset),
+            debug_event, locator),
         graph_ids=graph_ids,
         tensor_debug_mode=trace_proto.tensor_debug_mode,
         debug_tensor_value=debug_tensor_value,
@@ -1059,7 +1144,7 @@ class DebugDataReader(object):
     Returns:
       Stating wall time as seconds since the epoch, as a `float`.
     """
-    return self._starting_wall_time
+    return self._reader.starting_wall_time()
 
   def tensorflow_version(self):
     """TensorFlow version used in the debugged TensorFlow program.
@@ -1070,11 +1155,11 @@ class DebugDataReader(object):
     Returns:
       TensorFlow version used by the debugged program, as a `str`.
     """
-    return self._tensorflow_version
+    return self._reader.tensorflow_version()
 
   def tfdbg_run_id(self):
     """Get the debugger run ID of the debugged TensorFlow program."""
-    return self._tfdbg_run_id
+    return self._reader.tfdbg_run_id()
 
   def outermost_graphs(self):
     """Get the number of outer most graphs read so far."""
@@ -1171,9 +1256,9 @@ class DebugDataReader(object):
 
   def read_execution(self, execution_digest):
     """Read a detailed Execution object."""
-    debug_event = self._reader.read_execution_event(execution_digest.offset)
-    return _execution_from_debug_event_proto(
-        debug_event, execution_digest.offset)
+    debug_event = self._reader.read_execution_event(execution_digest.locator)
+    return _execution_from_debug_event_proto(debug_event,
+                                             execution_digest.locator)
 
   def read_graph_execution_trace(self, graph_execution_trace_digest):
     """Read the detailed graph execution trace.
@@ -1185,9 +1270,9 @@ class DebugDataReader(object):
       The corresponding `GraphExecutionTrace` object.
     """
     debug_event = self._reader.read_graph_execution_traces_event(
-        graph_execution_trace_digest.offset)
+        graph_execution_trace_digest.locator)
     return self._graph_execution_trace_from_debug_event_proto(
-        debug_event, graph_execution_trace_digest.offset)
+        debug_event, graph_execution_trace_digest.locator)
 
   def read_execution_stack_trace(self, execution):
     """Read the stack trace of a given Execution object.
@@ -1234,7 +1319,7 @@ class DebugDataReader(object):
       A list of numpy arrays representing the output tensor values of the
         execution event.
     """
-    debug_event = self._reader.read_execution_event(execution.offset)
+    debug_event = self._reader.read_execution_event(execution.locator)
     return [_parse_tensor_value(tensor_proto)
             for tensor_proto in debug_event.execution.tensor_protos]
 
@@ -1248,8 +1333,7 @@ class DebugDataReader(object):
       A numpy array representing the output tensor value of the intra-graph
         tensor execution event.
     """
-    debug_event = self._reader.read_graph_execution_traces_event(
-        trace.offset)
+    debug_event = self._reader.read_graph_execution_traces_event(trace.locator)
     return _parse_tensor_value(debug_event.graph_execution_trace.tensor_proto)
 
   def symbolic_tensor_id(self, graph_id, op_name, output_slot):
diff --git a/tensorflow/python/debug/lib/debug_events_writer_test.py b/tensorflow/python/debug/lib/debug_events_writer_test.py
index 7b06bf772be..3f3f9179e5d 100644
--- a/tensorflow/python/debug/lib/debug_events_writer_test.py
+++ b/tensorflow/python/debug/lib/debug_events_writer_test.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 import glob
 import json as json_lib
 import os
+import re
 import threading
 import time
 
@@ -264,14 +265,14 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
       writer.WriteGraphExecutionTrace(trace)
 
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
-      actuals = list(reader.graph_execution_traces_iterator())
+      actuals = list(reader.graph_execution_traces_iterators()[0])
       # Before FlushExecutionFiles() is called. No data should have been written
       # to the file.
       self.assertEmpty(actuals)
 
       writer.FlushExecutionFiles()
       actuals = list(item.debug_event.graph_execution_trace
-                     for item in reader.graph_execution_traces_iterator())
+                     for item in reader.graph_execution_traces_iterators()[0])
       self.assertLen(actuals, debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE)
       for i in range(debug_events_writer.DEFAULT_CIRCULAR_BUFFER_SIZE):
         self.assertEqual(
@@ -291,7 +292,7 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
 
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
       actuals = list(item.debug_event.graph_execution_trace
-                     for item in reader.graph_execution_traces_iterator())
+                     for item in reader.graph_execution_traces_iterators()[0])
     self.assertLen(actuals, num_execution_events)
     for i in range(num_execution_events):
       self.assertEqual(actuals[i].op_name, "Op%d" % i)
@@ -598,6 +599,86 @@ class DebugEventsWriterTest(dumping_callback_test_lib.DumpingCallbackTestBase,
     self.assertEqual(traces[-1].op_name, "Op_%d" % (expected_end - 1))
 
 
+class MultiSetReaderTest(dumping_callback_test_lib.DumpingCallbackTestBase):
+  """Test for DebugDataReader for multiple file sets under a dump root."""
+
+  def testReadingTwoFileSetsWithTheSameDumpRootSucceeds(self):
+    # To simulate a multi-host data dump, we first generate file sets in two
+    # different directories, with the same tfdbg_run_id, and then combine them.
+    tfdbg_run_id = "foo"
+    for i in range(2):
+      writer = debug_events_writer.DebugEventsWriter(
+          os.path.join(self.dump_root, str(i)),
+          tfdbg_run_id,
+          circular_buffer_size=-1)
+      if i == 0:
+        debugged_graph = debug_event_pb2.DebuggedGraph(
+            graph_id="graph1", graph_name="graph1")
+        writer.WriteDebuggedGraph(debugged_graph)
+        op_name = "Op_0"
+        graph_op_creation = debug_event_pb2.GraphOpCreation(
+            op_type="FooOp", op_name=op_name, graph_id="graph1")
+        writer.WriteGraphOpCreation(graph_op_creation)
+        op_name = "Op_1"
+        graph_op_creation = debug_event_pb2.GraphOpCreation(
+            op_type="FooOp", op_name=op_name, graph_id="graph1")
+        writer.WriteGraphOpCreation(graph_op_creation)
+      for _ in range(10):
+        trace = debug_event_pb2.GraphExecutionTrace(
+            op_name="Op_%d" % i, tfdbg_context_id="graph1")
+        writer.WriteGraphExecutionTrace(trace)
+        writer.FlushNonExecutionFiles()
+        writer.FlushExecutionFiles()
+
+    # Move all files from the subdirectory /1 to subdirectory /0.
+    dump_root_0 = os.path.join(self.dump_root, "0")
+    src_paths = glob.glob(os.path.join(self.dump_root, "1", "*"))
+    for src_path in src_paths:
+      dst_path = os.path.join(
+          dump_root_0,
+          # Rename the file set to avoid file name collision.
+          re.sub(r"(tfdbg_events\.\d+)", r"\g<1>1", os.path.basename(src_path)))
+      os.rename(src_path, dst_path)
+
+    with debug_events_reader.DebugDataReader(dump_root_0) as reader:
+      reader.update()
+      # Verify the content of the .graph_execution_traces file.
+      trace_digests = reader.graph_execution_traces(digest=True)
+      self.assertLen(trace_digests, 20)
+      for _ in range(10):
+        trace = reader.read_graph_execution_trace(trace_digests[i])
+        self.assertEqual(trace.op_name, "Op_0")
+      for _ in range(10):
+        trace = reader.read_graph_execution_trace(trace_digests[i + 10])
+        self.assertEqual(trace.op_name, "Op_1")
+
+  def testReadingTwoFileSetsWithTheDifferentRootsLeadsToError(self):
+    # To simulate a multi-host data dump, we first generate file sets in two
+    # different directories, with different tfdbg_run_ids, and then combine
+    # them.
+    for i in range(2):
+      writer = debug_events_writer.DebugEventsWriter(
+          os.path.join(self.dump_root, str(i)),
+          "run_id_%d" % i,
+          circular_buffer_size=-1)
+      writer.FlushNonExecutionFiles()
+      writer.FlushExecutionFiles()
+
+    # Move all files from the subdirectory /1 to subdirectory /0.
+    dump_root_0 = os.path.join(self.dump_root, "0")
+    src_paths = glob.glob(os.path.join(self.dump_root, "1", "*"))
+    for src_path in src_paths:
+      dst_path = os.path.join(
+          dump_root_0,
+          # Rename the file set to avoid file name collision.
+          re.sub(r"(tfdbg_events\.\d+)", r"\g<1>1", os.path.basename(src_path)))
+      os.rename(src_path, dst_path)
+
+    with self.assertRaisesRegexp(ValueError,
+                                 r"Found multiple \(2\) tfdbg2 runs"):
+      debug_events_reader.DebugDataReader(dump_root_0)
+
+
 class DataObjectsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
 
   def jsonRoundTripCheck(self, obj):
diff --git a/tensorflow/python/debug/lib/debug_v2_ops_test.py b/tensorflow/python/debug/lib/debug_v2_ops_test.py
index d715869f359..d70c505d3fc 100644
--- a/tensorflow/python/debug/lib/debug_v2_ops_test.py
+++ b/tensorflow/python/debug/lib/debug_v2_ops_test.py
@@ -92,16 +92,13 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
           write_debug_trace(x), [9.0 + np.sqrt(3.0), 16.0 + 2.0])
 
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
-      metadata_iter = reader.metadata_iterator()
       # Check that the .metadata DebugEvents data file has been created, even
       # before FlushExecutionFiles() is called.
-      debug_event = next(metadata_iter).debug_event
-      self.assertGreater(debug_event.wall_time, 0)
-      self.assertTrue(debug_event.debug_metadata.tensorflow_version)
-      self.assertTrue(
-          debug_event.debug_metadata.file_version.startswith("debug.Event:"))
+      self.assertGreater(reader.starting_wall_time(), 0)
+      self.assertTrue(reader.tensorflow_version())
+      self.assertTrue(reader.tfdbg_file_version().startswith("debug.Event"))
 
-      graph_trace_iter = reader.graph_execution_traces_iterator()
+      graph_trace_iter = reader.graph_execution_traces_iterators()[0]
       # Before FlushExecutionFiles() is called, the .graph_execution_traces file
       # ought to be empty.
       with self.assertRaises(StopIteration):
@@ -109,7 +106,7 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
 
       # Flush the circular buffer.
       self.writer.FlushExecutionFiles()
-      graph_trace_iter = reader.graph_execution_traces_iterator()
+      graph_trace_iter = reader.graph_execution_traces_iterators()[0]
 
       # The circular buffer has a size of 4. So only the data from the
       # last two iterations should have been written to self.dump_root.
@@ -167,7 +164,7 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
 
     self.writer.FlushExecutionFiles()
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
-      graph_trace_iter = reader.graph_execution_traces_iterator()
+      graph_trace_iter = reader.graph_execution_traces_iterators()[0]
       try:
         x_values = []
         timestamp = 0
@@ -216,7 +213,7 @@ class DebugIdentityV2OpTest(dumping_callback_test_lib.DumpingCallbackTestBase):
 
     for debug_root in (self.dump_root, another_dump_root):
       with debug_events_reader.DebugEventsReader(debug_root) as reader:
-        graph_trace_iter = reader.graph_execution_traces_iterator()
+        graph_trace_iter = reader.graph_execution_traces_iterators()[0]
 
         debug_event = next(graph_trace_iter).debug_event
         trace = debug_event.graph_execution_trace
@@ -272,7 +269,7 @@ class DebugIdentityV2OpUninitializedWriterTest(
     writer.FlushExecutionFiles()
 
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
-      graph_trace_iter = reader.graph_execution_traces_iterator()
+      graph_trace_iter = reader.graph_execution_traces_iterators()[0]
       graph_execution_traces = []
       while True:
         try:
diff --git a/tensorflow/python/debug/lib/dumping_callback_test_lib.py b/tensorflow/python/debug/lib/dumping_callback_test_lib.py
index 05bf3aeb6da..e58ffdbd79f 100644
--- a/tensorflow/python/debug/lib/dumping_callback_test_lib.py
+++ b/tensorflow/python/debug/lib/dumping_callback_test_lib.py
@@ -48,7 +48,6 @@ class DumpingCallbackTestBase(test_util.TensorFlowTestCase):
   def _readAndCheckMetadataFile(self):
     """Read and check the .metadata debug-events file."""
     with debug_events_reader.DebugEventsReader(self.dump_root) as reader:
-      metadata_iter = reader.metadata_iterator()
-      metadata = next(metadata_iter).debug_event.debug_metadata
-      self.assertEqual(metadata.tensorflow_version, versions.__version__)
-      self.assertTrue(metadata.file_version.startswith("debug.Event"))
+      self.assertTrue(reader.tfdbg_run_id())
+      self.assertEqual(reader.tensorflow_version(), versions.__version__)
+      self.assertTrue(reader.tfdbg_file_version().startswith("debug.Event"))

From c692c45daef9e4a1d5937f705fe8a3977b04fddd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 22 Jun 2020 17:10:15 -0700
Subject: [PATCH 0842/1390] tf numpy: some changes to ndarray constructor
 logic.

PiperOrigin-RevId: 317765968
Change-Id: Iea4338ad18707ff36fc49b450d0defad5c13a6a2
---
 tensorflow/python/ops/numpy_ops/np_arrays.py      |  7 +++----
 tensorflow/python/ops/numpy_ops/np_arrays_test.py | 14 ++++++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_arrays.py b/tensorflow/python/ops/numpy_ops/np_arrays.py
index eca84421d1b..77157544e8f 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays.py
@@ -141,13 +141,12 @@ class ndarray(composite_tensor.CompositeTensor):  # pylint: disable=invalid-name
         raise ValueError('Unexpected type for `buffer` {}. Must be an ndarray,'
                          ' Tensor or np.ndarray.'.format(type(buffer)))
 
-      if shape is not None and tuple(shape) != buffer._shape_tuple():  # pylint: disable=protected-access
-        # TODO(srbs): NumPy allows this. Investigate if/how to support this.
-        raise ValueError('shape arg must match buffer.shape.')
+      if shape is not None:
+        buffer.set_shape(shape)
 
     assert isinstance(buffer, ops.Tensor)
     if dtype and dtype != buffer.dtype:
-      buffer = array_ops.bitcast(buffer, dtype)
+      buffer = math_ops.cast(buffer, dtype)
     self._data = buffer
     self._type_spec_internal = None
 
diff --git a/tensorflow/python/ops/numpy_ops/np_arrays_test.py b/tensorflow/python/ops/numpy_ops/np_arrays_test.py
index 412addc0ad7..ab407d2bfcf 100644
--- a/tensorflow/python/ops/numpy_ops/np_arrays_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_arrays_test.py
@@ -22,6 +22,7 @@ import collections
 
 import numpy as np
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
@@ -51,6 +52,19 @@ class ArrayTest(test.TestCase):
     self.assertIs(a.dtype.type, np.bool_)
     self.assertAllEqual([False, True], a)
 
+  def testConstructor(self):
+    t = constant_op.constant([[1], [1]])
+    a = np_arrays.ndarray(shape=(2, 1), buffer=t)
+    self.assertAllEqual(t, a)
+    self.assertEqual(dtypes.float64, a.dtype)
+
+    a = np_arrays.ndarray(shape=(2, 1), dtype=dtypes.int32, buffer=t)
+    self.assertAllEqual(t, a)
+    self.assertEqual(dtypes.int32, a.dtype)
+
+    with self.assertRaises(ValueError):  # bad shape
+      _ = np_arrays.ndarray((2, 2), buffer=t)
+
   def testNeg(self):
     a = t2a(ops.convert_to_tensor(value=[1.0, 2.0]))
     self.assertAllEqual([-1.0, -2.0], -a)

From fe6e64b09853ac9dbb234ce9b7d4b10da11c7fe9 Mon Sep 17 00:00:00 2001
From: Xiao Yu 
Date: Mon, 22 Jun 2020 17:15:42 -0700
Subject: [PATCH 0843/1390] Refactor eager placement logic into three util
 methods: - MaybePinSmallOpsToCpu - MaybePinToResourceDevice -
 MaybePinToCustomDevice

We are going to reuse MaybePinSmallOpsToCpu in TFRT but not the other two. Because TFRT doesn't have native Resource neither Custom Device.

PiperOrigin-RevId: 317766813
Change-Id: I43241b5786120ddf39dc4bfff6071239afdfd785
---
 tensorflow/core/common_runtime/eager/BUILD    |  31 +++
 .../core/common_runtime/eager/context.h       |   2 +-
 tensorflow/core/common_runtime/eager/core.cc  |  22 ++
 .../common_runtime/eager/eager_operation.h    |   2 +-
 .../core/common_runtime/eager/execute.cc      | 169 -------------
 .../common_runtime/eager/placement_utils.cc   | 228 ++++++++++++++++++
 .../common_runtime/eager/placement_utils.h    |  55 +++++
 .../eager/eager_service_impl.cc               |   6 +-
 tensorflow/lite/delegates/flex/kernel.cc      |   9 +-
 9 files changed, 349 insertions(+), 175 deletions(-)
 create mode 100644 tensorflow/core/common_runtime/eager/placement_utils.cc
 create mode 100644 tensorflow/core/common_runtime/eager/placement_utils.h

diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index fb69bcb7ab5..911b59eed17 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -29,6 +29,7 @@ tf_cuda_library(
         ":context",
         ":eager_operation",
         ":execute",
+        ":placement_utils",
         ":tensor_handle",
         "//tensorflow/c:c_api_internal",
         "//tensorflow/c:tf_tensor_internal",
@@ -489,6 +490,7 @@ cc_library(
         ":eager_op_rewrite_registry",
         ":eager_operation",
         ":kernel_and_device",
+        ":placement_utils",
         ":tensor_handle",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
@@ -521,6 +523,35 @@ cc_library(
     }),
 )
 
+tf_cuda_library(
+    name = "placement_utils",
+    srcs = [
+        "placement_utils.cc",
+    ],
+    hdrs = [
+        "placement_utils.h",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":context",
+        ":attr_builder",
+        ":eager_operation",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+    ] + select({
+        "//tensorflow:android": [
+            "//tensorflow/core:portable_tensorflow_lib_lite",
+        ],
+        "//conditions:default": [
+            "//tensorflow/core:core_cpu_lib",
+            "//tensorflow/core:framework",
+            "//tensorflow/core:framework_internal",
+            "//tensorflow/core:lib",
+            "//tensorflow/core:lib_internal",
+            "//tensorflow/core:protos_all_cc",
+        ],
+    }),
+)
+
 tf_cuda_library(
     name = "attr_builder",
     srcs = ["attr_builder.cc"],
diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h
index c16e1f0f4ad..e6769279558 100644
--- a/tensorflow/core/common_runtime/eager/context.h
+++ b/tensorflow/core/common_runtime/eager/context.h
@@ -478,7 +478,7 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
   // On mobile, it just cleans the caches.
   void WaitForAndCloseRemoteContexts();
 
-  bool PinSmallOpsToCPU() { return pin_small_ops_to_cpu_; }
+  bool PinSmallOpsToCPU() const { return pin_small_ops_to_cpu_; }
 
   tensorflow::Env* TFEnv() const { return env_; }
 
diff --git a/tensorflow/core/common_runtime/eager/core.cc b/tensorflow/core/common_runtime/eager/core.cc
index 3d37250a4fe..77d2b665f5e 100644
--- a/tensorflow/core/common_runtime/eager/core.cc
+++ b/tensorflow/core/common_runtime/eager/core.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/context.h"
 #include "tensorflow/core/common_runtime/eager/eager_operation.h"
 #include "tensorflow/core/common_runtime/eager/execute.h"
+#include "tensorflow/core/common_runtime/eager/placement_utils.h"
 #include "tensorflow/core/common_runtime/eager/tensor_handle.h"
 #include "tensorflow/core/platform/errors.h"
 
@@ -187,6 +188,27 @@ Status EagerContext::RegisterFunction(AbstractFunction* f) {
 // eager_operation.cc we can avoid a circular dependency between them.
 Status EagerOperation::Execute(absl::Span retvals,
                                int* num_retvals) {
+  // Run eager placement logic.
+  VariantDevice device;
+  TF_RETURN_IF_ERROR(eager::MaybePinToCustomDevice(&device, *this));
+  if (device == kVariantDeviceNull) {
+    TF_RETURN_IF_ERROR(eager::MaybePinToResourceDevice(&device, *this));
+  }
+  if (device == kVariantDeviceNull) {
+    bool pin_to_cpu;
+    TF_RETURN_IF_ERROR(eager::MaybePinSmallOpsToCpu(
+        &pin_to_cpu, op_name(),
+        absl::MakeSpan(
+            reinterpret_cast(inputs_.data()),
+            inputs_.size()),
+        ctx_));
+    if (pin_to_cpu) {
+      device = ctx_.HostCPU();
+    }
+  }
+  if (device != kVariantDeviceNull) {
+    SetDevice(device);
+  }
   return EagerExecute(
       this, reinterpret_cast(retvals.data()),
       num_retvals);
diff --git a/tensorflow/core/common_runtime/eager/eager_operation.h b/tensorflow/core/common_runtime/eager/eager_operation.h
index 963aed25733..fa245649d5c 100644
--- a/tensorflow/core/common_runtime/eager/eager_operation.h
+++ b/tensorflow/core/common_runtime/eager/eager_operation.h
@@ -126,7 +126,7 @@ class EagerOperation : public ImmediateExecutionOperation {
   bool is_function() const { return is_function_; }
   bool colocation_exempt() const { return colocation_exempt_; }
 
-  tensorflow::EagerContext& EagerContext() { return ctx_; }
+  tensorflow::EagerContext& EagerContext() const { return ctx_; }
 
   AttrBuilder* MutableAttrs() { return &attrs_; }
   const AttrBuilder& Attrs() const { return attrs_; }
diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index a94c882b3b3..a030f4d0356 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -870,173 +870,6 @@ Status EagerRemoteExecute(EagerOperation* op, TensorHandle** retvals,
 }
 #endif  // IS_MOBILE_PLATFORM
 
-// These ops are not pinnable since they generate data. It can be slower to
-// generate and then copy the data instead of just generating the data on the
-// device directly.
-bool IsPinnableOp(const string& op_type) {
-  static const gtl::FlatSet* unpinnable_ops = new gtl::FlatSet({
-      "RandomUniform",
-      "RandomUniformInt",
-      "RandomStandardNormal",
-      "StatelessRandomUniform",
-      "StatelessRandomUniformInt",
-      "StatelessRandomUniformFullInt",
-      "StatelessRandomNormal",
-  });
-
-  // XRT ops refer to per-device handles that are not safe to move between
-  // devices.
-  return unpinnable_ops->find(op_type) == unpinnable_ops->end() &&
-         !absl::StartsWith(op_type, "XRT");
-}
-
-// Validate if the remote device with the given incarnation is valid in the
-// remote device manager of the current eager context.
-Status ValidateTensorHandleRemoteDevice(EagerContext* ctx,
-                                        int64 device_incarnation) {
-  if (ctx->remote_device_mgr()->ContainsDevice(device_incarnation)) {
-    return Status::OK();
-  }
-  return errors::InvalidArgument(
-      "Resource input tensor contains an invalid device. This might happen "
-      "when the client has connected to a different cluster, or some remote "
-      "workers have been restarted.");
-}
-
-// The Op device may be updated if:
-// - A resource touching input is specified: all resource-touching ops run in
-// the device the resource is, regardless of anything else that has been
-// specified. This is identical to the graph mode behavior.
-//
-// - All op inputs are on the CPU, small (<64 elements) and integers
-// (int32/int64). This can be disabled by setting the environment variable
-// "TF_EAGER_ENABLE_SMALL_TENSOR_CPU_PINNING" to "0" or "false".
-//
-// TODO(b/154234908): Unify placement logic.
-Status MaybeUpdateOpDevice(EagerOperation* op) {
-  // If operation was already placed on a custom device, use it.
-  if (VariantDeviceIsCustom(op->Device())) {
-    return Status::OK();
-  }
-
-  // If all the inputs are on the same custom device, use that custom
-  // device. Otherwise, it is an error to have a custom device as an input.
-  if (!op->Inputs().empty()) {
-    // We keep track of what we've seen with devices instead of booleans to be
-    // able to provide a meaningful error message below.
-    VariantDevice first = op->Inputs()[0]->device();
-    VariantDevice different = first;  // A different input device, if any.
-    VariantDevice custom = first;     // The first custom device seen, or an
-                                      // arbitrary non-custom device otherwise.
-    for (size_t i = 1; first == different && i < op->Inputs().size(); ++i) {
-      VariantDevice device = op->Inputs()[i]->device();
-      if (device != first) {
-        different = device;
-      }
-      if (!VariantDeviceIsCustom(custom) && VariantDeviceIsCustom(device)) {
-        custom = device;
-      }
-      if (different != first && VariantDeviceIsCustom(custom)) {
-        return errors::InvalidArgument(absl::StrCat(
-            "If an operation has one of its inputs in a custom device, then "
-            "all inputs should be on that same device. Operation ",
-            op->Name(), " has one input in custom device ",
-            VariantDeviceName(custom),
-            " and at least one input in a different device ",
-            VariantDeviceName(custom == first ? different : first)));
-      }
-    }
-    if (different == first && VariantDeviceIsCustom(custom)) {
-      op->SetDevice(first);
-      return Status::OK();
-    }
-  }
-
-  if (op->colocation_exempt()) {
-    return Status::OK();
-  }
-  EagerContext& ctx = op->EagerContext();
-  bool all_inputs_eligible_for_cpu_pinning =
-      ctx.PinSmallOpsToCPU() && !op->is_function() && IsPinnableOp(op->Name());
-  Device* op_device = op->Device() == kVariantDeviceNull
-                          ? ctx.HostCPU()
-                          : absl::get(op->Device());
-  for (int i = 0; i < op->Inputs().size(); ++i) {
-    TensorHandle* tensor_handle = op->Inputs()[i];
-    if (tensor_handle->dtype == DT_RESOURCE) {
-      if (tensor_handle->resource_remote_device_incarnation() != 0) {
-        TF_RETURN_IF_ERROR(ValidateTensorHandleRemoteDevice(
-            &ctx, tensor_handle->resource_remote_device_incarnation()));
-      }
-      Device* resource_device = tensor_handle->resource_device();
-      DVLOG(2) << "for op " << op->Name() << " input " << i << " "
-               << DataTypeString(tensor_handle->dtype)
-               << " input device = " << resource_device->name()
-               << ", op device = " << op_device->name();
-      // We check for `op->Device() == nullptr` because it can be later
-      // interpreted as unspecified device and a different device can
-      // be selected based on device priority. If any input to an op
-      // is a resource we must pin it to prevent different device selection.
-      // TODO(iga): null device can mean "unspecified" or "CPU". Clean this up.
-      if (resource_device != op_device || op->Device() == kVariantDeviceNull) {
-        DVLOG(1) << (resource_device != op_device ? "Changing " : "Setting ")
-                 << "device of operation " << op->Name() << " to "
-                 << resource_device->name() << " because input #" << i
-                 << " is a resource in this device.";
-        op->SetDevice(resource_device);
-      }
-      all_inputs_eligible_for_cpu_pinning = false;
-      // No point in looking at other inputs. If there are other resources,
-      // they must have the same device and we already declared the op to be
-      // ineligible for CPU pinning.
-      break;
-    } else if (all_inputs_eligible_for_cpu_pinning) {
-      auto input_device_variant = tensor_handle->DeviceOrHostCPU(ctx);
-      if (VariantDeviceIsCustom(input_device_variant)) {
-        all_inputs_eligible_for_cpu_pinning = false;
-        continue;
-      }
-      Device* input_device = absl::get(input_device_variant);
-      DVLOG(2) << "for op " << op->Name() << " input " << i << " "
-               << DataTypeString(tensor_handle->dtype)
-               << " input device = " << input_device->name()
-               << ", op device = " << op_device->name();
-
-      // Input is on CPU.
-      if (input_device != ctx.HostCPU()) {
-        all_inputs_eligible_for_cpu_pinning = false;
-        continue;
-      }
-
-      if (tensor_handle->dtype != DataType::DT_INT32 &&
-          tensor_handle->dtype != DataType::DT_INT64) {
-        all_inputs_eligible_for_cpu_pinning = false;
-        continue;
-      }
-
-      int64 num_elements;
-      TF_RETURN_IF_ERROR(tensor_handle->NumElements(&num_elements));
-      if (num_elements > 64) {
-        all_inputs_eligible_for_cpu_pinning = false;
-      }
-    }
-  }
-
-  // Ops without inputs are usually ops that generate a tensor in some way and
-  // usually require being present on whatever device they are scheduled on
-  // - for e.g. VarHandleOp or _Recv).
-  // TODO(nareshmodi): Is it possible there is no int32/int64 CPU kernel for
-  // an op, but there is a GPU kernel?
-  if (!op->Inputs().empty() && all_inputs_eligible_for_cpu_pinning) {
-    DVLOG(1) << "Forcing op " << op->Name()
-             << " to be on the CPU since all input tensors have an "
-                "int32/int64 dtype, and are small (less than 64 elements).";
-    op->SetDevice(ctx.HostCPU());
-  }
-
-  return Status::OK();
-}
-
 Status GetKernelOutputs(std::vector* outputs, int num_outputs,
                         TensorHandle** retvals, EagerContext* ctx,
                         KernelAndDevice* kernel) {
@@ -1099,8 +932,6 @@ Status EagerExecute(EagerOperation* op, TensorHandle** retvals,
       [&] { return absl::StrCat("EagerExecute: ", op->Name()); },
       profiler::TraceMeLevel::kInfo);
 
-  TF_RETURN_IF_ERROR(MaybeUpdateOpDevice(op));
-
   if (VariantDeviceIsCustom(op->Device())) {
     return absl::get(op->Device())
         ->Execute(op, retvals, num_retvals);
diff --git a/tensorflow/core/common_runtime/eager/placement_utils.cc b/tensorflow/core/common_runtime/eager/placement_utils.cc
new file mode 100644
index 00000000000..8898516612f
--- /dev/null
+++ b/tensorflow/core/common_runtime/eager/placement_utils.cc
@@ -0,0 +1,228 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#include "tensorflow/core/common_runtime/eager/placement_utils.h"
+
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/core/common_runtime/eager/attr_builder.h"
+#include "tensorflow/core/common_runtime/eager/eager_operation.h"
+#include "tensorflow/core/common_runtime/input_colocation_exemption_registry.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace eager {
+
+// These ops are not pinnable since they generate data. It can be slower to
+// generate and then copy the data instead of just generating the data on the
+// device directly.
+static bool IsPinnableOp(StringPiece op_name) {
+  static const gtl::FlatSet* unpinnable_ops = new gtl::FlatSet({
+      "RandomUniform",
+      "RandomUniformInt",
+      "RandomStandardNormal",
+      "StatelessRandomUniform",
+      "StatelessRandomUniformInt",
+      "StatelessRandomUniformFullInt",
+      "StatelessRandomNormal",
+  });
+
+  // XRT ops refer to per-device handles that are not safe to move between
+  // devices.
+  return unpinnable_ops->find(string(op_name)) == unpinnable_ops->end() &&
+         !absl::StartsWith(op_name, "XRT");
+}
+// Validate if the remote device with the given incarnation is valid in the
+// remote device manager of the current eager context.
+static Status ValidateTensorHandleRemoteDevice(EagerContext* ctx,
+                                               int64 device_incarnation) {
+  if (ctx->remote_device_mgr()->ContainsDevice(device_incarnation)) {
+    return Status::OK();
+  }
+  return errors::InvalidArgument(
+      "Resource input tensor contains an invalid device. This might happen "
+      "when the client has connected to a different cluster, or some remote "
+      "workers have been restarted.");
+}
+
+bool IsColocationExempt(StringPiece op_name) {
+  const auto& exempt_ops = InputColocationExemptionRegistry::Global()->Get();
+  return exempt_ops.find(string(op_name)) != exempt_ops.end();
+}
+
+bool IsFunction(StringPiece op_name) {
+  const OpDef* op_def = nullptr;
+  Status s = OpDefForOp(string(op_name), &op_def);
+  if (!s.ok()) {
+    if (!errors::IsNotFound(s)) {
+      LOG(WARNING) << "Looking up OpDef failed with error: " << s.ToString();
+    }
+    // Cannot find OpDef, it is a function.
+    return true;
+  }
+  return false;
+}
+
+bool IsCustomDevice(StringPiece device_name, const EagerContext& ctx) {
+  CustomDevice* custom_device;
+  return ctx.FindCustomDeviceFromName(string(device_name), &custom_device).ok();
+}
+
+Status MaybePinSmallOpsToCpu(bool* result, StringPiece op_name,
+                             absl::Span args,
+                             const EagerContext& ctx) {
+  if (!ctx.PinSmallOpsToCPU() || IsFunction(op_name) ||
+      IsColocationExempt(op_name) || !IsPinnableOp(op_name)) {
+    *result = false;
+    return Status::OK();
+  }
+
+  // Ops without inputs are usually ops that generate a tensor in some way and
+  // usually require being present on whatever device they are scheduled on
+  // - for e.g. VarHandleOp or _Recv).
+  if (args.empty()) {
+    *result = false;
+    return Status::OK();
+  }
+
+  int i = 0;
+  for (auto* arg : args) {
+    Status s;
+    const char* device_name = arg->DeviceName(&s);
+    DataType dtype = arg->DataType();
+    TF_RETURN_IF_ERROR(s);
+    if (IsCustomDevice(device_name, ctx)) {
+      *result = false;
+      return Status::OK();
+    }
+
+    DVLOG(2) << "for op " << op_name << " input " << i << " "
+             << DataTypeString(dtype) << " input device = " << device_name;
+
+    // Input is on CPU.
+    if (device_name != ctx.HostCPU()->name()) {
+      *result = false;
+      return Status::OK();
+    }
+
+    if (dtype != DataType::DT_INT32 && dtype != DataType::DT_INT64) {
+      *result = false;
+      return Status::OK();
+    }
+
+    int64 num_elements;
+    TF_RETURN_IF_ERROR(arg->NumElements(&num_elements));
+    if (num_elements > 64) {
+      *result = false;
+      return Status::OK();
+    }
+    i++;
+  }
+
+  // TODO(nareshmodi): Is it possible there is no int32/int64 CPU kernel for
+  // an op, but there is a GPU kernel?
+  DVLOG(1) << "Forcing op " << op_name
+           << " to be on the CPU since all input tensors have an "
+              "int32/int64 dtype, and are small (less than 64 elements).";
+  *result = true;
+  return Status::OK();
+}
+
+Status MaybePinToResourceDevice(VariantDevice* device,
+                                const EagerOperation& op) {
+  if (op.colocation_exempt()) {
+    return Status::OK();
+  }
+  EagerContext& ctx = op.EagerContext();
+  Device* op_device = op.Device() == kVariantDeviceNull
+                          ? ctx.HostCPU()
+                          : absl::get(op.Device());
+  for (int i = 0; i < op.Inputs().size(); ++i) {
+    TensorHandle* tensor_handle = op.Inputs()[i];
+    if (tensor_handle->dtype == DT_RESOURCE) {
+      if (tensor_handle->resource_remote_device_incarnation() != 0) {
+        TF_RETURN_IF_ERROR(ValidateTensorHandleRemoteDevice(
+            &ctx, tensor_handle->resource_remote_device_incarnation()));
+      }
+      Device* resource_device = tensor_handle->resource_device();
+      DVLOG(2) << "for op " << op.Name() << " input " << i << " "
+               << DataTypeString(tensor_handle->dtype)
+               << " input device = " << resource_device->name()
+               << ", op device = " << op_device->name();
+      // We check for `op->Device() == nullptr` because it can be later
+      // interpreted as unspecified device and a different device can
+      // be selected based on device priority. If any input to an op
+      // is a resource we must pin it to prevent different device selection.
+      // TODO(iga): null device can mean "unspecified" or "CPU". Clean this up.
+      if (resource_device != op_device || op.Device() == kVariantDeviceNull) {
+        DVLOG(1) << (resource_device != op_device ? "Changing " : "Setting ")
+                 << "device of operation " << op.Name() << " to "
+                 << resource_device->name() << " because input #" << i
+                 << " is a resource in this device.";
+        *device = resource_device;
+        return Status::OK();
+        // No point in looking at other inputs. If there are other resources,
+        // they must have the same device and we already declared the op to be
+        // ineligible for CPU pinning.
+      }
+    }
+  }
+  return Status::OK();
+}
+
+Status MaybePinToCustomDevice(VariantDevice* device, const EagerOperation& op) {
+  // If operation was already placed on a custom device, use it.
+  if (VariantDeviceIsCustom(op.Device())) {
+    *device = op.Device();
+    return Status::OK();
+  }
+
+  if (!op.Inputs().empty()) {
+    // We keep track of what we've seen with devices instead of booleans to be
+    // able to provide a meaningful error message below.
+    VariantDevice first = op.Inputs()[0]->device();
+    VariantDevice different = first;  // A different input device, if any.
+    VariantDevice custom = first;     // The first custom device seen, or an
+                                      // arbitrary non-custom device otherwise.
+    for (size_t i = 1; first == different && i < op.Inputs().size(); ++i) {
+      VariantDevice device = op.Inputs()[i]->device();
+      if (device != first) {
+        different = device;
+      }
+      if (!VariantDeviceIsCustom(custom) && VariantDeviceIsCustom(device)) {
+        custom = device;
+      }
+      if (different != first && VariantDeviceIsCustom(custom)) {
+        return errors::InvalidArgument(absl::StrCat(
+            "If an operation has one of its inputs in a custom device, then "
+            "all inputs should be on that same device. Operation ",
+            op.Name(), " has one input in custom device ",
+            VariantDeviceName(custom),
+            " and at least one input in a different device ",
+            VariantDeviceName(custom == first ? different : first)));
+      }
+    }
+    if (different == first && VariantDeviceIsCustom(custom)) {
+      *device = first;
+      return Status::OK();
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace eager
+}  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/placement_utils.h b/tensorflow/core/common_runtime/eager/placement_utils.h
new file mode 100644
index 00000000000..d58bd304b27
--- /dev/null
+++ b/tensorflow/core/common_runtime/eager/placement_utils.h
@@ -0,0 +1,55 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_PLACEMENT_UTILS_H_
+#define TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_PLACEMENT_UTILS_H_
+
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/stringpiece.h"
+
+namespace tensorflow {
+namespace eager {
+
+bool IsColocationExempt(StringPiece op_name);
+
+bool IsFunction(StringPiece op_name);
+
+bool IsCustomDevice(StringPiece device_name, const EagerContext& ctx);
+
+// TODO(b/154234908): Unify placement logic.
+// TODO(b/159647422): Add C++ unit tests for placement logic.
+
+// Pin the op to cpu if all op inputs are on the CPU, small (<64 elements) and
+// integers (int32/int64). This can be disabled by setting the environment
+// variable "TF_EAGER_ENABLE_SMALL_TENSOR_CPU_PINNING" to "0" or "false".
+Status MaybePinSmallOpsToCpu(bool* result, StringPiece op_name,
+                             absl::Span args,
+                             const EagerContext& ctx);
+
+// If a resource touching input is specified, all resource-touching ops run in
+// the device the resource is, regardless of anything else that has been
+// specified. This is identical to the graph mode behavior.
+Status MaybePinToResourceDevice(VariantDevice* device,
+                                const EagerOperation& op);
+
+// If all the inputs are on the same custom device, use that custom
+// device. Otherwise, it is an error to have a custom device as an input.
+Status MaybePinToCustomDevice(VariantDevice* device, const EagerOperation& op);
+
+}  // namespace eager
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_PLACEMENT_UTILS_H_
diff --git a/tensorflow/core/distributed_runtime/eager/eager_service_impl.cc b/tensorflow/core/distributed_runtime/eager/eager_service_impl.cc
index fd0606538c4..4735ff6eaf6 100644
--- a/tensorflow/core/distributed_runtime/eager/eager_service_impl.cc
+++ b/tensorflow/core/distributed_runtime/eager/eager_service_impl.cc
@@ -491,7 +491,11 @@ Status EagerServiceImpl::ExecuteOp(const Operation& operation,
 
   absl::FixedArray retvals(num_retvals);
   VLOG(3) << "ServerContext: Calling EagerExecute for op " << operation.id();
-  TF_RETURN_IF_ERROR(EagerExecute(&op, retvals.data(), &num_retvals));
+  TF_RETURN_IF_ERROR(op.Execute(
+      absl::MakeSpan(
+          reinterpret_cast(retvals.data()),
+          num_retvals),
+      &num_retvals));
 
   return AddOpRetvalsToResponse(
       eager_context, operation.id(), num_retvals, retvals.data(),
diff --git a/tensorflow/lite/delegates/flex/kernel.cc b/tensorflow/lite/delegates/flex/kernel.cc
index e7705ecf3ce..b3e978908bd 100644
--- a/tensorflow/lite/delegates/flex/kernel.cc
+++ b/tensorflow/lite/delegates/flex/kernel.cc
@@ -331,9 +331,12 @@ tensorflow::Status ExecuteFlexOp(TfLiteContext* context, BufferMap* buffer_map,
   node_data->mutable_outputs()->ResetTensorHandles();
   int num_retvals = node_data->NumOutputs();
   TF_RETURN_WITH_CONTEXT_IF_ERROR(
-      EagerExecute(node_data->op(),
-                   node_data->mutable_outputs()->GetTensorHandles()->data(),
-                   &num_retvals),
+      node_data->op()->Execute(
+          absl::MakeSpan(
+              reinterpret_cast(
+                  node_data->mutable_outputs()->GetTensorHandles()->data()),
+              num_retvals),
+          &num_retvals),
       " (while executing '", node_data->name(), "' via Eager)");
 
   if (num_retvals != node_data->NumOutputs()) {

From 34b4fab30a40efbdd22c5fb4532e818c4d8e441d Mon Sep 17 00:00:00 2001
From: Advait Jain 
Date: Mon, 22 Jun 2020 17:39:59 -0700
Subject: [PATCH 0844/1390] Remove deprecated AddBuiltin API from
 MicroMutableOpResolver.

* Added new API hooks for all the OPs currently supported in TFLM.
* These new APIs still need to be implemented with operator specific parse
  functions but this change allows us to remove the old API and incrementally
  update the implementations.

PiperOrigin-RevId: 317770205
Change-Id: Idaaa687401f2bac5fbf9925e27c04bf536b154ea
---
 tensorflow/lite/micro/all_ops_resolver.cc     | 120 +++----
 .../image_recognition_test.cc                 |  12 +-
 .../image_recognition_experimental/main.cc    |  12 +-
 .../examples/magic_wand/magic_wand_test.cc    |  16 +-
 .../examples/magic_wand/main_functions.cc     |  16 +-
 .../examples/micro_speech/main_functions.cc   |  16 +-
 .../micro_speech/micro_speech_test.cc         |  13 +-
 .../person_detection/main_functions.cc        |  10 +-
 .../person_detection/person_detection_test.cc |  10 +-
 .../main_functions.cc                         |  16 +-
 .../person_detection_test.cc                  |  16 +-
 .../lite/micro/micro_mutable_op_resolver.h    | 332 ++++++++++++++++--
 .../micro/micro_mutable_op_resolver_test.cc   |  33 +-
 13 files changed, 408 insertions(+), 214 deletions(-)

diff --git a/tensorflow/lite/micro/all_ops_resolver.cc b/tensorflow/lite/micro/all_ops_resolver.cc
index b0021a2e771..e728a95360a 100644
--- a/tensorflow/lite/micro/all_ops_resolver.cc
+++ b/tensorflow/lite/micro/all_ops_resolver.cc
@@ -26,74 +26,60 @@ const char* GetString_ETHOSU();
 
 AllOpsResolver::AllOpsResolver() {
   // Please keep this list of Builtin Operators in alphabetical order.
-  AddBuiltin(BuiltinOperator_ABS, tflite::ops::micro::Register_ABS());
-  AddBuiltin(BuiltinOperator_ADD, tflite::ops::micro::Register_ADD());
-  AddBuiltin(BuiltinOperator_ARG_MAX, tflite::ops::micro::Register_ARG_MAX());
-  AddBuiltin(BuiltinOperator_ARG_MIN, tflite::ops::micro::Register_ARG_MIN());
-  AddBuiltin(BuiltinOperator_AVERAGE_POOL_2D,
-             tflite::ops::micro::Register_AVERAGE_POOL_2D());
-  AddBuiltin(BuiltinOperator_CEIL, tflite::ops::micro::Register_CEIL());
-  AddBuiltin(BuiltinOperator_CONCATENATION,
-             tflite::ops::micro::Register_CONCATENATION());
-  AddBuiltin(BuiltinOperator_CONV_2D, tflite::ops::micro::Register_CONV_2D());
-  AddBuiltin(BuiltinOperator_COS, tflite::ops::micro::Register_COS());
-  AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D,
-             tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  AddBuiltin(BuiltinOperator_DEQUANTIZE,
-             tflite::ops::micro::Register_DEQUANTIZE());
-  AddBuiltin(BuiltinOperator_EQUAL, tflite::ops::micro::Register_EQUAL());
-  AddBuiltin(BuiltinOperator_FLOOR, tflite::ops::micro::Register_FLOOR());
-  AddBuiltin(BuiltinOperator_FULLY_CONNECTED,
-             tflite::ops::micro::Register_FULLY_CONNECTED());
-  AddBuiltin(BuiltinOperator_GREATER, tflite::ops::micro::Register_GREATER());
-  AddBuiltin(BuiltinOperator_GREATER_EQUAL,
-             tflite::ops::micro::Register_GREATER_EQUAL());
-  AddBuiltin(BuiltinOperator_L2_NORMALIZATION,
-             tflite::ops::micro::Register_L2_NORMALIZATION());
-  AddBuiltin(BuiltinOperator_LESS, tflite::ops::micro::Register_LESS());
-  AddBuiltin(BuiltinOperator_LESS_EQUAL,
-             tflite::ops::micro::Register_LESS_EQUAL());
-  AddBuiltin(BuiltinOperator_LOG, tflite::ops::micro::Register_LOG());
-  AddBuiltin(BuiltinOperator_LOGICAL_AND,
-             tflite::ops::micro::Register_LOGICAL_AND());
-  AddBuiltin(BuiltinOperator_LOGICAL_NOT,
-             tflite::ops::micro::Register_LOGICAL_NOT());
-  AddBuiltin(BuiltinOperator_LOGICAL_OR,
-             tflite::ops::micro::Register_LOGICAL_OR());
-  AddBuiltin(BuiltinOperator_LOGISTIC, tflite::ops::micro::Register_LOGISTIC());
-  AddBuiltin(BuiltinOperator_MAX_POOL_2D,
-             tflite::ops::micro::Register_MAX_POOL_2D());
-  AddBuiltin(BuiltinOperator_MAXIMUM, tflite::ops::micro::Register_MAXIMUM());
-  AddBuiltin(BuiltinOperator_MEAN, tflite::ops::micro::Register_MEAN());
-  AddBuiltin(BuiltinOperator_MINIMUM, tflite::ops::micro::Register_MINIMUM());
-  AddBuiltin(BuiltinOperator_MUL, tflite::ops::micro::Register_MUL());
-  AddBuiltin(BuiltinOperator_NEG, tflite::ops::micro::Register_NEG());
-  AddBuiltin(BuiltinOperator_NOT_EQUAL,
-             tflite::ops::micro::Register_NOT_EQUAL());
-  AddBuiltin(BuiltinOperator_PACK, tflite::ops::micro::Register_PACK());
-  AddBuiltin(BuiltinOperator_PAD, tflite::ops::micro::Register_PAD());
-  AddBuiltin(BuiltinOperator_PADV2, tflite::ops::micro::Register_PADV2());
-  AddBuiltin(BuiltinOperator_PRELU, tflite::ops::micro::Register_PRELU());
-  AddBuiltin(BuiltinOperator_QUANTIZE, tflite::ops::micro::Register_QUANTIZE());
-  AddBuiltin(BuiltinOperator_RELU, tflite::ops::micro::Register_RELU());
-  AddBuiltin(BuiltinOperator_RELU6, tflite::ops::micro::Register_RELU6());
-  AddBuiltin(BuiltinOperator_RESHAPE, tflite::ops::micro::Register_RESHAPE());
-  AddBuiltin(BuiltinOperator_RESIZE_NEAREST_NEIGHBOR,
-             tflite::ops::micro::Register_RESIZE_NEAREST_NEIGHBOR());
-  AddBuiltin(BuiltinOperator_ROUND, tflite::ops::micro::Register_ROUND());
-  AddBuiltin(BuiltinOperator_RSQRT, tflite::ops::micro::Register_RSQRT());
-  AddBuiltin(BuiltinOperator_SIN, tflite::ops::micro::Register_SIN());
-  AddBuiltin(BuiltinOperator_SOFTMAX, tflite::ops::micro::Register_SOFTMAX());
-  AddBuiltin(BuiltinOperator_SPLIT, tflite::ops::micro::Register_SPLIT());
-  AddBuiltin(BuiltinOperator_SQRT, tflite::ops::micro::Register_SQRT());
-  AddBuiltin(BuiltinOperator_SQUARE, tflite::ops::micro::Register_SQUARE());
-  AddBuiltin(BuiltinOperator_STRIDED_SLICE,
-             tflite::ops::micro::Register_STRIDED_SLICE());
-  AddBuiltin(BuiltinOperator_SUB, tflite::ops::micro::Register_SUB());
-  AddBuiltin(BuiltinOperator_SVDF, tflite::ops::micro::Register_SVDF());
-  AddBuiltin(BuiltinOperator_TANH, tflite::ops::micro::Register_TANH());
-  AddBuiltin(BuiltinOperator_UNPACK, tflite::ops::micro::Register_UNPACK());
+  AddAbs();
+  AddAdd();
+  AddArgMax();
+  AddArgMin();
+  AddAveragePool2D();
+  AddCeil();
+  AddConcatenation();
+  AddConv2D();
+  AddCos();
+  AddDepthwiseConv2D();
+  AddDequantize();
+  AddEqual();
+  AddFloor();
+  AddFullyConnected();
+  AddGreater();
+  AddGreaterEqual();
+  AddL2Normalization();
+  AddLess();
+  AddLessEqual();
+  AddLog();
+  AddLogicalAnd();
+  AddLogicalNot();
+  AddLogicalOr();
+  AddLogistic();
+  AddMaximum();
+  AddMaxPool2D();
+  AddMean();
+  AddMinimum();
+  AddMul();
+  AddNeg();
+  AddNotEqual();
+  AddPack();
+  AddPad();
+  AddPadV2();
+  AddPrelu();
+  AddQuantize();
+  AddRelu();
+  AddRelu6();
+  AddReshape();
+  AddResizeNearestNeighbor();
+  AddRound();
+  AddRsqrt();
+  AddSin();
+  AddSoftmax();
+  AddSplit();
+  AddSqrt();
+  AddSquare();
+  AddStridedSlice();
+  AddSub();
+  AddSvdf();
+  AddTanh();
+  AddUnpack();
 
+  // TODO(b/159644355): Figure out if custom Ops belong in AllOpsResolver.
   TfLiteRegistration* registration =
       tflite::ops::micro::custom::Register_ETHOSU();
   if (registration) {
diff --git a/tensorflow/lite/micro/examples/image_recognition_experimental/image_recognition_test.cc b/tensorflow/lite/micro/examples/image_recognition_experimental/image_recognition_test.cc
index ac4de118834..5ad2fb2acbe 100644
--- a/tensorflow/lite/micro/examples/image_recognition_experimental/image_recognition_test.cc
+++ b/tensorflow/lite/micro/examples/image_recognition_experimental/image_recognition_test.cc
@@ -44,14 +44,10 @@ TF_LITE_MICRO_TEST(TestImageRecognitionInvoke) {
 
   tflite::MicroMutableOpResolver<4> micro_op_resolver;
 
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_MAX_POOL_2D,
-                               tflite::ops::micro::Register_MAX_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
-                               tflite::ops::micro::Register_FULLY_CONNECTED());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddMaxPool2D();
+  micro_op_resolver.AddFullyConnected();
+  micro_op_resolver.AddSoftmax();
 
   const int tensor_arena_size = 50 * 1024;
   uint8_t tensor_arena[tensor_arena_size];
diff --git a/tensorflow/lite/micro/examples/image_recognition_experimental/main.cc b/tensorflow/lite/micro/examples/image_recognition_experimental/main.cc
index becdbdf1bd7..fcf7b41b827 100644
--- a/tensorflow/lite/micro/examples/image_recognition_experimental/main.cc
+++ b/tensorflow/lite/micro/examples/image_recognition_experimental/main.cc
@@ -58,14 +58,10 @@ int main(int argc, char** argv) {
 
   tflite::MicroMutableOpResolver<4> micro_op_resolver;
 
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_MAX_POOL_2D,
-                               tflite::ops::micro::Register_MAX_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
-                               tflite::ops::micro::Register_FULLY_CONNECTED());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddFullyConnected();
+  micro_op_resolver.AddMaxPool2D();
+  micro_op_resolver.AddSoftmax();
 
   constexpr int tensor_arena_size = 50 * 1024;
   uint8_t tensor_arena[tensor_arena_size];
diff --git a/tensorflow/lite/micro/examples/magic_wand/magic_wand_test.cc b/tensorflow/lite/micro/examples/magic_wand/magic_wand_test.cc
index 88bfad860e2..fb75afee309 100644
--- a/tensorflow/lite/micro/examples/magic_wand/magic_wand_test.cc
+++ b/tensorflow/lite/micro/examples/magic_wand/magic_wand_test.cc
@@ -47,17 +47,11 @@ TF_LITE_MICRO_TEST(LoadModelAndPerformInference) {
   // incur some penalty in code space for op implementations that are not
   // needed by this graph.
   static tflite::MicroMutableOpResolver<5> micro_op_resolver;  // NOLINT
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_MAX_POOL_2D,
-                               tflite::ops::micro::Register_MAX_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
-                               tflite::ops::micro::Register_FULLY_CONNECTED());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
+  micro_op_resolver.AddFullyConnected();
+  micro_op_resolver.AddMaxPool2D();
+  micro_op_resolver.AddSoftmax();
 
   // Create an area of memory to use for input, output, and intermediate arrays.
   // Finding the minimum value for your model may require some trial and error.
diff --git a/tensorflow/lite/micro/examples/magic_wand/main_functions.cc b/tensorflow/lite/micro/examples/magic_wand/main_functions.cc
index 26c2eb44747..8defeaad866 100644
--- a/tensorflow/lite/micro/examples/magic_wand/main_functions.cc
+++ b/tensorflow/lite/micro/examples/magic_wand/main_functions.cc
@@ -66,17 +66,11 @@ void setup() {
   // incur some penalty in code space for op implementations that are not
   // needed by this graph.
   static tflite::MicroMutableOpResolver<5> micro_op_resolver;  // NOLINT
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_MAX_POOL_2D,
-                               tflite::ops::micro::Register_MAX_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
-                               tflite::ops::micro::Register_FULLY_CONNECTED());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
+  micro_op_resolver.AddFullyConnected();
+  micro_op_resolver.AddMaxPool2D();
+  micro_op_resolver.AddSoftmax();
 
   // Build an interpreter to run the model with.
   static tflite::MicroInterpreter static_interpreter(
diff --git a/tensorflow/lite/micro/examples/micro_speech/main_functions.cc b/tensorflow/lite/micro/examples/micro_speech/main_functions.cc
index 30c5022b2d6..d09c4c7af06 100644
--- a/tensorflow/lite/micro/examples/micro_speech/main_functions.cc
+++ b/tensorflow/lite/micro/examples/micro_speech/main_functions.cc
@@ -75,24 +75,16 @@ void setup() {
   // tflite::AllOpsResolver resolver;
   // NOLINTNEXTLINE(runtime-global-variables)
   static tflite::MicroMutableOpResolver<4> micro_op_resolver(error_reporter);
-  if (micro_op_resolver.AddBuiltin(
-          tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-          tflite::ops::micro::Register_DEPTHWISE_CONV_2D()) != kTfLiteOk) {
+  if (micro_op_resolver.AddDepthwiseConv2D() != kTfLiteOk) {
     return;
   }
-  if (micro_op_resolver.AddBuiltin(
-          tflite::BuiltinOperator_FULLY_CONNECTED,
-          tflite::ops::micro::Register_FULLY_CONNECTED()) != kTfLiteOk) {
+  if (micro_op_resolver.AddFullyConnected() != kTfLiteOk) {
     return;
   }
-  if (micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                                   tflite::ops::micro::Register_SOFTMAX()) !=
-      kTfLiteOk) {
+  if (micro_op_resolver.AddSoftmax() != kTfLiteOk) {
     return;
   }
-  if (micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE,
-                                   tflite::ops::micro::Register_RESHAPE()) !=
-      kTfLiteOk) {
+  if (micro_op_resolver.AddReshape() != kTfLiteOk) {
     return;
   }
 
diff --git a/tensorflow/lite/micro/examples/micro_speech/micro_speech_test.cc b/tensorflow/lite/micro/examples/micro_speech/micro_speech_test.cc
index 2c442f955cc..0f6a2afd527 100644
--- a/tensorflow/lite/micro/examples/micro_speech/micro_speech_test.cc
+++ b/tensorflow/lite/micro/examples/micro_speech/micro_speech_test.cc
@@ -49,15 +49,10 @@ TF_LITE_MICRO_TEST(TestInvoke) {
   //
   // tflite::AllOpsResolver resolver;
   tflite::MicroMutableOpResolver<4> micro_op_resolver;
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_FULLY_CONNECTED,
-                               tflite::ops::micro::Register_FULLY_CONNECTED());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE,
-                               tflite::ops::micro::Register_RESHAPE());
+  micro_op_resolver.AddDepthwiseConv2D();
+  micro_op_resolver.AddFullyConnected();
+  micro_op_resolver.AddReshape();
+  micro_op_resolver.AddSoftmax();
 
   // Create an area of memory to use for input, output, and intermediate arrays.
   const int tensor_arena_size = 10 * 1024;
diff --git a/tensorflow/lite/micro/examples/person_detection/main_functions.cc b/tensorflow/lite/micro/examples/person_detection/main_functions.cc
index aa4d83a3334..d7e9f6826c4 100644
--- a/tensorflow/lite/micro/examples/person_detection/main_functions.cc
+++ b/tensorflow/lite/micro/examples/person_detection/main_functions.cc
@@ -66,13 +66,9 @@ void setup() {
   // tflite::AllOpsResolver resolver;
   // NOLINTNEXTLINE(runtime-global-variables)
   static tflite::MicroMutableOpResolver<3> micro_op_resolver;
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D,
-                               tflite::ops::micro::Register_AVERAGE_POOL_2D());
+  micro_op_resolver.AddAveragePool2D();
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
 
   // Build an interpreter to run the model with.
   static tflite::MicroInterpreter static_interpreter(
diff --git a/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc b/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc
index bc53a8410da..7e706d49fcc 100644
--- a/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc
+++ b/tensorflow/lite/micro/examples/person_detection/person_detection_test.cc
@@ -57,13 +57,9 @@ TF_LITE_MICRO_TEST(TestInvoke) {
   //
   // tflite::AllOpsResolver resolver;
   tflite::MicroMutableOpResolver<3> micro_op_resolver;
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D,
-                               tflite::ops::micro::Register_AVERAGE_POOL_2D());
+  micro_op_resolver.AddAveragePool2D();
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
 
   // Build an interpreter to run the model with.
   tflite::MicroInterpreter interpreter(model, micro_op_resolver, tensor_arena,
diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/main_functions.cc b/tensorflow/lite/micro/examples/person_detection_experimental/main_functions.cc
index ac47e36ff8f..09a9cb2c6c4 100644
--- a/tensorflow/lite/micro/examples/person_detection_experimental/main_functions.cc
+++ b/tensorflow/lite/micro/examples/person_detection_experimental/main_functions.cc
@@ -73,17 +73,11 @@ void setup() {
   // tflite::AllOpsResolver resolver;
   // NOLINTNEXTLINE(runtime-global-variables)
   static tflite::MicroMutableOpResolver<5> micro_op_resolver;
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D,
-                               tflite::ops::micro::Register_AVERAGE_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE,
-                               tflite::ops::micro::Register_RESHAPE());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddAveragePool2D();
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
+  micro_op_resolver.AddReshape();
+  micro_op_resolver.AddSoftmax();
 
   // Build an interpreter to run the model with.
   // NOLINTNEXTLINE(runtime-global-variables)
diff --git a/tensorflow/lite/micro/examples/person_detection_experimental/person_detection_test.cc b/tensorflow/lite/micro/examples/person_detection_experimental/person_detection_test.cc
index ddec8951596..270a427b1df 100644
--- a/tensorflow/lite/micro/examples/person_detection_experimental/person_detection_test.cc
+++ b/tensorflow/lite/micro/examples/person_detection_experimental/person_detection_test.cc
@@ -53,17 +53,11 @@ TF_LITE_MICRO_TEST(TestInvoke) {
   // incur some penalty in code space for op implementations that are not
   // needed by this graph.
   tflite::MicroMutableOpResolver<5> micro_op_resolver;
-  micro_op_resolver.AddBuiltin(
-      tflite::BuiltinOperator_DEPTHWISE_CONV_2D,
-      tflite::ops::micro::Register_DEPTHWISE_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_CONV_2D,
-                               tflite::ops::micro::Register_CONV_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_AVERAGE_POOL_2D,
-                               tflite::ops::micro::Register_AVERAGE_POOL_2D());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_RESHAPE,
-                               tflite::ops::micro::Register_RESHAPE());
-  micro_op_resolver.AddBuiltin(tflite::BuiltinOperator_SOFTMAX,
-                               tflite::ops::micro::Register_SOFTMAX());
+  micro_op_resolver.AddAveragePool2D();
+  micro_op_resolver.AddConv2D();
+  micro_op_resolver.AddDepthwiseConv2D();
+  micro_op_resolver.AddReshape();
+  micro_op_resolver.AddSoftmax();
 
   // Build an interpreter to run the model with.
   tflite::MicroInterpreter interpreter(model, micro_op_resolver, tensor_arena,
diff --git a/tensorflow/lite/micro/micro_mutable_op_resolver.h b/tensorflow/lite/micro/micro_mutable_op_resolver.h
index b9ce2bb4bba..1b76f440a61 100644
--- a/tensorflow/lite/micro/micro_mutable_op_resolver.h
+++ b/tensorflow/lite/micro/micro_mutable_op_resolver.h
@@ -104,39 +104,72 @@ class MicroMutableOpResolver : public MicroOpResolver {
     return kTfLiteOk;
   }
 
-  // Registers a Builtin Operator with the MicroOpResolver.
-  //
-  // Only the first call for a given BuiltinOperator enum will be successful.
-  // i.e. if this function is called again for a previously added
-  // BuiltinOperator, the MicroOpResolver will be unchanged and this function
-  // will return kTfLiteError.
-  //
-  // TODO(b/149408647): remove this API once the BuiltinOperator specific Add
-  // functions are fully implemented.
-  TfLiteStatus AddBuiltin(tflite::BuiltinOperator op,
-                          TfLiteRegistration* registration) {
-    TFLITE_DCHECK(registration != nullptr);
-    // For code that is not switched over to the new selective registration of
-    // the parse function, we pass in ParseOpData. This allows for backwards
-    // compatibility.
-    return AddBuiltin(op, *registration, ParseOpData);
-  }
-
   // The Add* functions below add the various Builtin operators to the
   // MicroMutableOpResolver object.
-  //
-  // This API is currently experimental (and only supported for a small subset
-  // of operators). It will soon be preferred over the AddBuiltin function for
-  // the following reason:
-  //  * If all calls to AddBuiltin for an application use this API, the code
-  //    size will be smaller by 5-8K (compared to the using the AddBuiltin
-  //    override).
+
+  TfLiteStatus AddAbs() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_ABS, *tflite::ops::micro::Register_ABS(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddAdd() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_ADD, *tflite::ops::micro::Register_ADD(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddArgMax() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_ARG_MAX,
+                      *tflite::ops::micro::Register_ARG_MAX(), ParseOpData);
+  }
+
+  TfLiteStatus AddArgMin() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_ARG_MIN,
+                      *tflite::ops::micro::Register_ARG_MIN(), ParseOpData);
+  }
+
+  TfLiteStatus AddAveragePool2D() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_AVERAGE_POOL_2D,
+                      *tflite::ops::micro::Register_AVERAGE_POOL_2D(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddCeil() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_CEIL,
+                      *tflite::ops::micro::Register_CEIL(), ParseOpData);
+  }
+
+  TfLiteStatus AddConcatenation() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_CONCATENATION,
+                      *tflite::ops::micro::Register_CONCATENATION(),
+                      ParseOpData);
+  }
 
   TfLiteStatus AddConv2D() {
     return AddBuiltin(BuiltinOperator_CONV_2D,
                       *tflite::ops::micro::Register_CONV_2D(), ParseConv2D);
   }
 
+  TfLiteStatus AddCos() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_COS, *tflite::ops::micro::Register_COS(),
+                      ParseOpData);
+  }
+
   TfLiteStatus AddDepthwiseConv2D() {
     return AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D,
                       *tflite::ops::micro::Register_DEPTHWISE_CONV_2D(),
@@ -149,12 +182,91 @@ class MicroMutableOpResolver : public MicroOpResolver {
                       ParseDequantize);
   }
 
+  TfLiteStatus AddEqual() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_EQUAL,
+                      *tflite::ops::micro::Register_EQUAL(), ParseOpData);
+  }
+
+  TfLiteStatus AddFloor() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_FLOOR,
+                      *tflite::ops::micro::Register_FLOOR(), ParseOpData);
+  }
+
   TfLiteStatus AddFullyConnected() {
     return AddBuiltin(BuiltinOperator_FULLY_CONNECTED,
                       *tflite::ops::micro::Register_FULLY_CONNECTED(),
                       ParseFullyConnected);
   }
 
+  TfLiteStatus AddGreater() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_GREATER,
+                      *tflite::ops::micro::Register_GREATER(), ParseOpData);
+  }
+
+  TfLiteStatus AddGreaterEqual() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_GREATER_EQUAL,
+                      *tflite::ops::micro::Register_GREATER_EQUAL(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddL2Normalization() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_L2_NORMALIZATION,
+                      *tflite::ops::micro::Register_L2_NORMALIZATION(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddLess() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LESS,
+                      *tflite::ops::micro::Register_LESS(), ParseOpData);
+  }
+
+  TfLiteStatus AddLessEqual() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LESS_EQUAL,
+                      *tflite::ops::micro::Register_LESS_EQUAL(), ParseOpData);
+  }
+
+  TfLiteStatus AddLog() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LOG, *tflite::ops::micro::Register_LOG(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddLogicalAnd() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LOGICAL_AND,
+                      *tflite::ops::micro::Register_LOGICAL_AND(), ParseOpData);
+  }
+
+  TfLiteStatus AddLogicalNot() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LOGICAL_NOT,
+                      *tflite::ops::micro::Register_LOGICAL_NOT(), ParseOpData);
+  }
+
+  TfLiteStatus AddLogicalOr() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_LOGICAL_OR,
+                      *tflite::ops::micro::Register_LOGICAL_OR(), ParseOpData);
+  }
+
   TfLiteStatus AddLogistic() {
     // TODO(b/149408647): Replace ParseOpData with the operator specific parse
     // function.
@@ -162,26 +274,196 @@ class MicroMutableOpResolver : public MicroOpResolver {
                       *tflite::ops::micro::Register_LOGISTIC(), ParseOpData);
   }
 
+  TfLiteStatus AddMaximum() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_MAXIMUM,
+                      *tflite::ops::micro::Register_MAXIMUM(), ParseOpData);
+  }
+
+  TfLiteStatus AddMaxPool2D() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_MAX_POOL_2D,
+                      *tflite::ops::micro::Register_MAX_POOL_2D(), ParseOpData);
+  }
+
+  TfLiteStatus AddMean() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_MEAN,
+                      *tflite::ops::micro::Register_MEAN(), ParseOpData);
+  }
+
+  TfLiteStatus AddMinimum() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_MINIMUM,
+                      *tflite::ops::micro::Register_MINIMUM(), ParseOpData);
+  }
+
+  TfLiteStatus AddMul() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_MUL, *tflite::ops::micro::Register_MUL(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddNeg() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_NEG, *tflite::ops::micro::Register_NEG(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddNotEqual() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_NOT_EQUAL,
+                      *tflite::ops::micro::Register_NOT_EQUAL(), ParseOpData);
+  }
+
+  TfLiteStatus AddPack() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_PACK,
+                      *tflite::ops::micro::Register_PACK(), ParseOpData);
+  }
+
+  TfLiteStatus AddPad() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_PAD, *tflite::ops::micro::Register_PAD(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddPadV2() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_PADV2,
+                      *tflite::ops::micro::Register_PADV2(), ParseOpData);
+  }
+
+  TfLiteStatus AddPrelu() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_PRELU,
+                      *tflite::ops::micro::Register_PRELU(), ParseOpData);
+  }
+
   TfLiteStatus AddQuantize() {
     return AddBuiltin(BuiltinOperator_QUANTIZE,
                       *tflite::ops::micro::Register_QUANTIZE(), ParseQuantize);
   }
 
+  TfLiteStatus AddRelu() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_RELU,
+                      *tflite::ops::micro::Register_RELU(), ParseOpData);
+  }
+
+  TfLiteStatus AddRelu6() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_RELU6,
+                      *tflite::ops::micro::Register_RELU6(), ParseOpData);
+  }
+
   TfLiteStatus AddReshape() {
     return AddBuiltin(BuiltinOperator_RESHAPE,
                       *tflite::ops::micro::Register_RESHAPE(), ParseReshape);
   }
 
+  TfLiteStatus AddResizeNearestNeighbor() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_RESIZE_NEAREST_NEIGHBOR,
+                      *tflite::ops::micro::Register_RESIZE_NEAREST_NEIGHBOR(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddRound() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_ROUND,
+                      *tflite::ops::micro::Register_ROUND(), ParseOpData);
+  }
+
+  TfLiteStatus AddRsqrt() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_RSQRT,
+                      *tflite::ops::micro::Register_RSQRT(), ParseOpData);
+  }
+
+  TfLiteStatus AddSin() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_SIN, *tflite::ops::micro::Register_SIN(),
+                      ParseOpData);
+  }
+
   TfLiteStatus AddSoftmax() {
     return AddBuiltin(BuiltinOperator_SOFTMAX,
                       *tflite::ops::micro::Register_SOFTMAX(), ParseSoftmax);
   }
 
+  TfLiteStatus AddSplit() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_SPLIT,
+                      *tflite::ops::micro::Register_SPLIT(), ParseOpData);
+  }
+
+  TfLiteStatus AddSqrt() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_SQRT,
+                      *tflite::ops::micro::Register_SQRT(), ParseOpData);
+  }
+
+  TfLiteStatus AddSquare() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_SQUARE,
+                      *tflite::ops::micro::Register_SQUARE(), ParseOpData);
+  }
+
+  TfLiteStatus AddStridedSlice() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_STRIDED_SLICE,
+                      *tflite::ops::micro::Register_STRIDED_SLICE(),
+                      ParseOpData);
+  }
+
+  TfLiteStatus AddSub() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_SUB, *tflite::ops::micro::Register_SUB(),
+                      ParseOpData);
+  }
+
   TfLiteStatus AddSvdf() {
     return AddBuiltin(BuiltinOperator_SVDF,
                       *tflite::ops::micro::Register_SVDF(), ParseSvdf);
   }
 
+  TfLiteStatus AddTanh() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_TANH,
+                      *tflite::ops::micro::Register_TANH(), ParseOpData);
+  }
+
+  TfLiteStatus AddUnpack() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_UNPACK,
+                      *tflite::ops::micro::Register_UNPACK(), ParseOpData);
+  }
+
   unsigned int GetRegistrationLength() { return registrations_len_; }
 
  private:
diff --git a/tensorflow/lite/micro/micro_mutable_op_resolver_test.cc b/tensorflow/lite/micro/micro_mutable_op_resolver_test.cc
index ff5dfdf3a9a..fe9c8de5959 100644
--- a/tensorflow/lite/micro/micro_mutable_op_resolver_test.cc
+++ b/tensorflow/lite/micro/micro_mutable_op_resolver_test.cc
@@ -68,14 +68,7 @@ TF_LITE_MICRO_TEST(TestOperations) {
   static TfLiteRegistration r = {tflite::MockInit, tflite::MockFree,
                                  tflite::MockPrepare, tflite::MockInvoke};
 
-  MicroMutableOpResolver<2> micro_op_resolver;
-  TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, micro_op_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &r));
-
-  // Only one AddBuiltin per operator should return kTfLiteOk.
-  TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteError, micro_op_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &r));
-
+  MicroMutableOpResolver<1> micro_op_resolver;
   TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk,
                           micro_op_resolver.AddCustom("mock_custom", &r));
 
@@ -85,16 +78,10 @@ TF_LITE_MICRO_TEST(TestOperations) {
 
   tflite::MicroOpResolver* resolver = µ_op_resolver;
 
+  TF_LITE_MICRO_EXPECT_EQ(1, micro_op_resolver.GetRegistrationLength());
+
   const TfLiteRegistration* registration =
-      resolver->FindOp(BuiltinOperator_CONV_2D);
-  TF_LITE_MICRO_EXPECT_NE(nullptr, registration);
-  TF_LITE_MICRO_EXPECT_EQ(nullptr, registration->init(nullptr, nullptr, 0));
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(nullptr, nullptr));
-  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(nullptr, nullptr));
-
-  TF_LITE_MICRO_EXPECT_EQ(2, micro_op_resolver.GetRegistrationLength());
-
-  registration = resolver->FindOp(BuiltinOperator_RELU);
+      resolver->FindOp(BuiltinOperator_RELU);
   TF_LITE_MICRO_EXPECT_EQ(nullptr, registration);
 
   registration = resolver->FindOp("mock_custom");
@@ -116,12 +103,7 @@ TF_LITE_MICRO_TEST(TestErrorReporting) {
                                  tflite::MockPrepare, tflite::MockInvoke};
 
   tflite::MockErrorReporter mock_reporter;
-  MicroMutableOpResolver<2> micro_op_resolver(&mock_reporter);
-  TF_LITE_MICRO_EXPECT_EQ(false, mock_reporter.HasBeenCalled());
-  mock_reporter.ResetState();
-
-  TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteOk, micro_op_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &r));
+  MicroMutableOpResolver<1> micro_op_resolver(&mock_reporter);
   TF_LITE_MICRO_EXPECT_EQ(false, mock_reporter.HasBeenCalled());
   mock_reporter.ResetState();
 
@@ -132,10 +114,7 @@ TF_LITE_MICRO_TEST(TestErrorReporting) {
 
   // Attempting to Add more operators than the class template parameter for
   // MicroMutableOpResolver should result in errors.
-  TF_LITE_MICRO_EXPECT_EQ(
-      kTfLiteError, micro_op_resolver.AddBuiltin(BuiltinOperator_RELU, &r));
-  TF_LITE_MICRO_EXPECT_EQ(true, mock_reporter.HasBeenCalled());
-  mock_reporter.ResetState();
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteError, micro_op_resolver.AddRelu());
 
   TF_LITE_MICRO_EXPECT_EQ(kTfLiteError,
                           micro_op_resolver.AddCustom("mock_custom_1", &r));

From e74a115bcf5cd27f476b46161a639e9ec599491d Mon Sep 17 00:00:00 2001
From: Peng Wang 
Date: Mon, 22 Jun 2020 17:47:47 -0700
Subject: [PATCH 0845/1390] [TF-numpy] Adds __rmatmul__ method to ndarray.

PiperOrigin-RevId: 317771125
Change-Id: I719c46d97ae1c68ac59dcd1cf8f65d067ddc7658
---
 tensorflow/python/ops/numpy_ops/np_math_ops.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 427aa96e5a4..03329bbdbf1 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -950,11 +950,12 @@ setattr(np_arrays.ndarray, '__sub__', _wrap(subtract))
 setattr(np_arrays.ndarray, '__rsub__', _wrap(subtract, True))
 setattr(np_arrays.ndarray, '__mul__', _wrap(multiply))
 setattr(np_arrays.ndarray, '__rmul__', _wrap(multiply, True))
+setattr(np_arrays.ndarray, '__matmul__', _wrap(matmul))
+setattr(np_arrays.ndarray, '__rmatmul__', _wrap(matmul, True))
 setattr(np_arrays.ndarray, '__pow__', _wrap(power))
 setattr(np_arrays.ndarray, '__rpow__', _wrap(power, True))
 setattr(np_arrays.ndarray, '__truediv__', _wrap(true_divide))
 setattr(np_arrays.ndarray, '__rtruediv__', _wrap(true_divide, True))
-setattr(np_arrays.ndarray, '__matmul__', _wrap(matmul))
 
 
 def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):

From 4d13d6416da49b82b375d631082607be3c58b57f Mon Sep 17 00:00:00 2001
From: Rick Chao 
Date: Mon, 22 Jun 2020 17:49:06 -0700
Subject: [PATCH 0846/1390] Make cluster_resolver standard property in
 tf.distribute strategies.

PiperOrigin-RevId: 317771299
Change-Id: I71b5c585cef7bd7ef80e66b75e30287fddcf89e2
---
 tensorflow/python/distribute/BUILD            |  5 +-
 .../collective_all_reduce_strategy.py         | 12 ++++
 .../collective_all_reduce_strategy_test.py    | 15 ++++-
 .../python/distribute/distribute_lib.py       | 59 +++++++++++++++++++
 .../python/distribute/distribute_lib_test.py  | 13 ++++
 .../python/distribute/strategy_common_test.py | 35 +++++++++++
 tensorflow/python/distribute/tpu_strategy.py  | 12 ++++
 .../python/distribute/tpu_strategy_test.py    |  7 +++
 ...orflow.distribute.-mirrored-strategy.pbtxt |  4 ++
 ...flow.distribute.-one-device-strategy.pbtxt |  4 ++
 .../v1/tensorflow.distribute.-strategy.pbtxt  |  4 ++
 ...perimental.-central-storage-strategy.pbtxt |  4 ++
 ...ntal.-multi-worker-mirrored-strategy.pbtxt |  4 ++
 ...erimental.-parameter-server-strategy.pbtxt |  4 ++
 ...tribute.experimental.-t-p-u-strategy.pbtxt |  4 ++
 ...orflow.distribute.-mirrored-strategy.pbtxt |  4 ++
 ...flow.distribute.-one-device-strategy.pbtxt |  4 ++
 .../v2/tensorflow.distribute.-strategy.pbtxt  |  4 ++
 ...ensorflow.distribute.-t-p-u-strategy.pbtxt |  4 ++
 ...perimental.-central-storage-strategy.pbtxt |  4 ++
 ...ntal.-multi-worker-mirrored-strategy.pbtxt |  4 ++
 ...erimental.-parameter-server-strategy.pbtxt |  4 ++
 ...tribute.experimental.-t-p-u-strategy.pbtxt |  4 ++
 23 files changed, 215 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 74d80b63e12..85ee8de5635 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -204,6 +204,7 @@ py_test(
         "//tensorflow/python:variables",
         "//tensorflow/python/autograph/core:test_lib",
         "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//third_party/py/numpy",
     ],
 )
@@ -1847,10 +1848,11 @@ py_test(
     ],
 )
 
-cuda_py_test(
+distribute_py_test(
     name = "strategy_common_test",
     srcs = ["strategy_common_test.py"],
     python_version = "PY3",
+    shard_count = 12,
     tags = [
         "multi_and_single_gpu",
         # TODO(b/155301154): Enable this test on multi-gpu guitar once multi process
@@ -1859,6 +1861,7 @@ cuda_py_test(
     ],
     xla_enable_strict_auto_jit = True,
     deps = [
+        ":collective_all_reduce_strategy",
         ":combinations",
         ":multi_worker_test_base",
         ":reduce_util",
diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy.py b/tensorflow/python/distribute/collective_all_reduce_strategy.py
index e2b039ceb23..e754cc43a41 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy.py
@@ -138,6 +138,18 @@ class CollectiveAllReduceStrategy(distribute_lib.Strategy):
     """
     return super(CollectiveAllReduceStrategy, self).scope()
 
+  @property
+  def cluster_resolver(self):
+    """Returns the cluster resolver associated with this strategy.
+
+    As a multi-worker strategy,
+    `tf.distribute.experimental.MultiWorkerMirroredStrategy` provides the
+    associated `tf.distribute.cluster_resolver.ClusterResolver`. If the user
+    provides one in `__init__`, that instance is returned; if the user does
+    not, a default `TFConfigClusterResolver` is provided.
+    """
+    return self.extended._cluster_resolver  # pylint: disable=protected-access
+
 
 @tf_export(v1=["distribute.experimental.MultiWorkerMirroredStrategy"])  # pylint: disable=missing-docstring
 class CollectiveAllReduceStrategyV1(distribute_lib.StrategyV1):
diff --git a/tensorflow/python/distribute/collective_all_reduce_strategy_test.py b/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
index 87212c85fc4..9d16c6f5a1f 100644
--- a/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
+++ b/tensorflow/python/distribute/collective_all_reduce_strategy_test.py
@@ -505,8 +505,7 @@ class DistributedCollectiveAllReduceStrategyTest(
     self.assertEqual(['CollectiveReduce'],
                      new_rewrite_options.scoped_allocator_opts.enable_op)
 
-  @combinations.generate(combinations.combine(mode=['eager']))
-  def testEnableCollectiveOps(self):
+  def _get_strategy_with_mocked_methods(self):
     mock_called = [False]
 
     # pylint: disable=dangerous-default-value
@@ -525,9 +524,21 @@ class DistributedCollectiveAllReduceStrategyTest(
                                 mock_configure_collective_ops):
       strategy, _, _ = self._get_test_object(
           task_type='worker', task_id=1, num_gpus=2)
+
+    return strategy, mock_called
+
+  @combinations.generate(combinations.combine(mode=['eager']))
+  def testEnableCollectiveOps(self):
+    strategy, mock_called = self._get_strategy_with_mocked_methods()
     self.assertTrue(strategy.extended._std_server_started)
     self.assertTrue(mock_called[0])
 
+  @combinations.generate(combinations.combine(mode=['eager']))
+  def testEnableCollectiveOpsAndClusterResolver(self):
+    strategy, _ = self._get_strategy_with_mocked_methods()
+    self.assertEqual(strategy.cluster_resolver.task_type, 'worker')
+    self.assertEqual(strategy.cluster_resolver.task_id, 1)
+
 
 class DistributedCollectiveAllReduceStrategyTestWithChief(
     CollectiveAllReduceStrategyTestBase, parameterized.TestCase):
diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index 5abfb6e1c09..f32427b88e0 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -1439,6 +1439,65 @@ class StrategyBase(object):
   def __copy__(self):
     raise RuntimeError("Must only deepcopy DistributionStrategy.")
 
+  @property
+  def cluster_resolver(self):
+    """Returns the cluster resolver associated with this strategy.
+
+    In general, when using a multi-worker `tf.distribute` strategy such as
+    `tf.distribute.experimental.MultiWorkerMirroredStrategy` or
+    `tf.distribute.experimental.TPUStrategy()`, there is a
+    `tf.distribute.cluster_resolver.ClusterResolver` associated with the
+    strategy used, and such an instance is returned by this property.
+
+    Strategies that intend to have an associated
+    `tf.distribute.cluster_resolver.ClusterResolver` must set the
+    relevant attribute, or override this property; otherwise, `None` is returned
+    by default. Those strategies should also provide information regarding what
+    is returned by this property.
+
+    Single-worker strategies usually do not have a
+    `tf.distribute.cluster_resolver.ClusterResolver`, and in those cases this
+    property will return `None`.
+
+    The `tf.distribute.cluster_resolver.ClusterResolver` may be useful when the
+    user needs to access information such as the cluster spec, task type or task
+    id. For example,
+
+    ```python
+
+    os.environ['TF_CONFIG'] = json.dumps({
+      'cluster': {
+          'worker': ["localhost:12345", "localhost:23456"],
+          'ps': ["localhost:34567"]
+      },
+      'task': {'type': 'worker', 'index': 0}
+    })
+
+    # This implicitly uses TF_CONFIG for the cluster and current task info.
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+
+    ...
+
+    if strategy.cluster_resolver.task_type == 'worker':
+      # Perform something that's only applicable on workers. Since we set this
+      # as a worker above, this block will run on this particular instance.
+    elif strategy.cluster_resolver.task_type == 'ps':
+      # Perform something that's only applicable on parameter servers. Since we
+      # set this as a worker above, this block will not run on this particular
+      # instance.
+    ```
+
+    For more information, please see
+    `tf.distribute.cluster_resolver.ClusterResolver`'s API docstring.
+
+    Returns:
+      The cluster resolver associated with this strategy. Returns `None` if a
+      cluster resolver is not applicable or available in this strategy.
+    """
+    if hasattr(self.extended, "_cluster_resolver"):
+      return self.extended._cluster_resolver  # pylint: disable=protected-access
+    return None
+
 
 @tf_export("distribute.Strategy", v1=[])  # pylint: disable=g-missing-docstring
 class Strategy(StrategyBase):
diff --git a/tensorflow/python/distribute/distribute_lib_test.py b/tensorflow/python/distribute/distribute_lib_test.py
index 8ea1cac6f02..b5924ec3b67 100644
--- a/tensorflow/python/distribute/distribute_lib_test.py
+++ b/tensorflow/python/distribute/distribute_lib_test.py
@@ -28,6 +28,7 @@ from tensorflow.python.distribute import distribute_lib
 from tensorflow.python.distribute import distribution_strategy_context as ds_context
 from tensorflow.python.distribute import input_lib
 from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
@@ -36,6 +37,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
 from tensorflow.python.util import nest
 
 
@@ -422,6 +424,17 @@ class TestStrategyTest(test.TestCase):
 
     test_fn()
 
+  def testClusterResolverDefaultNotImplemented(self):
+    dist = _TestStrategy()
+    self.assertIsNone(dist.cluster_resolver)
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+    cluster_resolver = SimpleClusterResolver(base_cluster_spec)
+    dist.extended._cluster_resolver = cluster_resolver
+    self.assertIs(dist.cluster_resolver, cluster_resolver)
+
 
 # _TestStrategy2 is like _TestStrategy, except it doesn't change variable
 # creation.
diff --git a/tensorflow/python/distribute/strategy_common_test.py b/tensorflow/python/distribute/strategy_common_test.py
index b823e2de331..d1a72da7e7c 100644
--- a/tensorflow/python/distribute/strategy_common_test.py
+++ b/tensorflow/python/distribute/strategy_common_test.py
@@ -27,6 +27,8 @@ from tensorflow.python.distribute import multi_worker_test_base
 from tensorflow.python.distribute import reduce_util
 from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.distribute import strategy_test_lib
+from tensorflow.python.distribute.collective_all_reduce_strategy import CollectiveAllReduceStrategy
+from tensorflow.python.distribute.tpu_strategy import TPUStrategy
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
@@ -184,5 +186,38 @@ class DistributedCollectiveAllReduceStrategyTest(
   # worker strategy combinations can run on a fixed number of GPUs.
 
 
+class StrategyClusterResolverTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          strategy=[strategy_combinations.multi_worker_mirrored_two_workers] +
+          strategy_combinations.all_strategies,
+          mode=['eager']))
+  def testClusterResolverProperty(self, strategy):
+    # CollectiveAllReduceStrategy and TPUStrategy must have a cluster resolver.
+    # `None` otherwise.
+    resolver = strategy.cluster_resolver
+    if not isinstance(strategy, CollectiveAllReduceStrategy) and not isinstance(
+        strategy, TPUStrategy):
+      self.assertIsNone(resolver)
+      return
+
+    with strategy.scope():
+      self.assertIs(strategy.cluster_resolver, resolver)
+    self.assertTrue(hasattr(resolver, 'cluster_spec'))
+    self.assertTrue(hasattr(resolver, 'environment'))
+    self.assertTrue(hasattr(resolver, 'master'))
+    self.assertTrue(hasattr(resolver, 'num_accelerators'))
+    self.assertIsNone(resolver.rpc_layer)
+    if isinstance(strategy, CollectiveAllReduceStrategy):
+      self.assertGreaterEqual(resolver.task_id, 0)
+      self.assertLessEqual(resolver.task_id, 1)
+      self.assertEqual(resolver.task_type, 'worker')
+    elif isinstance(strategy, TPUStrategy):
+      # TPUStrategy does not have task_id and task_type applicable.
+      self.assertIsNone(resolver.task_id)
+      self.assertIsNone(resolver.task_type)
+
+
 if __name__ == '__main__':
   combinations.main()
diff --git a/tensorflow/python/distribute/tpu_strategy.py b/tensorflow/python/distribute/tpu_strategy.py
index 4b3c4be0ccd..df393c61dbb 100644
--- a/tensorflow/python/distribute/tpu_strategy.py
+++ b/tensorflow/python/distribute/tpu_strategy.py
@@ -345,6 +345,18 @@ class TPUStrategy(distribute_lib.Strategy):
     options = options or distribute_lib.RunOptions()
     return self.extended.tpu_run(fn, args, kwargs, options)
 
+  @property
+  def cluster_resolver(self):
+    """Returns the cluster resolver associated with this strategy.
+
+    `tf.distribute.experimental.TPUStrategy` provides the
+    associated `tf.distribute.cluster_resolver.ClusterResolver`. If the user
+    provides one in `__init__`, that instance is returned; if the user does
+    not, a default
+    `tf.distribute.cluster_resolver.TPUClusterResolver` is provided.
+    """
+    return self.extended._tpu_cluster_resolver  # pylint: disable=protected-access
+
 
 @tf_export(v1=["distribute.experimental.TPUStrategy"])
 class TPUStrategyV1(distribute_lib.StrategyV1):
diff --git a/tensorflow/python/distribute/tpu_strategy_test.py b/tensorflow/python/distribute/tpu_strategy_test.py
index 5e47e750d87..142743a6ec2 100644
--- a/tensorflow/python/distribute/tpu_strategy_test.py
+++ b/tensorflow/python/distribute/tpu_strategy_test.py
@@ -555,6 +555,13 @@ class TPUStrategyTest(test.TestCase, parameterized.TestCase):
       update_variable.get_concrete_function()
       self.assertLen(strategy.extended.worker_devices, trace_count[0])
 
+  def test_cluster_resolver_available(self, enable_packed_var):
+    resolver = get_tpu_cluster_resolver()
+    remote.connect_to_cluster(resolver)
+    tpu_strategy_util.initialize_tpu_system(resolver)
+    strategy = tpu_lib.TPUStrategy(resolver)
+    self.assertIs(strategy.cluster_resolver, resolver)
+
 
 class TPUStrategyDataPrefetchTest(test.TestCase):
 
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
index 36c78c406b7..0c5db602029 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
index 09865ab02ee..ae62acffa44 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-one-device-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
index 0e6c10bd533..9285405ea4f 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.-strategy.pbtxt
@@ -3,6 +3,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
index fbc4c107a1a..3c3d785ac7c 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
index cd67e7d27c4..e1f8bea251b 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
index 0eff82474ff..6ae83d18589 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
index 2af9a5ad095..0e548eca9b5 100644
--- a/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v1/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
index be4c841aed7..8817f16d808 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-mirrored-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-one-device-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-one-device-strategy.pbtxt
index 4557fe1060b..992243ffe8a 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-one-device-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-one-device-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
index 9f6a2ac32be..8140088e701 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-strategy.pbtxt
@@ -3,6 +3,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-t-p-u-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-t-p-u-strategy.pbtxt
index f41a08454e2..29947a1c9c5 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.-t-p-u-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.-t-p-u-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
index dd61960c66f..695fb52358b 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-central-storage-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
index 500ae362e5f..43632e17b6d 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-multi-worker-mirrored-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
index d6dc9627d9a..39181625469 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-parameter-server-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
index 82a4362a597..855cdbfb175 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.distribute.experimental.-t-p-u-strategy.pbtxt
@@ -4,6 +4,10 @@ tf_class {
   is_instance: ""
   is_instance: ""
   is_instance: ""
+  member {
+    name: "cluster_resolver"
+    mtype: ""
+  }
   member {
     name: "extended"
     mtype: ""

From e60cf08994a837d1dc37ec2f85063e44b3987a81 Mon Sep 17 00:00:00 2001
From: Rahul Joshi 
Date: Mon, 22 Jun 2020 18:11:15 -0700
Subject: [PATCH 0847/1390] [NFC] Re-instate use of FuncOp::isPublic()

PiperOrigin-RevId: 317774876
Change-Id: I6a832236377d403b1ebd24ecbc26025c37dc1c13
---
 .../compiler/mlir/tensorflow/ir/tf_saved_model.cc   | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
index ef55761686e..5a7d81d4c0c 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
@@ -232,22 +232,19 @@ static LogicalResult VerifySavedModelModule(
   for (auto func : module.getOps()) {
     const bool is_exported = IsExported(func);
 
-    if (is_exported && func.getVisibility() != FuncOp::Visibility::Public) {
+    if (is_exported && !func.isPublic()) {
       return func.emitError()
              << "exported function @" << func.getName() << " should be public";
     }
 
-    if (!is_exported && func.getVisibility() == FuncOp::Visibility::Public) {
+    if (!is_exported && func.isPublic()) {
       return func.emitError() << "non-exported function @" << func.getName()
                               << " should be private";
     }
 
-    if (HasAnyTfSavedModelArgAttr(func)) {
-      if (!is_exported) {
-        return func.emitError()
-               << "can only apply 'tf_saved_model' argument attributes "
-                  "to exported functions";
-      }
+    if (!is_exported && HasAnyTfSavedModelArgAttr(func)) {
+      return func.emitError() << "can only apply 'tf_saved_model' argument "
+                                 "attributes to exported functions";
     }
   }
 

From 265de52331a10af793f733d21e2152123819a269 Mon Sep 17 00:00:00 2001
From: Scott Zhu 
Date: Mon, 22 Jun 2020 18:20:43 -0700
Subject: [PATCH 0848/1390] Fix input mapping issue when model is
 constructed/tested with dict input tensor.

The mapping of the dict input tensors was not correct since it was still using the tensor name, rather than the key of the tensor when build the model. This cause the issue down the stream when the inputs are provided with unknown keys.

We had some backup logic, which will probably do correct things, eg just flat the dict to keep the original order, which was correct most of the case, but not very reliable. In this change, we make the behavior change:

1. When model is build with dict input tensors, the key of the tensor, instead of the name, will be used to map the tensor with input data.
2. Unknown keys in the input data will result into a warning. We didn't throw error since user might do it intentionally, eg using part of the model to test with full input data.

PiperOrigin-RevId: 317776370
Change-Id: I91983443f2b770cb0b45ddb7726f52708cb91d61
---
 tensorflow/python/keras/engine/functional.py  | 24 ++++++++--
 .../python/keras/engine/functional_test.py    | 46 +++++++++++++++++++
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/keras/engine/functional.py b/tensorflow/python/keras/engine/functional.py
index 0612d70044d..fd80e7f8bb4 100644
--- a/tensorflow/python/keras/engine/functional.py
+++ b/tensorflow/python/keras/engine/functional.py
@@ -22,6 +22,7 @@ from __future__ import print_function
 import collections
 import copy
 import itertools
+import warnings
 
 from six.moves import zip  # pylint: disable=redefined-builtin
 
@@ -131,10 +132,10 @@ class Functional(training_lib.Model):
 
     # Models constructed with a single Tensor or list of Tensors can
     # be called with a dict, where the keys of the dict are the names
-    # of the `Input` objects. Extra keys are ignored.
+    # of the `Input` objects. Extra keys are ignored with warning.
     self._enable_dict_to_input_mapping = (
         not nest.is_sequence(self._nested_inputs) or
-        (isinstance(self._nested_inputs, (list, tuple)) and
+        (isinstance(self._nested_inputs, (list, tuple, dict)) and
          not any(nest.is_sequence(t) for t in self._nested_inputs)))
 
     if any(not hasattr(tensor, '_keras_history') for tensor in self.outputs):
@@ -524,10 +525,27 @@ class Functional(training_lib.Model):
       ref_inputs = self._nested_inputs
       if not nest.is_sequence(ref_inputs):
         ref_inputs = [self._nested_inputs]
+      if isinstance(ref_inputs, dict):
+        # In the case that the graph is constructed with dict input tensors,
+        # We will use the original dict key to map with the keys in the input
+        # data. Note that the model.inputs is using nest.flatten to process the
+        # input tensors, which means the dict input tensors are ordered by their
+        # keys.
+        ref_input_names = sorted(ref_inputs.keys())
+      else:
+        ref_input_names = [inp._keras_history.layer.name for inp in ref_inputs]
+
+      # Raise an warning if there are more input data comparing to input tensor
+      if len(tensors) > len(ref_input_names):
+        warnings.warn(
+            'Input dict contained keys {} which did not match any model input. '
+            'They will be ignored by the model.'.format(
+                [n for n in tensors.keys() if n not in ref_input_names])
+            )
 
       try:
         # Flatten in the order `Input`s were passed during Model construction.
-        return [tensors[inp._keras_history.layer.name] for inp in ref_inputs]
+        return [tensors[n] for n in ref_input_names]
       except KeyError:
         # TODO(b/151582614)
         return nest.flatten(tensors)
diff --git a/tensorflow/python/keras/engine/functional_test.py b/tensorflow/python/keras/engine/functional_test.py
index 3c14411deb9..0e82d95d3de 100644
--- a/tensorflow/python/keras/engine/functional_test.py
+++ b/tensorflow/python/keras/engine/functional_test.py
@@ -18,8 +18,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import warnings
+
 import numpy as np
 
+
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
@@ -43,6 +46,7 @@ from tensorflow.python.keras.utils import tf_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import string_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
 from tensorflow.python.platform import test
 from tensorflow.python.training.tracking.util import Checkpoint
@@ -1565,6 +1569,48 @@ class DefaultShapeInferenceBehaviorTest(keras_parameterized.TestCase):
     self.assertEqual(config['layers'][2]['inbound_nodes'],
                      [[['in1', 0, 0, {}], ['in2', 0, 0, {}]]])
 
+  @combinations.generate(combinations.combine(mode=['eager']))
+  def test_dict_inputs_tensors(self):
+    # Note that this test is running with v2 eager only, since the v1
+    # will behave differently wrt to dict input for training.
+    inputs = {
+        'sentence2': input_layer_lib.Input(
+            shape=(), name='a', dtype=dtypes.string),
+        'sentence1': input_layer_lib.Input(
+            shape=(), name='b', dtype=dtypes.string),
+    }
+    strlen = layers.Lambda(string_ops.string_length_v2)
+    diff = layers.Subtract()(
+        [strlen(inputs['sentence1']), strlen(inputs['sentence2'])])
+    diff = math_ops.cast(diff, dtypes.float32)
+    model = training_lib.Model(inputs, diff)
+
+    extra_keys = {
+        'sentence1': constant_op.constant(['brown fox', 'lazy dog']),
+        'sentence2': constant_op.constant(['owl', 'cheeky cat']),
+        'label': constant_op.constant([0, 1]),
+    }
+
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always')
+      model(extra_keys)
+      self.assertIn('ignored by the model', str(w[-1].message))
+
+    model.compile('sgd', 'mse')
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always')
+      model.fit(extra_keys, y=constant_op.constant([0, 1]), steps_per_epoch=1)
+      self.assertIn('ignored by the model', str(w[-1].message))
+
+    with warnings.catch_warnings(record=True) as w:
+      warnings.simplefilter('always')
+      model.evaluate(extra_keys, constant_op.constant([0, 1]))
+      self.assertIn('ignored by the model', str(w[-1].message))
+
+    # Make sure the model inputs are sorted with the dict keys.
+    self.assertEqual(model.inputs[0]._keras_history.layer.name, 'b')
+    self.assertEqual(model.inputs[1]._keras_history.layer.name, 'a')
+
 
 class GraphUtilsTest(test.TestCase):
 

From d8b90acf95bdff09e29b1065739e28e8bd7500d9 Mon Sep 17 00:00:00 2001
From: Dan Moldovan 
Date: Mon, 22 Jun 2020 18:46:12 -0700
Subject: [PATCH 0849/1390] Enable type annotations for python/ops.

PiperOrigin-RevId: 317779660
Change-Id: Ife6b7319ef394b611798f5d01b64ebdb3c0a01cc
---
 tensorflow/python/ops/logging_ops.py | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/tensorflow/python/ops/logging_ops.py b/tensorflow/python/ops/logging_ops.py
index 8ca63f55987..02fce277690 100644
--- a/tensorflow/python/ops/logging_ops.py
+++ b/tensorflow/python/ops/logging_ops.py
@@ -54,11 +54,9 @@ except NameError:
 # call relies on certain conditionals for its dependencies.  Use
 # control_flow_ops.Assert.
 
-# Assert and Print are special symbols in python, so we must
-# have an upper-case version of them.
-#
-# For users with Python 3 or Python 2.7
-# with `from __future__ import print_function`, we could also allow lowercase.
+# Assert and Print are special symbols in Python 2, so we must
+# have an upper-case version of them. When support for it is dropped,
+# we can allow lowercase.
 # See https://github.com/tensorflow/tensorflow/issues/18053
 
 
@@ -83,11 +81,6 @@ def Print(input_, data, message=None, first_n=None, summarize=None, name=None):
     with jupyter notebook (printing to the notebook *server's* output, not into
     the notebook).
 
-  Additionally, to use tf.print in python 2.7, users must make sure to import
-  the following:
-
-  `from __future__ import print_function`
-
   Args:
     input_: A tensor passed through this op.
     data: A list of tensors to print out when op is evaluated.
@@ -148,11 +141,6 @@ def print_v2(*inputs, **kwargs):
   Python objects. Printed tensors will recursively show the first and last
   elements of each dimension to summarize.
 
-  @compatibility(python2)
-  In python 2.7, make sure to import the following:
-  `from __future__ import print_function`
-  @end_compatibility
-
   Example:
     Single-input usage:
 

From 7c669e5f795a75bf6665c9596d730e438ab36511 Mon Sep 17 00:00:00 2001
From: Bixia Zheng 
Date: Mon, 22 Jun 2020 19:13:19 -0700
Subject: [PATCH 0850/1390] [TF:TRT] Add #if to logger_registry.h for
 consistency.

Add #if GOOGLE_TENSORRT to logger_registry.h to make it consistent with
logger_registry.cc.

Rewrite two lines of #if into #if GOOGLE_CUDA && GOOGLE_TENSORRT.

PiperOrigin-RevId: 317782636
Change-Id: Icaf7116b5013f57bb37b80511c7c195f8bd7e2d5
---
 tensorflow/compiler/tf2tensorrt/common/utils.h       |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_graph.cc    |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_graph.h     |  6 ++----
 .../tf2tensorrt/convert/convert_graph_test.cc        |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_nodes.cc    |  6 ++----
 .../compiler/tf2tensorrt/convert/convert_nodes.h     |  6 ++----
 .../tf2tensorrt/convert/convert_nodes_test.cc        |  6 ++----
 .../compiler/tf2tensorrt/convert/logger_registry.cc  |  6 ++----
 .../compiler/tf2tensorrt/convert/logger_registry.h   |  5 +++--
 .../tf2tensorrt/convert/trt_optimization_pass.cc     |  6 ++----
 .../tf2tensorrt/convert/trt_optimization_pass.h      |  6 ++----
 .../tf2tensorrt/kernels/get_calibration_data_op.cc   |  6 ++----
 .../compiler/tf2tensorrt/kernels/trt_engine_op.cc    |  6 ++----
 .../tf2tensorrt/kernels/trt_engine_op_test.cc        |  6 ++----
 .../tf2tensorrt/kernels/trt_engine_resource_ops.cc   |  6 ++----
 .../kernels/trt_engine_resource_ops_test.cc          |  6 ++----
 .../tf2tensorrt/ops/get_calibration_data_op.cc       |  6 ++----
 tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc |  6 ++----
 .../tf2tensorrt/ops/trt_engine_resource_ops.cc       |  6 ++----
 .../compiler/tf2tensorrt/plugin/plugin_cast.cu.cc    |  6 ++----
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc |  6 ++----
 tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h  |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/segment.cc   |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/segment.h    |  6 ++----
 .../compiler/tf2tensorrt/segment/segment_test.cc     |  6 ++----
 tensorflow/compiler/tf2tensorrt/segment/union_find.h |  6 ++----
 tensorflow/compiler/tf2tensorrt/tensorrt_test.cc     |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_allocator.cc      | 12 ++++--------
 .../compiler/tf2tensorrt/utils/trt_allocator.h       | 12 ++++--------
 .../compiler/tf2tensorrt/utils/trt_engine_utils.cc   |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_engine_utils.h    |  6 ++----
 .../tf2tensorrt/utils/trt_int8_calibrator.cc         |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_int8_calibrator.h |  6 ++----
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc  |  6 ++----
 tensorflow/compiler/tf2tensorrt/utils/trt_logger.h   |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_lru_cache.cc      |  6 ++----
 .../compiler/tf2tensorrt/utils/trt_lru_cache.h       |  6 ++----
 .../utils/trt_shape_optimization_profiles.h          |  6 ++----
 .../utils/trt_shape_optimization_profiles_test.cc    |  6 ++----
 39 files changed, 83 insertions(+), 162 deletions(-)

diff --git a/tensorflow/compiler/tf2tensorrt/common/utils.h b/tensorflow/compiler/tf2tensorrt/common/utils.h
index 9ab0145e1ec..b428733ecd4 100644
--- a/tensorflow/compiler/tf2tensorrt/common/utils.h
+++ b/tensorflow/compiler/tf2tensorrt/common/utils.h
@@ -16,8 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 #define TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/platform/logging.h"
 
@@ -29,7 +28,6 @@ namespace tensorrt {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_COMMON_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
index 1c51d51f1c9..5429aaf3362 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc
@@ -53,8 +53,7 @@ limitations under the License.
 #include "tensorflow/core/util/device_name_utils.h"
 #include "tensorflow/tools/graph_transforms/transform_utils.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 namespace tensorflow {
@@ -884,5 +883,4 @@ Status ConvertAfterShapes(const ConversionParams& params) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
index 53ab84a6fa9..d3897e864fa 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h
@@ -24,8 +24,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -66,7 +65,6 @@ Status RegisterGraphToFunctionLibrary(const GraphDef& segment_graph_def,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
index a1f523d6bfa..54fb1d56441 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc
@@ -34,8 +34,7 @@ limitations under the License.
 #include "tensorflow/core/protobuf/config.pb.h"  // NOLINT
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -231,5 +230,4 @@ TEST_F(ConvertAfterShapesTest, DirectlyConnectedEngines) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
index 96cec556942..2ec616ba621 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
@@ -59,8 +59,7 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/strided_slice_op.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 #include "third_party/tensorrt/NvInferPlugin.h"
 
@@ -6258,5 +6257,4 @@ bool OutputEdgeValidator::operator()(const Edge* out_edge) const {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
index 7a1276c645c..a621735fad1 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h
@@ -33,8 +33,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -694,7 +693,6 @@ BinaryOperationMap();
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
index c24b169f651..53ec9ee7ada 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc
@@ -21,8 +21,7 @@ limitations under the License.
 #include 
 #include 
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include 
 #include 
@@ -6636,5 +6635,4 @@ TEST_F(OpConverterTest, ConvertPad) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
index 82e68cbb28d..07c9c2f1ea0 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.cc
@@ -12,8 +12,7 @@ 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.
 ==============================================================================*/
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 
@@ -58,5 +57,4 @@ LoggerRegistry* GetLoggerRegistry() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
index 45b302742d0..2a265cf7caa 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/logger_registry.h
@@ -19,7 +19,8 @@ limitations under the License.
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -53,5 +54,5 @@ class RegisterLogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_LOGGER_REGISTRY_H_
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
index 72f4fe5ef9b..1cf98d135cb 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc
@@ -28,8 +28,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stacktrace.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
@@ -302,5 +301,4 @@ static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
index f79048bb5f6..e0aaa5500ab 100644
--- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h
@@ -23,8 +23,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -83,6 +82,5 @@ class TRTOptimizationPass : public grappler::CustomGraphOptimizer {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
index 3143b06817e..76fb40b9520 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
@@ -22,8 +22,7 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/lib/core/refcount.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -67,5 +66,4 @@ REGISTER_KERNEL_BUILDER(Name("GetCalibrationDataOp").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
index 98d199ca9ab..1094555a622 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
@@ -48,8 +48,7 @@ limitations under the License.
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/stream_executor/lib/statusor.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -1009,5 +1008,4 @@ REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
index a06010de1c7..71193dc24cf 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
@@ -50,8 +50,7 @@ limitations under the License.
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/public/version.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -306,5 +305,4 @@ TYPED_TEST(TRTEngineOpTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
index 2c5821df6ac..3b6e7e91d3b 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
@@ -33,8 +33,7 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -251,5 +250,4 @@ REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
index 4a24160569d..6a073ee24d0 100644
--- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
@@ -48,8 +48,7 @@ limitations under the License.
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -246,5 +245,4 @@ TEST_F(TRTEngineResourceOpsTest, Basic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
index 573172b92e6..2af3164c3e2 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/get_calibration_data_op.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -34,5 +33,4 @@ Returns calibration data for the given resource name
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
index bd3c2b299a9..2527fe9b910 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -59,5 +58,4 @@ REGISTER_OP("TRTEngineOp")
     .Attr("static_engine: bool = true");
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
index 01911de66ec..3141092de03 100644
--- a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
+++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_resource_ops.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/op.h"
@@ -46,5 +45,4 @@ REGISTER_OP("SerializeTRTResource")
 
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
index 4c0d8b0392a..141a7d1f462 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/plugin_cast.cu.cc
@@ -17,8 +17,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h"
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #define EIGEN_USE_GPU  // For definition of Eigen::GpuDevice.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
@@ -234,5 +233,4 @@ REGISTER_TFTRT_PLUGIN(CastPluginCreator);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
index 563ce724f43..83d5f9b5965 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc
@@ -17,8 +17,7 @@ limitations under the License.
 
 #include 
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -30,5 +29,4 @@ const char* kTfTrtPluginNamespace = "TF";
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
index bdb046e6c71..600ac6683da 100644
--- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
+++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h
@@ -20,8 +20,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -90,7 +89,6 @@ class TrtPluginRegistrar {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
index 32e30006f58..d9080b6f69a 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc
@@ -35,8 +35,7 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/env_var.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -1062,5 +1061,4 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h
index 7295c8f0d9d..3f79983cfd2 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment.h
@@ -25,8 +25,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -67,7 +66,6 @@ Status SegmentGraph(const Graph* tf_graph,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_
diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
index 2437481a9c4..f3bc5bfbee6 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc
@@ -26,8 +26,7 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/session.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -522,5 +521,4 @@ TEST_F(SegmentTest, IncompatibleBatchSizes) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/segment/union_find.h b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
index 70e83c12fca..b53615ec019 100644
--- a/tensorflow/compiler/tf2tensorrt/segment/union_find.h
+++ b/tensorflow/compiler/tf2tensorrt/segment/union_find.h
@@ -19,8 +19,7 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/types/optional.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -217,7 +216,6 @@ UnionFind* UnionFind::FindRoot() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_
diff --git a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
index 510591bfe00..e994d20df33 100644
--- a/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc
@@ -18,8 +18,7 @@ limitations under the License.
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/test.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -164,5 +163,4 @@ TEST(TensorrtTest, BasicFunctions) {
 }  // namespace
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
index 617ea7fad5c..d4f3a524577 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc
@@ -17,11 +17,9 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -52,8 +50,7 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -113,5 +110,4 @@ void TRTDeviceAllocator::free(void* memory) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
index 4ab8b52f523..d219a8a14e8 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
@@ -20,11 +20,9 @@ limitations under the License.
 
 #include "tensorflow/core/framework/allocator.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -33,8 +31,7 @@ void* Align(uint64_t alignment, uint64_t size, void*& ptr, uint64_t& space);
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 namespace tensorflow {
 namespace tensorrt {
@@ -69,6 +66,5 @@ class TRTDeviceAllocator : public TRTBaseAllocator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
index ed997b267b1..8ccfb8b06f0 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.cc
@@ -25,8 +25,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/errors.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -257,5 +256,4 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
index a471749877a..1ea4fe28cb4 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_engine_utils.h
@@ -24,8 +24,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/status.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -91,7 +90,6 @@ Status TrtEnqueue(nvinfer1::IExecutionContext* execution_context,
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ENGINE_UTILS_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
index 554c127fa37..24271e352a7 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc
@@ -20,8 +20,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/logging.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 
 namespace tensorflow {
@@ -147,5 +146,4 @@ TRTInt8Calibrator::~TRTInt8Calibrator() {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
index 06b39716490..4c670e85f52 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
@@ -22,8 +22,7 @@ limitations under the License.
 #include 
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/tensorrt/NvInfer.h"
@@ -101,6 +100,5 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
index 193687ebc8c..e34bf5e7397 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc
@@ -15,8 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "tensorflow/compiler/tf2tensorrt/common/utils.h"
 #include "tensorflow/compiler/tf2tensorrt/convert/logger_registry.h"
 #include "tensorflow/core/platform/logging.h"
@@ -68,5 +67,4 @@ REGISTER_TENSORRT_LOGGER("DefaultLogger", Logger::GetLogger());
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_CUDA
-#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
index 2ade1b48f47..ce6552e8fe9 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h
@@ -18,8 +18,7 @@ limitations under the License.
 
 #include "tensorflow/core/platform/types.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -40,7 +39,6 @@ class Logger : public nvinfer1::ILogger {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
index fbcdaad52c0..ee7e6272372 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.cc
@@ -23,8 +23,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/platform/mutex.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 #include "third_party/tensorrt/NvInfer.h"
 
 namespace tensorflow {
@@ -141,5 +140,4 @@ EngineContext* TRTEngineCacheResource::GetEngineContext(const int profile_id) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
index 8e345254f75..991b9a949e4 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h
@@ -115,8 +115,7 @@ class LRUCache {
   }
 };
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 struct EngineContext {
   EngineContext() {}  // Creates an empty context.
@@ -223,8 +222,7 @@ class TRTEngineCacheResource : public ResourceBase {
   TrtShapeOptimizationProfile profiles_;
 };
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
index 40c7f5dcf31..fc688b14139 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.h
@@ -29,8 +29,7 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include "third_party/tensorrt/NvInfer.h"
 
@@ -173,6 +172,5 @@ class TrtShapeOptimizationProfile {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
 #endif  // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_SHAPE_OPTIMIZATION_PROFILES_H_
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
index 501810587e0..32c2200fb71 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles_test.cc
@@ -13,8 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
 
 #include 
 
@@ -214,5 +213,4 @@ TEST_F(TrtShapeOptimizationProfileTest, Dynamic) {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif  // GOOGLE_TENSORRT
-#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT

From d8e73a2f303ddebb7da3694fe7ae20a752eea01b Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Mon, 22 Jun 2020 19:16:42 -0700
Subject: [PATCH 0851/1390] Add element tracing for prefetch.

PiperOrigin-RevId: 317782986
Change-Id: Ia6e6331a59559423f5a0eac76bf2e7ffd677f876
---
 tensorflow/core/kernels/data/BUILD            |  2 ++
 .../core/kernels/data/prefetch_dataset_op.cc  | 23 +++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 6d0351202df..0972dc83ccf 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -674,6 +674,8 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/core/profiler/lib:traceme_encode",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/prefetch_dataset_op.cc b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
index 0230bcd146d..20b78ba14ad 100644
--- a/tensorflow/core/kernels/data/prefetch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/prefetch_dataset_op.cc
@@ -28,6 +28,8 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/platform/stringprintf.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
+#include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
 
 namespace tensorflow {
@@ -303,6 +305,7 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
       // The buffered data element.
       std::vector value;
       int64 created_us;
+      int64 id;
     };
 
     int64 buffer_limit() const TF_EXCLUSIVE_LOCKS_REQUIRED(*mu_) {
@@ -339,6 +342,13 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
       // (if we successfully got an element) the output values.
       Status s = buffer_.front().status;
       if (s.ok()) {
+        int64 buffer_element_id = buffer_.front().id;
+        profiler::TraceMe traceme(
+            [&] {
+              return profiler::TraceMeEncode(
+                  "PrefetchConsume", {{"element_id", buffer_element_id}});
+            },
+            profiler::kInfo);
         if (dataset()->slack_period_ > 0 &&
             (num_elements() + 1) % dataset()->slack_period_ == 0) {
           // TODO(rachelim): Consider doing something more sophisticated
@@ -423,8 +433,16 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
         mutex_lock input_l(input_mu_);
         bool end_of_sequence;
         BufferElement buffer_element;
-        buffer_element.status = input_impl_->GetNext(
-            ctx.get(), &buffer_element.value, &end_of_sequence);
+        {
+          profiler::TraceMe traceme(
+              [&] {
+                return profiler::TraceMeEncode("PrefetchProduce",
+                                               {{"element_id", num_produced}});
+              },
+              profiler::kInfo);
+          buffer_element.status = input_impl_->GetNext(
+              ctx.get(), &buffer_element.value, &end_of_sequence);
+        }
         if (buffer_element.status.ok() && end_of_sequence) {
           mutex_lock l(*mu_);
           prefetch_thread_finished_ = true;
@@ -437,6 +455,7 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
           mutex_lock l(*mu_);
           RecordBufferEnqueue(ctx.get(), buffer_element.value);
           buffer_element.created_us = EnvTime::NowMicros();
+          buffer_element.id = num_produced;
           buffer_.push_back(std::move(buffer_element));
           cond_var_->notify_all();
         }

From 90078a19ee1ce5c7bc644bac1ceff7c16d12e4bb Mon Sep 17 00:00:00 2001
From: David Rim 
Date: Mon, 22 Jun 2020 19:16:48 -0700
Subject: [PATCH 0852/1390] Update quantize weights to use hybrid per-channel
 quantization and asymmetric quantization by default

PiperOrigin-RevId: 317783008
Change-Id: I3c0f026ded030407c434bca70fe15d5fe723c744
---
 tensorflow/lite/kernels/register.cc           | 18 ++--
 .../lite/tools/optimize/quantize_weights.cc   | 92 ++++++++++++++++---
 .../tools/optimize/quantize_weights_test.cc   |  2 +
 .../lite/tools/versioning/op_version.cc       | 63 ++++++++++++-
 tensorflow/lite/tools/versioning/op_version.h |  9 ++
 .../lite/tools/versioning/op_version_test.cc  | 27 ++++++
 .../lite/tools/versioning/runtime_version.cc  | 12 +++
 7 files changed, 195 insertions(+), 28 deletions(-)

diff --git a/tensorflow/lite/kernels/register.cc b/tensorflow/lite/kernels/register.cc
index c3a4aaad16d..452ce35ec78 100644
--- a/tensorflow/lite/kernels/register.cc
+++ b/tensorflow/lite/kernels/register.cc
@@ -54,24 +54,24 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_L2_POOL_2D, Register_L2_POOL_2D());
   AddBuiltin(BuiltinOperator_CONV_2D, Register_CONV_2D(),
              /* min_version = */ 1,
-             /* max_version = */ 4);
+             /* max_version = */ 5);
   AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D, Register_DEPTHWISE_CONV_2D(),
              /* min_version = */ 1,
-             /* max_version = */ 5);
+             /* max_version = */ 6);
   AddBuiltin(BuiltinOperator_SVDF, Register_SVDF(),
              /* min_version = */ 1,
-             /* max_version = */ 3);
+             /* max_version = */ 4);
   AddBuiltin(BuiltinOperator_RNN, Register_RNN(),
              /* min_version = */ 1,
-             /* max_version = */ 2);
+             /* max_version = */ 3);
   AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN,
              Register_BIDIRECTIONAL_SEQUENCE_RNN(),
              /* min_version = */ 1,
-             /* max_version = */ 2);
+             /* max_version = */ 3);
   AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN,
              Register_UNIDIRECTIONAL_SEQUENCE_RNN(),
              /* min_version = */ 1,
-             /* max_version = */ 2);
+             /* max_version = */ 3);
   AddBuiltin(BuiltinOperator_EMBEDDING_LOOKUP, Register_EMBEDDING_LOOKUP(),
              /* min_version = */ 1,
              /* max_version = */ 3);
@@ -79,7 +79,7 @@ BuiltinOpResolver::BuiltinOpResolver() {
              Register_EMBEDDING_LOOKUP_SPARSE());
   AddBuiltin(BuiltinOperator_FULLY_CONNECTED, Register_FULLY_CONNECTED(),
              /* min_version = */ 1,
-             /* max_version = */ 8);
+             /* max_version = */ 9);
   AddBuiltin(BuiltinOperator_LSH_PROJECTION, Register_LSH_PROJECTION());
   AddBuiltin(BuiltinOperator_HASHTABLE_LOOKUP, Register_HASHTABLE_LOOKUP());
   AddBuiltin(BuiltinOperator_SOFTMAX, Register_SOFTMAX(),
@@ -105,13 +105,13 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION,
              Register_LOCAL_RESPONSE_NORMALIZATION());
   AddBuiltin(BuiltinOperator_LSTM, Register_LSTM(), /* min_version = */ 1,
-             /* max_version = */ 3);
+             /* max_version = */ 4);
   AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM,
              Register_BIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version = */ 1,
              /* max_version = */ 3);
   AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM,
              Register_UNIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version = */ 1,
-             /* max_version = */ 2);
+             /* max_version = */ 3);
   AddBuiltin(BuiltinOperator_PAD, Register_PAD(), /* min_version = */ 1,
              /* max_version = */ 2);
   AddBuiltin(BuiltinOperator_PADV2, Register_PADV2(), /* min_version = */ 1,
diff --git a/tensorflow/lite/tools/optimize/quantize_weights.cc b/tensorflow/lite/tools/optimize/quantize_weights.cc
index 7e3853c645c..8bef019a83e 100644
--- a/tensorflow/lite/tools/optimize/quantize_weights.cc
+++ b/tensorflow/lite/tools/optimize/quantize_weights.cc
@@ -43,6 +43,12 @@ typedef struct {
   int32_t op_input_idx;
 } ConsumerOpInfo;
 
+typedef struct {
+  TensorT* t;
+  bool is_per_channel;
+  int channel_dim;
+} TensorPerChannel;
+
 // The default minimum number of elements a weights array must have to be
 // quantized by this transformation.
 const int kWeightsMinNumElementsDefault = 1024;
@@ -138,6 +144,7 @@ bool IsHybridEvaluationOp(const OperatorT* op, const OperatorCodeT* op_code,
     }
   } else if (builtin_op_code == BuiltinOperator_FULLY_CONNECTED ||
              builtin_op_code == BuiltinOperator_CONV_2D ||
+             builtin_op_code == BuiltinOperator_DEPTHWISE_CONV_2D ||
              builtin_op_code == BuiltinOperator_SVDF ||
              builtin_op_code == BuiltinOperator_RNN ||
              builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
@@ -181,9 +188,10 @@ bool CheckAllOpInputsQuantized(const SubGraphT* subgraph, const OperatorT* op,
 // Inserts Tensors for each input tensor of op that should be
 // quantized into tensor_map.
 TfLiteStatus InsertQuantizableInputTensorsFromOperator(
-    const ModelT* model, const OperatorT* op, uint64_t weights_min_num_elements,
+    const ModelT* model, OperatorT* op, uint64_t weights_min_num_elements,
     const CustomOpMap& custom_op_map,
-    absl::flat_hash_map* tensor_map, int subgraph_index) {
+    absl::flat_hash_map* tensor_map,
+    int subgraph_index) {
   SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
   const OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
 
@@ -222,7 +230,50 @@ TfLiteStatus InsertQuantizableInputTensorsFromOperator(
       continue;
     }
 
-    tensor_map->insert({tensor_idx, tensor});
+    if (op_code->builtin_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
+      tensor_map->insert(
+          {tensor_idx, {tensor, /*is_per_channel=*/true, /*dim=*/3}});
+    } else if (op_code->builtin_code == BuiltinOperator_CONV_2D) {
+      tensor_map->insert(
+          {tensor_idx, {tensor, /*is_per_channel=*/true, /*dim=*/0}});
+    } else {
+      switch (op_code->builtin_code) {
+        case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM:
+          op->builtin_options.AsBidirectionalSequenceLSTMOptions()
+              ->asymmetric_quantize_inputs = true;
+          break;
+        case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN:
+          op->builtin_options.AsBidirectionalSequenceRNNOptions()
+              ->asymmetric_quantize_inputs = true;
+          break;
+        case BuiltinOperator_FULLY_CONNECTED:
+          op->builtin_options.AsFullyConnectedOptions()
+              ->asymmetric_quantize_inputs = true;
+          break;
+        case BuiltinOperator_LSTM:
+          op->builtin_options.AsLSTMOptions()->asymmetric_quantize_inputs =
+              true;
+          break;
+        case BuiltinOperator_RNN:
+          op->builtin_options.AsRNNOptions()->asymmetric_quantize_inputs = true;
+          break;
+        case BuiltinOperator_SVDF:
+          op->builtin_options.AsSVDFOptions()->asymmetric_quantize_inputs =
+              true;
+          break;
+        case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM:
+          op->builtin_options.AsUnidirectionalSequenceLSTMOptions()
+              ->asymmetric_quantize_inputs = true;
+          break;
+        case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN:
+          op->builtin_options.AsSequenceRNNOptions()
+              ->asymmetric_quantize_inputs = true;
+          break;
+        default:
+          break;
+      }
+      tensor_map->insert({tensor_idx, {tensor, /*is_per_channel=*/false}});
+    }
   }
 
   return kTfLiteOk;
@@ -275,17 +326,22 @@ void MakeTensor(const string& name, const std::vector& shape,
 void UpdateInt8OperatorVersions(ModelT* model) {
   for (int i = 0; i < model->operator_codes.size(); ++i) {
     const BuiltinOperator& op_code = model->operator_codes[i]->builtin_code;
-    if (op_code == BuiltinOperator_CONV_2D || op_code == BuiltinOperator_SVDF ||
-        op_code == BuiltinOperator_RNN ||
+    if (op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
         op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN ||
+        op_code == BuiltinOperator_EMBEDDING_LOOKUP ||
+        op_code == BuiltinOperator_RNN ||
         op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM ||
         op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
-      model->operator_codes[i]->version = 2;
-    } else if (op_code == BuiltinOperator_FULLY_CONNECTED ||
-               op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
-               op_code == BuiltinOperator_EMBEDDING_LOOKUP ||
-               op_code == BuiltinOperator_LSTM) {
       model->operator_codes[i]->version = 3;
+    } else if (op_code == BuiltinOperator_LSTM ||
+               op_code == BuiltinOperator_SVDF) {
+      model->operator_codes[i]->version = 4;
+    } else if (op_code == BuiltinOperator_CONV_2D) {
+      model->operator_codes[i]->version = 5;
+    } else if (op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
+      model->operator_codes[i]->version = 6;
+    } else if (op_code == BuiltinOperator_FULLY_CONNECTED) {
+      model->operator_codes[i]->version = 9;
     }
   }
 }
@@ -354,7 +410,7 @@ TfLiteStatus QuantizeWeightsInt8(flatbuffers::FlatBufferBuilder* builder,
        ++subgraph_index) {
     SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
 
-    absl::flat_hash_map tensor_map;
+    absl::flat_hash_map tensor_map;
     for (int i = 0; i < subgraph->operators.size(); ++i) {
       OperatorT* op = subgraph->operators[i].get();
       TF_LITE_ENSURE_STATUS(InsertQuantizableInputTensorsFromOperator(
@@ -362,16 +418,22 @@ TfLiteStatus QuantizeWeightsInt8(flatbuffers::FlatBufferBuilder* builder,
           subgraph_index));
     }
 
-    for (std::pair tensor_pair : tensor_map) {
+    for (std::pair tensor_pair : tensor_map) {
       // Quantize the tensor.
-      TF_LITE_ENSURE_STATUS(
-          utils::SymmetricQuantizeTensor(model.get(), tensor_pair.second));
+      if (tensor_pair.second.is_per_channel) {
+        TF_LITE_ENSURE_STATUS(utils::SymmetricQuantizeTensorPerChannel(
+            model.get(), tensor_pair.second.t, tensor_pair.second.channel_dim,
+            nullptr));
+      } else {
+        TF_LITE_ENSURE_STATUS(
+            utils::SymmetricQuantizeTensor(model.get(), tensor_pair.second.t));
+      }
     }
 
     // Examine the tensor consumers to determine which require dequantize ops.
     for (const auto& tensor_pair : tensor_map) {
       int32_t tensor_idx = tensor_pair.first;
-      TensorT* tensor = tensor_pair.second;
+      TensorT* tensor = tensor_pair.second.t;
       std::vector consumer_op_infos =
           GetTensorConsumers(model.get(), subgraph, tensor_idx);
       if (IsQuantizationPassThroughOps(model.get(), consumer_op_infos)) {
diff --git a/tensorflow/lite/tools/optimize/quantize_weights_test.cc b/tensorflow/lite/tools/optimize/quantize_weights_test.cc
index 76f2815ef0b..2f92a9ad71c 100644
--- a/tensorflow/lite/tools/optimize/quantize_weights_test.cc
+++ b/tensorflow/lite/tools/optimize/quantize_weights_test.cc
@@ -215,6 +215,8 @@ TEST_F(QuantizeWeightsTest, HybridConv) {
       } else if (quant_tensor->buffer() != 0) {
         EXPECT_EQ(quant_tensor->type(), TensorType_INT8)
             << quant_tensor->name()->str();
+        auto shape = GetAsVector(quant_tensor->shape());
+        EXPECT_EQ(quant_tensor->quantization()->scale()->size(), shape[0]);
       } else {
         EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
       }
diff --git a/tensorflow/lite/tools/versioning/op_version.cc b/tensorflow/lite/tools/versioning/op_version.cc
index a97b9da47f1..a339976739b 100644
--- a/tensorflow/lite/tools/versioning/op_version.cc
+++ b/tensorflow/lite/tools/versioning/op_version.cc
@@ -70,10 +70,13 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
         return 3;
       }
       // If the op is a signed int8 hybrid operation, we need to return
-      // version 2.
+      // version 2 or 5 if per channel.
       if (op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(1) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.conv_2d.is_per_channel_quantized) {
+          return 5;
+        }
         return 2;
       }
       return 1;
@@ -87,10 +90,13 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
       }
 
       // If the op is a signed int8 hybrid operation, we need to return
-      // version 4.
+      // version 4 or 6 if per-channel.
       if (op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(1) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.depthwise_conv_2d.is_per_channel_quantized) {
+          return 6;
+        }
         return 4;
       }
       // If the op has signed int8 op_sig.inputs and op_sig.outputs, its
@@ -154,6 +160,10 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
       if (op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(1) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.fully_connected.asymmetric_quantize_inputs) {
+          // This is to use the updated quantization scheme.
+          return 9;
+        }
         return 3;
       }
       // For float and uint8 fixed point kernels, if the weight is
@@ -185,6 +195,10 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
       if (op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(1) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        // This is to use the updated quantization scheme
+        if (op_sig.options.input_quantization.asymmetric_quantize_inputs) {
+          return 4;
+        }
         return 2;
       }
       return 1;
@@ -251,6 +265,9 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
           op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(2) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.lstm.asymmetric_quantize_inputs) {
+          return 4;
+        }
         return 3;
       }
       // KERNEL_BASIC was added in version 2.
@@ -265,6 +282,9 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
       if (op_sig.input_types.at(0) == TensorType_FLOAT32 &&
           op_sig.input_types.at(2) == TensorType_INT8 &&
           op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.lstm.asymmetric_quantize_inputs) {
+          return 3;
+        }
         return 2;
       }
       return 1;
@@ -450,7 +470,6 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
         return 2;
       }
       return 1;
-
     case BuiltinOperator_TANH:
     case BuiltinOperator_LOGISTIC:
       if (op_sig.input_types.at(0) == TensorType_INT16 &&
@@ -500,6 +519,19 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
       }
       return 1;
 
+    case BuiltinOperator_RNN:
+    case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN:
+    case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN:
+    case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM:
+      if (op_sig.input_types.at(1) == TensorType_INT8 &&
+          op_sig.output_types.at(0) == TensorType_FLOAT32) {
+        if (op_sig.options.input_quantization.asymmetric_quantize_inputs) {
+          return 3;
+        } else {
+          return 2;
+        }
+      }
+      return 1;
     case BuiltinOperator_ADD:
     case BuiltinOperator_PAD:
     case BuiltinOperator_PADV2:
@@ -566,6 +598,16 @@ OpSignature GetOpSignature(const OperatorCode* op_code, const Operator* op,
         op_sig.options.depthwise_conv_2d.dilation_h_factor =
             conv_option->dilation_h_factor();
       }
+      const Tensor* filter_tensor =
+          subgraph->tensors()->Get(op->inputs()->Get(1));
+      const QuantizationParameters* filter_quant =
+          filter_tensor->quantization();
+      int num_channels = filter_tensor->shape()->Get(3);
+      if (filter_quant && filter_quant->scale() &&
+          filter_quant->scale()->Length() &&
+          filter_quant->scale()->Length() == num_channels) {
+        op_sig.options.depthwise_conv_2d.is_per_channel_quantized = true;
+      }
     } break;
 
     case BuiltinOperator_FAKE_QUANT: {
@@ -584,6 +626,8 @@ OpSignature GetOpSignature(const OperatorCode* op_code, const Operator* op,
             fully_connected_option->keep_num_dims();
         op_sig.options.fully_connected.weights_format =
             fully_connected_option->weights_format();
+        op_sig.options.fully_connected.asymmetric_quantize_inputs =
+            fully_connected_option->asymmetric_quantize_inputs();
       }
 
       const Tensor* weight_tensor =
@@ -644,6 +688,18 @@ OpSignature GetOpSignature(const OperatorCode* op_code, const Operator* op,
         op_sig.options.resize.align_corners = resize_nn_option->align_corners();
       }
     } break;
+    case BuiltinOperator_CONV_2D: {
+      const Tensor* filter_tensor =
+          subgraph->tensors()->Get(op->inputs()->Get(1));
+      const QuantizationParameters* filter_quant =
+          filter_tensor->quantization();
+      int num_channels = filter_tensor->shape()->Get(0);
+      if (filter_quant && filter_quant->scale() &&
+          filter_quant->scale()->Length() &&
+          filter_quant->scale()->Length() == num_channels) {
+        op_sig.options.conv_2d.is_per_channel_quantized = true;
+      }
+    } break;
     // TODO(b/150176627): Add tests for GetOpSignature.
     case BuiltinOperator_STRIDED_SLICE:
     case BuiltinOperator_SPACE_TO_BATCH_ND:
@@ -651,7 +707,6 @@ OpSignature GetOpSignature(const OperatorCode* op_code, const Operator* op,
     case BuiltinOperator_TRANSPOSE: {
       op_sig.options.single_input_op.num_dims = GetNumDims(subgraph, op, 0);
     } break;
-
     case BuiltinOperator_SUB:
     case BuiltinOperator_DIV:
     case BuiltinOperator_MAXIMUM:
diff --git a/tensorflow/lite/tools/versioning/op_version.h b/tensorflow/lite/tools/versioning/op_version.h
index df74ffaf6dd..71362001387 100644
--- a/tensorflow/lite/tools/versioning/op_version.h
+++ b/tensorflow/lite/tools/versioning/op_version.h
@@ -30,6 +30,7 @@ typedef struct {
     struct {
       int32_t dilation_w_factor;
       int32_t dilation_h_factor;
+      bool is_per_channel_quantized;
     } depthwise_conv_2d;
     struct {
       bool narrow_range;
@@ -40,6 +41,7 @@ typedef struct {
       // TODO(b/156530611): Make this global when more ops support sparse
       // computation.
       bool sparse_weight;
+      bool asymmetric_quantize_inputs;
     } fully_connected;
     struct {
       float input1_scale;
@@ -48,6 +50,7 @@ typedef struct {
     } mul;
     struct {
       LSTMKernelType kernel_type;
+      bool asymmetric_quantize_inputs;
     } lstm;
     struct {
       bool half_pixel_centers;
@@ -60,6 +63,12 @@ typedef struct {
       int32_t num_dims;
       bool need_broadcast;
     } broadcast;
+    struct {
+      bool is_per_channel_quantized;
+    } conv_2d;
+    struct {
+      bool asymmetric_quantize_inputs;
+    } input_quantization;
   } options;
 } OpSignature;
 
diff --git a/tensorflow/lite/tools/versioning/op_version_test.cc b/tensorflow/lite/tools/versioning/op_version_test.cc
index 4017fc3bff0..e9fd857a3f5 100644
--- a/tensorflow/lite/tools/versioning/op_version_test.cc
+++ b/tensorflow/lite/tools/versioning/op_version_test.cc
@@ -361,6 +361,19 @@ TEST(OpVersionTest, VersioningFullyConnectedTest) {
   fake_op_sig.options.fully_connected = {
       false, FullyConnectedOptionsWeightsFormat_DEFAULT, true};
   EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 8);
+
+  fake_op_sig = {
+      .op = BuiltinOperator_FULLY_CONNECTED,
+      .input_types =
+          std::vector{TensorType_FLOAT32, TensorType_INT8,
+                                  TensorType_FLOAT32},
+      .output_types = std::vector{TensorType_FLOAT32},
+  };
+  fake_op_sig.options.fully_connected = {
+      false, FullyConnectedOptionsWeightsFormat_DEFAULT, false, false};
+  EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 3);
+  fake_op_sig.options.fully_connected.asymmetric_quantize_inputs = true;
+  EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 9);
 }
 
 TEST(OpVersionTest, VersioningDequantizeTest) {
@@ -412,6 +425,15 @@ TEST(OpVersionTest, VersioningConv2DTest) {
       .output_types = std::vector{TensorType_FLOAT32},
   };
   EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 2);
+
+  fake_op_sig = {
+      .op = BuiltinOperator_CONV_2D,
+      .input_types =
+          std::vector{TensorType_FLOAT32, TensorType_INT8},
+      .output_types = std::vector{TensorType_FLOAT32},
+  };
+  fake_op_sig.options.conv_2d.is_per_channel_quantized = true;
+  EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 5);
 }
 
 TEST(OpVersionTest, VersioningFloorDivOperatorTest) {
@@ -479,6 +501,8 @@ TEST(OpVersionTest, VersioningSVDFOperatorTest) {
       .output_types = std::vector{TensorType_FLOAT32},
   };
   EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 2);
+  fake_op_sig.options.input_quantization.asymmetric_quantize_inputs = true;
+  EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 4);
 
   fake_op_sig = {
       .op = BuiltinOperator_SVDF,
@@ -489,6 +513,7 @@ TEST(OpVersionTest, VersioningSVDFOperatorTest) {
   };
   EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 3);
 }
+
 TEST(OpVersionTest, VersioningDepthwiseConv2DTest) {
   OpSignature fake_op_sig = {
       .op = BuiltinOperator_DEPTHWISE_CONV_2D,
@@ -497,6 +522,8 @@ TEST(OpVersionTest, VersioningDepthwiseConv2DTest) {
       .output_types = std::vector{TensorType_FLOAT32},
   };
   EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 4);
+  fake_op_sig.options.depthwise_conv_2d.is_per_channel_quantized = true;
+  EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 6);
 
   fake_op_sig = {
       .op = BuiltinOperator_DEPTHWISE_CONV_2D,
diff --git a/tensorflow/lite/tools/versioning/runtime_version.cc b/tensorflow/lite/tools/versioning/runtime_version.cc
index 36976354685..efec5a7da18 100644
--- a/tensorflow/lite/tools/versioning/runtime_version.cc
+++ b/tensorflow/lite/tools/versioning/runtime_version.cc
@@ -63,11 +63,13 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_CONV_2D, 2}, "1.14.0"},
               {{BuiltinOperator_CONV_2D, 3}, "1.14.0"},
               {{BuiltinOperator_CONV_2D, 4}, kPendingReleaseVersion},
+              {{BuiltinOperator_CONV_2D, 5}, kPendingReleaseVersion},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 1}, "1.5.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 2}, "1.12.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 3}, "1.14.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 4}, "2.2.0"},
               {{BuiltinOperator_DEPTHWISE_CONV_2D, 5}, kPendingReleaseVersion},
+              {{BuiltinOperator_DEPTHWISE_CONV_2D, 6}, kPendingReleaseVersion},
               {{BuiltinOperator_ADD, 1}, "1.5.0"},
               {{BuiltinOperator_ADD, 2}, "1.14.0"},
               {{BuiltinOperator_ADD_N, 1}, "1.14.0"},
@@ -102,6 +104,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_FULLY_CONNECTED, 6}, "2.1.0"},
               {{BuiltinOperator_FULLY_CONNECTED, 7}, kPendingReleaseVersion},
               {{BuiltinOperator_FULLY_CONNECTED, 8}, kPendingReleaseVersion},
+              {{BuiltinOperator_FULLY_CONNECTED, 9}, kPendingReleaseVersion},
               {{BuiltinOperator_GATHER, 1}, "1.6.0"},
               {{BuiltinOperator_GATHER, 2}, "1.14.0"},
               {{BuiltinOperator_GATHER, 3}, "1.15.0"},
@@ -111,6 +114,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_SVDF, 1}, "1.5.0"},
               {{BuiltinOperator_SVDF, 2}, "1.14.0"},
               {{BuiltinOperator_SVDF, 3}, "2.2.0"},
+              {{BuiltinOperator_SVDF, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_L2_NORMALIZATION, 1}, "1.5.0"},
               {{BuiltinOperator_L2_NORMALIZATION, 2}, "1.14.0"},
               {{BuiltinOperator_L2_POOL_2D, 1}, "1.5.0"},
@@ -151,13 +155,18 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_LSTM, 1}, "1.7.0"},
               {{BuiltinOperator_LSTM, 2}, "1.10.0"},
               {{BuiltinOperator_LSTM, 3}, "1.14.0"},
+              {{BuiltinOperator_LSTM, 4}, kPendingReleaseVersion},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 1}, "1.13.1"},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 2}, "1.14.0"},
+              {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, 3},
+               kPendingReleaseVersion},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 1}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 2}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, 3}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, 1}, "1.14.0"},
               {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, 2}, "1.14.0"},
+              {{BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN, 3},
+               kPendingReleaseVersion},
               {{BuiltinOperator_MEAN, 1}, "1.6.0"},
               {{BuiltinOperator_MEAN, 2}, "1.14.0"},
               {{BuiltinOperator_SUM, 1}, "1.10.0"},
@@ -179,6 +188,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
                kPendingReleaseVersion},
               {{BuiltinOperator_RNN, 1}, "1.5.0"},
               {{BuiltinOperator_RNN, 2}, "1.14.0"},
+              {{BuiltinOperator_RNN, 3}, kPendingReleaseVersion},
               {{BuiltinOperator_SKIP_GRAM, 1}, "1.5.0"},
               {{BuiltinOperator_SQUEEZE, 1}, "1.6.0"},
               {{BuiltinOperator_SPLIT, 1}, "1.5.0"},
@@ -233,6 +243,8 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code,
               {{BuiltinOperator_UNIQUE, 1}, "1.14.0"},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, 1}, "1.14.0"},
               {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, 2}, "1.14.0"},
+              {{BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN, 3},
+               kPendingReleaseVersion},
               {{BuiltinOperator_WHERE, 1}, "1.14.0"},
               {{BuiltinOperator_DEQUANTIZE, 1}, "1.13.1"},
               {{BuiltinOperator_DEQUANTIZE, 2}, "1.14.0"},

From 27d1a803165af2cab5025ebf2387aa4785349c5a Mon Sep 17 00:00:00 2001
From: Tare Gaskin 
Date: Tue, 23 Jun 2020 02:34:18 +0000
Subject: [PATCH 0853/1390] [-Wsign-compare] warning fixes batch 3

---
 indexed_warning_files.json                    |  1 +
 tensorflow/core/data/compression_utils.cc     |  2 +-
 .../optimizers/arithmetic_optimizer.cc        | 24 +++++++++--------
 .../grappler/optimizers/constant_folding.cc   | 27 ++++++++++---------
 .../generic_layout_optimizer_transposer.cc    | 24 ++++++++++-------
 .../generic_layout_optimizer_transposer.h     | 10 ++++---
 .../optimizers/graph_optimizer_stage.cc       |  2 +-
 .../optimizers/implementation_selector.cc     |  2 +-
 .../grappler/optimizers/memory_optimizer.cc   |  2 +-
 .../optimizers/scoped_allocator_optimizer.cc  | 11 ++++----
 .../grappler/optimizers/shape_optimizer.cc    |  3 ++-
 .../kernels/data/single_threaded_executor.cc  |  4 +--
 .../kernels/initializable_lookup_table.cc     |  2 +-
 .../remote_fused_graph_execute_utils.cc       | 15 ++++++-----
 tensorflow/core/kernels/stack.cc              |  3 ++-
 ...embedding_optimization_parameters_utils.cc |  4 +--
 tensorflow/core/util/bcast.h                  |  4 +--
 17 files changed, 79 insertions(+), 61 deletions(-)
 create mode 100644 indexed_warning_files.json

diff --git a/indexed_warning_files.json b/indexed_warning_files.json
new file mode 100644
index 00000000000..cbe0560ba18
--- /dev/null
+++ b/indexed_warning_files.json
@@ -0,0 +1 @@
+{"0": "tensorflow/lite/arena_planner.cc", "1": "tensorflow/core/platform/protobuf.cc", "2": "tensorflow/core/platform/protobuf.cc", "3": "tensorflow/core/platform/default/logging.cc", "4": "tensorflow/core/platform/default/logging.cc", "5": "tensorflow/core/lib/strings/proto_serialization.cc", "6": "tensorflow/core/lib/strings/proto_serialization.cc", "7": "tensorflow/core/platform/default/stacktrace_handler.cc", "8": "tensorflow/core/platform/default/stacktrace_handler.cc", "9": "tensorflow/core/framework/cpu_allocator_impl.cc", "10": "tensorflow/core/framework/allocator_registry.cc", "11": "tensorflow/core/framework/cpu_allocator_impl.cc", "12": "tensorflow/core/framework/allocator_registry.cc", "13": "tensorflow/lite/experimental/microfrontend/lib/fft.cc", "14": "tensorflow/python/util/tf_stack.cc", "15": "external/com_github_grpc_grpc/src/core/tsi/ssl_transport_security.cc", "16": "tensorflow/core/profiler/internal/parse_annotation.cc", "17": "tensorflow/compiler/mlir/xla/ir/chlo_ops.cc", "18": "tensorflow/core/platform/status.cc", "19": "tensorflow/core/platform/file_system_helper.cc", "20": "tensorflow/core/platform/file_system.cc", "21": "tensorflow/core/platform/env.cc", "22": "tensorflow/core/lib/io/random_inputstream.cc", "23": "tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc", "24": "tensorflow/core/lib/io/inputbuffer.cc", "25": "tensorflow/core/lib/io/zlib_outputbuffer.cc", "26": "tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc", "27": "tensorflow/core/framework/tensor_shape.cc", "28": "tensorflow/compiler/mlir/xla/ir/hlo_ops.cc", "29": "tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc", "30": "tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc", "31": "tensorflow/compiler/mlir/lite/quantization/quantization_config.cc", "32": "tensorflow/core/kernels/data/prefetch_autotuner.cc", "33": "tensorflow/core/kernels/quantization_utils.cc", "34": "tensorflow/core/profiler/utils/derived_timeline.cc", "35": "tensorflow/core/profiler/utils/xplane_utils.cc", "36": "tensorflow/core/profiler/lib/profiler_session.cc", "37": "tensorflow/core/platform/s3/s3_file_system.cc", "38": "tensorflow/lite/toco/model_cmdline_flags.cc", "39": "tensorflow/lite/toco/toco_cmdline_flags.cc", "40": "tensorflow/lite/toco/toco_cmdline_flags.cc", "41": "tensorflow/compiler/xla/window_util.cc", "42": "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc", "43": "tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc", "44": "tensorflow/core/kernels/batch_kernels.cc", "45": "tensorflow/lite/toco/graph_transformations/reorder_elementwise_unary.cc", "46": "tensorflow/core/kernels/range_sampler.cc", "47": "tensorflow/core/kernels/data/experimental/sql/sqlite_query_connection.cc", "48": "tensorflow/core/grappler/utils.cc", "49": "tensorflow/core/grappler/costs/op_level_cost_estimator.cc", "50": "tensorflow/lite/toco/graph_transformations/convert_trivial_tile_to_concat.cc", "51": "tensorflow/core/grappler/utils/topological_sort.cc", "52": "tensorflow/core/grappler/utils/frame.cc", "53": "tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc", "54": "tensorflow/core/grappler/optimizers/model_pruner.cc", "55": "tensorflow/python/grappler/model_analyzer.cc", "56": "tensorflow/core/grappler/optimizers/debug_stripper.cc", "57": "tensorflow/core/grappler/utils/graph_view.cc", "58": "tensorflow/core/grappler/utils/functions.cc", "59": "tensorflow/core/grappler/costs/graph_memory.cc", "60": "tensorflow/core/grappler/optimizers/generic_layout_optimizer.cc", "61": "tensorflow/core/grappler/optimizers/function_optimizer.cc", "62": "tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc", "63": "tensorflow/core/grappler/costs/virtual_scheduler.cc", "64": "tensorflow/core/grappler/optimizers/implementation_selector.cc", "65": "tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc", "66": "tensorflow/core/grappler/optimizers/shape_optimizer.cc", "67": "tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc", "68": "tensorflow/core/grappler/optimizers/memory_optimizer.cc", "69": "tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc", "70": "tensorflow/core/grappler/costs/graph_properties.cc", "71": "tensorflow/core/grappler/optimizers/constant_folding.cc", "72": "tensorflow/core/grappler/optimizers/loop_optimizer.cc", "73": "tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc", "74": "tensorflow/core/data/service/compression_utils.cc", "75": "tensorflow/core/kernels/stack.cc", "76": "tensorflow/core/kernels/remote_fused_graph_execute_utils.cc", "77": "tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc", "78": "tensorflow/core/kernels/initializable_lookup_table.cc", "79": "tensorflow/core/kernels/data/single_threaded_executor.cc", "80": "tensorflow/core/kernels/lookup_util.cc", "81": "tensorflow/lite/toco/graph_transformations/merge_reshape_into_preceding_transpose.cc", "82": "tensorflow/lite/toco/graph_transformations/identify_nearest_upsample.cc", "83": "tensorflow/lite/toco/graph_transformations/remove_successive_transpose.cc", "84": "tensorflow/lite/toco/graph_transformations/remove_trivial_passthrough.cc", "85": "tensorflow/lite/toco/graph_transformations/propagate_fixed_sizes.cc", "86": "tensorflow/lite/toco/graph_transformations/resolve_batch_normalization.cc", "87": "tensorflow/lite/toco/graph_transformations/resolve_constant_concatenation.cc", "88": "tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc", "89": "tensorflow/lite/toco/graph_transformations/resolve_constant_unary.cc", "90": "tensorflow/lite/toco/graph_transformations/drop_fake_quant.cc", "91": "tensorflow/lite/toco/graph_transformations/unpartition_embedding_lookup.cc", "92": "tensorflow/lite/toco/graph_transformations/dequantize.cc", "93": "tensorflow/lite/toco/graph_transformations/unroll_batch_matmul.cc", "94": "tensorflow/lite/toco/graph_transformations/hardcode_min_max.cc", "95": "tensorflow/lite/toco/graph_transformations/propagate_array_data_types.cc", "96": "tensorflow/lite/toco/graph_transformations/group_bidirectional_sequence_ops.cc", "97": "tensorflow/lite/toco/graph_transformations/propagate_fake_quant_num_bits.cc", "98": "tensorflow/lite/toco/graph_transformations/resolve_constant_slice.cc", "99": "tensorflow/lite/toco/graph_transformations/resolve_constant_transpose.cc", "100": "tensorflow/lite/toco/graph_transformations/resolve_constant_pack.cc", "101": "tensorflow/lite/toco/graph_transformations/fuse_broadcast_into_following_binary.cc", "102": "tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc", "103": "tensorflow/lite/toco/graph_transformations/convert_trivial_transpose_to_reshape.cc", "104": "tensorflow/core/profiler/convert/op_metrics_to_record.cc", "105": "tensorflow/core/profiler/utils/event_span.cc", "106": "tensorflow/python/framework/python_op_gen.cc", "107": "tensorflow/python/framework/python_op_gen_internal.cc", "108": "tensorflow/compiler/tf2xla/ops/xla_ops.cc", "109": "tensorflow/compiler/mlir/tensorflow/ir/tf_device.cc", "110": "tensorflow/compiler/mlir/tensorflow/transforms/unroll_batch_matmul.cc", "111": "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.cc", "112": "tensorflow/core/kernels/data/captured_function.cc", "113": "tensorflow/core/profiler/convert/xplane_to_tf_functions.cc", "114": "tensorflow/core/common_runtime/bfc_allocator.cc", "115": "tensorflow/core/util/padding.cc", "116": "tensorflow/core/framework/op_def_util.cc", "117": "tensorflow/core/framework/node_def_util.cc", "118": "tensorflow/core/framework/shape_inference.cc", "119": "tensorflow/core/framework/common_shape_fns.cc", "120": "tensorflow/core/common_runtime/lower_case_op.cc", "121": "tensorflow/core/common_runtime/gradients.cc", "122": "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.cc", "123": "tensorflow/stream_executor/device_description.cc", "124": "tensorflow/compiler/jit/shape_inference.cc", "125": "tensorflow/compiler/mlir/tensorflow/utils/export_utils.cc", "126": "tensorflow/compiler/xla/index_util.cc", "127": "tensorflow/compiler/xla/metric_table_report.cc", "128": "tensorflow/compiler/xla/layout.cc", "129": "tensorflow/stream_executor/stream_executor_pimpl.cc", "130": "tensorflow/compiler/xla/shape_util.cc", "131": "tensorflow/compiler/xla/service/hlo_lexer.cc", "132": "tensorflow/compiler/xla/service/cpu/shape_partition.cc", "133": "tensorflow/compiler/xla/util.cc", "134": "tensorflow/compiler/xla/service/name_uniquer.cc", "135": "tensorflow/compiler/xla/shape_layout.cc", "136": "tensorflow/compiler/xla/client/sharding_builder.cc", "137": "tensorflow/compiler/xla/service/computation_layout.cc", "138": "tensorflow/compiler/mlir/xla/type_to_shape.cc", "139": "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.cc", "140": "tensorflow/compiler/xla/service/shaped_buffer.cc", "141": "tensorflow/compiler/tf2xla/sharding_util.cc", "142": "tensorflow/stream_executor/stream.cc", "143": "tensorflow/compiler/xla/layout_util.cc", "144": "tensorflow/compiler/xla/shape.cc", "145": "tensorflow/core/profiler/convert/xplane_to_memory_profile.cc", "146": "tensorflow/core/profiler/convert/op_stats_to_overview_page.cc", "147": "tensorflow/compiler/mlir/xla/transforms/legalize_tf.cc", "148": "tensorflow/compiler/xla/client/xla_computation.cc", "149": "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.cc", "150": "tensorflow/compiler/tf2xla/rearrange_function_argument.cc", "151": "tensorflow/compiler/jit/encapsulate_util.cc", "152": "tensorflow/compiler/tf2xla/tf2xla_util.cc", "153": "tensorflow/compiler/tf2xla/functionalize_while.cc", "154": "tensorflow/compiler/tf2xla/functionalize_cond.cc", "155": "tensorflow/core/kernels/boosted_trees/resources.cc", "156": "tensorflow/python/client/session_ref.cc", "157": "tensorflow/core/distributed_runtime/rpc/grpc_state.cc", "158": "tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc", "159": "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.cc", "160": "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.cc", "161": "tensorflow/core/distributed_runtime/collective_rma_distributed.cc", "162": "tensorflow/core/platform/status.cc", "163": "tensorflow/core/framework/tensor_shape.cc", "164": "tensorflow/core/profiler/internal/tfprof_timeline.cc", "165": "tensorflow/core/platform/file_system.cc", "166": "tensorflow/core/platform/file_system_helper.cc", "167": "tensorflow/core/platform/env.cc", "168": "tensorflow/core/lib/io/snappy/snappy_outputbuffer.cc", "169": "tensorflow/core/lib/io/zlib_outputbuffer.cc", "170": "tensorflow/core/lib/io/snappy/snappy_inputbuffer.cc", "171": "tensorflow/core/lib/io/random_inputstream.cc", "172": "tensorflow/core/lib/io/inputbuffer.cc", "173": "tensorflow/core/util/padding.cc", "174": "tensorflow/core/framework/op_def_util.cc", "175": "tensorflow/core/platform/s3/s3_file_system.cc", "176": "tensorflow/stream_executor/device_description.cc", "177": "tensorflow/core/framework/shape_inference.cc", "178": "tensorflow/core/framework/node_def_util.cc", "179": "tensorflow/core/common_runtime/bfc_allocator.cc", "180": "tensorflow/core/framework/common_shape_fns.cc", "181": "tensorflow/stream_executor/stream_executor_pimpl.cc", "182": "tensorflow/core/profiler/utils/xplane_utils.cc", "183": "tensorflow/core/grappler/utils.cc", "184": "tensorflow/core/grappler/costs/op_level_cost_estimator.cc", "185": "tensorflow/core/grappler/utils/symbolic_shapes.cc", "186": "tensorflow/core/grappler/utils/frame.cc", "187": "tensorflow/core/grappler/utils/topological_sort.cc", "188": "tensorflow/stream_executor/stream.cc", "189": "tensorflow/core/kernels/initializable_lookup_table.cc", "190": "tensorflow/core/grappler/optimizers/common_subgraph_elimination.cc", "191": "tensorflow/core/grappler/utils/graph_view.cc", "192": "tensorflow/core/grappler/optimizers/model_pruner.cc", "193": "tensorflow/core/kernels/lookup_util.cc", "194": "tensorflow/compiler/tf2xla/ops/xla_ops.cc", "195": "tensorflow/python/framework/python_op_gen_internal.cc", "196": "tensorflow/python/framework/python_op_gen.cc", "197": "tensorflow/core/profiler/utils/derived_timeline.cc", "198": "tensorflow/core/grappler/optimizers/debug_stripper.cc", "199": "tensorflow/core/grappler/optimizers/implementation_selector.cc", "200": "tensorflow/core/grappler/utils/functions.cc", "201": "tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc", "202": "tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc", "203": "tensorflow/core/grappler/optimizers/shape_optimizer.cc", "204": "tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc", "205": "tensorflow/core/grappler/optimizers/function_optimizer.cc", "206": "tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc", "207": "tensorflow/core/grappler/optimizers/generic_layout_optimizer.cc", "208": "tensorflow/core/grappler/costs/virtual_scheduler.cc", "209": "tensorflow/core/grappler/optimizers/loop_optimizer.cc", "210": "tensorflow/core/grappler/optimizers/constant_folding.cc", "211": "tensorflow/core/profiler/lib/profiler_session.cc", "212": "tensorflow/core/grappler/costs/graph_properties.cc", "213": "tensorflow/core/grappler/costs/graph_memory.cc", "214": "tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc", "215": "tensorflow/core/grappler/optimizers/memory_optimizer.cc", "216": "tensorflow/core/common_runtime/gradients.cc", "217": "tensorflow/core/common_runtime/lower_case_op.cc", "218": "tensorflow/core/grappler/utils/symbolic_shapes.cc", "219": "tensorflow/compiler/jit/graphcycles/graphcycles.cc", "220": "tensorflow/compiler/jit/xla_cluster_util.cc", "221": "tensorflow/core/kernels/data/experimental/snapshot_util.cc", "222": "tensorflow/core/grappler/optimizers/data/vectorization/parse_single_example_vectorizer.cc", "223": "tensorflow/core/distributed_runtime/graph_mgr.cc", "224": "tensorflow/core/grappler/optimizers/data/vectorization_utils.cc", "225": "tensorflow/core/grappler/optimizers/data/map_vectorization.cc", "226": "tensorflow/cc/framework/while_gradients.cc", "227": "tensorflow/cc/framework/gradients.cc", "228": "tensorflow/core/grappler/graph_analyzer/subgraph.cc", "229": "tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc", "230": "tensorflow/core/grappler/graph_analyzer/sig_node.cc", "231": "tensorflow/compiler/tf2xla/shape_util.cc", "232": "tensorflow/compiler/xla/service/computation_placer.cc", "233": "tensorflow/compiler/xla/client/executable_build_options.cc", "234": "tensorflow/compiler/xla/service/hlo_module_config.cc", "235": "tensorflow/core/distributed_runtime/master_session.cc", "236": "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.cc", "237": "tensorflow/compiler/tf2xla/literal_util.cc", "238": "tensorflow/compiler/xla/service/cpu/cpu_options.cc", "239": "tensorflow/compiler/xla/client/client.cc", "240": "tensorflow/compiler/xla/literal_util.cc", "241": "tensorflow/core/distributed_runtime/rpc/grpc_session.cc", "242": "tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc", "243": "tensorflow/compiler/xla/service/call_graph.cc", "244": "tensorflow/compiler/xla/service/tuple_util.cc", "245": "tensorflow/compiler/xla/service/hlo_reachability.cc", "246": "tensorflow/compiler/xla/service/hlo_cost_analysis.cc", "247": "tensorflow/compiler/xla/service/hlo_execution_profile.cc", "248": "tensorflow/compiler/xla/service/hlo_module_group.cc", "249": "tensorflow/compiler/mlir/xla/hlo_utils.cc", "250": "tensorflow/compiler/xla/service/hlo_sharding_metadata.cc", "251": "tensorflow/compiler/xla/service/map_inliner.cc", "252": "tensorflow/compiler/xla/service/flatten_call_graph.cc", "253": "tensorflow/compiler/xla/service/hlo_domain_map.cc", "254": "tensorflow/compiler/xla/service/hlo_cse.cc", "255": "tensorflow/compiler/xla/service/batchnorm_expander.cc", "256": "tensorflow/compiler/xla/service/dynamic_index_splitter.cc", "257": "tensorflow/compiler/xla/service/dfs_hlo_visitor.cc", "258": "tensorflow/compiler/xla/service/hlo_subcomputation_unification.cc", "259": "tensorflow/compiler/xla/service/slice_sinker.cc", "260": "tensorflow/compiler/xla/service/dot_decomposer.cc", "261": "tensorflow/compiler/xla/service/sort_simplifier.cc", "262": "tensorflow/compiler/xla/service/reshape_mover.cc", "263": "tensorflow/compiler/xla/service/gpu/partition_assignment.cc", "264": "tensorflow/compiler/xla/service/fusion_node_indexing_evaluation.cc", "265": "tensorflow/compiler/xla/service/dynamic_parameter_binding.cc", "266": "tensorflow/compiler/xla/literal.cc", "267": "tensorflow/compiler/xla/service/hlo_schedule.cc", "268": "tensorflow/compiler/xla/service/buffer_value.cc", "269": "tensorflow/compiler/xla/service/hlo_dce.cc", "270": "tensorflow/compiler/xla/service/hlo_module.cc", "271": "tensorflow/compiler/xla/service/hlo_sharding.cc", "272": "tensorflow/compiler/xla/service/transpose_folding.cc", "273": "tensorflow/compiler/xla/service/logical_buffer.cc", "274": "tensorflow/compiler/xla/service/call_inliner.cc", "275": "tensorflow/compiler/xla/service/hlo_buffer.cc", "276": "tensorflow/compiler/xla/service/hlo_instructions.cc", "277": "tensorflow/compiler/xla/service/hlo_value.cc", "278": "tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc", "279": "tensorflow/compiler/xla/service/tuple_simplifier.cc", "280": "tensorflow/compiler/xla/service/hlo_computation.cc", "281": "tensorflow/compiler/xla/service/logical_buffer_analysis.cc", "282": "tensorflow/compiler/xla/service/hlo_phi_graph.cc", "283": "tensorflow/compiler/xla/service/conditional_simplifier.cc", "284": "tensorflow/compiler/xla/service/hlo_query.cc", "285": "tensorflow/compiler/xla/service/channel_tracker.cc", "286": "tensorflow/compiler/xla/client/lib/constants.cc", "287": "tensorflow/compiler/xla/service/collective_ops_utils.cc", "288": "tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc", "289": "tensorflow/compiler/tf2xla/lib/util.cc", "290": "tensorflow/compiler/xla/service/shape_inference.cc", "291": "tensorflow/compiler/tf2xla/lib/broadcast.cc", "292": "tensorflow/compiler/tf2xla/lib/data_format.cc", "293": "tensorflow/compiler/xla/service/hlo_parser.cc", "294": "tensorflow/compiler/xla/service/hlo_ordering.cc", "295": "tensorflow/compiler/xla/service/hlo_instruction.cc", "296": "tensorflow/compiler/xla/service/instruction_fusion.cc", "297": "tensorflow/compiler/xla/service/tuple_points_to_analysis.cc", "298": "tensorflow/compiler/xla/service/hlo_verifier.cc", "299": "tensorflow/compiler/xla/client/lib/comparators.cc", "300": "tensorflow/compiler/xla/client/lib/arithmetic.cc", "301": "tensorflow/compiler/xla/client/lib/sorting.cc", "302": "tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc", "303": "tensorflow/compiler/xla/client/lib/loops.cc", "304": "tensorflow/compiler/tf2xla/lib/scatter.cc", "305": "tensorflow/compiler/xla/client/lib/slicing.cc", "306": "tensorflow/compiler/xla/client/lib/prng.cc", "307": "tensorflow/compiler/xla/client/xla_builder.cc", "308": "tensorflow/compiler/xla/service/hlo_graph_dumper.cc", "309": "tensorflow/compiler/xla/service/hlo_alias_analysis.cc", "310": "tensorflow/compiler/xla/service/hlo_live_range.cc", "311": "tensorflow/compiler/xla/client/lib/pooling.cc", "312": "tensorflow/compiler/xla/client/lib/matrix.cc", "313": "tensorflow/compiler/xla/client/lib/tridiagonal.cc", "314": "tensorflow/compiler/xla/service/conditional_to_select.cc", "315": "tensorflow/compiler/xla/service/batch_dot_simplification.cc", "316": "tensorflow/compiler/xla/service/rng_expander.cc", "317": "tensorflow/compiler/xla/service/cholesky_expander.cc", "318": "tensorflow/compiler/xla/client/lib/svd.cc", "319": "tensorflow/compiler/xla/service/hlo_memory_scheduler.cc", "320": "tensorflow/compiler/xla/client/lib/qr.cc", "321": "tensorflow/compiler/xla/client/lib/self_adjoint_eig.cc", "322": "tensorflow/compiler/xla/service/convolution_group_converter.cc", "323": "tensorflow/compiler/xla/client/lib/math.cc", "324": "tensorflow/compiler/xla/service/hlo_creation_utils.cc", "325": "tensorflow/compiler/xla/service/layout_assignment.cc", "326": "tensorflow/compiler/xla/service/while_loop_constant_sinking.cc", "327": "tensorflow/compiler/xla/service/scatter_expander.cc", "328": "tensorflow/compiler/xla/service/heap_simulator.cc", "329": "tensorflow/compiler/xla/service/while_util.cc", "330": "tensorflow/compiler/xla/service/dynamic_dimension_inference.cc", "331": "tensorflow/compiler/xla/service/cpu/buffer_info_util.cc", "332": "tensorflow/compiler/xla/service/hlo_proto_util.cc", "333": "tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.cc", "334": "tensorflow/compiler/xla/service/executable.cc", "335": "tensorflow/compiler/xla/service/hlo_pass_pipeline.cc", "336": "tensorflow/compiler/xla/service/triangular_solve_expander.cc", "337": "tensorflow/compiler/xla/service/compilation_cache.cc", "338": "tensorflow/compiler/xla/service/platform_util.cc", "339": "tensorflow/compiler/xla/service/copy_insertion.cc", "340": "tensorflow/compiler/xla/service/compiler.cc", "341": "tensorflow/compiler/xla/service/hlo_constant_folding.cc", "342": "tensorflow/compiler/xla/service/memory_space_assignment.cc", "343": "tensorflow/compiler/xla/service/transfer_manager.cc", "344": "tensorflow/compiler/xla/service/generic_transfer_manager.cc", "345": "tensorflow/compiler/xla/service/hlo_element_type_converter.cc", "346": "tensorflow/compiler/xla/service/execution_tracker.cc", "347": "tensorflow/compiler/xla/service/buffer_assignment.cc", "348": "tensorflow/compiler/xla/service/backend.cc", "349": "tensorflow/compiler/xla/service/op_expander_pass.cc", "350": "tensorflow/compiler/xla/service/indexed_array_analysis.cc", "351": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex128.cc", "352": "tensorflow/compiler/xla/service/tree_reduction_rewriter.cc", "353": "tensorflow/compiler/xla/service/rng_bit_generator_expander.cc", "354": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex64.cc", "355": "tensorflow/compiler/xla/service/while_loop_analysis.cc", "356": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int32.cc", "357": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint16.cc", "358": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint8.cc", "359": "tensorflow/compiler/xla/service/algebraic_simplifier.cc", "360": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int8.cc", "361": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_half.cc", "362": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int64.cc", "363": "tensorflow/compiler/xla/service/allocation_tracker.cc", "364": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bfloat16.cc", "365": "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc", "366": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_float.cc", "367": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int16.cc", "368": "tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc", "369": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint32.cc", "370": "tensorflow/compiler/xla/service/dynamic_padder.cc", "371": "tensorflow/compiler/xla/service/dump.cc", "372": "tensorflow/compiler/xla/service/while_loop_simplifier.cc", "373": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint64.cc", "374": "tensorflow/compiler/xla/service/compile_only_service.cc", "375": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bool.cc", "376": "tensorflow/compiler/xla/service/local_service.cc", "377": "tensorflow/compiler/xla/client/compile_only_client.cc", "378": "tensorflow/compiler/tf2xla/xla_expression.cc", "379": "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_double.cc", "380": "tensorflow/compiler/xla/client/client_library.cc", "381": "tensorflow/compiler/xla/client/local_client.cc", "382": "tensorflow/compiler/jit/xla_tensor.cc", "383": "tensorflow/compiler/xla/service/service.cc", "384": "tensorflow/compiler/tf2xla/const_analysis.cc", "385": "tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc", "386": "tensorflow/compiler/jit/device_util.cc", "387": "tensorflow/compiler/tf2xla/xla_context.cc", "388": "tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc", "389": "tensorflow/compiler/tf2xla/xla_op_registry.cc", "390": "tensorflow/compiler/tf2xla/xla_helpers.cc", "391": "tensorflow/compiler/tf2xla/xla_compilation_device.cc", "392": "tensorflow/compiler/tf2xla/graph_compiler.cc", "393": "tensorflow/compiler/tf2xla/xla_resource.cc", "394": "tensorflow/compiler/tf2xla/xla_op_kernel.cc", "395": "tensorflow/compiler/tf2xla/lib/random.cc", "396": "tensorflow/compiler/jit/compilability_check_util.cc", "397": "tensorflow/compiler/xla/service/hlo_evaluator.cc", "398": "tensorflow/compiler/tf2xla/graph_compiler_util.cc", "399": "tensorflow/compiler/tf2xla/kernels/if_while_utils.cc", "400": "tensorflow/compiler/tf2xla/xla_compiler.cc", "401": "tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc", "402": "tensorflow/compiler/jit/build_xla_ops_pass.cc", "403": "tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc", "404": "tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc", "405": "tensorflow/compiler/jit/extract_outside_compilation_pass.cc", "406": "tensorflow/compiler/aot/aot_only_var_handle_op.cc", "407": "tensorflow/compiler/tf2xla/tf2xla.cc", "408": "tensorflow/compiler/mlir/tensorflow/transforms/tpu_merge_variables_with_execute.cc", "409": "tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc", "410": "tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc", "411": "tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc", "412": "tensorflow/compiler/mlir/tensorflow/transforms/tpu_cluster_formation.cc", "413": "tensorflow/compiler/mlir/tensorflow/transforms/tpu_sharding_identification_pass.cc", "414": "tensorflow/compiler/mlir/tensorflow/transforms/tpu_variable_runtime_reformatting.cc", "415": "tensorflow/compiler/jit/xla_compilation_cache.cc", "416": "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc", "417": "tensorflow/compiler/mlir/tensorflow/translate/import_model.cc", "418": "tensorflow/compiler/jit/xla_device_context.cc", "419": "tensorflow/compiler/jit/xla_launch_util.cc", "420": "tensorflow/compiler/mlir/tensorflow/transforms/resource_op_lifting.cc", "421": "tensorflow/compiler/jit/xla_compile_on_demand_op.cc", "422": "tensorflow/compiler/jit/xla_device_ops.cc", "423": "tensorflow/compiler/mlir/tensorflow/transforms/einsum.cc", "424": "tensorflow/compiler/mlir/tensorflow/transforms/batchmatmul_to_einsum.cc", "425": "tensorflow/compiler/jit/xla_device.cc", "426": "tensorflow/compiler/jit/kernels/xla_ops.cc", "427": "tensorflow/compiler/jit/xla_kernel_creator_util.cc", "428": "tensorflow/compiler/xla/service/llvm_compiler.cc", "429": "tensorflow/core/distributed_runtime/eager/remote_mgr.cc", "430": "tensorflow/core/common_runtime/eager/execute.cc", "431": "tensorflow/compiler/mlir/tensorflow/transforms/promote_resources_to_args.cc", "432": "tensorflow/core/distributed_runtime/eager/eager_service_impl.cc", "433": "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.cc", "434": "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc", "435": "tensorflow/compiler/xla/service/cpu/xfeed_manager.cc", "436": "tensorflow/compiler/xla/service/gpu/target_util.cc", "437": "tensorflow/compiler/aot/embedded_protocol_buffers.cc", "438": "tensorflow/compiler/xla/service/llvm_ir/math_ops.cc", "439": "tensorflow/compiler/xla/service/llvm_ir/tuple_ops.cc", "440": "tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc", "441": "tensorflow/compiler/xla/service/llvm_ir/ir_array.cc", "442": "tensorflow/compiler/xla/service/cpu/vector_support_library.cc", "443": "tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc", "444": "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc", "445": "tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc", "446": "tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc", "447": "tensorflow/python/tfcompile_wrapper.cc", "448": "tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc", "449": "tensorflow/compiler/xla/service/cpu/cpu_runtime.cc", "450": "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc", "451": "tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc", "452": "tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc", "453": "tensorflow/compiler/xla/service/cpu/ir_function.cc", "454": "tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc", "455": "tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc", "456": "tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.cc", "457": "tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc", "458": "tensorflow/compiler/xla/service/cpu/compiler_functor.cc", "459": "tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc", "460": "tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc", "461": "tensorflow/compiler/xla/service/elemental_ir_emitter.cc", "462": "tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc", "463": "tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc", "464": "tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc", "465": "tensorflow/compiler/xla/service/cpu/cpu_executable.cc", "466": "tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc", "467": "tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc", "468": "tensorflow/core/profiler/internal/tfprof_op.cc", "469": "tensorflow/compiler/jit/xla_cpu_device.cc", "470": "tensorflow/python/lib/core/ndarray_tensor.cc", "471": "tensorflow/compiler/xla/service/cpu/ir_emitter.cc", "472": "tensorflow/core/profiler/internal/tfprof_code.cc", "473": "tensorflow/compiler/xla/service/cpu/cpu_compiler.cc", "474": "tensorflow/core/profiler/internal/tfprof_stats.cc", "475": "tensorflow/core/profiler/internal/print_model_analysis.cc", "476": "tensorflow/compiler/aot/codegen.cc", "477": "tensorflow/compiler/aot/compile.cc", "478": "tensorflow/compiler/tf2xla/mlir_tf2xla.cc", "479": "tensorflow/python/eager/pywrap_tfe_src.cc", "480": "tensorflow/compiler/mlir/lite/utils/lstm_utils.cc", "481": "tensorflow/compiler/mlir/lite/transforms/prepare_composite_functions_tf.cc", "482": "tensorflow/compiler/mlir/lite/transforms/dilated_conv.cc", "483": "tensorflow/lite/delegates/nnapi/quant_lstm_sup.cc", "484": "tensorflow/compiler/mlir/lite/transforms/dense_to_sparse.cc", "485": "tensorflow/compiler/mlir/lite/transforms/lower_static_tensor_list.cc", "486": "tensorflow/lite/delegates/nnapi/nnapi_delegate.cc", "487": "tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc", "488": "tensorflow/compiler/mlir/lite/transforms/legalize_tf.cc", "489": "tensorflow/compiler/mlir/lite/transforms/optimize.cc", "490": "tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc", "491": "tensorflow/lite/tools/verifier.cc", "492": "tensorflow/lite/tools/optimize/model_utils.cc", "493": "tensorflow/lite/tools/optimize/quantization_utils.cc", "494": "tensorflow/lite/tools/versioning/op_version.cc", "495": "tensorflow/lite/tools/versioning/runtime_version.cc", "496": "tensorflow/lite/tools/optimize/quantize_model.cc", "497": "tensorflow/lite/tools/optimize/quantize_weights.cc", "498": "tensorflow/lite/python/optimize/calibration_wrapper.cc", "499": "tensorflow/lite/toco/tflite/import.cc", "500": "tensorflow/compiler/mlir/lite/flatbuffer_import.cc", "501": "tensorflow/compiler/mlir/lite/ir/tfl_ops.cc", "502": "tensorflow/compiler/mlir/lite/flatbuffer_export.cc", "503": "tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc", "504": "tensorflow/lite/experimental/microfrontend/lib/fft.cc", "505": "tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc", "506": "tensorflow/core/platform/default/logging.cc", "507": "tensorflow/core/platform/default/logging.cc", "508": "tensorflow/core/platform/protobuf.cc", "509": "tensorflow/core/framework/allocator_registry.cc", "510": "tensorflow/core/framework/cpu_allocator_impl.cc", "511": "tensorflow/core/kernels/lookup_util.cc", "512": "tensorflow/lite/delegates/nnapi/nnapi_delegate.cc", "513": "tensorflow/core/framework/allocator_registry.cc", "514": "tensorflow/core/common_runtime/lower_case_op.cc", "515": "tensorflow/core/data/service/compression_utils.cc", "516": "tensorflow/core/kernels/data/prefetch_autotuner.cc", "517": "tensorflow/lite/tools/versioning/runtime_version.cc", "518": "tensorflow/compiler/xla/service/hlo_evaluator.cc", "519": "tensorflow/compiler/tf2xla/xla_resource.cc", "520": "tensorflow/core/kernels/stack.cc", "521": "tensorflow/core/profiler/internal/parse_annotation.cc", "522": "tensorflow/core/framework/cpu_allocator_impl.cc", "523": "tensorflow/lite/experimental/microfrontend/lib/fft.cc", "524": "tensorflow/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc", "525": "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc", "526": "tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc"}
\ No newline at end of file
diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index d132bdca8da..2ab51712580 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -116,7 +116,7 @@ Status UncompressElement(const CompressedElement& compressed,
           compressed_data.data(), compressed_data.size(), &uncompressed_size)) {
     return errors::Internal("Could not get snappy uncompressed length");
   }
-  if (uncompressed_size != total_size) {
+  if (uncompressed_size != static_cast(total_size)) {
     return errors::Internal(
         "Uncompressed size mismatch. Snappy expects ", uncompressed_size,
         " whereas the tensor metadata suggests ", total_size);
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 520346b0166..c8015a6e50c 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -598,7 +598,7 @@ class AddOpsRewriteStage : public ArithmeticNodesGroupOptimizerStage {
     std::deque add_ops;
 
     // Prepare leaf AddN nodes for inputs of equal shape
-    for (int i = 0; i < shapes.size(); ++i) {
+    for (int i = 0, iter_limit = shapes.size(); i < iter_limit; ++i) {
       const auto node_name = leaf_node_name(i);
       const auto& inputs = shape_sig_to_inputs[ShapeSignature(shapes[i])];
       add_ops.push_back(AddInputsOfSymbolicallyEqualShape(*group.root_node,
@@ -750,7 +750,7 @@ class HoistCommonFactorOutOfAggregation : public ArithmeticOptimizerStage {
         ctx().node_map->AddOutput(new_add_node->name(), new_outer_node->name());
 
         // Hoist non-shared factors up into the new AddN node.
-        for (int i = 0; i < unique_factors.size(); ++i) {
+        for (int i = 0, iter_limit = unique_factors.size(); i < iter_limit; ++i) {
           const string& unique_factor_i = unique_factors[i];
           new_add_node->set_input(i, unique_factor_i);
           ctx().node_map->AddOutput(unique_factor_i, new_add_node->name());
@@ -1190,7 +1190,7 @@ class RemoveIdentityTranspose : public ArithmeticOptimizerStage {
     if (a.size() != b.size()) {
       return false;
     }
-    for (int i = 0; i < a.size(); ++i) {
+    for (int i = 0, iter_limit = a.size(); i < iter_limit; ++i) {
       if (a[b[i]] != i) {
         return false;
       }
@@ -1199,7 +1199,7 @@ class RemoveIdentityTranspose : public ArithmeticOptimizerStage {
   }
 
   bool IsIdentityPermutation(const std::vector& perm) {
-    for (int64 i = 0; i < perm.size(); ++i) {
+    for (int64 i = 0, iter_limit = perm.size(); i < iter_limit; ++i) {
       if (i != perm[i]) {
         return false;
       }
@@ -1500,7 +1500,8 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
     for (int i = start; i < end; ++i) {
       unique_inputs.insert(node.input(i));
     }
-    return unique_inputs.size() == n;
+    int unique_input_size = unique_inputs.size(); 
+    return unique_input_size == n;
   }
 
   // Returns the length of the common unary chain of ops that can be
@@ -3248,14 +3249,15 @@ class RemoveStackSliceSameAxis : public ArithmeticOptimizerStage {
                                      slice_begin_vec.size(), ") and size (",
                                      slice_size_vec.size(), ") vectors.");
     }
+    int slice_begin_vec_size = slice_begin_vec.size();
     if (!pack_output_shape.unknown_rank() &&
-        slice_begin_vec.size() != pack_output_shape.dims()) {
+        slice_begin_vec_size != pack_output_shape.dims()) {
       return Status::OK();
     }
-    if (pack_axis >= slice_begin_vec.size()) {
+    if (pack_axis >= slice_begin_vec_size) {
       return errors::InvalidArgument(
           "Input to node ", node->name(), " had pack_axis ", pack_axis,
-          " but rank was ", slice_begin_vec.size(), ".");
+          " but rank was ", slice_begin_vec_size, ".");
     }
 
     *slice_start_value = slice_begin_vec[pack_axis];
@@ -3264,7 +3266,7 @@ class RemoveStackSliceSameAxis : public ArithmeticOptimizerStage {
       return Status::OK();
     }
 
-    for (size_t i = 0; i < slice_begin_vec.size(); ++i) {
+    for (int i = 0; i < slice_begin_vec_size; ++i) {
       if (i != pack_axis) {
         if (slice_begin_vec[i] != 0 ||
             !(slice_size_vec[i] == -1 ||
@@ -3352,7 +3354,7 @@ class RemoveStackSliceSameAxis : public ArithmeticOptimizerStage {
 
     int begin_index = -1;
     int64 begin_value = 0;
-    for (int i = 0; i < slice_begin_vec.size(); ++i) {
+    for (int i = 0, iter_limit = slice_begin_vec.size(); i < iter_limit; ++i) {
       const int64 v = slice_begin_vec[i];
       if (v != 0) {
         if (begin_index != -1) {
@@ -3366,7 +3368,7 @@ class RemoveStackSliceSameAxis : public ArithmeticOptimizerStage {
 
     int end_index = -1;
     int64 end_value = 0;
-    for (int i = 0; i < slice_end_vec.size(); ++i) {
+    for (int i = 0, iter_limit = slice_begin_vec.size(); i < iter_limit; ++i) {
       const int64 v = slice_end_vec[i];
       if (v != pack_output_shape.dim_size(i)) {
         if (end_index != -1) {
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index d912eb7857b..f42340f9d09 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -479,7 +479,7 @@ Status ConstantFolding::MaterializeShapes(const GraphProperties& properties) {
     CHECK_EQ(op, "ShapeN");
     CHECK_EQ(input.size(), output.size());
     const NodeDef* const shape_n_node = node;
-    for (int port_idx = 0; port_idx < output.size(); ++port_idx) {
+    for (int port_idx = 0, idx_limit = output.size(); port_idx < idx_limit; ++port_idx) {
       const DataType type = output[port_idx].dtype();
       CHECK(type == DT_INT32 || type == DT_INT64);
       const PartialTensorShape shape(input[port_idx].shape());
@@ -641,12 +641,12 @@ Status ConstantFolding::MaterializeBroadcastGradientArgs(
   // These extra dims could be equal to 1, in which case there is no
   // broadcasting. It could also be greater than 1, in which case there would
   // be broadcasting. Since we don't know, we'll just punt.
-  for (int i = common_dims; i < shape1.size(); ++i) {
+  for (int i = common_dims, iter_limit = shape1.size(); i < iter_limit; ++i) {
     if (shape1[i] < 0) {
       return Status::OK();
     }
   }
-  for (int i = common_dims; i < shape2.size(); ++i) {
+  for (int i = common_dims, iter_limit = shape2.size(); i < iter_limit; ++i) {
     if (shape2[i] < 0) {
       return Status::OK();
     }
@@ -1165,7 +1165,7 @@ bool IsValidConstShapeForMulConvPushDown(
   // If the const is a scalar, or it has fewer or same number of dimensions
   // than the filter and it only has single element, the optimization should
   // work.
-  if (mul_const_input_shape.dim_size() <= data_format.size() &&
+  if (mul_const_input_shape.dim_size() <= static_cast(data_format.size()) &&
       TensorShape(mul_const_input_shape).num_elements() == 1) {
     return true;
   }
@@ -1461,7 +1461,7 @@ Status ConstantFolding::FoldNode(NodeDef* node, GraphDef* output_graph,
   VLOG(2) << "Folded node: " << SummarizeNodeDef(*node);
 
   NodeDef* constant_output = nullptr;
-  for (int i = 0; i < const_nodes.size(); i++) {
+  for (int i = 0, iter_limit = const_nodes.size(); i < iter_limit; i++) {
     NodeDef* const_node = &const_nodes[i];
     VLOG(3) << "Generated constant node: " << SummarizeNodeDef(*const_node);
     if (const_node->name().empty()) {
@@ -1549,7 +1549,7 @@ Status ConstantFolding::FoldNode(NodeDef* node, GraphDef* output_graph,
                                      constant_output->name());
               *output->mutable_input(i) = AsControlDependency(*constant_output);
             }
-          } else if (port < const_nodes.size() &&
+          } else if (port < static_cast(const_nodes.size()) &&
                      !const_nodes[port].name().empty()) {
             // Replace alive outputs with the corresponding constant.
             node_map_->UpdateInput(output->name(), NodeName(output->input(i)),
@@ -2068,7 +2068,8 @@ Status ConstantFolding::RemoveShuffleOrTranspose(
         permutation.push_back(permutation_tensor.vec()(j));
       }
     }
-    if (permutation.size() != shape.dim_size()) {
+    int permutation_size = permutation.size();
+    if (permutation_size != shape.dim_size()) {
       // Number of elements in perm should be same as dim_size. Skip if not.
       return Status::OK();
     }
@@ -2245,9 +2246,10 @@ Status ConstantFolding::SimplifyStridedSlice(const GraphProperties& properties,
       // as many as expanded_ellipsis_indices.size() axes during computation.
       // We need to subtract this number from j.
       int i = j;
+      int expanded_ellipsis_indices_size = expanded_ellipsis_indices.size();
       if (ellipsis_index != -1 &&
-          j >= ellipsis_index + expanded_ellipsis_indices.size()) {
-        i = j - expanded_ellipsis_indices.size();
+          j >= ellipsis_index + expanded_ellipsis_indices_size) {
+        i = j - expanded_ellipsis_indices_size;
       }
       int b = begin.dtype() == DT_INT32 ? begin.vec()(i)
                                         : begin.vec()(i);
@@ -3479,15 +3481,16 @@ bool ConstantFolding::PartialAssocOpConstFolding(GraphDef* optimized_graph,
   }
   // Promote AccumulateNV2 with all constant inputs to AddN, since it is
   // a fake node that cannot be constant folded by itself.
-  if (const_inputs.size() == num_non_control_inputs &&
+  int const_inputs_size = const_inputs.size();
+  if (const_inputs_size == num_non_control_inputs &&
       node->op() == "AccumulateNV2") {
     node->set_op("AddN");
     node->mutable_attr()->erase("shape");
     return true;
   }
   const string new_node_name = OptimizedNodeName(
-      *node, strings::StrCat("_partial_split_", const_inputs.size()));
-  if (const_inputs.size() > 1 && const_inputs.size() < num_non_control_inputs &&
+      *node, strings::StrCat("_partial_split_", const_inputs_size));
+  if (const_inputs_size > 1 && const_inputs_size < num_non_control_inputs &&
       !node_map_->NodeExists(new_node_name)) {
     NodeDef* added_node = optimized_graph->add_node();
     *added_node = *node;
diff --git a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc
index ab7d8fcd6cf..63239082134 100644
--- a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc
+++ b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc
@@ -242,7 +242,7 @@ Status Transposer::CreateConstPermNode(TransposeContext* context,
 
   AttrValue attr_tensor;
   Tensor tensor(DT_INT32, TensorShape({4}));
-  for (int i = 0; i < permutation.size(); i++) {
+  for (int i = 0, iter_limit = permutation.size(); i < iter_limit; i++) {
     tensor.flat()(i) = permutation[i];
   }
   tensor.AsProtoTensorContent(attr_tensor.mutable_tensor());
@@ -538,10 +538,11 @@ bool Transposer::IsFaninPortDimsNIfConst(const utils::MutableNodeView& node,
     if (!tensor.FromProto(value_attr->tensor())) {
       return false;
     }
-    if (tensor.dims() != dims.size()) {
+    int dims_size = dims.size(); 
+    if (tensor.dims() != dims_size) {
       return false;
     }
-    for (int i = 0; i < dims.size(); ++i) {
+    for (int i = 0; i < dims_size; ++i) {
       if (tensor.dim_size(i) != dims[i]) {
         return false;
       }
@@ -863,12 +864,13 @@ inline bool IsValidConstPermTransposeNode(const utils::MutableNodeView& node,
   if (!GetValueAttrFromConstInputNode(node, IsTranspose, 1, &tensor)) {
     return false;
   }
-  if (tensor.NumElements() != permutation.size()) {
+  int permutation_size = permutation.size();
+  if (tensor.NumElements() != permutation_size) {
     return false;
   }
 
   const auto& tensor_data = tensor.unaligned_flat();
-  for (int i = 0; i < permutation.size(); i++) {
+  for (int i = 0; i < permutation_size; i++) {
     if (permutation[i] != tensor_data(i)) {
       return false;
     }
@@ -1229,10 +1231,11 @@ bool ReduceTransposer::KeepDims(const utils::MutableNodeView& node) {
 
 bool ReduceTransposer::IsAlongAxis(const Tensor& tensor,
                                    absl::Span axis, int rank) {
-  if (tensor.dims() != 1 || tensor.dim_size(0) != axis.size()) {
+  int axis_size = axis.size();
+  if (tensor.dims() != 1 || tensor.dim_size(0) != axis_size) {
     return false;
   }
-  for (int i = 0; i < axis.size(); ++i) {
+  for (int i = 0; i < axis_size; ++i) {
     int local_axis = tensor.flat()(i);
     if (local_axis < 0) {
       local_axis += rank;
@@ -1444,12 +1447,13 @@ bool SqueezeTransposer::IsAlongAxis(const AttrValue& attr,
                                     int rank) const {
   const auto& list = attr.list();
   // If list is empty, Squeeze op will squeeze all dimensions of size 1.
+  int axis_size = axis.size();
   if (list.i_size() == 0) {
     return true;
-  } else if (list.i_size() != axis.size()) {
+  } else if (list.i_size() != axis_size) {
     return false;
   }
-  for (int i = 0; i < axis.size(); ++i) {
+  for (int i = 0; i < axis_size; ++i) {
     int local_axis = list.i(i);
     if (local_axis < 0) {
       local_axis += rank;
@@ -1563,7 +1567,7 @@ Status StridedSliceTransposer::PermuteMask(TransposeContext* context,
     return errors::InvalidArgument("invalid mask value: ", mask_i);
   }
   int result = 0;
-  for (int i = 0; i < context->src_to_dst.size(); i++) {
+  for (int i = 0, iter_limit = context->src_to_dst.size(); i < iter_limit; i++) {
     const int final_pos = context->src_to_dst[i];
     const int position_mask = 1 << final_pos;
     const int bit_i = (mask_i & position_mask) >> final_pos;
diff --git a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.h b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.h
index b518c32d8ec..bb00e965872 100644
--- a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.h
+++ b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.h
@@ -528,11 +528,12 @@ template 
 Status PermuteSingle(absl::string_view location,
                      absl::Span permutation, T* values) {
   DCHECK(values != nullptr);
-  if (values->size() != permutation.size()) {
+  int permutation_size = permutation.size(); 
+  if (values->size() != permutation_size) {
     return Status(tensorflow::error::Code::INVALID_ARGUMENT,
                   absl::StrCat("Size of values ", values->size(),
                                " does not match size of permutation ",
-                               permutation.size(), " @ ", location));
+                               permutation_size, " @ ", location));
   }
   typedef typename T::value_type V;
   std::vector elements(values->begin(), values->end());
@@ -549,11 +550,12 @@ template 
 Status PermuteDouble(absl::string_view location,
                      absl::Span permutation, T* values) {
   DCHECK(values != nullptr);
-  if (values->size() != permutation.size() * 2) {
+  int permutation_size = permutation.size(); 
+  if (values->size() != permutation_size * 2) {
     return Status(tensorflow::error::Code::INVALID_ARGUMENT,
                   absl::StrCat("Size of values ", values->size(),
                                " does not match twice the size of permutation ",
-                               permutation.size(), " @ ", location));
+                               permutation_size, " @ ", location));
   }
   typedef typename T::value_type V;
   std::vector elements(values->begin(), values->end());
diff --git a/tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc b/tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc
index 97033a180a6..4e955db2f5a 100644
--- a/tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc
+++ b/tensorflow/core/grappler/optimizers/graph_optimizer_stage.cc
@@ -58,7 +58,7 @@ Status GetTensorProperties(const GraphOptimizerContext& ctx,
 
   const auto& output_properties =
       ctx.graph_properties->GetOutputProperties(tensor_id.node());
-  auto num_outputs = output_properties.size();
+  int num_outputs = output_properties.size();
 
   if (num_outputs == 0 || tensor_id.index() > num_outputs - 1) {
     return errors::InvalidArgument(
diff --git a/tensorflow/core/grappler/optimizers/implementation_selector.cc b/tensorflow/core/grappler/optimizers/implementation_selector.cc
index 2b0a27aaa2d..51d61cfef2e 100644
--- a/tensorflow/core/grappler/optimizers/implementation_selector.cc
+++ b/tensorflow/core/grappler/optimizers/implementation_selector.cc
@@ -130,7 +130,7 @@ string FindForwardNode(utils::MutableNodeView* backward_node) {
 void UpdateForwardIdentityNodeDtype(utils::MutableNodeView* forward_node,
                                     const DataTypeVector& dtypes) {
   const auto& fanouts_vector = forward_node->GetRegularFanouts();
-  for (int pos = 0; pos < fanouts_vector.size(); ++pos) {
+  for (int pos = 0, pos_limit = fanouts_vector.size(); pos < pos_limit; ++pos) {
     const auto& fanouts_at_pos = fanouts_vector[pos];
     for (const auto& fanout : fanouts_at_pos) {
       if ("Identity" == fanout.node_view()->GetOp()) {
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
index 867433dcff5..c6fd2b48ac3 100644
--- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
@@ -468,7 +468,7 @@ void RecomputationRewritingPass(RewriterConfig::MemOptType optimization_level,
         // with "gradients/" or contains "/gradients/".
         return absl::StartsWith(node.name(),
                                 recomputation_targets_name_scope) ||
-               node.name().find("/" + recomputation_targets_name_scope) != -1;
+               static_cast(node.name().find("/" + recomputation_targets_name_scope)) != -1;
       };
 
   if (optimization_level == RewriterConfig::RECOMPUTATION_HEURISTICS ||
diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
index 464a2c17197..7a0079d2b4c 100644
--- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc
@@ -73,7 +73,8 @@ bool HasOpName(const string& node_name, const string& op_name) {
 Status GetOutputDataType(
     const std::vector& output_props, int output_index,
     DataType* dtype) {
-  if (output_index >= output_props.size()) {
+  int output_props_size = output_props.size();
+  if (output_index >= output_props_size) {
     return errors::Internal("Invalid output index ", output_index,
                             " size of output_props ", output_props.size());
   }
@@ -520,7 +521,7 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
 
     // Add control edges from the ScopedAllocatorOp to all of the
     // input nodes and mark them for allocation from backing tensor.
-    for (int i = 0; i < inputs.size(); ++i) {
+    for (int i = 0, iter_limit = inputs.size(); i < iter_limit; ++i) {
       auto& nd = inputs[i];
       if (IsArg(*nd.from_node_def)) {
         return errors::Internal(
@@ -547,7 +548,7 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
       std::vector inputs_to_first;
       LOG_WARNING_AND_RETURN_IF_ERROR(GetDataInputs(
           graph, sa_opti->node_map(), nd.from_node_def, &inputs_to_first));
-      for (int i = 0; i < inputs_to_first.size(); ++i) {
+      for (int i = 0, iter_limit = inputs_to_first.size(); i < iter_limit; ++i) {
         if (fanout.find(inputs_to_first[i].from_node_def) != fanout.end()) {
           VLOG(2) << "Found node " << inputs_to_first[i].from_node_def->name()
                   << " in the fanout of " << sa_name;
@@ -587,7 +588,7 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
     VLOG(2) << "BuildSAConcatNode " << sac_name;
     // control input: edge name -> source node name
     absl::flat_hash_map sac_ctl_inputs;
-    for (int i = 0; i < ops.size(); ++i) {
+    for (int i = 0, iter_limit = ops.size(); i < iter_limit; ++i) {
       NodeDef* old_op = ops[i];
       for (const string& old_op_input : old_op->input()) {
         int position = 0;
@@ -708,7 +709,7 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter {
                         const std::set& op_instance_names,
                         const string& op_name, const string& sas_name) {
     VLOG(2) << "RewireSubgraph";
-    for (int op_idx = 0; op_idx < ops.size(); ++op_idx) {
+    for (int op_idx = 0, idx_limit = ops.size(); op_idx < idx_limit; ++op_idx) {
       NodeDef* old_op = ops[op_idx];
       // Copy the output node set since we'll be modifying the version
       // maintained by NodeMap in the loop.
diff --git a/tensorflow/core/grappler/optimizers/shape_optimizer.cc b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
index 69de1cde4ca..656c1a1db1c 100644
--- a/tensorflow/core/grappler/optimizers/shape_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/shape_optimizer.cc
@@ -99,7 +99,8 @@ Status ShapeOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
         }
         const auto& prop =
             properties.GetOutputProperties(reduce_indices.node->name());
-        if (prop.size() <= reduce_indices.port_id) {
+        int prop_size = prop.size();
+        if (prop_size <= reduce_indices.port_id) {
           continue;
         }
         const TensorShapeProto& reduction_indices_shape =
diff --git a/tensorflow/core/kernels/data/single_threaded_executor.cc b/tensorflow/core/kernels/data/single_threaded_executor.cc
index 3a16f1018dd..eeb1ffd5ad0 100644
--- a/tensorflow/core/kernels/data/single_threaded_executor.cc
+++ b/tensorflow/core/kernels/data/single_threaded_executor.cc
@@ -51,8 +51,8 @@ class SingleThreadedExecutorImpl : public Executor {
     std::vector ordered_nodes;
     ordered_nodes.reserve(graph.num_nodes());
     GetReversePostOrder(graph, &ordered_nodes);
-
-    if (ordered_nodes.size() != graph.num_nodes()) {
+    int ordered_nodes_size = ordered_nodes.size();
+    if (ordered_nodes_size != graph.num_nodes()) {
       return errors::InvalidArgument("Graph had ", graph.num_nodes(),
                                      " but reverse post-order had ",
                                      ordered_nodes.size());
diff --git a/tensorflow/core/kernels/initializable_lookup_table.cc b/tensorflow/core/kernels/initializable_lookup_table.cc
index 196c2fe95a3..48041526022 100644
--- a/tensorflow/core/kernels/initializable_lookup_table.cc
+++ b/tensorflow/core/kernels/initializable_lookup_table.cc
@@ -74,7 +74,7 @@ Status InitializableLookupTable::Initialize(InitTableIterator& iter) {
 
 Status InitializableLookupTable::AreEntriesSame(const InitTableIterator& iter,
                                                 bool* result) {
-  *result = iter.total_size() == size();
+  *result = static_cast(iter.total_size()) == size();
   return Status::OK();
 }
 
diff --git a/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc b/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
index b57d6163cac..e3f220b9ff4 100644
--- a/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
+++ b/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
@@ -431,7 +431,8 @@ RemoteFusedGraphExecuteUtils::AddOutputTensorShapeTypeByTensorShapeMap(
   if (data_types.empty()) {
     return false;
   }
-  CHECK(data_types.size() > port);
+  int data_types_size = data_types.size();
+  CHECK(data_types_size > port);
   *data_type = data_types.at(port);
   *shape = shapes.at(port);
   return true;
@@ -788,7 +789,8 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
           ++input_count;
         }
       }
-      CHECK(input_count == 0 || input_count == node->in_edges().size())
+      int node_in_edges_size = node->in_edges().size();
+      CHECK(input_count == 0 || input_count == node_in_edges_size)
           << "Invalid input_count(" << input_count << ", "
           << node->in_edges().size() << ") " << node_name;
 
@@ -968,10 +970,10 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
       border_inputs, border_outputs, require_shape_type, &graph, &fused_node));
 
   for (const Node* node : graph.nodes()) {
-    for (int i = 0; i < node->num_inputs(); ++i) {
+    for (int i = 0, iter_limit = node->num_inputs(); i < iter_limit; ++i) {
       const Edge* edge = nullptr;
       TF_RETURN_IF_ERROR(node->input_edge(i, &edge));
-      for (int j = 0; j < border_outputs.size(); ++j) {
+      for (int j = 0, second_iter_limit = border_outputs.size(); j < second_iter_limit; ++j) {
         const string& output = border_outputs.at(j);
         const TensorId tid = ParseTensorName(output);
         const string output_name(tid.first);
@@ -1333,8 +1335,9 @@ RemoteFusedGraphExecuteUtils::FuseRemoteGraphByPlacedArguments(
 
 /* static */ Status RemoteFusedGraphExecuteUtils::CopyByteArrayToTensor(
     const void* src_ptr, const int src_size, Tensor* tensor) {
-  CHECK(tensor->TotalBytes() >= src_size)
-      << tensor->TotalBytes() << ", " << src_size;
+  int tensor_TotalBytes = tensor->TotalBytes();
+  CHECK(tensor_TotalBytes >= src_size)
+      << tensor_TotalBytes << ", " << src_size;
   void* dst_ptr;
   switch (tensor->dtype()) {
     case DT_FLOAT:
diff --git a/tensorflow/core/kernels/stack.cc b/tensorflow/core/kernels/stack.cc
index a30729c3b2f..dd20902a26d 100644
--- a/tensorflow/core/kernels/stack.cc
+++ b/tensorflow/core/kernels/stack.cc
@@ -57,7 +57,8 @@ class Stack : public ResourceBase {
   Status Push(const TensorAndAllocation& value) {
     mutex_lock l(mu_);
     TF_RETURN_IF_ERROR(CheckNotClosed());
-    if (max_size_ >= 0 && stack_.size() >= max_size_) {
+    int stack_size = stack_.size(); 
+    if (max_size_ >= 0 && stack_size >= max_size_) {
       return errors::InvalidArgument("Stack[", stack_name_, "] overflowed ",
                                      "its max_size (", max_size_, ")");
     }
diff --git a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
index 9daf9a2cef7..c40fea7c61c 100644
--- a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
+++ b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
@@ -336,7 +336,7 @@ Status LoadOpShapeFunction::operator()(
                     });
   std::vector inputs(user_param_count);
   int input_index = 0;
-  for (int i = 0; i < state_variable_specs.size(); ++i) {
+  for (int i = 0, iter_limit = state_variable_specs.size(); i < iter_limit; ++i) {
     if (state_variable_specs[i].has_user_defined() || is_debug_op_) {
       std::vector input_temp;
       TF_RETURN_IF_ERROR(c->input(state_variable_specs[i].name(), &input_temp));
@@ -388,7 +388,7 @@ Status RetrieveOpShapeFunction::operator()(
   TF_RETURN_IF_ERROR(c->GetAttr("num_shards", &num_shards));
   int shard_id;
   TF_RETURN_IF_ERROR(c->GetAttr("shard_id", &shard_id));
-  for (int j = 0; j < state_variable_specs.size(); ++j) {
+  for (int j = 0, iter_limit = state_variable_specs.size(); j < iter_limit; ++j) {
     if (state_variable_specs[j].has_user_defined() || is_debug_op_) {
       auto shape = c->MakeShape(
           std::vector(2, c->UnknownDim()));
diff --git a/tensorflow/core/util/bcast.h b/tensorflow/core/util/bcast.h
index 7bb8ea18ad3..a79a3f08622 100644
--- a/tensorflow/core/util/bcast.h
+++ b/tensorflow/core/util/bcast.h
@@ -139,7 +139,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N],
     if (x[i] != x[0]) {
       all_equal = false;
     }
-    if (x[i].size() > largest_rank) {
+    if ( static_cast(x[i].size()) > largest_rank) {
       largest_rank = x[i].size();
     }
   }
@@ -176,7 +176,7 @@ BCastList::BCastList(const BCastList::Vec (&x)[N],
 
   // 1-extend and align all vectors.
   for (int i = 0; i < N; ++i) {
-    if (copy[i].size() < largest_rank) {
+    if (static_cast(copy[i].size()) < largest_rank) {
       copy[i].resize(largest_rank, 1);
     }
   }

From 2399d37889c044fb6fb52fb3abc2bb5d41e64a28 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac 
Date: Mon, 22 Jun 2020 19:18:54 -0700
Subject: [PATCH 0854/1390] Create `RELEASE.md` stub.

Now, every PR and commit can add to the release notes before the branch is cut for a new release.

PiperOrigin-RevId: 317783211
Change-Id: If65da4b956f83940c7bd539d40f915cfd9af4db0
---
 RELEASE.md | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/RELEASE.md b/RELEASE.md
index f93626cc876..68d9399676a 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,3 +1,57 @@
+# Release 2.4.0
+
+
+
+## Breaking Changes
+
+* 
+* 
+
+## Known Caveats
+
+* 
+
+## Major Features and Improvements
+
+* 
+* 
+
+## Bug Fixes and Other Changes
+
+* 
+* 
+* 
+* TF Core:
+    * 
+* `tf.data`:
+    * 
+*   `tf.distribute`:
+    * 
+*   `tf.keras`:
+    * 
+*   `tf.function`/AutoGraph:
+    * 
+*   `tf.lite`:
+    * 
+*   `tf.random`:
+    * 
+*   Math and Linear Algebra:
+    * 
+*   TPU Enhancements:
+    * 
+*   XLA Support:
+    * 
+*   Tracing and Debugging:
+    * 
+*   Other:
+    * 
+
+## Thanks to our Contributors
+
+This release contains contributions from many people at Google, as well as:
+
+, , , , , 
+
 # Release 2.3.0
 
 ## Breaking Changes

From 821d1b087be221aa3c58be8e077d4553d6851ef2 Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Mon, 22 Jun 2020 19:30:34 -0700
Subject: [PATCH 0855/1390] Add element tracing for parallel_map to track
 exactly which upstream event produces an element for each parallel_map's
 GetNext call.

PiperOrigin-RevId: 317784362
Change-Id: I981692f8d152ab8b6c6c343d01563800a8f94600
---
 tensorflow/core/kernels/data/BUILD            |  2 ++
 .../kernels/data/parallel_map_dataset_op.cc   | 19 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 0972dc83ccf..74283b63b67 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -470,6 +470,8 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/core/profiler/lib:traceme_encode",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index bae90549841..1dc27ce6635 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -29,6 +29,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/platform/stringprintf.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
+#include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
 
 namespace tensorflow {
@@ -241,6 +243,10 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
       RecordStop(ctx);
       result->notification.WaitForNotification();
       RecordStart(ctx);
+      profiler::TraceMe traceme([&] {
+        return profiler::TraceMeEncode("ParallelMapConsume",
+                                       {{"element_id", result->id}});
+      });
       return ProcessResult(ctx, result, out_tensors, end_of_sequence);
     }
 
@@ -358,10 +364,14 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
 
    private:
     struct InvocationResult {
+      InvocationResult() {}
+      explicit InvocationResult(int64 id) : id(id) {}
+
       Notification notification;
       Status status;
       std::vector return_values;
       bool end_of_input;
+      int64 id;
     };
 
     void CancelThreads(bool wait) TF_LOCKS_EXCLUDED(mu_) {
@@ -402,6 +412,10 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
     void CallFunction(const std::shared_ptr& ctx,
                       const std::shared_ptr& result)
         TF_LOCKS_EXCLUDED(*mu_) {
+      profiler::TraceMe traceme([&] {
+        return profiler::TraceMeEncode("ParallelMapProduce",
+                                       {{"element_id", result->id}});
+      });
       // Get the next input element.
       std::vector input_element;
       result->status = input_impl_->GetNext(ctx.get(), &input_element,
@@ -490,6 +504,8 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
         return num_calls_ >= num_parallel_calls ||
                invocation_results_.size() >= num_parallel_calls;
       };
+      // Counts the total number of calls to use as an id of InvocationResult.
+      int64 num_total_calls = 0;
       while (true) {
         {
           mutex_lock l(*mu_);
@@ -502,7 +518,8 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
             return;
           }
           while (!busy()) {
-            invocation_results_.push_back(std::make_shared());
+            invocation_results_.push_back(
+                std::make_shared(num_total_calls++));
             new_calls.push_back(invocation_results_.back());
             num_calls_++;
           }

From 01fdbb866b4350acd74370d0a4b182feb4a056d8 Mon Sep 17 00:00:00 2001
From: rahul-kamat 
Date: Tue, 23 Jun 2020 02:46:05 +0000
Subject: [PATCH 0856/1390] Remove imports from python_op_gen_test

---
 tensorflow/python/framework/python_op_gen_test.cc | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tensorflow/python/framework/python_op_gen_test.cc b/tensorflow/python/framework/python_op_gen_test.cc
index 2561195c407..cf6566ea7ae 100644
--- a/tensorflow/python/framework/python_op_gen_test.cc
+++ b/tensorflow/python/framework/python_op_gen_test.cc
@@ -20,9 +20,6 @@ limitations under the License.
 #include "tensorflow/core/framework/op_gen_lib.h"
 #include "tensorflow/core/platform/test.h"
 
-#include "tensorflow/core/lib/io/path.h"
-#include "tensorflow/core/lib/strings/str_util.h"
-
 namespace tensorflow {
 namespace {
 

From b16994e2d646df982c2c8614a513742ac06c3cf0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Mon, 22 Jun 2020 20:27:46 -0700
Subject: [PATCH 0857/1390] Add element tracing for parallel_map to track
 exactly which upstream event produces an element for each parallel_map's
 GetNext call.

PiperOrigin-RevId: 317789815
Change-Id: Ic8a50a3a311051069ed25973b33cd2b0b9bbece5
---
 tensorflow/core/kernels/data/BUILD            |  2 --
 .../kernels/data/parallel_map_dataset_op.cc   | 19 +------------------
 2 files changed, 1 insertion(+), 20 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 74283b63b67..0972dc83ccf 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -470,8 +470,6 @@ tf_kernel_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
-        "//tensorflow/core/profiler/lib:traceme",
-        "//tensorflow/core/profiler/lib:traceme_encode",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index 1dc27ce6635..bae90549841 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -29,8 +29,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/platform/stringprintf.h"
-#include "tensorflow/core/profiler/lib/traceme.h"
-#include "tensorflow/core/profiler/lib/traceme_encode.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
 
 namespace tensorflow {
@@ -243,10 +241,6 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
       RecordStop(ctx);
       result->notification.WaitForNotification();
       RecordStart(ctx);
-      profiler::TraceMe traceme([&] {
-        return profiler::TraceMeEncode("ParallelMapConsume",
-                                       {{"element_id", result->id}});
-      });
       return ProcessResult(ctx, result, out_tensors, end_of_sequence);
     }
 
@@ -364,14 +358,10 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
 
    private:
     struct InvocationResult {
-      InvocationResult() {}
-      explicit InvocationResult(int64 id) : id(id) {}
-
       Notification notification;
       Status status;
       std::vector return_values;
       bool end_of_input;
-      int64 id;
     };
 
     void CancelThreads(bool wait) TF_LOCKS_EXCLUDED(mu_) {
@@ -412,10 +402,6 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
     void CallFunction(const std::shared_ptr& ctx,
                       const std::shared_ptr& result)
         TF_LOCKS_EXCLUDED(*mu_) {
-      profiler::TraceMe traceme([&] {
-        return profiler::TraceMeEncode("ParallelMapProduce",
-                                       {{"element_id", result->id}});
-      });
       // Get the next input element.
       std::vector input_element;
       result->status = input_impl_->GetNext(ctx.get(), &input_element,
@@ -504,8 +490,6 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
         return num_calls_ >= num_parallel_calls ||
                invocation_results_.size() >= num_parallel_calls;
       };
-      // Counts the total number of calls to use as an id of InvocationResult.
-      int64 num_total_calls = 0;
       while (true) {
         {
           mutex_lock l(*mu_);
@@ -518,8 +502,7 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
             return;
           }
           while (!busy()) {
-            invocation_results_.push_back(
-                std::make_shared(num_total_calls++));
+            invocation_results_.push_back(std::make_shared());
             new_calls.push_back(invocation_results_.back());
             num_calls_++;
           }

From 3ed1e3029e68ca8cb6306c8f31182306741dcf0c Mon Sep 17 00:00:00 2001
From: Nick Kreeger 
Date: Mon, 22 Jun 2020 20:37:38 -0700
Subject: [PATCH 0858/1390] Remove static_assert for type checking in
 FlatBufferVectorToTfLiteTypeArray.

It turns out that std::is_same() has dropped the non-string argument in c++17. This breaks internal users that are building against qualcomm.

PiperOrigin-RevId: 317790812
Change-Id: If56a61d20426670251b55f370a6b5fa886a49e21
---
 tensorflow/lite/micro/micro_allocator.cc | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc
index bf9e38d1050..239a23335a6 100644
--- a/tensorflow/lite/micro/micro_allocator.cc
+++ b/tensorflow/lite/micro/micro_allocator.cc
@@ -401,12 +401,9 @@ TfLiteStatus FlatBufferVectorToTfLiteTypeArray(
     kTfLiteArrayType** result) {
   TFLITE_DCHECK(error_reporter != nullptr);
   TFLITE_DCHECK(flatbuffer_array != nullptr);
-  // Only two conversions are supported - float and int32 - ensure that these
-  // match at compile time instead of duplicating functions here:
-  static_assert((std::is_same() &&
-                 std::is_same()) ||
-                (std::is_same() &&
-                 std::is_same()));
+  // TODO(b/159668691): Consider adding type assertion or breaking this function
+  // into multiple functions for each type. std::is_same is c++11 and has a
+  // special updated constructor in c++17 that requires a string argument.
   if (FLATBUFFERS_LITTLEENDIAN) {
     // On little-endian machines, TfLite*Array happens to have the same memory
     // layout as flatbuffers:Vector, so we can

From 0453c02f43b64e54caf1a3f001f0c01def38ba37 Mon Sep 17 00:00:00 2001
From: Andy Ly 
Date: Mon, 22 Jun 2020 21:36:57 -0700
Subject: [PATCH 0859/1390] Update tf.DeviceIndex description in TensorFlow
 MLIR ODS to match description in TensorFlow op registry (NFC).

PiperOrigin-RevId: 317796799
Change-Id: Ic4d6ca197cce7c56e491f59814cf913dcd64713c
---
 tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
index 7f250392cb2..65ca3ea4dbd 100644
--- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
+++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td
@@ -2378,6 +2378,13 @@ horizontal and vertices strides, `strides = [1, stride, stride, 1]`.
 def TF_DeviceIndexOp : TF_Op<"DeviceIndex", [NoSideEffect]> {
   let summary = "Return the index of device the op runs.";
 
+  let description = [{
+Given a list of device names, this operation returns the index of the device
+this op runs. The length of the list is returned in two cases:
+(1) Device does not exist in the given device list.
+(2) It is in XLA compilation.
+  }];
+
   let arguments = (ins
     StrArrayAttr:$device_names
   );

From 79fc3fd8fcc4b2936ca630cec6898a83ee92f59b Mon Sep 17 00:00:00 2001
From: Rick Chao 
Date: Mon, 22 Jun 2020 21:41:24 -0700
Subject: [PATCH 0860/1390] Skip data loading error in
 multi_worker_tutorial_test (the test does not aim to cover this).

PiperOrigin-RevId: 317797271
Change-Id: I8336d7ffeda0836beef0a2d04e633614a44e7fa4
---
 .../python/keras/distribute/multi_worker_tutorial_test.py      | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/python/keras/distribute/multi_worker_tutorial_test.py b/tensorflow/python/keras/distribute/multi_worker_tutorial_test.py
index 3f9ab18f89c..aafa00bcc9f 100644
--- a/tensorflow/python/keras/distribute/multi_worker_tutorial_test.py
+++ b/tensorflow/python/keras/distribute/multi_worker_tutorial_test.py
@@ -19,6 +19,7 @@ from __future__ import print_function
 import contextlib
 import os
 import re
+import zipfile
 from absl.testing import parameterized
 import numpy as np
 from tensorflow.python import keras
@@ -43,6 +44,8 @@ class MultiWorkerTutorialTest(parameterized.TestCase, test.TestCase):
   def skip_fetch_failure_exception(self):
     try:
       yield
+    except zipfile.BadZipfile as e:
+      self.skipTest('Data loading error: Bad magic number for file header.')
     except Exception as e:  # pylint: disable=broad-except
       if 'URL fetch failure' in str(e):
         self.skipTest('URL fetch error not considered failure of the test.')

From e3904ac3aef005326e126e8240eaf5bf23cb3b04 Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Mon, 22 Jun 2020 22:01:28 -0700
Subject: [PATCH 0861/1390] Internal TPU library change.

PiperOrigin-RevId: 317799072
Change-Id: Ia9b4fba3aa8292e5b089d5741f359eef9763df09
---
 .../core/tpu/kernels/tpu_compile_c_api.h      |  2 +-
 tensorflow/core/tpu/kernels/tpu_op_util.cc    | 37 ++++++++++++++-----
 tensorflow/core/tpu/kernels/tpu_op_util.h     |  7 ++++
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
index c101e489d56..eab53fe9da4 100644
--- a/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
+++ b/tensorflow/core/tpu/kernels/tpu_compile_c_api.h
@@ -35,7 +35,7 @@ struct CompilationCacheKeyProperty {
   const char* mlir_module;
   const int32_t* device_ids;
   size_t device_ids_size;
-  int32_t guaranteed_constants_size;
+  size_t guaranteed_constants_size;
   uint64_t function_library_fingerprint;
   int32_t num_cores_per_replica;
   int32_t num_replicas;
diff --git a/tensorflow/core/tpu/kernels/tpu_op_util.cc b/tensorflow/core/tpu/kernels/tpu_op_util.cc
index e2f717fea8b..31b0cc6c72d 100644
--- a/tensorflow/core/tpu/kernels/tpu_op_util.cc
+++ b/tensorflow/core/tpu/kernels/tpu_op_util.cc
@@ -24,12 +24,13 @@ namespace tpu {
 namespace {
 // Return fingerprint_in_metadata if it's not empty; otherwise read input tensor
 // data to compute the fingerprint.
-std::string GuaranteedConstFingerprint(
-    const string& fingerprint_in_metadata,
-    const OpInputList& guaranteed_constants) {
+std::string GuaranteedConstFingerprint(const string& fingerprint_in_metadata,
+                                       const Tensor* guaranteed_constants,
+                                       size_t guaranteed_constants_size) {
   if (fingerprint_in_metadata.empty()) {
     uint64_t fingerprint = 0;
-    for (const auto& constant : guaranteed_constants) {
+    for (size_t i = 0; i < guaranteed_constants_size; ++i) {
+      const Tensor& constant = guaranteed_constants[i];
       fingerprint = TpuCompile_CreateGuaranteedConstFingerprint(
           fingerprint, constant.tensor_data().data(),
           constant.tensor_data().size());
@@ -86,9 +87,12 @@ std::string CreateConfigPrefix(const TPUCompileMetadataProto& metadata) {
 }
 }  // namespace
 
+// The `guaranteed_constants` must be passed as reference due to the lazy
+// evaluation of `guaranteed_const_fingerprint()` callback.
 TpuCompilationCacheKey CreateCompilationCacheKey(
     absl::string_view function_name, uint64 function_library_fingerprint,
-    absl::string_view mlir_module, const OpInputList& guaranteed_constants,
+    absl::string_view mlir_module, const Tensor* guaranteed_constants,
+    size_t guaranteed_constants_size,
     const std::vector& dynamic_shapes,
     const TPUCompileMetadataProto& metadata,
     const TpuMeshStateInterface& mesh_state) {
@@ -114,7 +118,7 @@ TpuCompilationCacheKey CreateCompilationCacheKey(
           mlir_module.data(),
           flattened_device_ids.data(),
           flattened_device_ids.size(),
-          guaranteed_constants.size(),
+          guaranteed_constants_size,
           function_library_fingerprint,
           metadata.num_cores_per_replica(),
           metadata.num_replicas(),
@@ -128,7 +132,7 @@ TpuCompilationCacheKey CreateCompilationCacheKey(
 
   // Guaranteed constants can be different across sessions. Use session_handle
   // and guaranteed_const fingerprint to guarantee no collision.
-  if (guaranteed_constants.size() > 0) {
+  if (guaranteed_constants != nullptr && guaranteed_constants_size > 0) {
     key.has_guaranteed_const = true;
     key.session_handle = metadata.session_handle();
     // Both `metadata` and `guaranteed_constants` lifetime are captured by
@@ -136,16 +140,31 @@ TpuCompilationCacheKey CreateCompilationCacheKey(
     // managed through the `TPUCompileOpKernelImpl` that outlives the
     // lifetime of the compilation cache lookups.
     string fingerprint;
-    key.guaranteed_const_fingerprint = [&metadata, &guaranteed_constants,
+    key.guaranteed_const_fingerprint = [&metadata, guaranteed_constants,
+                                        guaranteed_constants_size,
                                         fingerprint]() mutable {
       if (fingerprint.empty()) {
         fingerprint = GuaranteedConstFingerprint(
-            metadata.guaranteed_const_fingerprint(), guaranteed_constants);
+            metadata.guaranteed_const_fingerprint(), guaranteed_constants,
+            guaranteed_constants_size);
       }
       return fingerprint;
     };
   }
   return key;
 }
+
+TpuCompilationCacheKey CreateCompilationCacheKey(
+    absl::string_view function_name, uint64 function_library_fingerprint,
+    absl::string_view mlir_module, const OpInputList& guaranteed_constants,
+    size_t guaranteed_constants_size,
+    const std::vector& dynamic_shapes,
+    const TPUCompileMetadataProto& metadata,
+    const TpuMeshStateInterface& mesh_state) {
+  return CreateCompilationCacheKey(
+      function_name, function_library_fingerprint, mlir_module,
+      (guaranteed_constants.size() > 0 ? &guaranteed_constants[0] : nullptr),
+      guaranteed_constants.size(), dynamic_shapes, metadata, mesh_state);
+}
 }  // namespace tpu
 }  // namespace tensorflow
diff --git a/tensorflow/core/tpu/kernels/tpu_op_util.h b/tensorflow/core/tpu/kernels/tpu_op_util.h
index 0a9657ca05e..bbaa05682e6 100644
--- a/tensorflow/core/tpu/kernels/tpu_op_util.h
+++ b/tensorflow/core/tpu/kernels/tpu_op_util.h
@@ -34,6 +34,13 @@ TpuCompilationCacheKey CreateCompilationCacheKey(
     const std::vector& dynamic_shapes,
     const TPUCompileMetadataProto& metadata,
     const TpuMeshStateInterface& mesh_state);
+TpuCompilationCacheKey CreateCompilationCacheKey(
+    absl::string_view function_name, uint64 function_library_fingerprint,
+    absl::string_view mlir_module, const Tensor* guaranteed_constants,
+    size_t guaranteed_constants_size,
+    const std::vector& dynamic_shapes,
+    const TPUCompileMetadataProto& metadata,
+    const TpuMeshStateInterface& mesh_state);
 }  // namespace tpu
 }  // namespace tensorflow
 

From ff4490f636c87bf9efacfca06bd1f70a61bcfdc9 Mon Sep 17 00:00:00 2001
From: Taehee Jeong 
Date: Mon, 22 Jun 2020 22:05:00 -0700
Subject: [PATCH 0862/1390] Config Core ML delegate provider correctly

PiperOrigin-RevId: 317799599
Change-Id: I6eef7b9de910b2401cf24909d1e245dc972103f9
---
 tensorflow/lite/tools/delegates/coreml_delegate_provider.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/lite/tools/delegates/coreml_delegate_provider.cc b/tensorflow/lite/tools/delegates/coreml_delegate_provider.cc
index c6509618aee..7d88f04c8b4 100644
--- a/tensorflow/lite/tools/delegates/coreml_delegate_provider.cc
+++ b/tensorflow/lite/tools/delegates/coreml_delegate_provider.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include "tensorflow/lite/tools/delegates/delegate_provider.h"
 #include "tensorflow/lite/tools/evaluation/utils.h"
 #if defined(__APPLE__)
+#include "TargetConditionals.h"
 #if TARGET_OS_IPHONE && !TARGET_IPHONE_SIMULATOR
 // Only enable metal delegate when using a real iPhone device.
 #define REAL_IPHONE_DEVICE

From 40e860b7a46cf3768bd177b25b1e35750860f73f Mon Sep 17 00:00:00 2001
From: Mehdi Amini 
Date: Mon, 22 Jun 2020 22:21:55 -0700
Subject: [PATCH 0863/1390] Add a doc on the current short-term plan for XLA
 GPU CodeGen with MLIR

PiperOrigin-RevId: 317801702
Change-Id: Idabeb7717b357622e9cbcc2d4352df6b4ce25373
---
 .../compiler/mlir/g3doc/xla_gpu_codegen.md    | 265 ++++++++++++++++++
 1 file changed, 265 insertions(+)
 create mode 100644 tensorflow/compiler/mlir/g3doc/xla_gpu_codegen.md

diff --git a/tensorflow/compiler/mlir/g3doc/xla_gpu_codegen.md b/tensorflow/compiler/mlir/g3doc/xla_gpu_codegen.md
new file mode 100644
index 00000000000..06c55abf1fa
--- /dev/null
+++ b/tensorflow/compiler/mlir/g3doc/xla_gpu_codegen.md
@@ -0,0 +1,265 @@
+# MLIR CodeGen for XLA
+
+
+
+XLA operates on `HloInstruction` and performs many optimizations on this
+representation, sharing a lot of these between targeted devices. As some point a
+linear schedule is computed and the memory buffer is assigned to each value
+statically. The device specific codegen operates by traversing this sequence and
+calling "emitters" to generate a representation suitable for the device (for
+example a single LLVM function per XLA computation on CPU, or a sequence of
+"thunks" encapsulating GPU operations and possibly generated PTX when targeting
+GPU).
+
+As a staging step, we're currently in the process of intercepting the process
+right after XLA completes the buffer-assignment phase and emit instead an MLIR
+module in the `lhlo` dialect. From there we perform the codegen using MLIR
+components (Linalg, affine, and GPU dialect mainly) depending on the device.
+
+Below is the plan of record to incrementally migrate XLA/GPU by using `lhlo` as
+the codegen input.
+
+## Tasks
+
+              | Host                     | Device
+------------- | ------------------------ | ------------------------
+Input format  | HloInstruction* (Task 1) | HloInstruction* (Task 1)
+Output format | xla::Thunk (Task 2)      | LLVM IR (Task 3)
+
+*   **Task 1** changes both host and device input format from HloInstruction* to
+    LHLO.
+*   **Task 2** changes output format of host from thunks to "some landing pad
+    for host" (see below).
+*   **Task 3** migrates device output from LLVM IR to some form of MLIR. It's
+    optional to this project, and see the section "Migrating Device LLVM IR" for
+    details.
+
+This project prioritizes having end-to-end runnable models with LHLO-emitters
+enabled as much as possible. This implies that the following order list of
+objectives by priority:
+
+*   Make XLA/GPU runnable with LHLO emitters, with existing Thunks and emitters
+    unmodified.
+*   Eliminate the references to HloInstruction\* in LHLO, case by case:
+    *   Switch a legacy emitter to an MLIR-based emitter (e.g. Linalg), or
+    *   Mechanically translate the existing emitter to take MLIR representation
+        (migrate to Standard with GPU Dialect).
+
+## Migrating Thunks (Task 2)
+
+xla::gpu::Thunk is a data structure that:
+
+*   Can be called into from the host (xla::gpu::Thunk::ExecuteOnStream()).
+*   Carries various data in its subclasses.
+*   Interacts with BufferAllocation::Slice and StreamExecutor.
+*   Launches kernels
+*   Calls into all runtime libraries.
+
+The cost of that includes:
+
+*   Representing op-specific configuration data (e.g. convolution configs).
+*   Migrating op shape and operand shapes.
+*   Representing a tree of thunks (while, condition, etc).
+
+The migration work is independent from LHLO / emitter migration. Under limited
+resources, it's prioritized behind LHLO / emitter migration.
+
+We have several choices on how to lower the host-side part from LHLO:
+
+*   TFRT
+    *   (Pro) great CUDA and HIP wrappers for use.
+    *   (Pro) easy to implement library calls (cuDNN, cuBLAS, cuFFT, etc), as
+        TFRT ops are interpreted by C++ code.
+    *   (Con) host side is under development and not tested.
+    *   (Con) the JAX integration isn’t clear from a runtime point of view
+*   Jitted CPU code
+    *   (Pro) great lower-ability. Create a few loops and conditions and it's
+        done.
+    *   (Con) GPUDialect doesn't yet model chains/streams/asynchronicity/device
+        allocation.
+    *   (Con) CUDA / HIP runtime support is minimal (toolkit path, version,
+        dynamic loading, etc).
+*   Existing (interpreting) XLA runtime
+
+Tentative conclusion: Use jitted CPU code during the transition, and optionally
+adopt TFRT in the end.
+
+## Migrating Device LLVM IR (Task 3)
+
+An elemental emitter generates target op by filling it element by element. Each
+output element depends on a set of elements from the operands. All elements are
+described by combining the buffer with dynamic indices. It's sufficient to
+describe almost all "math" ops, but for performance reasons only a large subset
+of "math" ops are implemented directly in (Cpu|Gpu)ElementalIrEmitter.
+
+ElementalIrEmitter is unique in that:
+
+*   A large portion of the code is shared between XLA/GPU and CPU.
+*   It represents a large portion of ops seen in models, including all
+    element-wise ops.
+*   Most fusions solely depend on ElementalIrEmitter.
+*   It's structurally simple, as it describes a data dependency DAG between op
+    elements and operand elements.
+*   It's mostly portable and high-level (e.g. unlike GPU kReduce and GPU kCopy).
+*   Dynamic shape support is easy for at least element-wise ops.
+
+Now, for all ops, elementally-emitted or not, there are several flavors of the
+end state of each XLA op:
+
+1.  Device code stays as LLVM IR.
+1.  Refactor the old emitter to be like LHLO -> MLIR LLVM Dialect:
+    *   (Cost) Will be throw-away work if we want to ultimately migrate to
+        Standard.
+    *   (Benefit) It is easy and mechanical. Can be done in a short period.
+    *   (Benefit) It doesn't benefit more compared to a).
+1.  Refactor old emitters to be like LHLO -> MLIR GPU + Standard + Loops:
+    *   (Cost) Lifting existing emitters to Standard introduces some challenges.
+        Pointers and GEPs need to be converted to MemRefs and SubViews. Ensuring
+        amdgpu completeness is another one.
+    *   (Cost) XLA/GPU heavily relies on LLVM metadata:
+        *   `range` for block/thread indices.
+        *   `align`, `dereferenceable`, `invariant.load`, `alias.scope`,
+            `noalias` for load/stores.
+        *   `llvm.loop.unroll.disable`, `llvm.loop.unroll.full`,
+            `llvm.loop.vectorize.enable` for sequential loops.
+    *   (Benefit) Can be long-term. More portable.
+1.  Refactor old emitters to be LHLO -> Linalg, and write new Linalg emitters
+    *   (Cost) This is case by case. Compared to previous options, a new
+        implementation that matches XLA's performance needs to go through the
+        benchmark <-> optimize workflow, which can be a significant cost for
+        some ops.
+    *   (Benefit) unified stack; community support; portability; more
+        optimization potentials.
+
+## Prioritization
+
+While all three tasks mentioned above are parallelizable, under limited
+resources they have to be serialized. The prioritization focuses on visible
+results for completion of each task.
+
+The prioritization is: Task1 (LHLO for legacy emitters) > Task 2 (Thunks) > Task
+3 (MLIR emitters).
+
+By the end of Task 1, users of XLA can generate an LHLO (e.g. kernel generator)
+and execute them. The compilation format will not be serializable MLIR.
+
+By the end of Task 2, LHLO lowers to proper, serializable MLIR. This enables
+offline compilation.
+
+By the end of Task 3, all XLA emitters are MLIR-based in its implementation.
+
+## Detailed Design
+
+### Step 1: (Task 1) Complete LHLO and Make Legacy Emitters Take LHLO
+
+This step makes all existing XLA/GPU emitters interact with MLIR ops. This step
+is pure refactoring and NFC.
+
+This step is mostly mechanical, but it's worth noticing the following
+discrepancies between an unnested HloComputation and LHLO:
+
+*   Each HloInstruction has direct access to its operands (a data-flow DAG). On
+    contrary, each LHLO op only has access to its operand buffers (a bipartite
+    between ops and buffers). LHLO ops have to go through use-def chains to
+    access their operand ops.
+*   Unnested legacy emitters empirically almost never access their operands. The
+    only exception is kReduce.
+*   Unnested legacy emitters access BufferAssignment only for getting slices,
+    not for accessing aux data structures like dataflow\_analysis() or
+    alias\_analysis(). llvm\_ir builds its own alias\_analysis() based on slice
+    information.
+
+The conclusion is that LHLO should fit right-in without major hassle.
+
+### Step 2: (Optional) Profiling Support
+
+**This step is only needed if we start to discard some of the XLA Thunk logic
+(see the next step).**
+
+Before actually turning on any MLIR-based emitters, we need profiling for
+MLIR-based emitters.
+
+Currently XLA performs its own profiling by calling into StreamExecutor's timer.
+The timer under the hood inserts two events before and after a kernel launch,
+and measures the sync time between these two events.
+
+There are roughly three approaches to support profiling in MLIR:
+
+*   Run a profiler end-to-end
+*   Add a profile op for each op in LHLO, using an injected profiler.
+
+The "end-to-end" approach is transparent to MLIR, but suffers the same problem
+that makes XLA not use it in the first place: library calls collected by a
+profiler (nvprof/...) can't easily relate to HLO ops. For example, cuDNN
+launches multiple kernels for each HLO, and it's hard to tell which kernels
+correspond to which HLO.
+
+The "injected profiler" approach requires:
+
+*   LHLO to take a profiler as a parameter.
+*   inserting profile.start / profile.end before and after each op.
+*   a pass from that lowers profile.{start,end} to a C++ implementation.
+
+The exact profiling can't be easily done for MLIR-generated ops, since:
+
+*   MLIR doesn't have a timer, nor it depends on TFRT / StreamExecutor.
+*   MLIR doesn't easily call into C functions with complicated parameters.
+
+### Step 3: (Task 2) Migrating Thunks
+
+This step migrates all host ops and library calls. This step will eliminate most
+of the thunks and produce serializable MLIR instead.
+
+There are roughly three kinds of thunks:
+
+*   KernelThunk, which launches a kernel.
+*   Control flow thunks, which has host control flow logic (conditional, while,
+    for, sequence) and launch body kernels.
+*   Library thunks: cuDNN, cuBLAS, cuFFT, NCCL, etc.
+
+The **bottom line** is to:
+
+*   Create a Thunk dialect that provides (de)serialize logic for all existing
+    C++-based Thunks.
+*   Change emitters to emit a graph of Thunk dialect.
+
+**Optionally**, we can relieve some thunks from C++ implementation. KernelThunk
+can lower to the GPU LaunchKernelOp. Control flow thunks can leverage the CFG
+Dialect for loops and conditions, combined with LaunchKernelOp. This optional
+step requires profiling and stream support.
+
+### Step 4: (Task 3) Migrated ElementalIrEmitter
+
+Once profiling is ready, we can complete and tune all ElementalIrEmitter-based
+emitters in MLIR. Then we turn them on by default, assuming that all of these
+MLIR-based emitters use a single stream.
+
+Notice that it's beneficial to migrate XLA/CPU's ElementalIrEmitter as well,
+since they share a large portion of the code.
+
+With all benchmarking and performance hunting done (TODO: define performance
+parity), we turn on the new MLIR-based elemental emitter, and delete the legacy
+ElementalIrEmitter.
+
+This step also provides easy fusion transitions (nested ops) for the later
+migration.
+
+### Step 5: Multi-Stream Support or Drop
+
+We can't delete
+[some of the emitters](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/gpu/stream_assignment.cc#L140)
+until we support it in MLIR, or we drop the feature. It's a relatively large
+amount of work in MLIR and a small amount of gain for XLA. We should investigate
+current users of multi-stream XLA/GPU users, and try to delete this feature if
+reasonable.
+
+### Step 6: (Task 3) Migrated Device Ops
+
+This step migrates all unnested ops, then we can delete all unnested emitters.
+
+This calls on a rewrite/refactor for kCopy and kReduce. kReduce is already
+worked on for plenty, so the actual amount of work that needs to be done remains
+to be seen.

From 934df8dcea0c176314a52d0062dcca08638bb52d Mon Sep 17 00:00:00 2001
From: Advait Jain 
Date: Mon, 22 Jun 2020 22:29:15 -0700
Subject: [PATCH 0864/1390] Remove portable_optimized depthwise_conv
 implementation.

PR #34999 surfaced inconsistencies between the reference and portable_optimized
implementations and we are also slowly transitioning to CMSIS-NN for the
optimized implementations for ARM Cortex-M targets.

There may still be an issue with the reference depthwise conv + dilation > 1 but
that would be a separate bug.

PiperOrigin-RevId: 317802639
Change-Id: Ic8c9acffb060bb3ec5f802a2045ac6ac8b9f2233
---
 tensorflow/lite/micro/BUILD                   |  24 -
 tensorflow/lite/micro/kernels/BUILD           |  81 ---
 .../portable_optimized/depthwise_conv.cc      | 515 ------------------
 .../lite/micro/tools/ci_build/test_mbed.sh    |   2 +-
 .../make/targets/apollo3evb_makefile.inc      |   3 -
 5 files changed, 1 insertion(+), 624 deletions(-)
 delete mode 100644 tensorflow/lite/micro/kernels/portable_optimized/depthwise_conv.cc

diff --git a/tensorflow/lite/micro/BUILD b/tensorflow/lite/micro/BUILD
index f63d9778634..bdfa0c909db 100644
--- a/tensorflow/lite/micro/BUILD
+++ b/tensorflow/lite/micro/BUILD
@@ -102,30 +102,6 @@ cc_library(
     ],
 )
 
-# TODO(b/144176795): This target should really be handled differently so that we
-# do not have a fork in the build graph. The bug has some initial ideas.
-cc_library(
-    name = "portable_optimized_op_resolver",
-    srcs = [
-        "all_ops_resolver.cc",
-        "micro_mutable_op_resolver.h",
-        "micro_op_resolver.h",
-    ],
-    hdrs = [
-        "all_ops_resolver.h",
-    ],
-    copts = micro_copts(),
-    deps = [
-        ":micro_compatibility",
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/core/api",
-        "//tensorflow/lite/kernels:op_macros",
-        "//tensorflow/lite/kernels/internal:compatibility",
-        "//tensorflow/lite/micro/kernels:portable_optimized_micro_ops",
-        "//tensorflow/lite/schema:schema_fbs",
-    ],
-)
-
 cc_library(
     name = "debug_log",
     srcs = [
diff --git a/tensorflow/lite/micro/kernels/BUILD b/tensorflow/lite/micro/kernels/BUILD
index c7fa19b8cea..0fd0be4e3a4 100644
--- a/tensorflow/lite/micro/kernels/BUILD
+++ b/tensorflow/lite/micro/kernels/BUILD
@@ -20,7 +20,6 @@ package_group(
     packages = ["//tensorflow/lite/micro"],
 )
 
-# LINT.IfChange(micro_ops)
 cc_library(
     name = "micro_ops",
     srcs = [
@@ -106,73 +105,6 @@ cc_library(
         ],
     }),
 )
-# LINT.ThenChange(//tensorflow/lite/micro/kernels/BUILD:portable_optimized_micro_ops)
-
-# LINT.IfChange(portable_optimized_micro_ops)
-cc_library(
-    name = "portable_optimized_micro_ops",
-    srcs = [
-        "activations.cc",
-        "add.cc",
-        "arg_min_max.cc",
-        "ceil.cc",
-        "circular_buffer.cc",
-        "comparisons.cc",
-        "concatenation.cc",
-        "conv.cc",
-        "dequantize.cc",
-        "elementwise.cc",
-        "ethosu.cc",
-        "floor.cc",
-        "fully_connected.cc",
-        "l2norm.cc",
-        "logical.cc",
-        "logistic.cc",
-        "maximum_minimum.cc",
-        "mul.cc",
-        "neg.cc",
-        "pack.cc",
-        "pad.cc",
-        "pooling.cc",
-        "portable_optimized/depthwise_conv.cc",
-        "prelu.cc",
-        "quantize.cc",
-        "reduce.cc",
-        "reshape.cc",
-        "resize_nearest_neighbor.cc",
-        "round.cc",
-        "softmax.cc",
-        "split.cc",
-        "strided_slice.cc",
-        "sub.cc",
-        "svdf.cc",
-        "tanh.cc",
-        "unpack.cc",
-    ],
-    hdrs = ["micro_ops.h"],
-    copts = micro_copts(),
-    visibility = [
-        # Needed for micro:portable_optimized_ops_resolver but visibility can not be
-        # finer-grained than a package.
-        ":micro_top_level",
-    ],
-    deps = [
-        ":activation_utils",
-        ":micro_utils",
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/kernels:kernel_util",
-        "//tensorflow/lite/kernels:op_macros",
-        "//tensorflow/lite/kernels:padding",
-        "//tensorflow/lite/kernels/internal:common",
-        "//tensorflow/lite/kernels/internal:compatibility",
-        "//tensorflow/lite/kernels/internal:quantization_util",
-        "//tensorflow/lite/kernels/internal:reference_base",
-        "//tensorflow/lite/kernels/internal:tensor",
-        "//tensorflow/lite/kernels/internal:types",
-        "//tensorflow/lite/micro:micro_utils",
-    ],
-)
-# LINT.ThenChange(//tensorflow/lite/micro/kernels/BUILD:micro_ops)
 
 test_suite(
     name = "all_tests",
@@ -214,19 +146,6 @@ tflite_micro_cc_test(
     ],
 )
 
-tflite_micro_cc_test(
-    name = "portable_optimized_depthwise_conv_test",
-    srcs = [
-        "depthwise_conv_test.cc",
-    ],
-    deps = [
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/kernels/internal:tensor",
-        "//tensorflow/lite/micro:portable_optimized_op_resolver",
-        "//tensorflow/lite/micro/testing:micro_test",
-    ],
-)
-
 tflite_micro_cc_test(
     name = "fully_connected_test",
     srcs = [
diff --git a/tensorflow/lite/micro/kernels/portable_optimized/depthwise_conv.cc b/tensorflow/lite/micro/kernels/portable_optimized/depthwise_conv.cc
deleted file mode 100644
index 9fb8f2e32cc..00000000000
--- a/tensorflow/lite/micro/kernels/portable_optimized/depthwise_conv.cc
+++ /dev/null
@@ -1,515 +0,0 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "tensorflow/lite/kernels/internal/reference/integer_ops/depthwise_conv.h"
-
-#include "tensorflow/lite/c/builtin_op_data.h"
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/kernels/internal/common.h"
-#include "tensorflow/lite/kernels/internal/quantization_util.h"
-#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_float.h"
-#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
-#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
-#include "tensorflow/lite/kernels/kernel_util.h"
-#include "tensorflow/lite/kernels/padding.h"
-
-namespace tflite {
-namespace ops {
-namespace micro {
-namespace depthwise_conv {
-namespace {
-
-constexpr int kInputTensor = 0;
-constexpr int kFilterTensor = 1;
-constexpr int kBiasTensor = 2;
-constexpr int kOutputTensor = 0;
-constexpr int kMaxChannels = 256;
-
-// Depthwise conv is quantized along dimension 3:
-// https://www.tensorflow.org/lite/performance/quantization_spec
-constexpr int kDepthwiseConvQuantizedDimension = 3;
-
-// Size of the cached buffer we'll be using to hold reordered weights.
-constexpr int kReshapedFilterDataSize = 1 * 1024;
-
-struct OpData {
-  TfLitePaddingValues padding;
-  // The scaling factor from input to output (aka the 'real multiplier') can
-  // be represented as a fixed point multiplier plus a left shift.
-  int32_t output_multiplier;
-  int output_shift;
-
-  // Per channel output multiplier and shift.
-  int32_t per_channel_output_multiplier[kMaxChannels];
-  int32_t per_channel_output_shift[kMaxChannels];
-
-  // The range of the fused activation layer. For example for kNone and
-  // uint8_t these would be 0 and 255.
-  int32_t output_activation_min;
-  int32_t output_activation_max;
-};
-
-TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node,
-                             TfLiteDepthwiseConvParams* params, int width,
-                             int height, int filter_width, int filter_height,
-                             int out_width, int out_height,
-                             const TfLiteType data_type, OpData* data) {
-  bool has_bias = node->inputs->size == 3;
-  // Check number of inputs/outputs
-  TF_LITE_ENSURE(context, has_bias || node->inputs->size == 2);
-  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
-
-  // Matching GetWindowedOutputSize in TensorFlow.
-  auto padding = params->padding;
-  data->padding = ComputePaddingHeightWidth(
-      params->stride_height, params->stride_width,
-      params->dilation_height_factor, params->dilation_width_factor, height,
-      width, filter_height, filter_width, padding, &out_height, &out_width);
-
-  // Note that quantized inference requires that all tensors have their
-  // parameters set. This is usually done during quantized training.
-  if (data_type != kTfLiteFloat32) {
-    const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-    const TfLiteTensor* filter = GetInput(context, node, kFilterTensor);
-    const TfLiteTensor* bias =
-        GetOptionalInputTensor(context, node, kBiasTensor);
-    TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-    int num_channels = filter->dims->data[kDepthwiseConvQuantizedDimension];
-
-    TF_LITE_ENSURE_STATUS(tflite::PopulateConvolutionQuantizationParams(
-        context, input, filter, bias, output, params->activation,
-        &data->output_multiplier, &data->output_shift,
-        &data->output_activation_min, &data->output_activation_max,
-        data->per_channel_output_multiplier,
-        reinterpret_cast(data->per_channel_output_shift), num_channels));
-  }
-  return kTfLiteOk;
-}
-
-// Specialized implementation of the depthwise convolution operation designed to
-// work with the particular filter width of eight used by the default micro
-// speech sample code. It uses 1KB of RAM to hold reordered weight parameters,
-// converted from TFLite's NHWC format to NCHW format, and expressed as signed
-// eight bit integers, rather than unsigned. Care must be taken when calling
-// this not to use it for more than one node since there's only a single static
-// buffer holding the weights. You should use this implementation if depthwise
-// convolutions are a performance bottleneck, you have a layer that meets the
-// parameter requirements, and the extra RAM usage and additional code size are
-// not an issue.
-static inline void DepthwiseConvOptimizedForFilterWidthEight(
-    TfLiteContext* context, const DepthwiseParams& params,
-    const RuntimeShape& input_shape, const uint8* input_data,
-    const RuntimeShape& filter_shape, const uint8* filter_data,
-    const RuntimeShape& bias_shape, const int32* bias_data,
-    const RuntimeShape& output_shape, uint8* output_data) {
-  const int stride_width = params.stride_width;
-  const int stride_height = params.stride_height;
-  const int pad_width = params.padding_values.width;
-  const int pad_height = params.padding_values.height;
-  const int depth_multiplier = params.depth_multiplier;
-  const int32 output_activation_min = params.quantized_activation_min;
-  const int32 output_activation_max = params.quantized_activation_max;
-  const int32 input_offset = params.input_offset;
-  const int32 filter_offset = params.weights_offset;
-  const int32 output_offset = params.output_offset;
-  const int32 output_multiplier = params.output_multiplier;
-  const int output_shift = params.output_shift;
-  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
-
-  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
-  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
-  const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
-  const int input_height = input_shape.Dims(1);
-  const int input_width = input_shape.Dims(2);
-  const int input_depth = input_shape.Dims(3);
-  const int filter_height = filter_shape.Dims(1);
-  const int filter_width = filter_shape.Dims(2);
-  const int output_height = output_shape.Dims(1);
-  const int output_width = output_shape.Dims(2);
-  TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier);
-  TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
-
-  static int16_t reshaped_filter_data[kReshapedFilterDataSize];
-  const int needed_size =
-      output_depth * filter_width * filter_height * input_depth;
-  if (needed_size > kReshapedFilterDataSize) {
-    TF_LITE_KERNEL_LOG(
-        context,
-        "Size too large for reshaped weight buffer (%d needed, %d available)",
-        needed_size, kReshapedFilterDataSize);
-    return;
-  }
-
-  RuntimeShape reshaped_filter_shape;
-  reshaped_filter_shape.BuildFrom(
-      {1, output_depth, filter_height, filter_width});
-
-  // If this is the first time through, repack the weights into a cached buffer
-  // so that they can be accessed sequentially.
-  static bool is_reshaped_filter_initialized = false;
-  if (!is_reshaped_filter_initialized) {
-    for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
-      for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
-        for (int oc = 0; oc < output_depth; ++oc) {
-          const uint8* current_filter =
-              filter_data + Offset(filter_shape, 0, filter_y, filter_x, oc);
-          int16_t* reshaped_filter =
-              reshaped_filter_data +
-              Offset(reshaped_filter_shape, 0, oc, filter_y, filter_x);
-          *reshaped_filter =
-              static_cast(*current_filter) + filter_offset;
-        }
-      }
-    }
-    is_reshaped_filter_initialized = true;
-  }
-
-  for (int b = 0; b < batches; ++b) {
-    for (int out_y = 0; out_y < output_height; ++out_y) {
-      for (int out_x = 0; out_x < output_width; ++out_x) {
-        for (int ic = 0; ic < input_depth; ++ic) {
-          for (int m = 0; m < depth_multiplier; m++) {
-            const int oc = m + ic * depth_multiplier;
-            const int in_x_origin = (out_x * stride_width) - pad_width;
-            const int in_y_origin = (out_y * stride_height) - pad_height;
-            int32 acc = 0;
-            int in_y_start = in_y_origin;
-            int filter_y_start = 0;
-            if (in_y_origin < 0) {
-              in_y_start = 0;
-              filter_y_start = 0 - in_y_origin;
-            }
-            int filter_y_end = filter_height;
-            if ((in_y_origin + filter_height) >= input_height) {
-              filter_y_end -= (in_y_origin + filter_height) - input_height;
-            }
-            int in_y = in_y_start;
-            int in_x_start = in_x_origin;
-            int filter_x_start = 0;
-            bool is_out_of_x_bounds = false;
-            if (in_x_origin < 0) {
-              in_x_start = 0;
-              filter_x_start = 0 - in_x_origin;
-              is_out_of_x_bounds = true;
-            }
-            int filter_x_end = filter_width;
-            if ((in_x_origin + filter_width) >= input_width) {
-              filter_x_end -= (in_x_origin + filter_width) - input_width;
-              is_out_of_x_bounds = true;
-            }
-            for (int filter_y = filter_y_start; filter_y < filter_y_end;
-                 ++filter_y, ++in_y) {
-              const uint8* current_input =
-                  input_data + Offset(input_shape, b, in_y, in_x_start, ic);
-              if ((filter_width == 8) && !is_out_of_x_bounds) {
-                int16* current_filter =
-                    reshaped_filter_data + Offset(reshaped_filter_shape, 0, oc,
-                                                  filter_y, filter_x_start);
-                const uint32_t input_vals0 =
-                    *reinterpret_cast(current_input);
-                current_input += 4;
-                const int32_t filter_vals0 =
-                    *reinterpret_cast(current_filter);
-                current_filter += 2;
-                const uint8 input_val0 = input_vals0 & 0xff;
-                const int16 filter_val0 = filter_vals0 & 0xffff;
-                acc += filter_val0 * input_val0;
-                const uint8 input_val1 = (input_vals0 >> 8) & 0xff;
-                const int16 filter_val1 = (filter_vals0 >> 16) & 0xffff;
-                acc += filter_val1 * input_val1;
-
-                const int32_t filter_vals1 =
-                    *reinterpret_cast(current_filter);
-                current_filter += 2;
-                const uint8 input_val2 = (input_vals0 >> 16) & 0xff;
-                const int16 filter_val2 = filter_vals1 & 0xffff;
-                acc += filter_val2 * input_val2;
-                const uint8 input_val3 = (input_vals0 >> 24) & 0xff;
-                const int16 filter_val3 = (filter_vals1 >> 16) & 0xffff;
-                acc += filter_val3 * input_val3;
-
-                const uint32_t input_vals1 =
-                    *reinterpret_cast(current_input);
-                const int32_t filter_vals2 =
-                    *reinterpret_cast(current_filter);
-                current_filter += 2;
-                const uint8 input_val4 = input_vals1 & 0xff;
-                const int16 filter_val4 = filter_vals2 & 0xffff;
-                acc += filter_val4 * input_val4;
-                const uint8 input_val5 = (input_vals1 >> 8) & 0xff;
-                const int16 filter_val5 = (filter_vals2 >> 16) & 0xffff;
-                acc += filter_val5 * input_val5;
-
-                const int32_t filter_vals3 =
-                    *reinterpret_cast(current_filter);
-                const uint8 input_val6 = (input_vals1 >> 16) & 0xff;
-                const int16 filter_val6 = filter_vals3 & 0xffff;
-                acc += filter_val6 * input_val6;
-                const uint8 input_val7 = (input_vals1 >> 24) & 0xff;
-                const int16 filter_val7 = (filter_vals3 >> 16) & 0xffff;
-                acc += filter_val7 * input_val7;
-              } else {
-                const uint8* current_filter =
-                    filter_data +
-                    Offset(filter_shape, 0, filter_y, filter_x_start, oc);
-                for (int filter_x = filter_x_start; filter_x < filter_x_end;
-                     ++filter_x) {
-                  int32 input_val = *current_input;
-                  current_input += input_depth;
-                  int32 filter_val = *current_filter;
-                  current_filter += output_depth;
-                  acc +=
-                      (filter_val + filter_offset) * (input_val + input_offset);
-                }
-              }
-            }
-            if (bias_data) {
-              acc += bias_data[oc];
-            }
-            acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-                DepthwiseConvOutputRounding::kAwayFromZero>(
-                acc, output_multiplier, output_shift);
-            acc += output_offset;
-            acc = std::max(acc, output_activation_min);
-            acc = std::min(acc, output_activation_max);
-            output_data[Offset(output_shape, b, out_y, out_x, oc)] =
-                static_cast(acc);
-          }
-        }
-      }
-    }
-  }
-}  // namespace
-
-}  // namespace
-
-void EvalFloat(TfLiteContext* context, TfLiteNode* node,
-               TfLiteDepthwiseConvParams* params, OpData* data,
-               const TfLiteTensor* input, const TfLiteTensor* filter,
-               const TfLiteTensor* bias, TfLiteTensor* output) {
-  float output_activation_min, output_activation_max;
-  CalculateActivationRange(params->activation, &output_activation_min,
-                           &output_activation_max);
-
-  tflite::DepthwiseParams op_params;
-  // Padding type is ignored, but still set.
-  op_params.padding_type = PaddingType::kSame;
-  op_params.padding_values.width = data->padding.width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.stride_width = params->stride_width;
-  op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = 1;
-  op_params.dilation_height_factor = 1;
-  op_params.depth_multiplier = params->depth_multiplier;
-  op_params.float_activation_min = output_activation_min;
-  op_params.float_activation_max = output_activation_max;
-
-  tflite::reference_ops::DepthwiseConv(
-      op_params, GetTensorShape(input), GetTensorData(input),
-      GetTensorShape(filter), GetTensorData(filter),
-      GetTensorShape(bias), GetTensorData(bias), GetTensorShape(output),
-      GetTensorData(output));
-}
-
-// TODO(njeff): Optimize for int8 like we do for uint8.
-
-void EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node,
-                             TfLiteDepthwiseConvParams* params, OpData* data,
-                             const TfLiteTensor* input,
-                             const TfLiteTensor* filter,
-                             const TfLiteTensor* bias, TfLiteTensor* output) {
-  DepthwiseParams op_params;
-  op_params.padding_type = PaddingType::kSame;
-  op_params.padding_values.width = data->padding.width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.stride_width = params->stride_width;
-  op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = params->dilation_width_factor;
-  op_params.dilation_height_factor = params->dilation_height_factor;
-  op_params.depth_multiplier = params->depth_multiplier;
-  op_params.input_offset = -input->params.zero_point;
-  op_params.weights_offset = 0;
-  op_params.output_offset = output->params.zero_point;
-  // TODO(b/130439627): Use calculated value for clamping.
-  op_params.quantized_activation_min = std::numeric_limits::min();
-  op_params.quantized_activation_max = std::numeric_limits::max();
-
-  reference_integer_ops::DepthwiseConvPerChannel(
-      op_params, data->per_channel_output_multiplier,
-      data->per_channel_output_shift, GetTensorShape(input),
-      GetTensorData(input), GetTensorShape(filter),
-      GetTensorData(filter), GetTensorShape(bias),
-      GetTensorData(bias), GetTensorShape(output),
-      GetTensorData(output));
-}
-
-void EvalQuantized(TfLiteContext* context, TfLiteNode* node,
-                   TfLiteDepthwiseConvParams* params, OpData* data,
-                   const TfLiteTensor* input, const TfLiteTensor* filter,
-                   const TfLiteTensor* bias, TfLiteTensor* output) {
-  const int32_t input_offset = -input->params.zero_point;
-  const int32_t filter_offset = -filter->params.zero_point;
-  const int32_t output_offset = output->params.zero_point;
-
-  tflite::DepthwiseParams op_params;
-  // Padding type is ignored, but still set.
-  op_params.padding_type = PaddingType::kSame;
-  op_params.padding_values.width = data->padding.width;
-  op_params.padding_values.height = data->padding.height;
-  op_params.stride_width = params->stride_width;
-  op_params.stride_height = params->stride_height;
-  op_params.dilation_width_factor = 1;
-  op_params.dilation_height_factor = 1;
-  op_params.depth_multiplier = params->depth_multiplier;
-  op_params.quantized_activation_min = data->output_activation_min;
-  op_params.quantized_activation_max = data->output_activation_max;
-  op_params.input_offset = input_offset;
-  op_params.weights_offset = filter_offset;
-  op_params.output_offset = output_offset;
-  op_params.output_multiplier = data->output_multiplier;
-  // Legacy ops used mixed left and right shifts. Now all are +ve-means-left.
-  op_params.output_shift = -data->output_shift;
-
-  // Figure out if we can use the optimized path for this set of parameters.
-  const int filter_width = GetTensorShape(filter).Dims(2);
-  const int input_depth = GetTensorShape(input).Dims(3);
-  const int output_depth = GetTensorShape(filter).Dims(3);
-  const int filter_height = GetTensorShape(filter).Dims(1);
-  const int needed_size =
-      output_depth * filter_width * filter_height * input_depth;
-  bool use_optimized_path = false;
-  if ((filter_width == 8) && (input_offset == 0) && (input_depth == 1) &&
-      (needed_size <= kReshapedFilterDataSize)) {
-    // FIXME(petewarden) - We need a more robust way of handling this, ideally
-    // with an allocation mechanism available through the context API.
-    // Use the address of the node as a proxy for its identity, since we need
-    // to ensure the weight values are consistent between calls, and there's
-    // no easy way to do that quickly other than relying on the identity of
-    // the owning node.
-    static TfLiteNode* initialized_node_address = node;
-    if (initialized_node_address == node) {
-      use_optimized_path = true;
-    } else {
-      static bool has_warned = false;
-      if (!has_warned) {
-        TF_LITE_KERNEL_LOG(
-            context,
-            "Multiple depthwise conv ops match optimization parameters, but "
-            "only the first will use the fast path, because there's only one "
-            "RAM cache available");
-        has_warned = true;
-      }
-    }
-  }
-  if (use_optimized_path) {
-    DepthwiseConvOptimizedForFilterWidthEight(
-        context, op_params, GetTensorShape(input),
-        GetTensorData(input), GetTensorShape(filter),
-        GetTensorData(filter), GetTensorShape(bias),
-        GetTensorData(bias), GetTensorShape(output),
-        GetTensorData(output));
-  } else {
-    tflite::reference_ops::DepthwiseConv(
-        op_params, GetTensorShape(input), GetTensorData(input),
-        GetTensorShape(filter), GetTensorData(filter),
-        GetTensorShape(bias), GetTensorData(bias),
-        GetTensorShape(output), GetTensorData(output));
-  }
-}
-
-TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
-  auto* params =
-      reinterpret_cast(node->builtin_data);
-
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  const TfLiteTensor* filter = GetInput(context, node, kFilterTensor);
-  const TfLiteTensor* bias =
-      (NumInputs(node) == 3) ? GetInput(context, node, kBiasTensor) : nullptr;
-
-  const TfLiteType data_type = input->type;
-  int width = SizeOfDimension(input, 2);
-  int height = SizeOfDimension(input, 1);
-  int filter_width = SizeOfDimension(filter, 2);
-  int filter_height = SizeOfDimension(filter, 1);
-  int out_width = ComputeOutSize(params->padding, width, filter_width,
-                                 params->stride_width);
-  int out_height = ComputeOutSize(params->padding, height, filter_height,
-                                  params->stride_height);
-  OpData data;
-
-  // All per-channel quantized tensors need valid zero point and scale arrays.
-  if (input->type == kTfLiteInt8) {
-    TF_LITE_ENSURE_EQ(context, filter->quantization.type,
-                      kTfLiteAffineQuantization);
-
-    const auto* affine_quantization =
-        reinterpret_cast(
-            filter->quantization.params);
-    TF_LITE_ENSURE(context, affine_quantization);
-    TF_LITE_ENSURE(context, affine_quantization->scale);
-    TF_LITE_ENSURE(context, affine_quantization->zero_point);
-    TF_LITE_ENSURE(
-        context, affine_quantization->scale->size == 1 ||
-                     affine_quantization->scale->size ==
-                         filter->dims->data[kDepthwiseConvQuantizedDimension]);
-    TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size,
-                      affine_quantization->zero_point->size);
-  }
-
-  TF_LITE_ENSURE_STATUS(CalculateOpData(context, node, params, width, height,
-                                        filter_width, filter_height, out_width,
-                                        out_height, data_type, &data));
-
-  // TODO(aselle): Consider whether float conv and quantized conv should be
-  // separate ops to avoid dispatch overhead here.
-  switch (input->type) {  // Already know in/out types are same.
-    case kTfLiteFloat32:
-      EvalFloat(context, node, params, &data, input, filter, bias, output);
-      break;
-    case kTfLiteInt8:
-      EvalQuantizedPerChannel(context, node, params, &data, input, filter, bias,
-                              output);
-      break;
-    case kTfLiteUInt8:
-      EvalQuantized(context, node, params, &data, input, filter, bias, output);
-      break;
-    default:
-      TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.",
-                         TfLiteTypeGetName(input->type), input->type);
-      return kTfLiteError;
-  }
-  return kTfLiteOk;
-}
-
-}  // namespace depthwise_conv
-
-TfLiteRegistration* Register_DEPTHWISE_CONV_2D() {
-  static TfLiteRegistration r = {/*init=*/nullptr,
-                                 /*free=*/nullptr,
-                                 /*prepare=*/nullptr,
-                                 /*invoke=*/depthwise_conv::Eval,
-                                 /*profiling_string=*/nullptr,
-                                 /*builtin_code=*/0,
-                                 /*custom_name=*/nullptr,
-                                 /*version=*/0};
-  return &r;
-}
-
-}  // namespace micro
-}  // namespace ops
-}  // namespace tflite
diff --git a/tensorflow/lite/micro/tools/ci_build/test_mbed.sh b/tensorflow/lite/micro/tools/ci_build/test_mbed.sh
index a4d47009c93..fa4506fa6b8 100755
--- a/tensorflow/lite/micro/tools/ci_build/test_mbed.sh
+++ b/tensorflow/lite/micro/tools/ci_build/test_mbed.sh
@@ -49,7 +49,7 @@ fi
 
 make -f tensorflow/lite/micro/tools/make/Makefile \
   TARGET=${TARGET} \
-  TAGS="portable_optimized disco_f746ng" \
+  TAGS="disco_f746ng" \
   ${PROJECTS}
 
 readable_run tensorflow/lite/micro/tools/ci_build/install_mbed_cli.sh
diff --git a/tensorflow/lite/micro/tools/make/targets/apollo3evb_makefile.inc b/tensorflow/lite/micro/tools/make/targets/apollo3evb_makefile.inc
index 7d2a0e65b97..dc7a689daed 100644
--- a/tensorflow/lite/micro/tools/make/targets/apollo3evb_makefile.inc
+++ b/tensorflow/lite/micro/tools/make/targets/apollo3evb_makefile.inc
@@ -24,9 +24,6 @@ ifeq ($(TARGET),$(filter $(TARGET),\
 $(MAKEFILE_DIR)/downloads/$(AM_SDK_DEST)/$(SF_BSPS_DEST): $(MAKEFILE_DIR)/downloads/$(AM_SDK_DEST)
   endif
 
-  # Use the faster depthwise conv implementation.
-  ALL_TAGS += portable_optimized
-
   PLATFORM_FLAGS = \
     -DPART_apollo3 \
     -DAM_PACKAGE_BGA \

From 088d4376945c0ee6f3def04284ff1a835ea57097 Mon Sep 17 00:00:00 2001
From: Hye Soo Yang 
Date: Mon, 22 Jun 2020 22:38:05 -0700
Subject: [PATCH 0865/1390] Fix BMP header bytes.

PiperOrigin-RevId: 317803857
Change-Id: Iafa8e4e44400ec3c119af8897ba3fa693138501f
---
 tensorflow/python/kernel_tests/decode_bmp_op_test.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/kernel_tests/decode_bmp_op_test.py b/tensorflow/python/kernel_tests/decode_bmp_op_test.py
index 5e7991382ed..b0a8f9ffbbb 100644
--- a/tensorflow/python/kernel_tests/decode_bmp_op_test.py
+++ b/tensorflow/python/kernel_tests/decode_bmp_op_test.py
@@ -25,14 +25,14 @@ from tensorflow.python.ops import image_ops
 from tensorflow.python.platform import test
 
 
-
 class DecodeBmpOpTest(test.TestCase):
 
   def testex1(self):
     img_bytes = [[[0, 0, 255], [0, 255, 0]], [[255, 0, 0], [255, 255, 255]]]
     # Encoded BMP bytes from Wikipedia
+    # BMP header bytes: https://en.wikipedia.org/wiki/List_of_file_signatures
     encoded_bytes = [
-        0x42, 0x40,
+        0x42, 0x4d,
         0x46, 0, 0, 0,
         0, 0,
         0, 0,
@@ -66,9 +66,10 @@ class DecodeBmpOpTest(test.TestCase):
 
   def testGrayscale(self):
     img_bytes = [[[255], [0]], [[255], [0]]]
+    # BMP header bytes: https://en.wikipedia.org/wiki/List_of_file_signatures
     encoded_bytes = [
         0x42,
-        0x40,
+        0x4d,
         0x3d,
         0,
         0,
@@ -133,6 +134,8 @@ class DecodeBmpOpTest(test.TestCase):
 
     byte_string = bytes(bytearray(encoded_bytes))
     img_in = constant_op.constant(byte_string, dtype=dtypes.string)
+    # TODO(b/159600494): Currently, `decode_bmp` op does not validate input
+    # magic bytes.
     decode = image_ops.decode_bmp(img_in)
 
     with self.cached_session():

From dce31a431df80aa55f2cfd87d60952c413d6f690 Mon Sep 17 00:00:00 2001
From: chuanqiw 
Date: Tue, 23 Jun 2020 14:08:27 +0800
Subject: [PATCH 0866/1390] Update sqlite version to 3.32.1 to fix the
 CVE-2020-13630 and CVE-2020-11656

---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 0f591ba8b90..29c6ef99397 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -409,12 +409,12 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     tf_http_archive(
         name = "org_sqlite",
         build_file = clean_dep("//third_party:sqlite.BUILD"),
-        sha256 = "f3c79bc9f4162d0b06fa9fe09ee6ccd23bb99ce310b792c5145f87fbcc30efca",
-        strip_prefix = "sqlite-amalgamation-3310100",
+        sha256 = "8d46ef69b96628bedb781bd8309210f2a1f4a353792097302f6b754044e6540f",
+        strip_prefix = "sqlite-amalgamation-3320100",
         system_build_file = clean_dep("//third_party/systemlibs:sqlite.BUILD"),
         urls = [
-            "https://storage.googleapis.com/mirror.tensorflow.org/www.sqlite.org/2020/sqlite-amalgamation-3310100.zip",
-            "https://www.sqlite.org/2020/sqlite-amalgamation-3310100.zip",
+            "https://storage.googleapis.com/mirror.tensorflow.org/www.sqlite.org/2020/sqlite-amalgamation-3320100.zip",
+            "https://www.sqlite.org/2020/sqlite-amalgamation-3320100.zip",
         ],
     )
 

From 51ee63247d1d4e4c6db089875b605963f31de2d9 Mon Sep 17 00:00:00 2001
From: Gaurav Jain 
Date: Mon, 22 Jun 2020 23:04:46 -0700
Subject: [PATCH 0867/1390] Rollforward: Add uint32 & uint64 to
 TF_CALL_INTEGRAL_TYPES

PiperOrigin-RevId: 317806755
Change-Id: I3752198bf98d5659bf163455a5dfe1fbe39ffd1c
---
 tensorflow/core/framework/register_types.h    | 21 +++++---------
 tensorflow/core/framework/types.cc            |  5 ----
 tensorflow/core/kernels/BUILD                 |  2 ++
 tensorflow/core/kernels/concat_lib_cpu.cc     |  2 --
 tensorflow/core/kernels/concat_op.cc          |  2 --
 tensorflow/core/kernels/constant_op.cc        |  1 -
 tensorflow/core/kernels/control_flow_ops.cc   |  5 ----
 .../core/kernels/data/dataset_test_base.cc    |  2 --
 tensorflow/core/kernels/dense_update_ops.cc   |  1 -
 .../core/kernels/dynamic_partition_op.cc      |  2 --
 tensorflow/core/kernels/fill_functor.cc       |  5 +++-
 tensorflow/core/kernels/gather_op.cc          |  2 --
 tensorflow/core/kernels/identity_op.cc        |  1 -
 tensorflow/core/kernels/ragged_gather_op.cc   |  2 --
 .../kernels/ragged_tensor_from_variant_op.cc  |  2 --
 .../kernels/ragged_tensor_to_tensor_op.cc     |  2 --
 .../kernels/ragged_tensor_to_variant_op.cc    |  2 --
 .../core/kernels/resource_variable_ops.cc     |  1 -
 tensorflow/core/kernels/split_lib_cpu.cc      |  1 -
 tensorflow/core/kernels/split_op.cc           |  1 -
 tensorflow/core/kernels/strided_slice_op.cc   |  2 --
 .../core/kernels/strided_slice_op_impl.h      |  2 --
 tensorflow/core/kernels/topk_op.cc            |  2 --
 .../core/kernels/topk_op_gpu_uint32.cu.cc     | 28 +++++++++++++++++++
 .../core/kernels/topk_op_gpu_uint64.cu.cc     | 28 +++++++++++++++++++
 tensorflow/core/util/batch_util.cc            |  8 ------
 .../core/util/saved_tensor_slice_util.h       |  2 ++
 27 files changed, 71 insertions(+), 63 deletions(-)
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
 create mode 100644 tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc

diff --git a/tensorflow/core/framework/register_types.h b/tensorflow/core/framework/register_types.h
index bc3e5e1743b..0cf6536e8c2 100644
--- a/tensorflow/core/framework/register_types.h
+++ b/tensorflow/core/framework/register_types.h
@@ -153,16 +153,9 @@ limitations under the License.
 #endif  // defined(IS_MOBILE_PLATFORM)  - end of TF_CALL_type defines
 
 // Defines for sets of types.
-
-// TODO(b/111604096): Add uint32 and uint64 to TF_CALL_INTEGRAL_TYPES.
-//
-// The uint32 and uint64 types were introduced in 10/2017 to be used via XLA and
-// thus were not included in TF_CALL_INTEGRAL_TYPES. Including them in
-// TF_CALL_INTEGRAL_TYPES should only happen after evaluating the effect on the
-// TF binary size and performance.
-#define TF_CALL_INTEGRAL_TYPES(m)                                      \
-  TF_CALL_int64(m) TF_CALL_int32(m) TF_CALL_uint16(m) TF_CALL_int16(m) \
-      TF_CALL_uint8(m) TF_CALL_int8(m)
+#define TF_CALL_INTEGRAL_TYPES(m)                                       \
+  TF_CALL_uint64(m) TF_CALL_int64(m) TF_CALL_uint32(m) TF_CALL_int32(m) \
+      TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)
 
 #define TF_CALL_FLOAT_TYPES(m) \
   TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)
@@ -174,10 +167,10 @@ limitations under the License.
 #define TF_CALL_REAL_NUMBER_TYPES_NO_BFLOAT16(m) \
   TF_CALL_INTEGRAL_TYPES(m) TF_CALL_half(m) TF_CALL_float(m) TF_CALL_double(m)
 
-#define TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m)                              \
-  TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)   \
-      TF_CALL_int64(m) TF_CALL_uint16(m) TF_CALL_int16(m) TF_CALL_uint8(m) \
-          TF_CALL_int8(m)
+#define TF_CALL_REAL_NUMBER_TYPES_NO_INT32(m)                                \
+  TF_CALL_half(m) TF_CALL_bfloat16(m) TF_CALL_float(m) TF_CALL_double(m)     \
+      TF_CALL_uint64(m) TF_CALL_int64(m) TF_CALL_uint32(m) TF_CALL_uint16(m) \
+          TF_CALL_int16(m) TF_CALL_uint8(m) TF_CALL_int8(m)
 
 #define TF_CALL_COMPLEX_TYPES(m) TF_CALL_complex64(m) TF_CALL_complex128(m)
 
diff --git a/tensorflow/core/framework/types.cc b/tensorflow/core/framework/types.cc
index 97eaec98ffe..d6455e012d0 100644
--- a/tensorflow/core/framework/types.cc
+++ b/tensorflow/core/framework/types.cc
@@ -238,11 +238,6 @@ int DataTypeSize(DataType dt) {
     TF_CALL_qint16(CASE);
     TF_CALL_quint16(CASE);
 
-    // uint32 and uint64 aren't included in TF_CALL_POD_TYPES because we
-    // don't want to define kernels for them at this stage to avoid binary
-    // bloat.
-    TF_CALL_uint32(CASE);
-    TF_CALL_uint64(CASE);
     default:
       return 0;
   }
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 7da864a6027..e2ff5aed283 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -4919,7 +4919,9 @@ tf_kernel_library(
         "topk_op_gpu_double.cu.cc",
         "topk_op_gpu_float.cu.cc",
         "topk_op_gpu_half.cu.cc",
+        "topk_op_gpu_uint64.cu.cc",
         "topk_op_gpu_int64.cu.cc",
+        "topk_op_gpu_uint32.cu.cc",
         "topk_op_gpu_int32.cu.cc",
         "topk_op_gpu_int16.cu.cc",
         "topk_op_gpu_uint16.cu.cc",
diff --git a/tensorflow/core/kernels/concat_lib_cpu.cc b/tensorflow/core/kernels/concat_lib_cpu.cc
index da73d3d2c56..1dec589d3ff 100644
--- a/tensorflow/core/kernels/concat_lib_cpu.cc
+++ b/tensorflow/core/kernels/concat_lib_cpu.cc
@@ -116,8 +116,6 @@ REGISTER(qint8)
 REGISTER(quint16)
 REGISTER(qint16)
 REGISTER(qint32)
-REGISTER(uint32)
-REGISTER(uint64)
 
 #if defined(IS_MOBILE_PLATFORM) && !defined(SUPPORT_SELECTIVE_REGISTRATION) && \
     !defined(__ANDROID_TYPES_FULL__)
diff --git a/tensorflow/core/kernels/concat_op.cc b/tensorflow/core/kernels/concat_op.cc
index be3e9a67c5f..d3f3a04f33b 100644
--- a/tensorflow/core/kernels/concat_op.cc
+++ b/tensorflow/core/kernels/concat_op.cc
@@ -208,8 +208,6 @@ REGISTER_CONCAT(qint8);
 REGISTER_CONCAT(quint16);
 REGISTER_CONCAT(qint16);
 REGISTER_CONCAT(qint32);
-REGISTER_CONCAT(uint32);
-REGISTER_CONCAT(uint64);
 
 #undef REGISTER_CONCAT
 
diff --git a/tensorflow/core/kernels/constant_op.cc b/tensorflow/core/kernels/constant_op.cc
index 4bcbc076446..dc178d17d49 100644
--- a/tensorflow/core/kernels/constant_op.cc
+++ b/tensorflow/core/kernels/constant_op.cc
@@ -211,7 +211,6 @@ TF_CALL_ALL_TYPES(REGISTER_CPU_KERNEL);
 // the conversion from uint8 to quint8.
 REGISTER_KERNEL(CPU, quint8);
 REGISTER_KERNEL(CPU, quint16);
-REGISTER_KERNEL(CPU, uint32);
 #undef REGISTER_CPU_KERNEL
 
 #ifdef TENSORFLOW_USE_SYCL
diff --git a/tensorflow/core/kernels/control_flow_ops.cc b/tensorflow/core/kernels/control_flow_ops.cc
index 435de3c5954..1a0082c6a3b 100644
--- a/tensorflow/core/kernels/control_flow_ops.cc
+++ b/tensorflow/core/kernels/control_flow_ops.cc
@@ -101,16 +101,12 @@ TF_CALL_ALL_TYPES(REGISTER_CPU_SWITCH);
 TF_CALL_ALL_TYPES(REGISTER_CPU_REF_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_REF_SWITCH);
-REGISTER_CPU_SWITCH(uint64);
 
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_SWITCH);
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_REF_SWITCH);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_REF_SWITCH);
-REGISTER_GPU_SWITCH(uint64);
 TF_CALL_variant(REGISTER_GPU_SWITCH);
-TF_CALL_uint32(REGISTER_GPU_SWITCH);
-TF_CALL_uint32(REGISTER_GPU_REF_SWITCH);
 TF_CALL_bool(REGISTER_GPU_SWITCH);
 TF_CALL_bool(REGISTER_GPU_REF_SWITCH);
 
@@ -311,7 +307,6 @@ TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GPU_REF_KERNEL);
 REGISTER_GPU_KERNEL(bool);
 REGISTER_GPU_REF_KERNEL(bool);
-REGISTER_GPU_KERNEL(uint64);
 TF_CALL_variant(REGISTER_GPU_KERNEL);
 
 #undef REGISTER_GPU_KERNEL
diff --git a/tensorflow/core/kernels/data/dataset_test_base.cc b/tensorflow/core/kernels/data/dataset_test_base.cc
index b91ab9b733c..e41e35be1e9 100644
--- a/tensorflow/core/kernels/data/dataset_test_base.cc
+++ b/tensorflow/core/kernels/data/dataset_test_base.cc
@@ -220,8 +220,6 @@ Status DatasetOpsTestBase::ExpectEqual(const Tensor& a, const Tensor& b) {
     break;
     TF_CALL_NUMBER_TYPES(CASE);
     TF_CALL_tstring(CASE);
-    TF_CALL_uint32(CASE);
-    TF_CALL_uint64(CASE);
     // TODO(feihugis): figure out how to support variant tensors.
 #undef CASE
     default:
diff --git a/tensorflow/core/kernels/dense_update_ops.cc b/tensorflow/core/kernels/dense_update_ops.cc
index 55e4cd7606a..71235fca143 100644
--- a/tensorflow/core/kernels/dense_update_ops.cc
+++ b/tensorflow/core/kernels/dense_update_ops.cc
@@ -98,7 +98,6 @@ typedef Eigen::SyclDevice SYCLDevice;
 
 TF_CALL_ALL_TYPES(REGISTER_KERNELS);
 // uint32 not included in ALL_TYPES
-TF_CALL_uint32(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 // quint16 not included in QUANTIZIED_TYPES
 TF_CALL_quint16(REGISTER_KERNELS);
diff --git a/tensorflow/core/kernels/dynamic_partition_op.cc b/tensorflow/core/kernels/dynamic_partition_op.cc
index 90ed71dccce..95af19c4c48 100644
--- a/tensorflow/core/kernels/dynamic_partition_op.cc
+++ b/tensorflow/core/kernels/dynamic_partition_op.cc
@@ -164,8 +164,6 @@ class DynamicPartitionOp : public DynamicPartitionOp_Shared {
       DynamicPartitionOp)
 
 TF_CALL_ALL_TYPES(REGISTER_DYNAMIC_PARTITION);
-// For partitioning fingerprints.
-TF_CALL_uint64(REGISTER_DYNAMIC_PARTITION);
 #undef REGISTER_DYNAMIC_PARTITION
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/fill_functor.cc b/tensorflow/core/kernels/fill_functor.cc
index 10dd3df1915..174a4e45a79 100644
--- a/tensorflow/core/kernels/fill_functor.cc
+++ b/tensorflow/core/kernels/fill_functor.cc
@@ -45,6 +45,8 @@ DEFINE_SETZERO_CPU(Eigen::half);
 DEFINE_SETZERO_CPU(bfloat16);
 DEFINE_SETZERO_CPU(float);
 DEFINE_SETZERO_CPU(double);
+DEFINE_SETZERO_CPU(uint32);
+DEFINE_SETZERO_CPU(uint64);
 DEFINE_SETZERO_CPU(uint8);
 DEFINE_SETZERO_CPU(int8);
 DEFINE_SETZERO_CPU(uint16);
@@ -96,6 +98,8 @@ DEFINE_SETONE_CPU(Eigen::half);
 DEFINE_SETONE_CPU(bfloat16);
 DEFINE_SETONE_CPU(float);
 DEFINE_SETONE_CPU(double);
+DEFINE_SETONE_CPU(uint32);
+DEFINE_SETONE_CPU(uint64);
 DEFINE_SETONE_CPU(uint8);
 DEFINE_SETONE_CPU(int8);
 DEFINE_SETONE_CPU(uint16);
@@ -137,7 +141,6 @@ struct FillFunctor {
 TF_CALL_ALL_TYPES(DEFINE_FILL_CPU);
 DEFINE_FILL_CPU(quint8);
 DEFINE_FILL_CPU(quint16);
-DEFINE_FILL_CPU(uint32);
 #undef DEFINE_FILL_CPU
 
 #ifdef TENSORFLOW_USE_SYCL
diff --git a/tensorflow/core/kernels/gather_op.cc b/tensorflow/core/kernels/gather_op.cc
index 6d493a5f2ea..948567e019a 100644
--- a/tensorflow/core/kernels/gather_op.cc
+++ b/tensorflow/core/kernels/gather_op.cc
@@ -211,8 +211,6 @@ TF_CALL_ALL_TYPES(REGISTER_GATHER_CPU);
 TF_CALL_QUANTIZED_TYPES(REGISTER_GATHER_CPU);
 TF_CALL_quint16(REGISTER_GATHER_CPU);
 TF_CALL_qint16(REGISTER_GATHER_CPU);
-TF_CALL_uint32(REGISTER_GATHER_CPU);
-TF_CALL_uint64(REGISTER_GATHER_CPU);
 
 #undef REGISTER_GATHER_CPU
 
diff --git a/tensorflow/core/kernels/identity_op.cc b/tensorflow/core/kernels/identity_op.cc
index d15b64597f5..4b226dd72d4 100644
--- a/tensorflow/core/kernels/identity_op.cc
+++ b/tensorflow/core/kernels/identity_op.cc
@@ -122,7 +122,6 @@ REGISTER_SYCL_HOST_KERNEL(bool);
 
 TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL);
 REGISTER_GPU_KERNEL(Variant);
-TF_CALL_uint32(REGISTER_GPU_KERNEL);
 REGISTER_GPU_KERNEL(bool);
 
 #undef REGISTER_GPU_KERNEL
diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc
index 88c0d1ebd69..3bf82cba050 100644
--- a/tensorflow/core/kernels/ragged_gather_op.cc
+++ b/tensorflow/core/kernels/ragged_gather_op.cc
@@ -296,8 +296,6 @@ TF_CALL_tstring(REGISTER_CPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_KERNEL);
 TF_CALL_quint16(REGISTER_CPU_KERNEL);
 TF_CALL_qint16(REGISTER_CPU_KERNEL);
-TF_CALL_uint32(REGISTER_CPU_KERNEL);
-TF_CALL_uint64(REGISTER_CPU_KERNEL);
 #undef REGISTER_CPU_KERNEL
 #undef REGISTER_CPU_KERNEL_WITH_INDEX_TYPE
 
diff --git a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
index f83bcb38c6c..ad0712e6fd0 100644
--- a/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_from_variant_op.cc
@@ -308,8 +308,6 @@ TF_CALL_tstring(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 TF_CALL_quint16(REGISTER_KERNELS);
 TF_CALL_qint16(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
-TF_CALL_uint64(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 #undef REGISTER_KERNELS_WITH_SPLIT_TYPE
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
index d729c43f25a..9ae5d7ffbdc 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_tensor_op.cc
@@ -561,8 +561,6 @@ TF_CALL_string(REGISTER_CPU_KERNEL);
 TF_CALL_QUANTIZED_TYPES(REGISTER_CPU_KERNEL);
 TF_CALL_quint16(REGISTER_CPU_KERNEL);
 TF_CALL_qint16(REGISTER_CPU_KERNEL);
-TF_CALL_uint32(REGISTER_CPU_KERNEL);
-TF_CALL_uint64(REGISTER_CPU_KERNEL);
 
 #undef REGISTER_CPU_KERNEL
 
diff --git a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
index 7a5ae1c6240..64c372b005e 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_variant_op.cc
@@ -213,8 +213,6 @@ TF_CALL_tstring(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
 TF_CALL_quint16(REGISTER_KERNELS);
 TF_CALL_qint16(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
-TF_CALL_uint64(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 #undef REGISTER_KERNELS_WITH_SPLIT_TYPE
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index b9c883c7e2f..510e95ca606 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -508,7 +508,6 @@ class AssignVariableOp : public OpKernel {
 
 TF_CALL_ALL_TYPES(REGISTER_KERNELS);
 TF_CALL_QUANTIZED_TYPES(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS);
 #undef REGISTER_KERNELS
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
diff --git a/tensorflow/core/kernels/split_lib_cpu.cc b/tensorflow/core/kernels/split_lib_cpu.cc
index 0cb0a94d498..a3060e4e90d 100644
--- a/tensorflow/core/kernels/split_lib_cpu.cc
+++ b/tensorflow/core/kernels/split_lib_cpu.cc
@@ -43,7 +43,6 @@ void Split::operator()(
 
 TF_CALL_ALL_TYPES(DEFINE_CPU_KERNELS)
 DEFINE_CPU_KERNELS(quint8)
-DEFINE_CPU_KERNELS(uint64)
 
 #ifdef TENSORFLOW_USE_SYCL
 template 
diff --git a/tensorflow/core/kernels/split_op.cc b/tensorflow/core/kernels/split_op.cc
index f09740c6198..08575f01f67 100644
--- a/tensorflow/core/kernels/split_op.cc
+++ b/tensorflow/core/kernels/split_op.cc
@@ -404,7 +404,6 @@ class SplitOpSYCL : public SplitOpBase {
 
 TF_CALL_ALL_TYPES(REGISTER_SPLIT);
 REGISTER_SPLIT(quint8);
-REGISTER_SPLIT(uint64);
 
 #undef REGISTER_SPLIT
 
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index ccc1984bb98..b4099213303 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -440,8 +440,6 @@ class StridedSliceAssignOp : public OpKernel {
                           StridedSliceAssignOp)
 
 TF_CALL_ALL_TYPES(REGISTER_STRIDED_SLICE);
-TF_CALL_uint32(REGISTER_STRIDED_SLICE);
-TF_CALL_uint64(REGISTER_STRIDED_SLICE);
 
 #undef REGISTER_STRIDED_SLICE
 
diff --git a/tensorflow/core/kernels/strided_slice_op_impl.h b/tensorflow/core/kernels/strided_slice_op_impl.h
index 1ae959b7b3f..5ce1d773e33 100644
--- a/tensorflow/core/kernels/strided_slice_op_impl.h
+++ b/tensorflow/core/kernels/strided_slice_op_impl.h
@@ -287,8 +287,6 @@ TF_CALL_GPU_ALL_TYPES(DECLARE_FOR_N_GPU);
 #endif  // END GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
 TF_CALL_ALL_TYPES(DECLARE_FOR_N_CPU);
-TF_CALL_uint32(DECLARE_FOR_N_CPU);
-TF_CALL_uint64(DECLARE_FOR_N_CPU);
 
 #ifdef TENSORFLOW_USE_SYCL
 #define PREVENT_FOR_N_SYCL(T) \
diff --git a/tensorflow/core/kernels/topk_op.cc b/tensorflow/core/kernels/topk_op.cc
index c555b42f005..50325b7bcfe 100644
--- a/tensorflow/core/kernels/topk_op.cc
+++ b/tensorflow/core/kernels/topk_op.cc
@@ -258,7 +258,6 @@ namespace functor {
 
 TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
 TF_CALL_INTEGRAL_TYPES(DECLARE_GPU_SPEC);
-TF_CALL_uint32(DECLARE_GPU_SPEC);
 
 #undef DECLARE_GPU_SPEC
 
@@ -276,7 +275,6 @@ TF_CALL_uint32(DECLARE_GPU_SPEC);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER_KERNELS);
 TF_CALL_INTEGRAL_TYPES(REGISTER_KERNELS);
-TF_CALL_uint32(REGISTER_KERNELS)
 #undef REGISTER_KERNELS
 
 #endif  // end GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc b/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
new file mode 100644
index 00000000000..16e2e0e9420
--- /dev/null
+++ b/tensorflow/core/kernels/topk_op_gpu_uint32.cu.cc
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include "tensorflow/core/kernels/topk_op.h"
+#include "tensorflow/core/kernels/topk_op_gpu.h"
+
+namespace tensorflow {
+using Eigen::GpuDevice;
+
+template struct functor::TopKFunctor;
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc b/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc
new file mode 100644
index 00000000000..895247a63a2
--- /dev/null
+++ b/tensorflow/core/kernels/topk_op_gpu_uint64.cu.cc
@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+
+#include "tensorflow/core/kernels/topk_op.h"
+#include "tensorflow/core/kernels/topk_op_gpu.h"
+
+namespace tensorflow {
+using Eigen::GpuDevice;
+
+template struct functor::TopKFunctor;
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
diff --git a/tensorflow/core/util/batch_util.cc b/tensorflow/core/util/batch_util.cc
index b88c365ced0..e03188b04da 100644
--- a/tensorflow/core/util/batch_util.cc
+++ b/tensorflow/core/util/batch_util.cc
@@ -182,8 +182,6 @@ Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index) {
   switch (element.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopyElementToSlice Unhandled data type: ",
@@ -207,8 +205,6 @@ Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index) {
   switch (parent.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopySliceToElement Unhandled data type: ",
@@ -280,8 +276,6 @@ Status CopyContiguousSlices(const Tensor& src, int64 src_offset,
   switch (src.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopyContiguousSlices unhandled data type: ",
@@ -308,8 +302,6 @@ Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element, int64 index) {
   switch (parent->dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
-    TF_CALL_uint32(HANDLE_TYPE);
-    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented(
diff --git a/tensorflow/core/util/saved_tensor_slice_util.h b/tensorflow/core/util/saved_tensor_slice_util.h
index 09b9235b711..1f9768f5163 100644
--- a/tensorflow/core/util/saved_tensor_slice_util.h
+++ b/tensorflow/core/util/saved_tensor_slice_util.h
@@ -116,7 +116,9 @@ TENSOR_PROTO_EXTRACT_TYPE(double, double, double);
 TENSOR_PROTO_EXTRACT_TYPE_COMPLEX(complex64, scomplex, float);
 TENSOR_PROTO_EXTRACT_TYPE_COMPLEX(complex128, dcomplex, double);
 TENSOR_PROTO_EXTRACT_TYPE(int32, int, int32);
+TENSOR_PROTO_EXTRACT_TYPE(uint32, uint32, uint32);
 TENSOR_PROTO_EXTRACT_TYPE(int64, int64, protobuf_int64);
+TENSOR_PROTO_EXTRACT_TYPE(uint64, uint64, protobuf_uint64);
 TENSOR_PROTO_EXTRACT_TYPE(uint16, int, int32);
 TENSOR_PROTO_EXTRACT_TYPE(uint8, int, int32);
 TENSOR_PROTO_EXTRACT_TYPE(int8, int, int32);

From e5023a1738cce7efcdf9d87863b85c80ab2f8c9e Mon Sep 17 00:00:00 2001
From: Karim Nosir 
Date: Mon, 22 Jun 2020 23:10:17 -0700
Subject: [PATCH 0868/1390] - Fix type checking in elementwise.cc - Update
 error messages for some Abs Cast Ceil Cos Sin Not Square Sqrt RSqrt Log

PiperOrigin-RevId: 317807251
Change-Id: I2a4f359f04346551eda5a382b25f34bab2c73dc7
---
 tensorflow/lite/kernels/cast.cc        | 24 ++++++++-------
 tensorflow/lite/kernels/ceil.cc        |  4 +++
 tensorflow/lite/kernels/elementwise.cc | 42 ++++++++++++++++++--------
 tensorflow/lite/kernels/op_macros.h    |  9 ++++++
 4 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/tensorflow/lite/kernels/cast.cc b/tensorflow/lite/kernels/cast.cc
index 415f1270328..ab95afa979f 100644
--- a/tensorflow/lite/kernels/cast.cc
+++ b/tensorflow/lite/kernels/cast.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
@@ -67,8 +68,8 @@ void copyCast(const std::complex* in, std::complex* out,
 }
 
 template 
-TfLiteStatus copyToTensor(const FromT* in, TfLiteTensor* out,
-                          int num_elements) {
+TfLiteStatus copyToTensor(TfLiteContext* context, const FromT* in,
+                          TfLiteTensor* out, int num_elements) {
   switch (out->type) {
     case kTfLiteInt64:
       copyCast(in, out->data.i64, num_elements);
@@ -91,7 +92,7 @@ TfLiteStatus copyToTensor(const FromT* in, TfLiteTensor* out,
       break;
     default:
       // Unsupported type.
-      return kTfLiteError;
+      TF_LITE_UNSUPPORTED_TYPE(context, out->type, "Cast");
   }
   return kTfLiteOk;
 }
@@ -103,22 +104,23 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, num_elements, NumElements(output));
   switch (input->type) {
     case kTfLiteInt64:
-      return copyToTensor(input->data.i64, output, num_elements);
+      return copyToTensor(context, input->data.i64, output, num_elements);
     case kTfLiteInt32:
-      return copyToTensor(input->data.i32, output, num_elements);
+      return copyToTensor(context, input->data.i32, output, num_elements);
     case kTfLiteUInt8:
-      return copyToTensor(input->data.uint8, output, num_elements);
+      return copyToTensor(context, input->data.uint8, output, num_elements);
     case kTfLiteFloat32:
-      return copyToTensor(GetTensorData(input), output, num_elements);
+      return copyToTensor(context, GetTensorData(input), output,
+                          num_elements);
     case kTfLiteBool:
-      return copyToTensor(input->data.b, output, num_elements);
+      return copyToTensor(context, input->data.b, output, num_elements);
     case kTfLiteComplex64:
       return copyToTensor(
-          reinterpret_cast*>(input->data.c64), output,
-          num_elements);
+          context, reinterpret_cast*>(input->data.c64),
+          output, num_elements);
     default:
       // Unsupported type.
-      return kTfLiteError;
+      TF_LITE_UNSUPPORTED_TYPE(context, input->type, "Cast");
   }
   return kTfLiteOk;
 }
diff --git a/tensorflow/lite/kernels/ceil.cc b/tensorflow/lite/kernels/ceil.cc
index d8c6eaad7a4..95c660f3376 100644
--- a/tensorflow/lite/kernels/ceil.cc
+++ b/tensorflow/lite/kernels/ceil.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
@@ -41,6 +42,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  if (input->type != kTfLiteFloat32) {
+    TF_LITE_UNSUPPORTED_TYPE(context, input->type, "Ceil");
+  }
 
   optimized_ops::Ceil(GetTensorShape(input), GetTensorData(input),
                       GetTensorShape(output), GetTensorData(output));
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 1b91244af33..61c6aeaa811 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
 
 namespace tflite {
 namespace ops {
@@ -39,17 +40,15 @@ bool IsLogicalSupportedType(const TfLiteType type) {
 }
 
 typedef bool (*IsSupportedType)(TfLiteType);
-template 
+template 
 TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
   const TfLiteTensor* input = GetInput(context, node, 0);
   TfLiteTensor* output = GetOutput(context, node, 0);
   TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
-  if (!IsSupportedType(input->type)) {
-    context->ReportError(context, "Current data type %d is not supported.",
-                         input->type);
-    return kTfLiteError;
+  if (!is_supported_type(input->type)) {
+    TF_LITE_UNSUPPORTED_TYPE(context, input->type, op_name);
   }
   return context->ResizeTensor(context, output,
                                TfLiteIntArrayCopy(input->dims));
@@ -112,13 +111,23 @@ TfLiteStatus LogicalNotEval(TfLiteContext* context, TfLiteNode* node) {
   return EvalLogical(context, node, [](bool v) { return !v; });
 }
 
+constexpr char kAbsName[] = "Abs";
+constexpr char kSinName[] = "Sin";
+constexpr char kCosName[] = "Cos";
+constexpr char kLogName[] = "Log";
+constexpr char kSqrtName[] = "Sqrt";
+constexpr char kRsqrtName[] = "Rsqrt";
+constexpr char kSquareName[] = "Square";
+constexpr char kNotName[] = "Not";
+
 }  // namespace
 }  // namespace elementwise
 
 TfLiteRegistration* Register_ABS() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::AbsEval};
   return &r;
 }
@@ -126,7 +135,8 @@ TfLiteRegistration* Register_ABS() {
 TfLiteRegistration* Register_SIN() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::SinEval};
   return &r;
 }
@@ -134,7 +144,8 @@ TfLiteRegistration* Register_SIN() {
 TfLiteRegistration* Register_COS() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::CosEval};
   return &r;
 }
@@ -142,7 +153,8 @@ TfLiteRegistration* Register_COS() {
 TfLiteRegistration* Register_LOG() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::LogEval};
   return &r;
 }
@@ -150,7 +162,8 @@ TfLiteRegistration* Register_LOG() {
 TfLiteRegistration* Register_SQRT() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::SqrtEval};
   return &r;
 }
@@ -158,7 +171,8 @@ TfLiteRegistration* Register_SQRT() {
 TfLiteRegistration* Register_RSQRT() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::RsqrtEval};
   return &r;
 }
@@ -166,7 +180,8 @@ TfLiteRegistration* Register_RSQRT() {
 TfLiteRegistration* Register_SQUARE() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::SquareEval};
   return &r;
 }
@@ -174,7 +189,8 @@ TfLiteRegistration* Register_SQUARE() {
 TfLiteRegistration* Register_LOGICAL_NOT() {
   static TfLiteRegistration r = {
       /*init=*/nullptr, /*free=*/nullptr,
-      elementwise::GenericPrepare,
+      elementwise::GenericPrepare,
       elementwise::LogicalNotEval};
   return &r;
 }
diff --git a/tensorflow/lite/kernels/op_macros.h b/tensorflow/lite/kernels/op_macros.h
index 5c190f1c595..5786756f408 100644
--- a/tensorflow/lite/kernels/op_macros.h
+++ b/tensorflow/lite/kernels/op_macros.h
@@ -44,6 +44,15 @@ inline void InfiniteLoop() {
     fprintf(stderr, "%s", (x)); \
   } while (0)
 
+// Report Error for unsupported type by op 'op_name' and returns kTfLiteError.
+#define TF_LITE_UNSUPPORTED_TYPE(context, type, op_name)                    \
+  do {                                                                      \
+    TF_LITE_KERNEL_LOG((context), "%s:%d Type %s is unsupported by op %s.", \
+                       __FILE__, __LINE__, TfLiteTypeGetName(type),         \
+                       (op_name));                                          \
+    return kTfLiteError;                                                    \
+  } while (0)
+
 #define TFLITE_ABORT abort()
 
 #endif  // TF_LITE_MCU_DEBUG_LOG

From daa3c52aa1c301fc698574a2943f2a6eed2de08f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= 
Date: Fri, 13 Mar 2020 17:05:54 +0100
Subject: [PATCH 0869/1390] TFL: Port HARD_SWISH operator from TFLite to TFLu

Change-Id: I18092fa30dee33df4577a4abaf5321e7ba96a04c
---
 tensorflow/lite/kernels/activations.cc        |  16 -
 tensorflow/lite/kernels/internal/common.h     |  47 +++
 .../internal/reference/reference_ops.h        |  19 -
 tensorflow/lite/micro/all_ops_resolver.cc     |   1 +
 tensorflow/lite/micro/kernels/activations.cc  | 193 +++++++++-
 .../lite/micro/kernels/activations_test.cc    | 354 ++++++++++++++++++
 tensorflow/lite/micro/kernels/micro_ops.h     |   1 +
 .../lite/micro/micro_mutable_op_resolver.h    |   8 +
 tensorflow/lite/micro/micro_utils.h           |  10 +
 9 files changed, 612 insertions(+), 37 deletions(-)

diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 7ad33973b38..b6181973c7d 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -298,22 +298,6 @@ void HardSwishFree(TfLiteContext* context, void* buffer) {
   delete static_cast(buffer);
 }
 
-void DownScaleInt32ToInt16Multiplier(int32_t multiplier_int32,
-                                     int16_t* multiplier_int16) {
-  TFLITE_DCHECK_GE(multiplier_int32, 0);
-  static constexpr int32_t kRoundingOffset = 1 << 15;
-  if (multiplier_int32 >=
-      std::numeric_limits::max() - kRoundingOffset) {
-    *multiplier_int16 = std::numeric_limits::max();
-    return;
-  }
-  const int32_t result = (multiplier_int32 + kRoundingOffset) >> 16;
-  TFLITE_DCHECK_LE(result << 16, multiplier_int32 + kRoundingOffset);
-  TFLITE_DCHECK_GT(result << 16, multiplier_int32 - kRoundingOffset);
-  *multiplier_int16 = result;
-  TFLITE_DCHECK_EQ(*multiplier_int16, result);
-}
-
 TfLiteStatus HardSwishPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_STATUS(GenericPrepare(context, node));
   TfLiteTensor* output = GetOutput(context, node, 0);
diff --git a/tensorflow/lite/kernels/internal/common.h b/tensorflow/lite/kernels/internal/common.h
index c1db3587415..9bc65029f34 100644
--- a/tensorflow/lite/kernels/internal/common.h
+++ b/tensorflow/lite/kernels/internal/common.h
@@ -932,6 +932,53 @@ void optimized_ops_prefetch_write_l1_keep(const T* ptr) {
 #endif
 }
 
+// Similar to ARM instruction SQDMULH.
+// Similar to gemmlowp::SaturatingRoundingDoublingHighMul except
+// rounding to zero instead of to nearest (SQRDMULH).
+inline std::int16_t SaturatingDoublingHighMul(std::int16_t a, std::int16_t b) {
+  bool overflow = a == b && a == std::numeric_limits::min();
+  std::int32_t a_32(a);
+  std::int32_t b_32(b);
+  std::int32_t ab_32 = a_32 * b_32;
+  std::int16_t ab_x2_high16 = static_cast((ab_32) / (1 << 15));
+  return overflow ? std::numeric_limits::max() : ab_x2_high16;
+}
+
+// Similar to gemmlowp::SaturatingRoundingDoublingHighMul
+inline std::int16_t SaturatingRoundingDoublingHighMul(std::int16_t a, std::int16_t b) {
+  bool overflow = a == b && a == std::numeric_limits::min();
+  std::int32_t a_32(a);
+  std::int32_t b_32(b);
+  std::int32_t ab_32 = a_32 * b_32;
+  std::int16_t nudge = ab_32 >= 0 ? (1 << 14) : (1 - (1 << 14));
+  std::int16_t ab_x2_high16 =
+      static_cast((ab_32 + nudge) / (1 << 15));
+  return overflow ? std::numeric_limits::max() : ab_x2_high16;
+}
+
+inline int16_t SaturatingLeftShift(int16_t value, int amount) {
+  int32_t result = static_cast(value) * (1 << amount);
+  result = std::min(result, std::numeric_limits::max());
+  result = std::max(result, std::numeric_limits::min());
+  return result;
+}
+
+inline void DownScaleInt32ToInt16Multiplier(int32_t multiplier_int32,
+                                            int16_t* multiplier_int16) {
+  TFLITE_DCHECK_GE(multiplier_int32, 0);
+  static constexpr int32_t kRoundingOffset = 1 << 15;
+  if (multiplier_int32 >=
+      std::numeric_limits::max() - kRoundingOffset) {
+    *multiplier_int16 = std::numeric_limits::max();
+    return;
+  }
+  const int32_t result = (multiplier_int32 + kRoundingOffset) >> 16;
+  TFLITE_DCHECK_LE(result << 16, multiplier_int32 + kRoundingOffset);
+  TFLITE_DCHECK_GT(result << 16, multiplier_int32 - kRoundingOffset);
+  *multiplier_int16 = result;
+  TFLITE_DCHECK_EQ(*multiplier_int16, result);
+}
+
 }  // namespace tflite
 
 #endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_COMMON_H_
diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index 5208b21eb4d..b94c2060b2d 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -2615,25 +2615,6 @@ inline void HardSwish(const RuntimeShape& input_shape, const T* input_data,
   }
 }
 
-inline int16_t SaturatingLeftShift(int16_t value, int amount) {
-  int32_t result = static_cast(value) * (1 << amount);
-  result = std::min(result, std::numeric_limits::max());
-  result = std::max(result, std::numeric_limits::min());
-  return result;
-}
-
-// Similar to ARM instruction SQDMULH.
-// Similar to gemmlowp::SaturatingRoundingDoublingHighMul except
-// rounding to zero instead of to nearest (SQRDMULH).
-inline std::int16_t SaturatingDoublingHighMul(std::int16_t a, std::int16_t b) {
-  bool overflow = a == b && a == std::numeric_limits::min();
-  std::int32_t a_32(a);
-  std::int32_t b_32(b);
-  std::int32_t ab_32 = a_32 * b_32;
-  std::int16_t ab_x2_high16 = static_cast((ab_32) / (1 << 15));
-  return overflow ? std::numeric_limits::max() : ab_x2_high16;
-}
-
 template 
 inline void HardSwish(const HardSwishParams& params,
                       const RuntimeShape& input_shape, const T* input_data,
diff --git a/tensorflow/lite/micro/all_ops_resolver.cc b/tensorflow/lite/micro/all_ops_resolver.cc
index e728a95360a..ff461cb947e 100644
--- a/tensorflow/lite/micro/all_ops_resolver.cc
+++ b/tensorflow/lite/micro/all_ops_resolver.cc
@@ -42,6 +42,7 @@ AllOpsResolver::AllOpsResolver() {
   AddFullyConnected();
   AddGreater();
   AddGreaterEqual();
+  AddHardSwish();
   AddL2Normalization();
   AddLess();
   AddLessEqual();
diff --git a/tensorflow/lite/micro/kernels/activations.cc b/tensorflow/lite/micro/kernels/activations.cc
index 4a9b8ce5d8e..75b0339002d 100644
--- a/tensorflow/lite/micro/kernels/activations.cc
+++ b/tensorflow/lite/micro/kernels/activations.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/micro/micro_utils.h"
@@ -77,6 +78,157 @@ inline void Relu6Quantized(Q lower, Q upper, const RuntimeShape& input_shape,
   }
 }
 
+inline std::int32_t RoundingDivideByPOT(std::int32_t numerator, int exponent) {
+  std::int32_t sign = numerator >= 0 ? 1 : -1;
+  std::int32_t abs_numerator = std::abs(numerator);
+  std::int32_t mask = (1LL << exponent) - 1;
+  std::int32_t remainder = abs_numerator & mask;
+  std::int32_t threshold = mask >> 1;
+  std::int32_t abs_result =
+      (abs_numerator >> exponent) + (remainder > threshold ? 1 : 0);
+  return sign * abs_result;
+}
+
+inline void HardSwishFloatOp(const RuntimeShape& input_shape, const float* input_data,
+                           const RuntimeShape& output_shape, float* output_data) {
+  auto matching_size = MatchingFlatSize(input_shape, output_shape);
+  const float* in_end = input_data + matching_size;
+  for (; input_data < in_end; input_data++, output_data++) {
+    const float in = *input_data;
+    *output_data =
+        in * std::min(static_cast(6), std::max(static_cast(0), in + 3)) /
+        6;
+  }
+}
+
+template 
+void HardSwishOp(HardSwishParams& params,
+                      const RuntimeShape& input_shape, const T* input_data,
+                      const RuntimeShape& output_shape, T* output_data) {
+  const int flat_size = MatchingFlatSize(input_shape, output_shape);
+
+  for (int i = 0; i < flat_size; i++) {
+    const int16_t input_value = input_data[i] - params.input_zero_point;
+    // Left-shift as much as we can without overflow/saturation to put
+    // significant bits in the high bits of our 16-bit fixedpoint values, so
+    // that fixed-point approximate computations below are as accurate as
+    // possible.
+    const int16_t input_value_on_hires_input_scale = input_value << 7;
+    // Compute the input value on essentially the output scale, just not
+    // right-shifted yet. This is the value that we'll use in the (x >= +3)
+    // case, and that in the general case we'll multiply against the "relu-ish"
+    // fixed-point multiplier in [0, 1].
+    const int16_t input_value_on_preshift_output_scale =
+        SaturatingRoundingDoublingHighMul(input_value_on_hires_input_scale,
+                                          params.output_multiplier_fixedpoint_int16);
+    // Now compute the "relu-ish multiplier". In the (-3 <= x <= +3) case, that
+    // is just an affine rescaling of x from [-3, 3] to [0, 1]. In the general
+    // case, it is just that plus saturation at the boundaries of [-3, 3].
+    // First, we rescale from [-3, 3] to [-1, 1], saturating.
+    // That is done by rescaling the input value with a fixed-point multiplier
+    // (reluish_multiplier_fixedpoint) and bit-shift such that we represent
+    // that input value on the scale where the real value 3.0f is represented
+    // by the quantized value 32768.  (+32768 is actually not representable as
+    // int16, so this saturates at +32767, and that is seen empirically to be
+    // a negligible contribution to numerical error/bias).
+    //
+    // This code is careful to correctly implement any magnitude of multiplier,
+    // involving either a right shift or a left shift, with correct saturation
+    // behavior in the left-shift case. This forces this code to be more
+    // complicated, but is necessary for real applications: a partially
+    // trained quantized MobileNet v3-small model that motivated this code
+    // exhibits some large [min, max] range boundaries, of the order of
+    // magnitude of 10 or 100 depending on layers.
+    //
+    // The next few lines are basically just an ordinary
+    // MultiplyByQuantizedMultiplier, except that we are more careful here
+    // about the fine details of saturation when left-shifting, because here
+    // overflow in left-shift is a common case, not an anomaly as
+    // MultiplyByQuantizedMultiplier assumes.
+    int16_t reluish_value = input_value_on_hires_input_scale;
+    // Shift left, saturating, as much as we can while ensuring that this
+    // saturation will not contribute to the result. That is, left shift amount
+    // reduced by 1.
+    if (params.reluish_multiplier_exponent > 0) {
+      reluish_value = SaturatingLeftShift(
+          reluish_value, params.reluish_multiplier_exponent - 1);
+    }
+    // Apply the fixed-point multiplier, dividing the value by a divisor
+    // ranging in [1, 2].
+    reluish_value = SaturatingRoundingDoublingHighMul(reluish_value, params.reluish_multiplier_fixedpoint_int16);
+    // Apply the last bit of left-shift. Thus, in the left-shifting case, if
+    // any saturation affects the result, it is happening here --- any
+    // saturation having occurred above is overwritten here, not affecting the
+    // result.
+    if (params.reluish_multiplier_exponent > 0) {
+      reluish_value = SaturatingLeftShift(reluish_value, 1);
+    }
+    // Shift right, in the right-shifting case.
+    if (params.reluish_multiplier_exponent < 0) {
+      reluish_value = RoundingDivideByPOT(
+          reluish_value, -params.reluish_multiplier_exponent);
+    }
+    // At this point we have rescaled the value into a 16bit fixedpoint
+    // reluish_value in [-1, 1].
+    // We now convert that to a 16bit fixedpoint value in [0, 1].
+    reluish_value = (reluish_value + (1 << 15)) >> 1;
+    // Use of SaturatingDoublingHighMul here is important to cancel the biases
+    // from the above SaturatingRoundingDoublingHighMul.
+    //
+    const int16_t preshift_output_value = SaturatingDoublingHighMul(
+        reluish_value, input_value_on_preshift_output_scale);
+    // We were so far operating on the pre-shift output scale. Now we finally
+    // apply that output shift, arriving at the final output scale.
+    int16_t output_value = RoundingDivideByPOT(
+        preshift_output_value, -params.output_multiplier_exponent);
+    output_value += params.output_zero_point;
+    output_value =
+        std::min(output_value, std::numeric_limits::max());
+    output_value =
+        std::max(output_value, std::numeric_limits::min());
+    output_data[i] = output_value;
+  }
+}
+
+
+template 
+TfLiteStatus HardSwishQuantized(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  HardSwishParams params;
+
+   params.input_zero_point = input->params.zero_point;
+   params.output_zero_point = output->params.zero_point;
+
+   const float input_scale = input->params.scale;
+   const float hires_input_scale = (1.0f / 128.0f) * input_scale;
+   const float reluish_scale = 3.0f / 32768.0f;
+   const float output_scale = output->params.scale;
+
+   const double output_multiplier = static_cast(hires_input_scale / output_scale);
+   int32_t output_multiplier_fixedpoint_int32;
+   QuantizeMultiplier(output_multiplier, &output_multiplier_fixedpoint_int32,
+                      ¶ms.output_multiplier_exponent);
+   DownScaleInt32ToInt16Multiplier(
+       output_multiplier_fixedpoint_int32,
+       ¶ms.output_multiplier_fixedpoint_int16);
+
+   TF_LITE_ENSURE(context, params.output_multiplier_exponent <= 0);
+
+   const double reluish_multiplier = static_cast(hires_input_scale / reluish_scale);
+   int32_t reluish_multiplier_fixedpoint_int32;
+   QuantizeMultiplier(reluish_multiplier, &reluish_multiplier_fixedpoint_int32,
+                      ¶ms.reluish_multiplier_exponent);
+   DownScaleInt32ToInt16Multiplier(
+       reluish_multiplier_fixedpoint_int32,
+       ¶ms.reluish_multiplier_fixedpoint_int16);
+
+   HardSwishOp(params, GetTensorShape(input),
+                GetTensorData(input), GetTensorShape(output), GetTensorData(output));
+   return kTfLiteOk;
+}
+
 TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
   return kTfLiteOk;
 }
@@ -107,7 +259,7 @@ TfLiteStatus ReluEval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteOk;
     }
     default: {
-      TF_LITE_KERNEL_LOG(context, "Only float32 is supported currently, got %s",
+      TF_LITE_KERNEL_LOG(context, "Only float32/int8/uint8 are supported currently, got %s",
                          TfLiteTypeGetName(input->type));
       return kTfLiteError;
     }
@@ -148,13 +300,43 @@ TfLiteStatus Relu6Eval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteOk;
     }
     default: {
-      TF_LITE_KERNEL_LOG(context, "Only float32 is supported currently, got %s",
+      TF_LITE_KERNEL_LOG(context, "Only float32/int8/uint8 are supported currently, got %s",
                          TfLiteTypeGetName(input->type));
       return kTfLiteError;
     }
   }
 }
 
+TfLiteStatus HardSwishPrepare(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus HardSwishEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  switch (input->type) {
+    case kTfLiteFloat32: {
+      HardSwishFloatOp(
+          GetTensorShape(input),
+          GetTensorData(input),
+          GetTensorShape(output),
+          GetTensorData(output));
+      return kTfLiteOk;
+    } break;
+    case kTfLiteUInt8: {
+      return HardSwishQuantized(context, node);
+    } break;
+    case kTfLiteInt8: {
+      return HardSwishQuantized(context, node);
+    } break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Only float32/int8/uint8 are supported currently, got %s",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+  }
+}
+
 }  // namespace activations
 
 TfLiteRegistration* Register_RELU() {
@@ -181,6 +363,13 @@ TfLiteRegistration* Register_RELU6() {
   return &r;
 }
 
+TfLiteRegistration* Register_HARD_SWISH() {
+  static TfLiteRegistration r = {};
+  r.prepare = activations::HardSwishPrepare;
+  r.invoke = activations::HardSwishEval;
+  return &r;
+}
+
 }  // namespace micro
 }  // namespace ops
 }  // namespace tflite
diff --git a/tensorflow/lite/micro/kernels/activations_test.cc b/tensorflow/lite/micro/kernels/activations_test.cc
index 221f8f66d58..0a7db1e6589 100644
--- a/tensorflow/lite/micro/kernels/activations_test.cc
+++ b/tensorflow/lite/micro/kernels/activations_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include 
+
 #include "tensorflow/lite/c/builtin_op_data.h"
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/micro/all_ops_resolver.h"
@@ -135,6 +137,318 @@ void TestRelu6Float(const int* input_dims_data, const float* input_data,
   }
 }
 
+void GenerateUniformRandomVector(int size, float min, float max,
+                                 std::minstd_rand* random_engine,
+                                 std::vector* result) {
+  // Never use std::uniform_*_distribution in tests, it's
+  // implementation-defined. Likewise, don't use std::default_random_engine,
+  // implementation-defined. Implementation-defined is bad because it means that
+  // any toolchain update or new platform may run into test failures.
+  // std::minstd_rand is a standard instantiation of
+  // std::linear_congruential_engine, the cheapest generator in c++11 stdlib,
+  // it's good enough here.
+  result->resize(size);
+  for (int i = 0; i < size; i++) {
+    // We don't care whether the `max` value may ever be produced exactly.
+    // It may actually be thanks to rounding, as std::minstd_rand::modulus
+    // is 2^31 - 1 is greater than the inverse float epsilon.
+    float random_value_scaled_0_1 =
+        (*random_engine)() *
+        (1.0f / static_cast(std::minstd_rand::modulus));
+    (*result)[i] = min + (max - min) * random_value_scaled_0_1;
+  }
+}
+
+void EvalTestReferenceHardSwish(int size, const std::vector& input,
+                                std::vector* result) {
+  result->resize(size);
+  for (int i = 0; i < size; i++) {
+    const float in = input[i];
+    (*result)[i] = in * std::min(6.0f, std::max(0.0f, in + 3)) * (1.0f / 6.0f);
+  }
+}
+
+template 
+void TestHardSwishQuantized(int size, float input_min,
+                            float input_max, float output_min,
+                            float output_max, std::minstd_rand* random_engine) {
+  T output_data[size];
+  T input_data_quantized[size];
+  const int input_dims_data[] = {2, 1, size};
+  const int output_dims_data[] = {2, 1, size};
+  const float input_scale = ScaleFromMinMax(input_min, input_max);
+  const int input_zero_point = ZeroPointFromMinMax(input_min, input_max);
+  const float output_scale = ScaleFromMinMax(output_min, output_max);
+  const int output_zero_point = ZeroPointFromMinMax(output_min, output_max);
+
+  // The numerical error for any 8bit quantized function is at least one half
+  // times the quantization step: 0.5 * (kOutMax - kOutMin) / 256.
+  // To that we add again the quantization step (kOutMax - kOutMin) / 256
+  // to allow for an off-by-one rounding error.
+  const float kTolerance = std::max(input_max - input_min, output_max - output_min) * (1.5f / 256.f);
+
+  TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
+  TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+  const int output_elements_count = ElementCount(*output_dims);
+
+  TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
+
+  float dequantized_output[output_elements_count];
+
+  std::vector float_input_values;
+  std::vector float_ref_output_values;
+  GenerateUniformRandomVector(size, input_min, input_max, random_engine,
+                              &float_input_values);
+  EvalTestReferenceHardSwish(size, float_input_values,
+                             &float_ref_output_values);
+  for (float& val : float_ref_output_values) {
+    val = std::min(output_max, std::max(output_min, val));
+  }
+
+  constexpr int inputs_size = 1;
+  constexpr int outputs_size = 1;
+  constexpr int tensors_size = inputs_size + outputs_size;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateQuantizedTensor(float_input_values.data(), input_data_quantized, input_dims,
+                            input_scale, input_zero_point, "input_tensor"),
+      CreateQuantizedTensor(output_data, output_dims, output_scale,
+                            output_zero_point, "output_tensor"),
+  };
+
+  TfLiteContext context;
+  PopulateContext(tensors, tensors_size, micro_test::reporter, &context);
+
+  ::tflite::AllOpsResolver resolver;
+  const TfLiteRegistration* registration =
+      resolver.FindOp(tflite::BuiltinOperator_HARD_SWISH);
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration);
+
+  const char* init_data = nullptr;
+  size_t init_data_size = 0;
+  void* user_data = nullptr;
+  if (registration->init) {
+    user_data = registration->init(&context, init_data, init_data_size);
+  }
+  int inputs_array_data[] = {1, 0};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 1};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+  TfLiteNode node;
+  node.inputs = inputs_array;
+  node.outputs = outputs_array;
+  node.temporaries = nullptr;
+  node.user_data = user_data;
+  node.builtin_data = nullptr;
+  node.custom_initial_data = nullptr;
+  node.custom_initial_data_size = 0;
+  node.delegate = nullptr;
+  if (registration->prepare) {
+    TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node));
+  }
+
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node));
+
+  if (registration->free) {
+    registration->free(&context, user_data);
+  }
+
+  AsymmetricDequantize(output_data, output_elements_count, output_scale, output_zero_point, dequantized_output);
+
+  for (int i = 0; i < output_elements_count; ++i) {
+    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values.data()[i], dequantized_output[i], kTolerance);
+  }
+}
+
+template 
+void TestHardSwishQuantizedBias(float input_min, float input_max, float output_min,
+                                float output_max, float tolerated_bias) {
+  const float quantized_type_range =
+      static_cast(std::numeric_limits::max()) -
+      static_cast(std::numeric_limits::min());
+
+
+  const float input_scale = ScaleFromMinMax(input_min, input_max);
+  const float output_scale = ScaleFromMinMax(output_min, output_max);
+
+  const int input_zero_point = ZeroPointFromMinMax(input_min, input_max);
+  const int output_zero_point = ZeroPointFromMinMax(output_min, output_max);
+
+  const float max_scale = std::max(output_scale, input_scale);
+
+  // In this bias-focused test case, no need for randomly generated input
+  // values.
+  TF_LITE_MICRO_EXPECT_LE(input_min, -3.0f);
+  TF_LITE_MICRO_EXPECT_GE(input_max, 3.0f);
+  const int quantized_input_negative_three =
+      std::round(std::numeric_limits::min() +
+                 (-3.0f - input_min) / input_scale);
+  const int quantized_input_positive_three =
+      std::round(std::numeric_limits::min() +
+                 (3.0f - input_min) / input_scale);
+  std::vector float_input_values;
+  for (int i = quantized_input_negative_three;
+       i <= quantized_input_positive_three; i++) {
+    float_input_values.push_back(
+        input_min +
+        (i - std::numeric_limits::min()) * input_scale);
+  }
+  const int size = float_input_values.size();
+  std::vector float_ref_output_values;
+  EvalTestReferenceHardSwish(size, float_input_values,
+                             &float_ref_output_values);
+  for (float& val : float_ref_output_values) {
+    val = std::min(output_max, std::max(output_min, val));
+  }
+
+  T output_data[size];
+  T input_data_quantized[size];
+  const int input_dims_data[] = {2, 1, size};
+  const int output_dims_data[] = {2, 1, size};
+
+  // The numerical error for any 8bit quantized function is at least one half
+  // times the quantization step: 0.5 * (kOutMax - kOutMin) / 256.
+  // To that we add again the quantization step (kOutMax - kOutMin) / 256
+  // to allow for an off-by-one rounding error.
+  const float kTolerance = std::max(input_max - input_min, output_max - output_min) * (1.5f / 256.f);
+
+  TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
+  TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+  const int output_elements_count = ElementCount(*output_dims);
+
+  TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
+
+  float dequantized_output[output_elements_count];
+
+  constexpr int inputs_size = 1;
+  constexpr int outputs_size = 1;
+  constexpr int tensors_size = inputs_size + outputs_size;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateQuantizedTensor(float_input_values.data(), input_data_quantized, input_dims,
+                            input_scale, input_zero_point, "input_tensor"),
+      CreateQuantizedTensor(output_data, output_dims, output_scale,
+                            output_zero_point, "output_tensor"),
+  };
+
+  TfLiteContext context;
+  PopulateContext(tensors, tensors_size, micro_test::reporter, &context);
+
+  ::tflite::AllOpsResolver resolver;
+  const TfLiteRegistration* registration =
+      resolver.FindOp(tflite::BuiltinOperator_HARD_SWISH);
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration);
+
+  const char* init_data = nullptr;
+  size_t init_data_size = 0;
+  void* user_data = nullptr;
+  if (registration->init) {
+    user_data = registration->init(&context, init_data, init_data_size);
+  }
+  int inputs_array_data[] = {1, 0};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 1};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+  TfLiteNode node;
+  node.inputs = inputs_array;
+  node.outputs = outputs_array;
+  node.temporaries = nullptr;
+  node.user_data = user_data;
+  node.builtin_data = nullptr;
+  node.custom_initial_data = nullptr;
+  node.custom_initial_data_size = 0;
+  node.delegate = nullptr;
+  if (registration->prepare) {
+    TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node));
+  }
+
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node));
+
+  if (registration->free) {
+    registration->free(&context, user_data);
+  }
+
+  AsymmetricDequantize(output_data, output_elements_count, output_scale, output_zero_point, dequantized_output);
+
+  float sum_diff = 0;
+  for (int i = 0; i < size; i++) {
+    sum_diff += dequantized_output[i] - float_ref_output_values[i];
+  }
+  const float bias = sum_diff / (size * max_scale);
+  TF_LITE_MICRO_EXPECT_LE(std::abs(bias), tolerated_bias);
+}
+
+void TestHardSwishFloat(int size, std::minstd_rand* random_engine) {
+  std::vector float_input_values;
+  const float kMin = -10.0f;
+  const float kMax = 10.0f;
+  GenerateUniformRandomVector(size, kMin, kMax, random_engine,
+                              &float_input_values);
+  std::vector float_ref_output_values;
+  EvalTestReferenceHardSwish(size, float_input_values,
+                             &float_ref_output_values);
+
+  float output_data[size];
+  const int input_dims_data[] = {1, size};
+  const int output_dims_data[] = {1, size};
+
+  TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
+  TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+  const int output_elements_count = ElementCount(*output_dims);
+
+  TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
+
+  constexpr int inputs_size = 1;
+  constexpr int outputs_size = 1;
+  constexpr int tensors_size = inputs_size + outputs_size;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateFloatTensor(float_input_values.data(), input_dims, "input_tensor"),
+      CreateFloatTensor(output_data, output_dims, "output_tensor"),
+  };
+  TfLiteContext context;
+  PopulateContext(tensors, tensors_size, micro_test::reporter, &context);
+
+  ::tflite::AllOpsResolver resolver;
+  const TfLiteRegistration* registration =
+      resolver.FindOp(tflite::BuiltinOperator_HARD_SWISH);
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration);
+
+  const char* init_data = nullptr;
+  size_t init_data_size = 0;
+  void* user_data = nullptr;
+  if (registration->init) {
+    user_data = registration->init(&context, init_data, init_data_size);
+  }
+  int inputs_array_data[] = {1, 0};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 1};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+  TfLiteNode node;
+  node.inputs = inputs_array;
+  node.outputs = outputs_array;
+  node.temporaries = nullptr;
+  node.user_data = user_data;
+  node.builtin_data = nullptr;
+  node.custom_initial_data = nullptr;
+  node.custom_initial_data_size = 0;
+  node.delegate = nullptr;
+  if (registration->prepare) {
+    TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node));
+  }
+  TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke);
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node));
+
+  if (registration->free) {
+    registration->free(&context, user_data);
+  }
+
+  for (int i = 0; i < output_elements_count; ++i) {
+    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values.data()[i], output_data[i], 1e-5f);
+  }
+}
+
 void TestReluUint8(const int* input_dims_data, const float* input_data,
                    uint8_t* input_data_quantized, const float input_scale,
                    const int input_zero_point, const float* golden,
@@ -431,6 +745,46 @@ TF_LITE_MICRO_TEST(SimpleRelu6TestFloat) {
                                   output_data);
 }
 
+TF_LITE_MICRO_TEST(SimpleHardSwishTestFloat) {
+  std::minstd_rand random_engine;
+  for (int size : {1, 2, 3, 4, 10, 20, 30, 40, 100}) {
+    tflite::testing::TestHardSwishFloat(size, &random_engine);
+  }
+}
+
+TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantized) {
+  std::minstd_rand random_engine;
+  std::vector> minmax_pairs{
+      {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
+  for (const auto& input_minmax : minmax_pairs) {
+    for (const auto& output_minmax : minmax_pairs) {
+      float input_min = input_minmax.first;
+      float input_max = input_minmax.second;
+      float output_min = output_minmax.first;
+      float output_max = output_minmax.second;
+      for (int size : {1, 3, 10, 100}) {
+        tflite::testing::TestHardSwishQuantized(size, input_min, input_max,
+                                                        output_min, output_max,
+                                                        &random_engine);
+        tflite::testing::TestHardSwishQuantized(size, input_min, input_max,
+                                                         output_min, output_max,
+                                                         &random_engine);
+      }
+    }
+  }
+}
+
+// See the comment in the reference implementation of quantized HardSwish:
+// A numerical issue significantly affecting ImageNet classification accuracy
+// with MobileNet v3 is only observable at the scale of HardSwish unit tests
+// if we monitor specifically bias. This testcase is extracted from one of the
+// HardSwish nodes in that MobileNet v3 that exhibited this issue.
+TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantizedBias) {
+  tflite::testing::TestHardSwishQuantizedBias(-11.654928f, 25.036512f,
+                                                       -0.3905796f, 24.50887f, 0.035);
+}
+
+
 TF_LITE_MICRO_TEST(SimpleReluTestUint8) {
   const int elements_count = 10;
 
diff --git a/tensorflow/lite/micro/kernels/micro_ops.h b/tensorflow/lite/micro/kernels/micro_ops.h
index 24180aab8c5..25f9bd41d63 100644
--- a/tensorflow/lite/micro/kernels/micro_ops.h
+++ b/tensorflow/lite/micro/kernels/micro_ops.h
@@ -46,6 +46,7 @@ TfLiteRegistration* Register_FLOOR();
 TfLiteRegistration* Register_FULLY_CONNECTED();
 TfLiteRegistration* Register_GREATER();
 TfLiteRegistration* Register_GREATER_EQUAL();
+TfLiteRegistration* Register_HARD_SWISH();
 TfLiteRegistration* Register_LESS();
 TfLiteRegistration* Register_LESS_EQUAL();
 TfLiteRegistration* Register_LOG();
diff --git a/tensorflow/lite/micro/micro_mutable_op_resolver.h b/tensorflow/lite/micro/micro_mutable_op_resolver.h
index 1b76f440a61..d77dc9577fd 100644
--- a/tensorflow/lite/micro/micro_mutable_op_resolver.h
+++ b/tensorflow/lite/micro/micro_mutable_op_resolver.h
@@ -217,6 +217,14 @@ class MicroMutableOpResolver : public MicroOpResolver {
                       ParseOpData);
   }
 
+  TfLiteStatus AddHardSwish() {
+    // TODO(b/149408647): Replace ParseOpData with the operator specific parse
+    // function.
+    return AddBuiltin(BuiltinOperator_HARD_SWISH,
+                      *tflite::ops::micro::Register_HARD_SWISH(),
+                      ParseOpData);
+  }
+
   TfLiteStatus AddL2Normalization() {
     // TODO(b/149408647): Replace ParseOpData with the operator specific parse
     // function.
diff --git a/tensorflow/lite/micro/micro_utils.h b/tensorflow/lite/micro/micro_utils.h
index 4f8689b943e..0a4ab1ef09c 100644
--- a/tensorflow/lite/micro/micro_utils.h
+++ b/tensorflow/lite/micro/micro_utils.h
@@ -94,6 +94,16 @@ void SymmetricDequantize(const int8_t* values, const int size,
                          const float dequantization_scale,
                          float* dequantized_values);
 
+template 
+void AsymmetricDequantize(const T* values, const int size,
+                          const float dequantization_scale,
+                          int dequantization_zero_point,
+                          float* dequantized_values) {
+  for (int i = 0; i < size; ++i) {
+    dequantized_values[i] = (values[i] - dequantization_zero_point) * dequantization_scale;
+  }
+}
+
 }  // namespace tflite
 
 #endif  // TENSORFLOW_LITE_MICRO_MICRO_UTILS_H_

From 820257026fa125ee102948986f0e6826b6bee883 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= 
Date: Tue, 2 Jun 2020 09:35:02 +0200
Subject: [PATCH 0870/1390] TFLu: Fix variable length array error in hard_swish
 unit test

---
 .../lite/micro/kernels/activations_test.cc    | 87 ++++++++++++-------
 1 file changed, 57 insertions(+), 30 deletions(-)

diff --git a/tensorflow/lite/micro/kernels/activations_test.cc b/tensorflow/lite/micro/kernels/activations_test.cc
index 0a7db1e6589..03a6b698a86 100644
--- a/tensorflow/lite/micro/kernels/activations_test.cc
+++ b/tensorflow/lite/micro/kernels/activations_test.cc
@@ -169,11 +169,9 @@ void EvalTestReferenceHardSwish(int size, const std::vector& input,
 }
 
 template 
-void TestHardSwishQuantized(int size, float input_min,
-                            float input_max, float output_min,
+void TestHardSwishQuantized(int size, const T* output_data, T* input_data_quantized, float* dequantized_output,
+                            float input_min, float input_max, float output_min,
                             float output_max, std::minstd_rand* random_engine) {
-  T output_data[size];
-  T input_data_quantized[size];
   const int input_dims_data[] = {2, 1, size};
   const int output_dims_data[] = {2, 1, size};
   const float input_scale = ScaleFromMinMax(input_min, input_max);
@@ -193,8 +191,6 @@ void TestHardSwishQuantized(int size, float input_min,
 
   TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
 
-  float dequantized_output[output_elements_count];
-
   std::vector float_input_values;
   std::vector float_ref_output_values;
   GenerateUniformRandomVector(size, input_min, input_max, random_engine,
@@ -262,8 +258,9 @@ void TestHardSwishQuantized(int size, float input_min,
 }
 
 template 
-void TestHardSwishQuantizedBias(float input_min, float input_max, float output_min,
-                                float output_max, float tolerated_bias) {
+void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_data_quantized,
+                                float* dequantized_output, float input_min, float input_max,
+                                float output_min, float output_max, float tolerated_bias) {
   const float quantized_type_range =
       static_cast(std::numeric_limits::max()) -
       static_cast(std::numeric_limits::min());
@@ -294,7 +291,8 @@ void TestHardSwishQuantizedBias(float input_min, float input_max, float output_m
         input_min +
         (i - std::numeric_limits::min()) * input_scale);
   }
-  const int size = float_input_values.size();
+  TF_LITE_MICRO_EXPECT_EQ(float_input_values.size(), size);
+
   std::vector float_ref_output_values;
   EvalTestReferenceHardSwish(size, float_input_values,
                              &float_ref_output_values);
@@ -302,8 +300,6 @@ void TestHardSwishQuantizedBias(float input_min, float input_max, float output_m
     val = std::min(output_max, std::max(output_min, val));
   }
 
-  T output_data[size];
-  T input_data_quantized[size];
   const int input_dims_data[] = {2, 1, size};
   const int output_dims_data[] = {2, 1, size};
 
@@ -319,8 +315,6 @@ void TestHardSwishQuantizedBias(float input_min, float input_max, float output_m
 
   TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
 
-  float dequantized_output[output_elements_count];
-
   constexpr int inputs_size = 1;
   constexpr int outputs_size = 1;
   constexpr int tensors_size = inputs_size + outputs_size;
@@ -370,7 +364,8 @@ void TestHardSwishQuantizedBias(float input_min, float input_max, float output_m
     registration->free(&context, user_data);
   }
 
-  AsymmetricDequantize(output_data, output_elements_count, output_scale, output_zero_point, dequantized_output);
+  AsymmetricDequantize(output_data, output_elements_count, output_scale,
+                          output_zero_point, dequantized_output);
 
   float sum_diff = 0;
   for (int i = 0; i < size; i++) {
@@ -380,7 +375,7 @@ void TestHardSwishQuantizedBias(float input_min, float input_max, float output_m
   TF_LITE_MICRO_EXPECT_LE(std::abs(bias), tolerated_bias);
 }
 
-void TestHardSwishFloat(int size, std::minstd_rand* random_engine) {
+void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* random_engine) {
   std::vector float_input_values;
   const float kMin = -10.0f;
   const float kMax = 10.0f;
@@ -390,7 +385,6 @@ void TestHardSwishFloat(int size, std::minstd_rand* random_engine) {
   EvalTestReferenceHardSwish(size, float_input_values,
                              &float_ref_output_values);
 
-  float output_data[size];
   const int input_dims_data[] = {1, size};
   const int output_dims_data[] = {1, size};
 
@@ -747,29 +741,57 @@ TF_LITE_MICRO_TEST(SimpleRelu6TestFloat) {
 
 TF_LITE_MICRO_TEST(SimpleHardSwishTestFloat) {
   std::minstd_rand random_engine;
-  for (int size : {1, 2, 3, 4, 10, 20, 30, 40, 100}) {
-    tflite::testing::TestHardSwishFloat(size, &random_engine);
-  }
+  constexpr int size = 100;
+  float output_data[size] = {0.f};
+
+  tflite::testing::TestHardSwishFloat(size, output_data, &random_engine);
 }
 
-TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantized) {
+
+TF_LITE_MICRO_TEST(SimpleHardSwishTestInt8) {
   std::minstd_rand random_engine;
   std::vector> minmax_pairs{
       {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
+  constexpr int size = 101;
+  constexpr int8_t output_data[size] = {0};
+  int8_t input_data_quantized[size] = {0};
+  float dequantized_output[size] = {0.f};
+
   for (const auto& input_minmax : minmax_pairs) {
     for (const auto& output_minmax : minmax_pairs) {
       float input_min = input_minmax.first;
       float input_max = input_minmax.second;
       float output_min = output_minmax.first;
       float output_max = output_minmax.second;
-      for (int size : {1, 3, 10, 100}) {
-        tflite::testing::TestHardSwishQuantized(size, input_min, input_max,
-                                                        output_min, output_max,
-                                                        &random_engine);
-        tflite::testing::TestHardSwishQuantized(size, input_min, input_max,
-                                                         output_min, output_max,
-                                                         &random_engine);
-      }
+
+      tflite::testing::TestHardSwishQuantized(size, output_data, input_data_quantized, dequantized_output,
+                                                      input_min, input_max, output_min, output_max,
+                                                      &random_engine);
+
+    }
+  }
+}
+
+TF_LITE_MICRO_TEST(SimpleHardSwishTestUint8) {
+  std::minstd_rand random_engine;
+  std::vector> minmax_pairs{
+      {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
+  constexpr int size = 99;
+  constexpr uint8_t output_data[size] = {0};
+  uint8_t input_data_quantized[size] = {0};
+  float dequantized_output[size] = {0.f};
+
+  for (const auto& input_minmax : minmax_pairs) {
+    for (const auto& output_minmax : minmax_pairs) {
+      float input_min = input_minmax.first;
+      float input_max = input_minmax.second;
+      float output_min = output_minmax.first;
+      float output_max = output_minmax.second;
+
+      tflite::testing::TestHardSwishQuantized(size, output_data, input_data_quantized, dequantized_output,
+                                                       input_min, input_max, output_min, output_max,
+                                                       &random_engine);
+
     }
   }
 }
@@ -780,8 +802,13 @@ TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantized) {
 // if we monitor specifically bias. This testcase is extracted from one of the
 // HardSwish nodes in that MobileNet v3 that exhibited this issue.
 TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantizedBias) {
-  tflite::testing::TestHardSwishQuantizedBias(-11.654928f, 25.036512f,
-                                                       -0.3905796f, 24.50887f, 0.035);
+  constexpr int size = 43;
+  constexpr uint8_t output_data[size] = {0};
+  uint8_t input_data_quantized[size] = {0};
+  float dequantized_output[size] = {0.f};
+
+  tflite::testing::TestHardSwishQuantizedBias(size, output_data, input_data_quantized, dequantized_output,
+                                                       -11.654928f, 25.036512f, -0.3905796f, 24.50887f, 0.035);
 }
 
 

From c2fc5dccfa021fa8f7ec48da874c0d91a2d211c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= 
Date: Wed, 10 Jun 2020 11:14:27 +0200
Subject: [PATCH 0871/1390] TFLu: Update hard swish op as per review comments

---
 tensorflow/lite/kernels/internal/BUILD        |   2 +
 .../kernels/internal/reference/activations.h  |  54 +++++++
 .../internal/reference/reference_ops.h        |  15 +-
 tensorflow/lite/micro/kernels/activations.cc  | 150 +++++++++---------
 .../lite/micro/kernels/activations_test.cc    | 121 +++++++-------
 5 files changed, 196 insertions(+), 146 deletions(-)
 create mode 100644 tensorflow/lite/kernels/internal/reference/activations.h

diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD
index a02a5bf3981..31a6a74eb5e 100644
--- a/tensorflow/lite/kernels/internal/BUILD
+++ b/tensorflow/lite/kernels/internal/BUILD
@@ -437,6 +437,7 @@ cc_library(
     name = "reference_base",
     srcs = [],
     hdrs = [
+        "reference/activations.h",
         "reference/add.h",
         "reference/arg_min_max.h",
         "reference/batch_matmul.h",
@@ -525,6 +526,7 @@ cc_library(
     name = "legacy_reference_base",
     srcs = [],
     hdrs = [
+        "reference/activations.h",
         "reference/add.h",
         "reference/arg_min_max.h",
         "reference/binary_function.h",
diff --git a/tensorflow/lite/kernels/internal/reference/activations.h b/tensorflow/lite/kernels/internal/reference/activations.h
new file mode 100644
index 00000000000..2110f5e3cbc
--- /dev/null
+++ b/tensorflow/lite/kernels/internal/reference/activations.h
@@ -0,0 +1,54 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ACTIVATIONS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ACTIVATIONS_H_
+
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+
+
+namespace tflite {
+namespace reference_ops {
+
+inline std::int32_t RoundingDivideByPOT(std::int32_t numerator, int exponent) {
+  std::int32_t sign = numerator >= 0 ? 1 : -1;
+  std::int32_t abs_numerator = std::abs(numerator);
+  std::int32_t mask = (1LL << exponent) - 1;
+  std::int32_t remainder = abs_numerator & mask;
+  std::int32_t threshold = mask >> 1;
+  std::int32_t abs_result =
+      (abs_numerator >> exponent) + (remainder > threshold ? 1 : 0);
+  return sign * abs_result;
+}
+
+template 
+inline void HardSwish(const RuntimeShape& input_shape, const T* input_data,
+                      const RuntimeShape& output_shape, T* output_data) {
+  ruy::profiler::ScopeLabel label("ReferenceHardSwish/Float");
+  auto matching_size = MatchingFlatSize(input_shape, output_shape);
+  const T* in_end = input_data + matching_size;
+  for (; input_data < in_end; input_data++, output_data++) {
+    const float in = *input_data;
+    *output_data =
+        in * std::min(static_cast(6), std::max(static_cast(0), in + 3)) /
+        6;
+  }
+}
+
+}  // namespace reference_ops
+}  // namespace tflite
+
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_CONV_H_
diff --git a/tensorflow/lite/kernels/internal/reference/reference_ops.h b/tensorflow/lite/kernels/internal/reference/reference_ops.h
index b94c2060b2d..3e9ebff69ad 100644
--- a/tensorflow/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/lite/kernels/internal/reference/reference_ops.h
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/activations.h"
 #include "tensorflow/lite/kernels/internal/reference/add.h"
 #include "tensorflow/lite/kernels/internal/reference/arg_min_max.h"
 #include "tensorflow/lite/kernels/internal/reference/binary_function.h"
@@ -2601,20 +2602,6 @@ void ReverseSequence(const TS* seq_lengths, const int seq_dim,
   }
 }
 
-template 
-inline void HardSwish(const RuntimeShape& input_shape, const T* input_data,
-                      const RuntimeShape& output_shape, T* output_data) {
-  ruy::profiler::ScopeLabel label("ReferenceHardSwish/Float");
-  auto matching_size = MatchingFlatSize(input_shape, output_shape);
-  const T* in_end = input_data + matching_size;
-  for (; input_data < in_end; input_data++, output_data++) {
-    const float in = *input_data;
-    *output_data =
-        in * std::min(static_cast(6), std::max(static_cast(0), in + 3)) /
-        6;
-  }
-}
-
 template 
 inline void HardSwish(const HardSwishParams& params,
                       const RuntimeShape& input_shape, const T* input_data,
diff --git a/tensorflow/lite/micro/kernels/activations.cc b/tensorflow/lite/micro/kernels/activations.cc
index 75b0339002d..934f029ac6a 100644
--- a/tensorflow/lite/micro/kernels/activations.cc
+++ b/tensorflow/lite/micro/kernels/activations.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/micro/micro_utils.h"
@@ -78,33 +79,10 @@ inline void Relu6Quantized(Q lower, Q upper, const RuntimeShape& input_shape,
   }
 }
 
-inline std::int32_t RoundingDivideByPOT(std::int32_t numerator, int exponent) {
-  std::int32_t sign = numerator >= 0 ? 1 : -1;
-  std::int32_t abs_numerator = std::abs(numerator);
-  std::int32_t mask = (1LL << exponent) - 1;
-  std::int32_t remainder = abs_numerator & mask;
-  std::int32_t threshold = mask >> 1;
-  std::int32_t abs_result =
-      (abs_numerator >> exponent) + (remainder > threshold ? 1 : 0);
-  return sign * abs_result;
-}
-
-inline void HardSwishFloatOp(const RuntimeShape& input_shape, const float* input_data,
-                           const RuntimeShape& output_shape, float* output_data) {
-  auto matching_size = MatchingFlatSize(input_shape, output_shape);
-  const float* in_end = input_data + matching_size;
-  for (; input_data < in_end; input_data++, output_data++) {
-    const float in = *input_data;
-    *output_data =
-        in * std::min(static_cast(6), std::max(static_cast(0), in + 3)) /
-        6;
-  }
-}
-
 template 
 void HardSwishOp(HardSwishParams& params,
-                      const RuntimeShape& input_shape, const T* input_data,
-                      const RuntimeShape& output_shape, T* output_data) {
+                 const RuntimeShape& input_shape, const T* input_data,
+                 const RuntimeShape& output_shape, T* output_data) {
   const int flat_size = MatchingFlatSize(input_shape, output_shape);
 
   for (int i = 0; i < flat_size; i++) {
@@ -165,7 +143,7 @@ void HardSwishOp(HardSwishParams& params,
     }
     // Shift right, in the right-shifting case.
     if (params.reluish_multiplier_exponent < 0) {
-      reluish_value = RoundingDivideByPOT(
+      reluish_value = tflite::reference_ops::RoundingDivideByPOT(
           reluish_value, -params.reluish_multiplier_exponent);
     }
     // At this point we have rescaled the value into a 16bit fixedpoint
@@ -179,7 +157,7 @@ void HardSwishOp(HardSwishParams& params,
         reluish_value, input_value_on_preshift_output_scale);
     // We were so far operating on the pre-shift output scale. Now we finally
     // apply that output shift, arriving at the final output scale.
-    int16_t output_value = RoundingDivideByPOT(
+    int16_t output_value = tflite::reference_ops::RoundingDivideByPOT(
         preshift_output_value, -params.output_multiplier_exponent);
     output_value += params.output_zero_point;
     output_value =
@@ -190,46 +168,11 @@ void HardSwishOp(HardSwishParams& params,
   }
 }
 
-
-template 
-TfLiteStatus HardSwishQuantized(TfLiteContext* context, TfLiteNode* node) {
-  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
-  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
-
-  HardSwishParams params;
-
-   params.input_zero_point = input->params.zero_point;
-   params.output_zero_point = output->params.zero_point;
-
-   const float input_scale = input->params.scale;
-   const float hires_input_scale = (1.0f / 128.0f) * input_scale;
-   const float reluish_scale = 3.0f / 32768.0f;
-   const float output_scale = output->params.scale;
-
-   const double output_multiplier = static_cast(hires_input_scale / output_scale);
-   int32_t output_multiplier_fixedpoint_int32;
-   QuantizeMultiplier(output_multiplier, &output_multiplier_fixedpoint_int32,
-                      ¶ms.output_multiplier_exponent);
-   DownScaleInt32ToInt16Multiplier(
-       output_multiplier_fixedpoint_int32,
-       ¶ms.output_multiplier_fixedpoint_int16);
-
-   TF_LITE_ENSURE(context, params.output_multiplier_exponent <= 0);
-
-   const double reluish_multiplier = static_cast(hires_input_scale / reluish_scale);
-   int32_t reluish_multiplier_fixedpoint_int32;
-   QuantizeMultiplier(reluish_multiplier, &reluish_multiplier_fixedpoint_int32,
-                      ¶ms.reluish_multiplier_exponent);
-   DownScaleInt32ToInt16Multiplier(
-       reluish_multiplier_fixedpoint_int32,
-       ¶ms.reluish_multiplier_fixedpoint_int16);
-
-   HardSwishOp(params, GetTensorShape(input),
-                GetTensorData(input), GetTensorShape(output), GetTensorData(output));
-   return kTfLiteOk;
-}
-
 TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
+  // Validate number of inputs and outputs
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
   return kTfLiteOk;
 }
 
@@ -267,6 +210,10 @@ TfLiteStatus ReluEval(TfLiteContext* context, TfLiteNode* node) {
 }
 
 TfLiteStatus Relu6Prepare(TfLiteContext* context, TfLiteNode* node) {
+  // Validate number of inputs and outputs
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
   return kTfLiteOk;
 }
 
@@ -307,34 +254,84 @@ TfLiteStatus Relu6Eval(TfLiteContext* context, TfLiteNode* node) {
   }
 }
 
+void* HardSwishInit(TfLiteContext* context, const char* buffer, size_t length) {
+  void* data = nullptr;
+  if (context->AllocatePersistentBuffer(context, sizeof(HardSwishParams), &data) ==
+      kTfLiteError) {
+    return nullptr;
+  }
+  return data;
+}
+
 TfLiteStatus HardSwishPrepare(TfLiteContext* context, TfLiteNode* node) {
+  TFLITE_DCHECK(node->user_data != nullptr);
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
+    HardSwishParams* params = static_cast(node->user_data);
+
+    params->input_zero_point = input->params.zero_point;
+    params->output_zero_point = output->params.zero_point;
+
+    const float input_scale = input->params.scale;
+    const float hires_input_scale = (1.0f / 128.0f) * input_scale;
+    const float reluish_scale = 3.0f / 32768.0f;
+    const float output_scale = output->params.scale;
+
+    const double output_multiplier = static_cast(hires_input_scale / output_scale);
+    int32_t output_multiplier_fixedpoint_int32;
+    QuantizeMultiplier(output_multiplier, &output_multiplier_fixedpoint_int32,
+                       ¶ms->output_multiplier_exponent);
+    DownScaleInt32ToInt16Multiplier(
+        output_multiplier_fixedpoint_int32,
+        ¶ms->output_multiplier_fixedpoint_int16);
+
+    TF_LITE_ENSURE(context, params->output_multiplier_exponent <= 0);
+
+    const double reluish_multiplier = static_cast(hires_input_scale / reluish_scale);
+    int32_t reluish_multiplier_fixedpoint_int32;
+    QuantizeMultiplier(reluish_multiplier, &reluish_multiplier_fixedpoint_int32,
+                       ¶ms->reluish_multiplier_exponent);
+    DownScaleInt32ToInt16Multiplier(
+        reluish_multiplier_fixedpoint_int32,
+        ¶ms->reluish_multiplier_fixedpoint_int16);
+  }
+
   return kTfLiteOk;
 }
 
 TfLiteStatus HardSwishEval(TfLiteContext* context, TfLiteNode* node) {
   const TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  HardSwishParams* params = static_cast(node->user_data);
 
   switch (input->type) {
     case kTfLiteFloat32: {
-      HardSwishFloatOp(
+      tflite::reference_ops::HardSwish(
           GetTensorShape(input),
           GetTensorData(input),
           GetTensorShape(output),
           GetTensorData(output));
-      return kTfLiteOk;
     } break;
     case kTfLiteUInt8: {
-      return HardSwishQuantized(context, node);
+      HardSwishOp(*params, GetTensorShape(input),
+                           GetTensorData(input), GetTensorShape(output), GetTensorData(output));
     } break;
     case kTfLiteInt8: {
-      return HardSwishQuantized(context, node);
+      HardSwishOp(*params, GetTensorShape(input),
+                          GetTensorData(input), GetTensorShape(output), GetTensorData(output));
     } break;
-    default:
+    default: {
       TF_LITE_KERNEL_LOG(context, "Only float32/int8/uint8 are supported currently, got %s",
                          TfLiteTypeGetName(input->type));
       return kTfLiteError;
+    }
   }
+  return kTfLiteOk;
 }
 
 }  // namespace activations
@@ -364,9 +361,14 @@ TfLiteRegistration* Register_RELU6() {
 }
 
 TfLiteRegistration* Register_HARD_SWISH() {
-  static TfLiteRegistration r = {};
-  r.prepare = activations::HardSwishPrepare;
-  r.invoke = activations::HardSwishEval;
+  static TfLiteRegistration r = {/*init=*/activations::HardSwishInit,
+                                 /*free=*/nullptr,
+                                 /*prepare=*/activations::HardSwishPrepare,
+                                 /*invoke=*/activations::HardSwishEval,
+                                 /*profiling_string=*/nullptr,
+                                 /*builtin_code=*/0,
+                                 /*custom_name=*/nullptr,
+                                 /*version=*/0};
   return &r;
 }
 
diff --git a/tensorflow/lite/micro/kernels/activations_test.cc b/tensorflow/lite/micro/kernels/activations_test.cc
index 03a6b698a86..b1d8e429025 100644
--- a/tensorflow/lite/micro/kernels/activations_test.cc
+++ b/tensorflow/lite/micro/kernels/activations_test.cc
@@ -139,7 +139,7 @@ void TestRelu6Float(const int* input_dims_data, const float* input_data,
 
 void GenerateUniformRandomVector(int size, float min, float max,
                                  std::minstd_rand* random_engine,
-                                 std::vector* result) {
+                                 float* result) {
   // Never use std::uniform_*_distribution in tests, it's
   // implementation-defined. Likewise, don't use std::default_random_engine,
   // implementation-defined. Implementation-defined is bad because it means that
@@ -147,7 +147,6 @@ void GenerateUniformRandomVector(int size, float min, float max,
   // std::minstd_rand is a standard instantiation of
   // std::linear_congruential_engine, the cheapest generator in c++11 stdlib,
   // it's good enough here.
-  result->resize(size);
   for (int i = 0; i < size; i++) {
     // We don't care whether the `max` value may ever be produced exactly.
     // It may actually be thanks to rounding, as std::minstd_rand::modulus
@@ -155,23 +154,23 @@ void GenerateUniformRandomVector(int size, float min, float max,
     float random_value_scaled_0_1 =
         (*random_engine)() *
         (1.0f / static_cast(std::minstd_rand::modulus));
-    (*result)[i] = min + (max - min) * random_value_scaled_0_1;
+    result[i] = min + (max - min) * random_value_scaled_0_1;
   }
 }
 
-void EvalTestReferenceHardSwish(int size, const std::vector& input,
-                                std::vector* result) {
-  result->resize(size);
+void EvalTestReferenceHardSwish(int size, float* input,
+                                float* result) {
   for (int i = 0; i < size; i++) {
     const float in = input[i];
-    (*result)[i] = in * std::min(6.0f, std::max(0.0f, in + 3)) * (1.0f / 6.0f);
+    result[i] = in * std::min(6.0f, std::max(0.0f, in + 3)) * (1.0f / 6.0f);
   }
 }
 
 template 
 void TestHardSwishQuantized(int size, const T* output_data, T* input_data_quantized, float* dequantized_output,
                             float input_min, float input_max, float output_min,
-                            float output_max, std::minstd_rand* random_engine) {
+                            float output_max, std::minstd_rand* random_engine,
+                            float *float_input_values, float* float_ref_output_values) {
   const int input_dims_data[] = {2, 1, size};
   const int output_dims_data[] = {2, 1, size};
   const float input_scale = ScaleFromMinMax(input_min, input_max);
@@ -191,21 +190,20 @@ void TestHardSwishQuantized(int size, const T* output_data, T* input_data_quanti
 
   TF_LITE_MICRO_EXPECT_EQ(output_elements_count, size);
 
-  std::vector float_input_values;
-  std::vector float_ref_output_values;
   GenerateUniformRandomVector(size, input_min, input_max, random_engine,
-                              &float_input_values);
+                              float_input_values);
   EvalTestReferenceHardSwish(size, float_input_values,
-                             &float_ref_output_values);
-  for (float& val : float_ref_output_values) {
-    val = std::min(output_max, std::max(output_min, val));
+                             float_ref_output_values);
+  for (int i = 0; i < size; i++) {
+    float val = float_ref_output_values[i];
+    float_ref_output_values[i] = std::min(output_max, std::max(output_min, val));
   }
 
   constexpr int inputs_size = 1;
   constexpr int outputs_size = 1;
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
-      CreateQuantizedTensor(float_input_values.data(), input_data_quantized, input_dims,
+      CreateQuantizedTensor(float_input_values, input_data_quantized, input_dims,
                             input_scale, input_zero_point, "input_tensor"),
       CreateQuantizedTensor(output_data, output_dims, output_scale,
                             output_zero_point, "output_tensor"),
@@ -253,19 +251,19 @@ void TestHardSwishQuantized(int size, const T* output_data, T* input_data_quanti
   AsymmetricDequantize(output_data, output_elements_count, output_scale, output_zero_point, dequantized_output);
 
   for (int i = 0; i < output_elements_count; ++i) {
-    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values.data()[i], dequantized_output[i], kTolerance);
+    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values[i], dequantized_output[i], kTolerance);
   }
 }
 
 template 
 void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_data_quantized,
                                 float* dequantized_output, float input_min, float input_max,
-                                float output_min, float output_max, float tolerated_bias) {
+                                float output_min, float output_max, float tolerated_bias,
+                                float* float_input_values, float* float_ref_output_values) {
   const float quantized_type_range =
       static_cast(std::numeric_limits::max()) -
       static_cast(std::numeric_limits::min());
 
-
   const float input_scale = ScaleFromMinMax(input_min, input_max);
   const float output_scale = ScaleFromMinMax(output_min, output_max);
 
@@ -284,20 +282,18 @@ void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_d
   const int quantized_input_positive_three =
       std::round(std::numeric_limits::min() +
                  (3.0f - input_min) / input_scale);
-  std::vector float_input_values;
+
   for (int i = quantized_input_negative_three;
        i <= quantized_input_positive_three; i++) {
-    float_input_values.push_back(
-        input_min +
-        (i - std::numeric_limits::min()) * input_scale);
+    float_input_values[i] = input_min +
+                            (i - std::numeric_limits::min()) * input_scale;
   }
-  TF_LITE_MICRO_EXPECT_EQ(float_input_values.size(), size);
 
-  std::vector float_ref_output_values;
   EvalTestReferenceHardSwish(size, float_input_values,
-                             &float_ref_output_values);
-  for (float& val : float_ref_output_values) {
-    val = std::min(output_max, std::max(output_min, val));
+                             float_ref_output_values);
+  for (int i = 0; i < size; i++) {
+    float val = float_ref_output_values[i];
+    float_ref_output_values[i] = std::min(output_max, std::max(output_min, val));
   }
 
   const int input_dims_data[] = {2, 1, size};
@@ -319,7 +315,7 @@ void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_d
   constexpr int outputs_size = 1;
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
-      CreateQuantizedTensor(float_input_values.data(), input_data_quantized, input_dims,
+      CreateQuantizedTensor(float_input_values, input_data_quantized, input_dims,
                             input_scale, input_zero_point, "input_tensor"),
       CreateQuantizedTensor(output_data, output_dims, output_scale,
                             output_zero_point, "output_tensor"),
@@ -375,15 +371,15 @@ void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_d
   TF_LITE_MICRO_EXPECT_LE(std::abs(bias), tolerated_bias);
 }
 
-void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* random_engine) {
-  std::vector float_input_values;
+void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* random_engine,
+                        float* float_input_values, float* float_ref_output_values) {
   const float kMin = -10.0f;
   const float kMax = 10.0f;
   GenerateUniformRandomVector(size, kMin, kMax, random_engine,
-                              &float_input_values);
-  std::vector float_ref_output_values;
+                              float_input_values);
+
   EvalTestReferenceHardSwish(size, float_input_values,
-                             &float_ref_output_values);
+                             float_ref_output_values);
 
   const int input_dims_data[] = {1, size};
   const int output_dims_data[] = {1, size};
@@ -398,7 +394,7 @@ void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* ra
   constexpr int outputs_size = 1;
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
-      CreateFloatTensor(float_input_values.data(), input_dims, "input_tensor"),
+      CreateFloatTensor(float_input_values, input_dims, "input_tensor"),
       CreateFloatTensor(output_data, output_dims, "output_tensor"),
   };
   TfLiteContext context;
@@ -439,7 +435,7 @@ void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* ra
   }
 
   for (int i = 0; i < output_elements_count; ++i) {
-    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values.data()[i], output_data[i], 1e-5f);
+    TF_LITE_MICRO_EXPECT_NEAR(float_ref_output_values[i], output_data[i], 1e-5f);
   }
 }
 
@@ -743,55 +739,62 @@ TF_LITE_MICRO_TEST(SimpleHardSwishTestFloat) {
   std::minstd_rand random_engine;
   constexpr int size = 100;
   float output_data[size] = {0.f};
+  float input_values[size] = {0.f};
+  float output_values[size] = {0.f};
 
-  tflite::testing::TestHardSwishFloat(size, output_data, &random_engine);
+  tflite::testing::TestHardSwishFloat(size, output_data, &random_engine,
+                                      input_values, output_values);
 }
 
 
 TF_LITE_MICRO_TEST(SimpleHardSwishTestInt8) {
   std::minstd_rand random_engine;
-  std::vector> minmax_pairs{
-      {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
+  constexpr int pairs = 4, one_pair = 2;
   constexpr int size = 101;
+  constexpr float minmax_pairs[pairs][one_pair] = {
+    {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
   constexpr int8_t output_data[size] = {0};
   int8_t input_data_quantized[size] = {0};
   float dequantized_output[size] = {0.f};
+  float input_values[size] = {0.f};
+  float output_values[size] = {0.f};
 
-  for (const auto& input_minmax : minmax_pairs) {
-    for (const auto& output_minmax : minmax_pairs) {
-      float input_min = input_minmax.first;
-      float input_max = input_minmax.second;
-      float output_min = output_minmax.first;
-      float output_max = output_minmax.second;
+  for (int x = 0; x < pairs; x++) {
+    for (int y = 0; y < pairs; y++) {
+      float input_min = minmax_pairs[x][0];
+      float input_max = minmax_pairs[x][1];
+      float output_min = minmax_pairs[y][0];
+      float output_max = minmax_pairs[y][1];
 
       tflite::testing::TestHardSwishQuantized(size, output_data, input_data_quantized, dequantized_output,
                                                       input_min, input_max, output_min, output_max,
-                                                      &random_engine);
-
+                                                      &random_engine, input_values, output_values);
     }
   }
 }
 
 TF_LITE_MICRO_TEST(SimpleHardSwishTestUint8) {
   std::minstd_rand random_engine;
-  std::vector> minmax_pairs{
-      {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
   constexpr int size = 99;
+  constexpr int pairs = 4, one_pair = 2;
+  constexpr float minmax_pairs[pairs][one_pair] = {
+    {0.f, 1.f}, {-2.f, 1.f}, {-5.f, 10.f}, {-40.f, 60.f}};
   constexpr uint8_t output_data[size] = {0};
   uint8_t input_data_quantized[size] = {0};
   float dequantized_output[size] = {0.f};
+  float input_values[size] = {0.f};
+  float output_values[size] = {0.f};
 
-  for (const auto& input_minmax : minmax_pairs) {
-    for (const auto& output_minmax : minmax_pairs) {
-      float input_min = input_minmax.first;
-      float input_max = input_minmax.second;
-      float output_min = output_minmax.first;
-      float output_max = output_minmax.second;
+  for (int x = 0; x < pairs; x++) {
+    for (int y = 0; y < pairs; y++) {
+      float input_min = minmax_pairs[x][0];
+      float input_max = minmax_pairs[x][1];
+      float output_min = minmax_pairs[y][0];
+      float output_max = minmax_pairs[y][1];
 
       tflite::testing::TestHardSwishQuantized(size, output_data, input_data_quantized, dequantized_output,
                                                        input_min, input_max, output_min, output_max,
-                                                       &random_engine);
-
+                                                       &random_engine, input_values, output_values);
     }
   }
 }
@@ -806,12 +809,14 @@ TF_LITE_MICRO_TEST(SimpleHardSwishTestQuantizedBias) {
   constexpr uint8_t output_data[size] = {0};
   uint8_t input_data_quantized[size] = {0};
   float dequantized_output[size] = {0.f};
+  float input_values[size] = {0.f};
+  float output_values[size] = {0.f};
 
   tflite::testing::TestHardSwishQuantizedBias(size, output_data, input_data_quantized, dequantized_output,
-                                                       -11.654928f, 25.036512f, -0.3905796f, 24.50887f, 0.035);
+                                                       -11.654928f, 25.036512f, -0.3905796f, 24.50887f, 0.035,
+                                                       input_values, output_values);
 }
 
-
 TF_LITE_MICRO_TEST(SimpleReluTestUint8) {
   const int elements_count = 10;
 

From 801779e9231ec58497b6ce4d921abe459045a87d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= 
Date: Thu, 11 Jun 2020 15:56:58 +0200
Subject: [PATCH 0872/1390] TFLu: do not directly include reference_ops.h

---
 tensorflow/lite/kernels/internal/reference/activations.h | 2 +-
 tensorflow/lite/micro/kernels/activations.cc             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/kernels/internal/reference/activations.h b/tensorflow/lite/kernels/internal/reference/activations.h
index 2110f5e3cbc..c043f1e0845 100644
--- a/tensorflow/lite/kernels/internal/reference/activations.h
+++ b/tensorflow/lite/kernels/internal/reference/activations.h
@@ -15,10 +15,10 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ACTIVATIONS_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ACTIVATIONS_H_
 
+#include "ruy/profiler/instrumentation.h"  // from @ruy
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/internal/common.h"
 
-
 namespace tflite {
 namespace reference_ops {
 
diff --git a/tensorflow/lite/micro/kernels/activations.cc b/tensorflow/lite/micro/kernels/activations.cc
index 934f029ac6a..a536878e4cb 100644
--- a/tensorflow/lite/micro/kernels/activations.cc
+++ b/tensorflow/lite/micro/kernels/activations.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/quantization_util.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
-#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/activations.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/kernels/op_macros.h"
 #include "tensorflow/lite/micro/micro_utils.h"

From 7473be170522cf994f541de07f867648b25c450e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= 
Date: Tue, 23 Jun 2020 10:09:18 +0200
Subject: [PATCH 0873/1390] TFLu: remove parameter to CreateTensor calls as it
 was removed

---
 tensorflow/lite/micro/kernels/activations_test.cc | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/micro/kernels/activations_test.cc b/tensorflow/lite/micro/kernels/activations_test.cc
index b1d8e429025..d8fa1c56a5d 100644
--- a/tensorflow/lite/micro/kernels/activations_test.cc
+++ b/tensorflow/lite/micro/kernels/activations_test.cc
@@ -204,9 +204,9 @@ void TestHardSwishQuantized(int size, const T* output_data, T* input_data_quanti
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
       CreateQuantizedTensor(float_input_values, input_data_quantized, input_dims,
-                            input_scale, input_zero_point, "input_tensor"),
+                            input_scale, input_zero_point),
       CreateQuantizedTensor(output_data, output_dims, output_scale,
-                            output_zero_point, "output_tensor"),
+                            output_zero_point),
   };
 
   TfLiteContext context;
@@ -316,9 +316,9 @@ void TestHardSwishQuantizedBias(const int size, const T* output_data, T* input_d
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
       CreateQuantizedTensor(float_input_values, input_data_quantized, input_dims,
-                            input_scale, input_zero_point, "input_tensor"),
+                            input_scale, input_zero_point),
       CreateQuantizedTensor(output_data, output_dims, output_scale,
-                            output_zero_point, "output_tensor"),
+                            output_zero_point),
   };
 
   TfLiteContext context;
@@ -394,8 +394,8 @@ void TestHardSwishFloat(const int size, float* output_data, std::minstd_rand* ra
   constexpr int outputs_size = 1;
   constexpr int tensors_size = inputs_size + outputs_size;
   TfLiteTensor tensors[tensors_size] = {
-      CreateFloatTensor(float_input_values, input_dims, "input_tensor"),
-      CreateFloatTensor(output_data, output_dims, "output_tensor"),
+      CreateFloatTensor(float_input_values, input_dims),
+      CreateFloatTensor(output_data, output_dims),
   };
   TfLiteContext context;
   PopulateContext(tensors, tensors_size, micro_test::reporter, &context);

From f5e1a27293823b799a080c384d854aff8f3ef133 Mon Sep 17 00:00:00 2001
From: Tres Popp 
Date: Tue, 23 Jun 2020 02:01:33 -0700
Subject: [PATCH 0874/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/f570d5810485

PiperOrigin-RevId: 317824506
Change-Id: I5da78511de1e20d9b0a62fd3a8092b791a35b9da
---
 .../compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc     | 2 +-
 tensorflow/workspace.bzl                                      | 4 ++--
 third_party/mlir/BUILD                                        | 1 +
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
index 99d2c08aa98..78a77dc3b4d 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
@@ -129,7 +129,7 @@ struct DynamicMemRefCastOpConverter
 void PopulateLhloToLLVMConversionPatterns(LLVMTypeConverter *converter,
                                           OwningRewritePatternList *patterns) {
   patterns->insert(
-      *converter);
+      *converter, LowerToLLVMOptions());
 }
 
 }  // namespace xla_lhlo
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 0f591ba8b90..a55dcd17b07 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "7e825abd5704ce28b166f9463d4bd304348fd2a9"
-    LLVM_SHA256 = "a21b752ee1866e195f3f72c7931c79f8c4ecc0f14861488284bdc2bdf14d6fe9"
+    LLVM_COMMIT = "f570d5810485fa6fb2e1009f795a899d79bd429f"
+    LLVM_SHA256 = "e154a1a97c3b6bead73a32bcb6fc37aac2d80628abee2961315304a4964f04fe"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD
index 2be59fc44b1..8fd0a94bf64 100644
--- a/third_party/mlir/BUILD
+++ b/third_party/mlir/BUILD
@@ -1973,6 +1973,7 @@ cc_library(
         ":StandardOps",
         ":Support",
         ":Transforms",
+        ":VectorOps",
         "@llvm-project//llvm:Support",
     ],
 )

From 7a124ecab8070e7affc2fabf019b5b57588d85ee Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 23 Jun 2020 02:01:48 -0700
Subject: [PATCH 0875/1390] Update GraphDef version to 441.

PiperOrigin-RevId: 317824531
Change-Id: I7fbf95f3a662d2d8d823656e9ddf5ae2dfd487f6
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index e63ca48863f..52a926c8d8b 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 440  // Updated: 2020/6/22
+#define TF_GRAPH_DEF_VERSION 441  // Updated: 2020/6/23
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 5aa8be781eebe1f3544d4c19f83ca0fa6351b4a3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 23 Jun 2020 02:01:50 -0700
Subject: [PATCH 0876/1390] compat: Update forward compatibility horizon to
 2020-06-23

PiperOrigin-RevId: 317824535
Change-Id: I12a1d7a71a420b996f67cf3c536a9c5fd3365154
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index 301f4608aba..521d7eaf30f 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 22)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 23)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From 775b9cb5acffaf86c1ce09d4d1460f40b1392c2a Mon Sep 17 00:00:00 2001
From: Alexander Belyaev 
Date: Tue, 23 Jun 2020 03:19:37 -0700
Subject: [PATCH 0877/1390] [XLA][MLIR] Use callback-based builders for
 (indexed_)generic in (L)HLO->Linalg.

PiperOrigin-RevId: 317833249
Change-Id: Ifb60bd3a4797ce360c2a2685e3e96a19dcd6b161
---
 .../xla/transforms/xla_legalize_to_linalg.cc  | 124 +++++++-----------
 1 file changed, 44 insertions(+), 80 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
index 8a2f8ce7d04..e7bb5df8233 100644
--- a/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/xla_legalize_to_linalg.cc
@@ -125,32 +125,19 @@ class PointwiseToLinalgConverter : public OpConversionPattern {
       opResultTypes.push_back(shapedType);
     }
 
+    int64_t args_count = bodyArgTypes.size();
+    int64_t results_count = bodyResultTypes.size();
     auto linalgOp = rewriter.create(
-        loc, opResultTypes, args,
-        /*inputCount=*/bodyArgTypes.size(),
-        /*outputCount=*/bodyResultTypes.size(), indexing_maps,
-        GetNParallelLoopsAttrs(nloops));
-
-    // Add a block to the region.
-    auto* region = &linalgOp.region();
-    auto* block = rewriter.createBlock(region, region->end());
-    block->addArguments(bodyArgTypes);
-    if (isLHLO) block->addArguments(bodyResultTypes);
-
-    SmallVector bodyArgs;
-    for (int i = 0, e = bodyArgTypes.size(); i < e; ++i) {
-      bodyArgs.push_back(block->getArgument(i));
-    }
-
-    rewriter.setInsertionPointToEnd(block);
-    // TODO(ravishankarm) : For now use the method in xla_lhlo namespace. That
-    // method needs to be moved out of there.
-    Value opResult = xla_lhlo::XlaOpToStdScalarOp::map(
-        op, bodyResultTypes, bodyArgs, &rewriter);
-    if (!opResult) {
-      return failure();
-    }
-    rewriter.create(loc, opResult);
+        loc, opResultTypes, args, args_count, results_count, indexing_maps,
+        GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+          // TODO(ravishankarm) : For now use the method in xla_lhlo namespace.
+          // That method needs to be moved out of there.
+          Value opResult = xla_lhlo::XlaOpToStdScalarOp::map(
+              op, bodyResultTypes,
+              llvm::to_vector<2>(args.take_front(args_count)), &rewriter);
+          nestedBuilder.create(loc, opResult);
+        });
     rewriter.replaceOp(op, linalgOp.getOperation()->getResults());
     return success();
   }
@@ -301,27 +288,20 @@ class DataMovementOpConverter : public OpConversionPattern {
       OpTy op, ArrayRef args,
       ConversionPatternRewriter& rewriter) const final {
     if (!verifyXLAOpBufferOrTensorSemantics(op)) return failure();
-    auto operandType = op.operand().getType().template cast();
     auto resultType = getXLAOpResultType(op);
 
     SmallVector indexing_maps =
         Derived::getIndexingMaps(op, &rewriter);
     if (indexing_maps.empty()) return failure();
 
-    OpBuilder::InsertionGuard linalgOpGuard(rewriter);
     auto nloops = resultType.getRank();
     auto loc = op.getLoc();
     auto linalgOp = rewriter.create(
         loc, isLHLO ? ArrayRef{} : resultType, args, /*inputCount=*/1,
-        /*outputCount=*/1, indexing_maps, GetNParallelLoopsAttrs(nloops));
-
-    auto* region = &linalgOp.region();
-    auto* block = rewriter.createBlock(region, region->end());
-    block->addArguments(operandType.getElementType());
-    if (isLHLO) block->addArgument(resultType.getElementType());
-
-    rewriter.setInsertionPointToEnd(block);
-    rewriter.create(loc, block->getArgument(0));
+        /*outputCount=*/1, indexing_maps, GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+          nestedBuilder.create(loc, *args.begin());
+        });
 
     rewriter.replaceOp(op, linalgOp.getOperation()->getResults());
     return success();
@@ -437,36 +417,26 @@ class LhloBroadcastInDimConverter
       Value zero = rewriter.create(loc, 0);
       Value val =
           rewriter.create(loc, operand, llvm::makeArrayRef({zero}));
-      auto linalgOp = rewriter.create(
+      rewriter.create(
           loc, llvm::None, llvm::makeArrayRef(operand_adaptor.output()),
           /*inputCount=*/0, /*outputCount=*/1,
           llvm::makeArrayRef(rewriter.getMultiDimIdentityMap(nloops)),
-          GetNParallelLoopsAttrs(nloops));
+          GetNParallelLoopsAttrs(nloops),
+          [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+            nestedBuilder.create(loc, val);
+          });
 
-      auto* region = &linalgOp.region();
-      auto* block = rewriter.createBlock(region, region->end());
-      block->addArgument(result_type.getElementType());
-
-      rewriter.setInsertionPointToEnd(block);
-      rewriter.create(loc, val);
     } else {
       auto indexing_maps = getIndexingMaps(op, broadcast_dims, result_shape,
                                            operand_type, &rewriter);
-
-      OpBuilder::InsertionGuard linalgOpGuard(rewriter);
-      auto linalgOp = rewriter.create(
+      rewriter.create(
           loc, llvm::None,
           llvm::makeArrayRef({operand, operand_adaptor.output()}),
           /*inputCount=*/1, /*outputCount=*/1, indexing_maps,
-          GetNParallelLoopsAttrs(nloops));
-
-      auto* region = &linalgOp.region();
-      auto* block = rewriter.createBlock(region, region->end());
-      block->addArguments(operand_type.getElementType());
-      block->addArgument(result_type.getElementType());
-
-      rewriter.setInsertionPointToEnd(block);
-      rewriter.create(loc, block->getArgument(0));
+          GetNParallelLoopsAttrs(nloops),
+          [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange args) {
+            nestedBuilder.create(loc, *args.begin());
+          });
     }
     rewriter.replaceOp(op, llvm::None);
     return success();
@@ -686,32 +656,26 @@ class IotaConverter : public OpConversionPattern {
     // Construct the indexing maps needed for linalg.generic ops.
     unsigned nloops = resultMemrefType.getRank();
 
-    auto loc = iotaOp.getLoc();
-    auto linalgOp = rewriter.create(
-        loc, ArrayRef{}, args,
+    rewriter.create(
+        iotaOp.getLoc(), ArrayRef{}, args,
         0,  // args_in
         1,  // args_out
         llvm::makeArrayRef(rewriter.getMultiDimIdentityMap(nloops)),
-        GetNParallelLoopsAttrs(nloops));
+        GetNParallelLoopsAttrs(nloops),
+        [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange ivs,
+            ValueRange args) {
+          Value castOp = nestedBuilder.create(
+              nestedLoc, ivs[iotaOp.iota_dimension().getZExtValue()],
+              nestedBuilder.getIntegerType(
+                  resultElementType.getIntOrFloatBitWidth()));
+          if (resultElementType.isa()) {
+            castOp = nestedBuilder.create(nestedLoc, castOp,
+                                                    resultElementType);
+          }
+          nestedBuilder.create(nestedLoc, castOp);
+        });
 
-    // Add a block to the region.
-    auto* region = &linalgOp.region();
-    auto* block = rewriter.createBlock(region, region->end());
-    for (unsigned i = 0; i < nloops; ++i) {
-      block->addArgument(rewriter.getIndexType());
-    }
-    block->addArguments(llvm::makeArrayRef(resultElementType));
-
-    rewriter.setInsertionPointToEnd(block);
-    Operation* castOp = rewriter.create(
-        loc, block->getArgument(iotaOp.iota_dimension().getZExtValue()),
-        rewriter.getIntegerType(resultElementType.getIntOrFloatBitWidth()));
-    if (resultElementType.isa()) {
-      castOp = rewriter.create(loc, castOp->getResult(0),
-                                         resultElementType);
-    }
-    rewriter.create(loc, castOp->getResult(0));
-    rewriter.eraseOp(iotaOp);
+    rewriter.replaceOp(iotaOp, llvm::None);
     return success();
   }
 };
@@ -867,7 +831,7 @@ struct LhloLegalizeToLinalg
 
     auto func = getFunction();
     populateLHLOToLinalgConversionPattern(func.getContext(), &patterns);
-    if (failed(applyPartialConversion(func, target, patterns))) {
+    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
       signalPassFailure();
     }
   }
@@ -882,7 +846,7 @@ struct HloLegalizeToLinalg
 
     auto func = getFunction();
     xla_hlo::populateHLOToLinalgConversionPattern(func.getContext(), &patterns);
-    if (failed(applyPartialConversion(func, target, patterns))) {
+    if (failed(applyPartialConversion(func, target, patterns, nullptr))) {
       signalPassFailure();
     }
   }

From 8eea0658d48c4110076b478ac507445e2b2e47b3 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel 
Date: Tue, 23 Jun 2020 04:51:06 -0700
Subject: [PATCH 0878/1390] Replace tensorflow command line flags with LLVM
 command line flags.

This avoids a dependency on the framework_internal library, and also LLVM
directly supports comma-separated value lists.

PiperOrigin-RevId: 317842748
Change-Id: I9207375c753fbab33f6211736d57248c8c3f558a
---
 .../compiler/mlir/tools/kernel_gen/BUILD      |  3 +-
 .../mlir/tools/kernel_gen/tf_to_cubin.cc      | 88 +++++--------------
 2 files changed, 26 insertions(+), 65 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/BUILD b/tensorflow/compiler/mlir/tools/kernel_gen/BUILD
index cebfa7cd9d4..80b597d962d 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/BUILD
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/BUILD
@@ -44,8 +44,9 @@ tf_cc_binary(
     visibility = ["//tensorflow/core/kernels/cubin_headers:__pkg__"],
     deps = [
         ":cubin_creator",
-        "//tensorflow/core:framework_internal",
+        "//tensorflow/compiler/mlir:init_mlir",
         "//tensorflow/core:lib",
         "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
     ],
 )
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_cubin.cc b/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_cubin.cc
index 66fcabde0ac..96831689600 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_cubin.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_cubin.cc
@@ -21,77 +21,37 @@
 #include 
 #include 
 
-#include "absl/strings/numbers.h"
-#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "llvm/Support/CommandLine.h"
+#include "tensorflow/compiler/mlir/init_mlir.h"
 #include "tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.h"
 #include "tensorflow/core/platform/env.h"
-#include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/util/command_line_flags.h"
-
-namespace {
-bool ParseStringList(std::string string_list, std::vector* result) {
-  result->clear();
-  uint32_t item;
-  auto items = absl::StrSplit(string_list, ',');
-  for (const auto& item_str : items) {
-    if (!absl::SimpleAtoi(item_str, &item)) {
-      LOG(ERROR) << "Expected token " << item_str << " to be an integer";
-      return false;
-    }
-    result->push_back(item);
-  }
-  return true;
-}
-}  // namespace
 
 int main(int argc, char** argv) {
-  std::string input_file = "foo.mlir";
-  std::string output_file = "foo.bin";
-  int32_t architecture = 50;
-  std::vector tile_sizes;
-  std::vector unroll_factors;
-  std::vector same_shape;
+  llvm::cl::opt input_file("input", llvm::cl::desc("input file"),
+                                        llvm::cl::value_desc("filename"),
+                                        llvm::cl::init("foo.mlir"));
+  llvm::cl::opt output_file(
+      "output", llvm::cl::desc("output file"), llvm::cl::value_desc("filename"),
+      llvm::cl::init("foo.bin"));
+  llvm::cl::opt architecture(
+      "arch", llvm::cl::desc("target architecture (e.g. 50 for sm_50)"),
+      llvm::cl::init(50));
+  llvm::cl::list tile_sizes(
+      "tile_sizes", llvm::cl::desc("tile sizes to use"), llvm::cl::ZeroOrMore,
+      llvm::cl::CommaSeparated);
+  llvm::cl::list unroll_factors(
+      "unroll_factors",
+      llvm::cl::desc("factors to unroll by, separated by commas"),
+      llvm::cl::ZeroOrMore, llvm::cl::CommaSeparated);
+  llvm::cl::list same_shape(
+      "same_shape",
+      llvm::cl::desc("arguments with same shape, separated by commas"),
+      llvm::cl::ZeroOrMore, llvm::cl::CommaSeparated);
 
-  auto parse_tile_sizes = [&tile_sizes](std::string tile_sizes_str) {
-    if (!ParseStringList(tile_sizes_str, &tile_sizes)) {
-      return false;
-    }
-    // Initialize with the default.
-    if (tile_sizes.empty()) {
-      tile_sizes.push_back(16);
-      tile_sizes.push_back(64);
-    }
-    return true;
-  };
-
-  auto parse_unroll_factors =
-      [&unroll_factors](std::string unroll_factors_str) {
-        return ParseStringList(unroll_factors_str, &unroll_factors);
-      };
-
-  auto parse_same_shape = [&same_shape](std::string same_shape_str) {
-    return ParseStringList(same_shape_str, &same_shape);
-  };
-
-  std::vector flag_list = {
-      tensorflow::Flag("input", &input_file, "input file"),
-      tensorflow::Flag("output", &output_file, "output file"),
-      tensorflow::Flag("arch", &architecture,
-                       "target architecture (e.g. 50 for sm_50)"),
-      tensorflow::Flag("tile_sizes", parse_tile_sizes, "16,64",
-                       "tile sizes to use"),
-      tensorflow::Flag("unroll_factors", parse_unroll_factors, "",
-                       "factors to unroll by, separated by commas"),
-      tensorflow::Flag("same_shape", parse_same_shape, "",
-                       "arguments with same shape, separated by commas"),
-  };
-  bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
-  tensorflow::port::InitMain("usage", &argc, &argv);
-  if (!parse_ok) {
-    return 1;
-  }
+  tensorflow::InitMlir y(&argc, &argv);
+  llvm::cl::ParseCommandLineOptions(argc, argv, "TF op GPU kernel generator\n");
 
   std::pair compute_capability(architecture / 10,
                                                  architecture % 10);

From af94e801cf90fe8fb70d930d8658c20913d09ca5 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 23 Jun 2020 04:56:42 -0700
Subject: [PATCH 0879/1390] Use ForLoopSpecializationPass to help with
 vectorization of kernels.

PiperOrigin-RevId: 317843378
Change-Id: I4ead02c24f957269888af5491934567cd3e311fb
---
 tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index 3f99d40c717..196ea218ef3 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -505,6 +505,11 @@ Status LowerLHLOToGPU(mlir::ModuleOp module, LowerLHLOToGPUOptions options) {
   // Some basic cleanup.
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCanonicalizerPass());
   pm.addNestedPass<::mlir::FuncOp>(::mlir::createCSEPass());
+  // Make loops with min bounds into a conditional plus static bounds.
+  // Only do this if we unrolled in the first place.
+  if (!options.unroll_factors.empty()) {
+    pm.addNestedPass<::mlir::FuncOp>(mlir::createForLoopSpecializationPass());
+  }
   // Approximate of requested.
   if (options.use_approximations) {
     pm.addNestedPass<::mlir::FuncOp>(

From 7198070f4d032493eeae9741ba9fd868f3d3fb7c Mon Sep 17 00:00:00 2001
From: Tres Popp 
Date: Tue, 23 Jun 2020 04:57:39 -0700
Subject: [PATCH 0880/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/7a55d984971c

PiperOrigin-RevId: 317843475
Change-Id: I2824384ac6a2535fd779ee40301d83004e55b5f3
---
 tensorflow/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index a55dcd17b07..7c47628818b 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "f570d5810485fa6fb2e1009f795a899d79bd429f"
-    LLVM_SHA256 = "e154a1a97c3b6bead73a32bcb6fc37aac2d80628abee2961315304a4964f04fe"
+    LLVM_COMMIT = "7a55d984971c11daa55e9423934f98bdc9c04f2f"
+    LLVM_SHA256 = "6ef6d1f92f51936ed3027433ea26875ac4d2ac8eed88b1fbce472018c4fb7720"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),

From e071e66f03eb6ae234212400f34061f25e79a699 Mon Sep 17 00:00:00 2001
From: Lev Proleev 
Date: Tue, 23 Jun 2020 05:29:56 -0700
Subject: [PATCH 0881/1390] Add support for TENSOR_QUANT8_ASYMM_SIGNED in NNAPI
 delegate

PiperOrigin-RevId: 317846923
Change-Id: I1c61f53e89228cd2482435e9255e390864bd83e3
---
 tensorflow/lite/delegates/nnapi/BUILD         |  26 +
 .../delegates/nnapi/acceleration_test_list.cc |  14 +-
 .../lite/delegates/nnapi/nnapi_delegate.cc    | 246 +++--
 .../delegates/nnapi/nnapi_delegate_kernel.h   |   3 +
 .../nnapi/nnapi_delegate_mock_test.h          |   2 +
 ...nnapi_delegate_signed_quantization_test.cc | 920 ++++++++++++++++++
 tensorflow/lite/nnapi/NeuralNetworksTypes.h   |   1 +
 tensorflow/lite/nnapi/nnapi_handler.h         |  22 +
 8 files changed, 1154 insertions(+), 80 deletions(-)
 create mode 100644 tensorflow/lite/delegates/nnapi/nnapi_delegate_signed_quantization_test.cc

diff --git a/tensorflow/lite/delegates/nnapi/BUILD b/tensorflow/lite/delegates/nnapi/BUILD
index ec9f6907f21..beeaff1b99d 100644
--- a/tensorflow/lite/delegates/nnapi/BUILD
+++ b/tensorflow/lite/delegates/nnapi/BUILD
@@ -190,6 +190,32 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "nnapi_delegate_signed_quantization_test",
+    size = "small",
+    srcs = [
+        "nnapi_delegate_signed_quantization_test.cc",
+    ],
+    tags = [
+        "no_mac",
+        "no_windows",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":nnapi_delegate",
+        ":nnapi_delegate_mock_test",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:kernel_api",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/nnapi:nnapi_implementation",
+        "//tensorflow/lite/nnapi:nnapi_lib",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 cc_test(
     name = "quant_lstm_sup_test",
     size = "small",
diff --git a/tensorflow/lite/delegates/nnapi/acceleration_test_list.cc b/tensorflow/lite/delegates/nnapi/acceleration_test_list.cc
index b20628016f0..31bdc5f8b99 100644
--- a/tensorflow/lite/delegates/nnapi/acceleration_test_list.cc
+++ b/tensorflow/lite/delegates/nnapi/acceleration_test_list.cc
@@ -60,6 +60,10 @@ FloatActivationsOpTest/Elu,30
 FloatActivationsOpTest/HardSwish
 QuantizedActivationsOpTest/HardSwish
 QuantizedActivationsOpTest/HardSwishBias
+QuantizedActivationsOpTest/Relu*
+QuantizedActivationsOpTest/PRelu,29
+QuantizedActivationsOpTest/PReluSameShapes,29
+QuantizedActivationsOpTest/PReluInt8.+,30
 
 # add_test
 FloatAddOpModel/.+
@@ -145,6 +149,7 @@ ConvolutionOpTest/ConvolutionOpTest/.+/\d+
 
 # dequantize_test
 DequantizeOpTest/Uint8
+DequantizeOpTest/Int8,30
 
 # depth_to_space_test
 DepthToSpaceOpModel/Float32
@@ -190,6 +195,7 @@ QuantizedFullyConnectedOpTest/SimpleTestQuantizedOutputMultiplierGreaterThan1Uin
 QuantizedFullyConnectedOpTest/SimpleTestQuantizedOutputMultiplierGreaterThan1Int8/\d+,29
 HybridFullyConnectedOpTest/SimpleTestQuantizedUint8,29
 HybridFullyConnectedOpTest/SimpleTestQuantizedInt8,29
+HybridAsymmetricInputFullyConnectedOpTest.SimpleTestQuantizedUint8,29
 FloatFullyConnectedOpTest/FloatFullyConnectedOpTest/SimpleTest4DInput/\d+
 QuantizedFullyConnectedOpTest/QuantizedFullyConnectedOpTest/SimpleTest4dInputQuantizedUint8/\d+
 QuantizedFullyConnectedOpTest/QuantizedFullyConnectedOpTest/SimpleTest4dInputQuantizedOutputMultiplierGreaterThan1Uint8/\d+,29
@@ -207,6 +213,7 @@ FloatGatherOpTest/LastAxis,29
 TypesGatherOpTest/Float32Int32,29
 TypesGatherOpTest/Int32Int32,29
 TypesGatherOpTest/Uint8Int32,29
+TypesGatherOpTest/Int8Int32,29
 
 # hashtable_lookup_test
 # All test excepted the string one should be accelerated
@@ -286,13 +293,18 @@ QuantizedLstmTest/BasicQuantizedLstmTest/29
 
 # quantize_test
 QuantizeOpTest/UINT8,29
+QuantizeOpTest/INT8,30
+
+# rank
 
 # reduce_test
 -Dynamic.+(Mean|Sum|Prod|Max|Min)OpTest/.+
 -ConstUint8(Mean|Sum)OpTest/.+
+-ConstInt8MeanOpTest.NonSpecialAxisNonSameScale
+-ConstInt8MeanOpTest.QuantizedDifferentScale
 ConstUint8(Max|Min)OpTest/.+,29
 ConstUint8(Mean)OpTest/.+
-Constint8(Mean|Max|Min)OpTest/.+
+ConstInt8(Mean|Max|Min)OpTest/.+,29
 ConstFloat(Sum|Prod|Max|Min)OpTest/NotKeepDims,29
 ConstFloat(Sum|Prod|Max|Min)OpTest/KeepDims,29
 ConstFloat(Mean|Any)OpTest/NotKeepDims
diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
index 1c35ee370c2..58ab13ab657 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc
@@ -201,6 +201,7 @@ bool NeedInt8Conversion(const TfLiteContext* context, int builtin_code,
     case kTfLiteBuiltinConcatenation:
     case kTfLiteBuiltinEqual:
     case kTfLiteBuiltinExpandDims:
+    case kTfLiteBuiltinGather:
     case kTfLiteBuiltinGreater:
     case kTfLiteBuiltinGreaterEqual:
     case kTfLiteBuiltinHardSwish:
@@ -377,6 +378,7 @@ bool HasZeroes(TfLiteIntArrayView array) {
 enum {
   NN_TENSOR_FLAG_SCALAR_AS_TENSOR = 1U << 0,
   NN_TENSOR_FLAG_INT8_CONVERSION = 1U << 1,
+  NN_TENSOR_FLAG_USE_INT8_ASYMM_SIGNED = 1U << 2,
 };
 
 // Returns the SDK level to target when delegating to the given devices.
@@ -1065,6 +1067,8 @@ class NNAPIOpBuilder {
         tensor_flags & NN_TENSOR_FLAG_SCALAR_AS_TENSOR;
     const bool need_int8_conversion =
         tensor_flags & NN_TENSOR_FLAG_INT8_CONVERSION;
+    const bool use_int8_asymm_signed =
+        tensor_flags & NN_TENSOR_FLAG_USE_INT8_ASYMM_SIGNED;
     int ann_tensor_index = operand_mapping_->lite_index_to_ann(tensor_index);
     if (ann_tensor_index != -1) {
       indices->push_back(ann_tensor_index);
@@ -1095,12 +1099,25 @@ class NNAPIOpBuilder {
         nn_type = ANEURALNETWORKS_TENSOR_FLOAT32;
         break;
       case kTfLiteUInt8:
+        nn_type = ANEURALNETWORKS_TENSOR_QUANT8_ASYMM;
+        scale = tensor->params.scale;
+        zeroPoint = tensor->params.zero_point;
+        if (scale == 0) {
+          // ANEURALNETWORKS_TENSOR_QUANT8_ASYMM with zero scale is not valid in
+          // NNAPI.
+          scale = 1;
+        }
+        break;
       case kTfLiteInt8:
         // If explicit int8 conversion is needed, we still need
         // ANEURALNETWORKS_TENSOR_QUANT8_ASYMM type.
-        nn_type = (tensor_type == kTfLiteUInt8 || need_int8_conversion)
-                      ? ANEURALNETWORKS_TENSOR_QUANT8_ASYMM
-                      : ANEURALNETWORKS_TENSOR_QUANT8_SYMM;
+        if (use_int8_asymm_signed) {
+          nn_type = ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED;
+        } else if (need_int8_conversion) {
+          nn_type = ANEURALNETWORKS_TENSOR_QUANT8_ASYMM;
+        } else {
+          nn_type = ANEURALNETWORKS_TENSOR_QUANT8_SYMM;
+        }
         scale = tensor->params.scale;
         zeroPoint = tensor->params.zero_point;
         if (tensor->quantization.type == kTfLiteAffineQuantization) {
@@ -1130,8 +1147,7 @@ class NNAPIOpBuilder {
             operand_mapping_->add_type_conversion(tensor_index, kTfLiteUInt8);
           }
           if (scale == 0) {
-            // TENSOR_QUANT8_ASYMM and ANEURALNETWORKS_TENSOR_QUANT8_ASYMM
-            // with zero scale are not valid in NNAPI.
+            // QUANT8 tensors with zero scale are not valid in NNAPI.
             scale = 1;
           }
         }
@@ -1248,7 +1264,6 @@ class NNAPIOpBuilder {
             "setting new operand value", nnapi_errno_);
       }
     }
-
     indices->push_back(ann_tensor_index);
     return kTfLiteOk;
   }
@@ -1437,7 +1452,6 @@ bool NNAPIDelegateKernel::Validate(
     bool is_accelerator_specified,
     std::vector* map_failures) {
   OpValidationContext val_ctx{true, map_failures};
-
   switch (builtin_code) {
     case kTfLiteBuiltinAdd: {
       ExpectMaxOpVersion(version, 2, &val_ctx);
@@ -1789,18 +1803,21 @@ bool NNAPIDelegateKernel::Validate(
              "Supported op versions are 1 and 2 only", &val_ctx);
 
       const auto& input = context->tensors[node->inputs->data[0]];
-      Expect(input.type != kTfLiteFloat16,
-             NNAPIValidationFailureType::kUnsupportedInputType,
-             "kTfLiteFloat16 not supported as input", &val_ctx);
+      if (android_sdk_version < kMinSdkVersionForNNAPI12) {
+        EXPECT_INPUT_TYPE_IN(input.type, kTfLiteUInt8);
+      } else {
+        EXPECT_INPUT_TYPE_IN(input.type, kTfLiteUInt8, kTfLiteInt8);
 
-      const auto zero_point = input.params.zero_point;
-      Expect(input.type != kTfLiteInt8 ||
-                 (zero_point == 0 &&
-                  android_sdk_version >= kMinSdkVersionForNNAPI12),
-             NNAPIValidationFailureType::kUnsupportedInputType,
-             "NN API supports int8 type since version 1.2 but only for "
-             "symmetric quantization.",
-             &val_ctx);
+        if (android_sdk_version == kMinSdkVersionForNNAPI12 &&
+            input.type == kTfLiteInt8) {
+          const auto zero_point = input.params.zero_point;
+          Expect(zero_point == 0,
+                 NNAPIValidationFailureType::kUnsupportedInputType,
+                 "NN API supports int8 type since version 1.2 but only for "
+                 "symmetric quantization.",
+                 &val_ctx);
+        }
+      }
     } break;
     case kTfLiteBuiltinFloor: {
       ExpectOpVersion(version, 1, &val_ctx);
@@ -2150,21 +2167,38 @@ bool NNAPIDelegateKernel::Validate(
                                  &val_ctx);
       const TfLiteType input_type =
           context->tensors[node->inputs->data[0]].type;
-      EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteInt32,
-                           kTfLiteUInt8);
       const TfLiteType output_type =
           context->tensors[node->outputs->data[0]].type;
-      ExpectTypeIn(output_type, {kTfLiteFloat32, kTfLiteInt32, kTfLiteUInt8},
-                   NNAPIValidationFailureType::kUnsupportedOutputType,
-                   "Output type should be one of kTfLiteFloat32, kTfLiteInt32, "
-                   "kTfLiteUInt8.",
-                   &val_ctx);
+      if (android_sdk_version >= kMinSdkVersionForNNAPI13) {
+        EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteInt32,
+                             kTfLiteUInt8, kTfLiteInt8);
+
+        ExpectTypeIn(
+            output_type,
+            {kTfLiteFloat32, kTfLiteInt32, kTfLiteUInt8, kTfLiteInt8},
+            NNAPIValidationFailureType::kUnsupportedOutputType,
+            "Output type should be one of kTfLiteFloat32, kTfLiteInt32, "
+            "kTfLiteUInt8, kTfLiteInt8.",
+            &val_ctx);
+      } else {
+        EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteInt32,
+                             kTfLiteUInt8);
+
+        ExpectTypeIn(
+            output_type, {kTfLiteFloat32, kTfLiteInt32, kTfLiteUInt8},
+            NNAPIValidationFailureType::kUnsupportedOutputType,
+            "Output type should be one of kTfLiteFloat32, kTfLiteInt32, "
+            "kTfLiteUInt8.",
+            &val_ctx);
+      }
     } break;
     case kTfLiteBuiltinPrelu: {
       ExpectOpVersion(version, 1, &val_ctx);
       ExpectMinAndroidSdkVersion(android_sdk_version, kMinSdkVersionForNNAPI12,
                                  &val_ctx);
-      ExpectIsFloatOrUint8Operator(context, node, &val_ctx);
+      const auto input_type = context->tensors[node->inputs->data[0]].type;
+      EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteUInt8,
+                           kTfLiteInt8);
     } break;
     case kTfLiteBuiltinTile: {
       ExpectOpVersion(version, 1, &val_ctx);
@@ -2240,19 +2274,18 @@ bool NNAPIDelegateKernel::Validate(
              &val_ctx);
     } break;
     case kTfLiteBuiltinGather: {
-      ExpectOpVersion(version, 1, &val_ctx);
+      ExpectOpVersion(version, 2, &val_ctx);
       ExpectMinAndroidSdkVersion(android_sdk_version, kMinSdkVersionForNNAPI12,
                                  &val_ctx);
       const auto input_type = context->tensors[node->inputs->data[0]].type;
       const auto& positions = context->tensors[node->inputs->data[1]];
+
       EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteFloat16,
-                           kTfLiteInt32, kTfLiteUInt8);
-      ExpectTypeIn(positions.type,
-                   {kTfLiteFloat32, kTfLiteFloat16, kTfLiteInt32, kTfLiteUInt8},
-                   NNAPIValidationFailureType::kUnsupportedInputType,
-                   "Positions type should be one of kTfLiteFloat32, "
-                   "kTfLiteFloat16, kTfLiteInt32, kTfLiteUInt8",
-                   &val_ctx);
+                           kTfLiteInt32, kTfLiteUInt8, kTfLiteInt8);
+
+      Expect(positions.type == kTfLiteInt32,
+             NNAPIValidationFailureType::kUnsupportedInputType,
+             "Positions type should be one of kTfLiteInt32", &val_ctx);
       Expect(positions.dims->size != 0,
              NNAPIValidationFailureType::kUnsupportedOperandRank,
              "0-dimension args are not supported by NNAPI.", &val_ctx);
@@ -2283,8 +2316,13 @@ bool NNAPIDelegateKernel::Validate(
                                  &val_ctx);
       // Tensor indices: split_dim: 0, value: 1
       const TfLiteTensor& input = context->tensors[node->inputs->data[1]];
-      EXPECT_INPUT_TYPE_IN(input.type, kTfLiteFloat32, kTfLiteUInt8,
-                           kTfLiteInt32);
+      if (android_sdk_version >= kMinSdkVersionForNNAPI13) {
+        EXPECT_INPUT_TYPE_IN(input.type, kTfLiteFloat32, kTfLiteUInt8,
+                             kTfLiteInt8, kTfLiteInt32);
+      } else {
+        EXPECT_INPUT_TYPE_IN(input.type, kTfLiteFloat32, kTfLiteUInt8,
+                             kTfLiteInt32);
+      }
       const TfLiteTensor& axis = context->tensors[node->inputs->data[0]];
       Expect(axis.type == kTfLiteInt32 && axis.allocation_type == kTfLiteMmapRo,
              NNAPIValidationFailureType::kUnsupportedInputType,
@@ -2308,30 +2346,41 @@ bool NNAPIDelegateKernel::Validate(
              NNAPIValidationFailureType::kUnsupportedInputType,
              "Value should be Float32.", &val_ctx);
       const auto output_type = context->tensors[node->outputs->data[0]].type;
-      Expect(output_type == kTfLiteUInt8,
-             NNAPIValidationFailureType::kUnsupportedOutputType,
-             "Output should be kTfLiteUInt8.", &val_ctx);
+      if (android_sdk_version < kMinSdkVersionForNNAPI13) {
+        Expect(output_type == kTfLiteUInt8,
+               NNAPIValidationFailureType::kUnsupportedOutputType,
+               "Output should be kTfLiteUInt8.", &val_ctx);
+      } else {
+        ExpectTypeIn(output_type, {kTfLiteUInt8, kTfLiteInt8},
+                     NNAPIValidationFailureType::kUnsupportedOutputType,
+                     "Output should be kTfLiteUInt8.", &val_ctx);
+      }
       const auto quantization_params =
           context->tensors[node->outputs->data[0]].params;
       Expect(quantization_params.scale > 0.f,
              NNAPIValidationFailureType::kUnsupportedQuantizationParameters,
              "Quantization scale should be > 0.", &val_ctx);
     } break;
-    case kTfLiteBuiltinReduceAny:
-    case kTfLiteBuiltinReduceMin:
-    case kTfLiteBuiltinReduceMax: {
-      ExpectOpVersion(version, 1, &val_ctx);
+    case kTfLiteBuiltinReduceAny: {
+      ExpectOpVersion(version, 2, &val_ctx);
       ExpectMinAndroidSdkVersion(android_sdk_version, kMinSdkVersionForNNAPI12,
                                  &val_ctx);
       Expect(context->tensors[node->outputs->data[0]].dims->size != 0,
              NNAPIValidationFailureType::kUnsupportedOutputType,
              "NNAPI does not support generating a scalar as output.", &val_ctx);
-      if (builtin_code == kTfLiteBuiltinReduceProd) {
-        const auto input_type = context->tensors[node->inputs->data[0]].type;
-        Expect(input_type == kTfLiteFloat32,
-               NNAPIValidationFailureType::kUnsupportedInputType,
-               "NNAPI only supports floating point REDUCE_PROD.", &val_ctx);
-      }
+    } break;
+    case kTfLiteBuiltinReduceMin:
+    case kTfLiteBuiltinReduceMax: {
+      ExpectMaxOpVersion(version, 2, &val_ctx);
+      ExpectMinAndroidSdkVersion(android_sdk_version, kMinSdkVersionForNNAPI12,
+                                 &val_ctx);
+      const auto input_tensor = context->tensors[node->inputs->data[0]];
+      const auto input_type = input_tensor.type;
+      EXPECT_INPUT_TYPE_IN(input_type, kTfLiteFloat32, kTfLiteUInt8,
+                           kTfLiteInt8);
+      Expect(input_tensor.dims->size != 0,
+             NNAPIValidationFailureType::kUnsupportedOutputType,
+             "NNAPI does not support generating a scalar as output.", &val_ctx);
     } break;
     case kTfLiteBuiltinDepthToSpace: {
       const TfLiteType input_type =
@@ -3093,16 +3142,10 @@ TfLiteStatus NNAPIDelegateKernel::Map(
     case kTfLiteBuiltinGather: {
       auto builtin = reinterpret_cast(
           mapping_args.node->builtin_data);
-      mapping_args.builder->AddTensorInput(mapping_args.node->inputs->data[0],
-                                           /* hybrid_op */ false,
-                                           /* scalar_as_tensor */ false);
-
       mapping_args.builder->AddScalarInt32Operand(builtin->axis);
-
       mapping_args.builder->AddTensorInput(mapping_args.node->inputs->data[1],
                                            /* hybrid_op */ false,
-                                           /* scalar_as_tensor */ false);
-
+                                           /* tensor_flags */ 0);
       *nn_op_type = ANEURALNETWORKS_GATHER;
     } break;
     case kTfLiteBuiltinBidirectionalSequenceLstm: {
@@ -3430,6 +3473,9 @@ TfLiteStatus NNAPIDelegateKernel::Invoke(TfLiteContext* context,
   // absolute indices but NN api indices inputs by relative indices.
   int relative_input_index = 0;
 
+  const bool use_int8_asymm_signed =
+      target_sdk_version_ >= kMinSdkVersionForNNAPI13;
+
   size_t input_offset = 0;
   for (auto absolute_input_index : TfLiteIntArrayView(node->inputs)) {
     if (absolute_input_index == kTfLiteOptionalTensor) {
@@ -3472,9 +3518,16 @@ TfLiteStatus NNAPIDelegateKernel::Invoke(TfLiteContext* context,
           }
         } else if (tensor->type == kTfLiteInt8 &&
                    ann_type_equivalent == kTfLiteInt32) {
-          for (int i = 0; i < num_elements; ++i) {
-            reinterpret_cast(input_ptr)[i] =
-                static_cast(tensor->data.int8[i]) + 128;
+          if (use_int8_asymm_signed) {
+            for (int i = 0; i < num_elements; ++i) {
+              reinterpret_cast(input_ptr)[i] =
+                  static_cast(tensor->data.int8[i]);
+            }
+          } else {
+            for (int i = 0; i < num_elements; ++i) {
+              reinterpret_cast(input_ptr)[i] =
+                  static_cast(tensor->data.int8[i]) + 128;
+            }
           }
         } else {
           context->ReportError(
@@ -3685,6 +3738,15 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
                          &dequantize_mapping, &allocation_memory_mapping_,
                          &nnapi_to_tflite_op_mapping_, nn_model_.get(),
                          nnapi_errno);
+
+  // If we have target accelerators the target SDK version might be
+  // different than the current android version.
+  target_sdk_version_ = nnapi_->android_sdk_version;
+  if (!nnapi_devices_.empty()) {
+    TF_LITE_ENSURE_STATUS(GetTargetSdkVersion(
+        context, nnapi_, nnapi_devices_, &target_sdk_version_, nnapi_errno));
+  }
+
   // Add Tensors.
   for (auto node_index : nodes_) {
     // Obtain the op and registration.
@@ -3696,11 +3758,18 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
     const bool hybrid_op = IsHybridOperator(context, reg->builtin_code, node);
     const bool scalar_as_tensor = IsScalarInputSupported(reg->builtin_code);
     const bool need_int8_conversion =
+        target_sdk_version_ < kMinSdkVersionForNNAPI13 &&
         NeedInt8Conversion(context, reg->builtin_code, node);
+    const bool use_int8_asymm_signed =
+        target_sdk_version_ >= kMinSdkVersionForNNAPI13 && !hybrid_op;
+
     int input_tensor_flags = 0;
     if (scalar_as_tensor) {
       input_tensor_flags |= NN_TENSOR_FLAG_SCALAR_AS_TENSOR;
     }
+    if (use_int8_asymm_signed) {
+      input_tensor_flags |= NN_TENSOR_FLAG_USE_INT8_ASYMM_SIGNED;
+    }
 
     // On SDK level less than 30, h_swish will be lowered into supported NNAPI
     // operations. Since SDK level 30, h_swish is supported as a single
@@ -3807,8 +3876,12 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
             break;
           case kTfLiteInt8:
             if (constant_value.allocation_type == kTfLiteMmapRo) {
-              builder.AddScalarInt32Operand(
-                  static_cast(*constant_value.data.int8) + 128);
+              if (need_int8_conversion) {
+                builder.AddScalarInt32Operand(
+                    static_cast(*constant_value.data.int8) + 128);
+              } else {
+                builder.AddScalarInt32Operand(*constant_value.data.int8);
+              }
             } else {
               builder.AddSingleValueTensorAsScalarOperand(
                   constant_value_id, ANEURALNETWORKS_INT32);
@@ -3836,7 +3909,8 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
           // specifying the output height and width, is not added and
           // instead the height and width will be added individually as
           // scalars by the mapping function returned by Map().
-          TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op));
+          TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op,
+                                                       input_tensor_flags));
         }
       } else if (reg->builtin_code == kTfLiteBuiltinTopkV2 && input_pos > 0) {
         // The K parameter tensor is not handled here but by the functor
@@ -3844,8 +3918,12 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
         // the else clause below
         continue;
       } else if (reg->builtin_code == kTfLiteBuiltinGather) {
-        // Everything is added during Map since input tensors
+        // Everything else is added during Map since input tensors
         // have different order.
+        if (input_pos == 0) {
+          TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op,
+                                                       input_tensor_flags));
+        }
         continue;
       } else if (reg->builtin_code == kTfLiteBuiltinExpandDims &&
                  input_pos == 1) {
@@ -3862,7 +3940,8 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
         // the axis, needs to be converted to a scalar since TFLite uses a
         // tensor but NNAPI uses a scalar as the axis.
         if (input_pos == 0) {
-          TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op));
+          TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index, hybrid_op,
+                                                       input_tensor_flags));
         } else {
           const int axis_id = node->inputs->data[1];
           const TfLiteTensor& axis_tensor = context->tensors[axis_id];
@@ -3908,12 +3987,26 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
                   std::vector(1, operand_tensor.data.uint8[0]),
                   operand_tensor.params, &tensor_index));
               break;
-            case kTfLiteInt8:
-              TF_LITE_ENSURE_STATUS(builder.AddNewInputConstantTensor(
-                  ANEURALNETWORKS_TENSOR_QUANT8_SYMM, operand_tensor.type, {1},
-                  std::vector(1, operand_tensor.data.int8[0]),
-                  operand_tensor.params, &tensor_index));
-              break;
+            case kTfLiteInt8: {
+              auto params = operand_tensor.params;
+              if (params.scale == 0.0) {
+                params.scale = 1.0;
+              }
+
+              if (use_int8_asymm_signed) {
+                TF_LITE_ENSURE_STATUS(builder.AddNewInputConstantTensor(
+                    ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED,
+                    operand_tensor.type, {1},
+                    std::vector(1, operand_tensor.data.int8[0]), params,
+                    &tensor_index));
+              } else {
+                TF_LITE_ENSURE_STATUS(builder.AddNewInputConstantTensor(
+                    ANEURALNETWORKS_TENSOR_QUANT8_ASYMM, operand_tensor.type,
+                    {1},
+                    std::vector(1, operand_tensor.data.int8[0] + 128),
+                    params, &tensor_index));
+              }
+            } break;
             case kTfLiteInt32:
               TF_LITE_ENSURE_STATUS(builder.AddNewInputConstantTensor(
                   ANEURALNETWORKS_TENSOR_INT32, operand_tensor.type, {1},
@@ -3995,19 +4088,11 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
       }
     }
 
-    // If we have target accelerators the target SDK version might be
-    // different than the current android version.
-    int target_sdk_version = nnapi_->android_sdk_version;
-    if (!nnapi_devices_.empty()) {
-      TF_LITE_ENSURE_STATUS(GetTargetSdkVersion(
-          context, nnapi_, nnapi_devices_, &target_sdk_version, nnapi_errno));
-    }
-
     // Get op type and operands
     // Fails if the Validate function failed
     int nn_op_type;
     TF_LITE_ENSURE_STATUS(
-        Map(context, reg->builtin_code, reg->version, target_sdk_version,
+        Map(context, reg->builtin_code, reg->version, target_sdk_version_,
             {context, &builder, node, &model_state_outputs_,
              &model_state_tfl_inputs_, &feedback_loops_, nnapi_errno},
             &nn_op_type));
@@ -4017,6 +4102,9 @@ TfLiteStatus NNAPIDelegateKernel::AddOpsAndTensors(TfLiteContext* context,
     if (need_int8_conversion) {
       output_tensor_flags |= NN_TENSOR_FLAG_INT8_CONVERSION;
     }
+    if (use_int8_asymm_signed) {
+      output_tensor_flags |= NN_TENSOR_FLAG_USE_INT8_ASYMM_SIGNED;
+    }
     for (int output_pos = 0; output_pos < node->outputs->size; ++output_pos) {
       const auto output_index = node->outputs->data[output_pos];
 
diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate_kernel.h b/tensorflow/lite/delegates/nnapi/nnapi_delegate_kernel.h
index 26822c011e3..9aa0f303cc2 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate_kernel.h
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate_kernel.h
@@ -341,6 +341,9 @@ class NNAPIDelegateKernel {
 
   std::vector nnapi_to_tflite_op_mapping_;
 
+  // Fully initialized in NNAPIDelegateKernel::AddOpsAndTensors
+  int target_sdk_version_ = 27;  // kMinSdkVersionForNNAPI13
+
   void AddDequantizeOperatorsWhereNeeded(
       const TfLiteContext* context, int builtin_code, const TfLiteNode* node,
       int tflite_node_index, NNAPIOpBuilder* builder, int* nnapi_errno);
diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h b/tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h
index fa7ff9dd1f1..5dbe4110131 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h
@@ -71,6 +71,8 @@ class NnApiMock : public ::tflite::nnapi::NnApiHandler {
     ExecutionComputeReturns();
     ExecutionStartComputeReturns();
     EventWaitReturns();
+    SetPriorityReturns();
+    SetOperandSymmPerChannelQuantParamsReturns();
     SetNnapiSupportedDevice("test-device", android_sdk_version);
   }
 
diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate_signed_quantization_test.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate_signed_quantization_test.cc
new file mode 100644
index 00000000000..b9d702015c2
--- /dev/null
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate_signed_quantization_test.cc
@@ -0,0 +1,920 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include 
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
+#include "tensorflow/lite/delegates/nnapi/nnapi_delegate_mock_test.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/kernels/fully_connected.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/nnapi/NeuralNetworksTypes.h"
+#include "tensorflow/lite/nnapi/nnapi_implementation.h"
+
+namespace tflite {
+
+namespace ops {
+namespace builtin {
+
+TfLiteRegistration* Register_CONVOLUTION_REF();
+TfLiteRegistration* Register_DEQUANTIZE();
+
+}  // namespace builtin
+}  // namespace ops
+
+namespace {
+
+class SingleOpModelWithNNAPI : public SingleOpModel {
+ public:
+  SingleOpModelWithNNAPI() = default;
+  void Init(const NnApi* nnapi) {
+    stateful_delegate_.reset(new StatefulNnApiDelegate(nnapi));
+    SetDelegate(stateful_delegate_.get());
+  }
+
+  StatefulNnApiDelegate* GetDelegate() { return stateful_delegate_.get(); }
+
+  void SetBufferHandle(int index, TfLiteBufferHandle handle) {
+    interpreter_->SetBufferHandle(index, handle, stateful_delegate_.get());
+  }
+  TfLiteStatus GetCompilationStatus() { return compilation_status_; }
+
+ protected:
+  std::unique_ptr stateful_delegate_;
+  TfLiteStatus compilation_status_;
+};
+
+class HybridFullyConnectedOpModel : public SingleOpModelWithNNAPI {
+ public:
+  HybridFullyConnectedOpModel(const NnApi* nnapi, int units, int batches,
+                              const TensorData& input,
+                              const TensorData& weights,
+                              const TensorData& output = {TensorType_FLOAT32},
+                              bool asymmetric_inputs = false)
+      : batches_(batches), units_(units) {
+    SingleOpModelWithNNAPI::Init(nnapi);
+    int total_input_size = 1;
+    for (size_t i = 0; i < input.shape.size(); ++i) {
+      total_input_size *= input.shape[i];
+    }
+    input_size_ = total_input_size / batches_;
+
+    input_ = AddInput(input);
+    weights_ = AddInput(weights);
+
+    TensorData bias{TensorType_FLOAT32, {units_}};
+    bias_ = AddInput(bias);
+
+    output_ = AddOutput(output);
+
+    auto options = CreateFullyConnectedOptions(
+                       builder_, ActivationFunctionType_RELU,
+                       tflite::FullyConnectedOptionsWeightsFormat_DEFAULT,
+                       false, asymmetric_inputs)
+                       .Union();
+    SetBuiltinOp(BuiltinOperator_FULLY_CONNECTED,
+                 BuiltinOptions_FullyConnectedOptions, options);
+    resolver_ = absl::make_unique(
+        BuiltinOperator_FULLY_CONNECTED,
+        ops::builtin::Register_FULLY_CONNECTED_PIE());
+    BuildInterpreter({GetShape(input_), GetShape(weights_), GetShape(bias_)},
+                     /*num_threads=*/-1,
+                     /* allow_fp32_relax_to_fp16 */ false,
+                     /*apply_delegate=*/false);
+    compilation_status_ = ApplyDelegate();
+  }
+  void SetBias(const std::vector& f) { PopulateTensor(bias_, f); }
+  void SetWeights(const std::vector& data) {
+    SymmetricQuantizeAndPopulate(weights_, data);
+  }
+  void SetSignedWeights(std::initializer_list f) {
+    SignedSymmetricQuantizeAndPopulate(weights_, f);
+  }
+
+  void SetInput(const std::vector& f) { PopulateTensor(input_, f); }
+  std::vector GetOutput() { return ExtractVector(output_); }
+  std::vector GetOutputShape() { return GetTensorShape(output_); }
+
+  int input_size() { return input_size_; }
+  int num_units() { return units_; }
+  int num_batches() { return batches_; }
+
+ protected:
+  int input_;
+  int weights_;
+  int bias_;
+  int output_;
+
+  int batches_;
+  int units_;
+  int input_size_;
+};
+
+struct NnApiSignedQuantizationTest
+    : ::tflite::delegate::nnapi::NnApiDelegateMockTest {
+  static void SetUpTestSuite() { tensors_count = new std::map(); }
+  void SetUp() override {
+    ::tflite::delegate::nnapi::NnApiDelegateMockTest::SetUp();
+    nnapi_mock_->StubAddOperandWith(
+        [](ANeuralNetworksModel* model,
+           const ANeuralNetworksOperandType* type) -> int {
+          const auto nn_tensor_type = type->type;
+          if (tensors_count->find(nn_tensor_type) == tensors_count->end()) {
+            tensors_count->insert({nn_tensor_type, 0});
+          }
+          tensors_count->at(nn_tensor_type)++;
+          return ANEURALNETWORKS_NO_ERROR;
+        });
+  }
+  void TearDown() override { tensors_count->clear(); }
+  static void TearDownTestSuite() {
+    delete tensors_count;
+    tensors_count = nullptr;
+  }
+  static std::map* tensors_count;
+};
+std::map* NnApiSignedQuantizationTest::tensors_count = nullptr;
+
+TEST_F(NnApiSignedQuantizationTest,
+       HybridFullyConnectedMapsToSignedSymmOnSdk29) {
+  nnapi_mock_->SetAndroidSdkVersion(29);
+
+  HybridFullyConnectedOpModel m(
+      nnapi_mock_->GetNnApi(), /*units=*/3, /*batches=*/2,
+      /*input=*/{TensorType_FLOAT32, {2, 10}},
+      /*weights=*/{TensorType_INT8, {3, 10}, 0, 0, 10.0 / 127.0, 0});
+  m.SetSignedWeights({
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 0
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 1
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 2
+  });
+  m.SetBias({1, 2, 3});
+  m.SetInput({
+      1, 2, 3, 4, 5, 6, 7, 8,  -9, -10,  // b = 0
+      1, 2, 3, 4, 5, 6, 7, -8, 9,  -10,  // b = 1
+  });
+
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            4);  // fc_input, fc_weights, fc_bias, fc_output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32), 1);  // activation
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            1);  // dequantize_weights_input
+}
+
+TEST_F(NnApiSignedQuantizationTest,
+       HybridFullyConnectedMapsToSignedSymmOnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+
+  HybridFullyConnectedOpModel m(
+      nnapi_mock_->GetNnApi(), /*units=*/3, /*batches=*/2,
+      /*input=*/{TensorType_FLOAT32, {2, 10}},
+      /*weights=*/{TensorType_INT8, {3, 10}, 0, 0, 10.0 / 127.0, 0});
+  m.SetSignedWeights({
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 0
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 1
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 2
+  });
+  m.SetBias({1, 2, 3});
+  m.SetInput({
+      1, 2, 3, 4, 5, 6, 7, 8,  -9, -10,  // b = 0
+      1, 2, 3, 4, 5, 6, 7, -8, 9,  -10,  // b = 1
+  });
+
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            4);  // fc_input, fc_weights, fc_bias, fc_output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32), 1);  // activation
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            1);  // dequantize_weights_input
+}
+
+template 
+class BaseConvolutionOpModel : public SingleOpModelWithNNAPI {
+ public:
+  BaseConvolutionOpModel(
+      const NnApi* nnapi, TfLiteRegistration* registration,
+      const TensorData& input, const TensorData& filter,
+      const TensorData& output, int stride_width = 2, int stride_height = 2,
+      enum Padding padding = Padding_VALID,
+      enum ActivationFunctionType activation = ActivationFunctionType_NONE,
+      int dilation_width_factor = 1, int dilation_height_factor = 1,
+      std::initializer_list filter_data = {}) {
+    SingleOpModelWithNNAPI::Init(nnapi);
+
+    input_ = AddInput(input);
+
+    if (filter_data.size()) {
+      filter_ = AddConstInput(filter, filter_data);
+    } else {
+      filter_ = AddInput(filter);
+    }
+
+    int bias_size = GetShape(filter_)[0];
+    if (input.type == TensorType_FLOAT32) {
+      bias_ = AddInput({TensorType_FLOAT32, {bias_size}});
+    } else {
+      // This is a quantized version. The scale of 'bias' depends on the scales
+      // of input and filter. Supposedly this is correctly set during quantized
+      // training.
+      if (filter.per_channel_quantization) {
+        // per channel quantization.
+        std::vector bias_scale(
+            filter.per_channel_quantization_scales.size());
+        std::vector bias_zero_points(
+            filter.per_channel_quantization_scales.size());
+        for (size_t i = 0; i < filter.per_channel_quantization_scales.size();
+             ++i) {
+          bias_scale[i] =
+              input.scale * filter.per_channel_quantization_scales[i];
+          bias_zero_points[i] = 0;
+        }
+        tflite::TensorType bias_type = TensorType_INT32;
+        if (input.type == TensorType_INT16) {
+          // In case of 16-bit, the bias type is set to be int 64.
+          bias_type = TensorType_INT64;
+        }
+        TensorData bias{bias_type,
+                        {bias_size},
+                        /*min=*/0,
+                        /*max=*/0,
+                        /*scale=*/0,
+                        /*zero_point=*/0,
+                        true,
+                        /*per_channel_quantization_scales=*/bias_scale,
+                        /*per_channel_quantization_offsets=*/bias_zero_points,
+                        /*channel_index==*/0};
+        bias_ = AddInput(bias);
+      } else {
+        // per tensor quantization.
+        auto bias_scale = GetScale(input_) * GetScale(filter_);
+        TensorData bias{TensorType_INT32, {bias_size}, 0, 0, bias_scale};
+        bias_ = AddInput(bias);
+      }
+    }
+
+    output_ = AddOutput(output);
+
+    SetBuiltinOp(BuiltinOperator_CONV_2D, BuiltinOptions_Conv2DOptions,
+                 CreateConv2DOptions(
+                     builder_, padding, stride_width, stride_height, activation,
+                     dilation_width_factor, dilation_height_factor)
+                     .Union());
+
+    resolver_ = absl::make_unique(BuiltinOperator_CONV_2D,
+                                                    registration);
+    BuildInterpreter({GetShape(input_), GetShape(filter_), GetShape(bias_)},
+                     /*num_threads=*/-1,
+                     /* allow_fp32_relax_to_fp16 */ false,
+                     /*apply_delegate=*/false);
+    compilation_status_ = ApplyDelegate();
+  }
+
+ protected:
+  int input_;
+  int filter_;
+  int bias_;
+  int output_;
+};
+
+class QuantizedConvolutionOpModel : public BaseConvolutionOpModel {
+ public:
+  using BaseConvolutionOpModel::BaseConvolutionOpModel;
+
+  void SetInput(std::initializer_list data) {
+    QuantizeAndPopulate(input_, data);
+  }
+
+  void SetFilter(std::initializer_list data) {
+    QuantizeAndPopulate(filter_, data);
+  }
+
+  void SetBias(std::initializer_list data) {
+    QuantizeAndPopulate(bias_, data);
+  }
+
+  std::vector GetOutput() { return ExtractVector(output_); }
+  std::vector GetDequantizedOutput() {
+    return Dequantize(ExtractVector(output_),
+                               GetScale(output_), GetZeroPoint(output_));
+  }
+};
+
+TEST_F(NnApiSignedQuantizationTest,
+       Conv2DUnsignedPerTensorMapsToUnsignedOnSdk29) {
+  QuantizedConvolutionOpModel m(nnapi_mock_->GetNnApi(),
+                                ops::builtin::Register_CONVOLUTION_REF(),
+                                {TensorType_UINT8, {2, 2, 4, 1}, -63.5, 64},
+                                {TensorType_UINT8, {3, 2, 2, 1}, -63.5, 64},
+                                {TensorType_UINT8, {}, -127, 128});
+  m.SetInput({
+      // First batch
+      1, 1, 1, 1,  // row = 1
+      2, 2, 2, 2,  // row = 2
+      // Second batch
+      1, 2, 3, 4,  // row = 1
+      1, 2, 3, 4,  // row = 2
+  });
+  m.SetFilter({
+      1, 2, 3, 4,    // first 2x2 filter
+      -1, 1, -1, 1,  // second 2x2 filter
+      -1, -1, 1, 1,  // third 2x2 filter
+  });
+  m.SetBias({1, 2, 3});
+
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            3);  // input, filter, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+TEST_F(NnApiSignedQuantizationTest,
+       Conv2dUnsignedPerTensorMapsToUnsignedOnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  QuantizedConvolutionOpModel m(nnapi_mock_->GetNnApi(),
+                                ops::builtin::Register_CONVOLUTION_REF(),
+                                {TensorType_UINT8, {2, 2, 4, 1}, -63.5, 64},
+                                {TensorType_UINT8, {3, 2, 2, 1}, -63.5, 64},
+                                {TensorType_UINT8, {}, -127, 128});
+  m.SetInput({
+      // First batch
+      1, 1, 1, 1,  // row = 1
+      2, 2, 2, 2,  // row = 2
+      // Second batch
+      1, 2, 3, 4,  // row = 1
+      1, 2, 3, 4,  // row = 2
+  });
+  m.SetFilter({
+      1, 2, 3, 4,    // first 2x2 filter
+      -1, 1, -1, 1,  // second 2x2 filter
+      -1, -1, 1, 1,  // third 2x2 filter
+  });
+  m.SetBias({1, 2, 3});
+
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            3);  // input, filter, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+class PerChannelQuantizedConvolutionOpModel
+    : public BaseConvolutionOpModel {
+ public:
+  using BaseConvolutionOpModel::BaseConvolutionOpModel;
+
+  void SetInput(std::initializer_list data) {
+    QuantizeAndPopulate(input_, data);
+  }
+
+  void SetFilter(std::initializer_list data) {
+    PerChannelSymmetricQuantizeAndPopulate(filter_, data);
+  }
+
+  void SetBias(std::initializer_list data) {
+    PerChannelQuantizeBias(bias_, data);
+  }
+
+  std::vector GetOutput() { return ExtractVector(output_); }
+  std::vector GetDequantizedOutput() {
+    return Dequantize(ExtractVector(output_), GetScale(output_),
+                              GetZeroPoint(output_));
+  }
+};
+
+TEST_F(NnApiSignedQuantizationTest,
+       Conv2dSignedPerTensorMapsToUnsignedOnSdk29) {
+  nnapi_mock_->SetAndroidSdkVersion(29);
+  PerChannelQuantizedConvolutionOpModel m(
+      nnapi_mock_->GetNnApi(), ops::builtin::Register_CONVOLUTION_REF(),
+      {TensorType_INT8, {1, 2, 3, 2}, -63.5, 64, 0.5, -1},
+      {TensorType_INT8,
+       // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+       {2, 2, 2, 2},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{1},
+       /*per_channel_quantization_offsets=*/{0},
+       /*channel_index=*/0},
+      {TensorType_INT8, {}, -63.5, 64, 0.5, -1},
+      /*stride_width=*/1, /*stride_height=*/1);
+  m.SetInput({
+      // [1 * 2 * 3 * 2] as [batch, y, x, input_channel]
+      3, 2,    // batch = 0, y = 0, x = 0
+      1, -1,   // batch = 0, y = 0, x = 1
+      -2, -3,  // batch = 0, y = 0, x = 2
+      4, 3,    // batch = 0, y = 1, x = 0
+      2, -2,   // batch = 0, y = 1, x = 1
+      -3, -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetFilter(
+      // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+      {
+          1, 2,  // out channel = 0, y = 0, x = 0
+          3, 4,  // out channel = 0, y = 0, x = 1
+          3, 4,  // out channel = 0, y = 1, x = 0
+          5, 6,  // out channel = 0, y = 1, x = 1
+          7, 8,  // out channel = 1, y = 0, x = 0
+          5, 6,  // out channel = 1, y = 0, x = 1
+          3, 4,  // out channel = 1, y = 1, x = 0
+          1, 2,  // out channel = 1, y = 1, x = 1
+      });
+  m.SetBias({3, -2});
+
+  // Invoke and verify output.
+  // output has dimension [1 * 1 * 2 * 2] as [batch, y, x, output_channel]
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            3);  // input, filter, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+TEST_F(NnApiSignedQuantizationTest,
+       Conv2dSignedPerTensorMapsToUnsignedOnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  PerChannelQuantizedConvolutionOpModel m(
+      nnapi_mock_->GetNnApi(), ops::builtin::Register_CONVOLUTION_REF(),
+      {TensorType_INT8, {1, 2, 3, 2}, -63.5, 64, 0.5, -1},
+      {TensorType_INT8,
+       // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+       {2, 2, 2, 2},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{1},
+       /*per_channel_quantization_offsets=*/{0},
+       /*channel_index=*/0},
+      {TensorType_INT8, {}, -63.5, 64, 0.5, -1},
+      /*stride_width=*/1, /*stride_height=*/1);
+  m.SetInput({
+      // [1 * 2 * 3 * 2] as [batch, y, x, input_channel]
+      3, 2,    // batch = 0, y = 0, x = 0
+      1, -1,   // batch = 0, y = 0, x = 1
+      -2, -3,  // batch = 0, y = 0, x = 2
+      4, 3,    // batch = 0, y = 1, x = 0
+      2, -2,   // batch = 0, y = 1, x = 1
+      -3, -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetFilter(
+      // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+      {
+          1, 2,  // out channel = 0, y = 0, x = 0
+          3, 4,  // out channel = 0, y = 0, x = 1
+          3, 4,  // out channel = 0, y = 1, x = 0
+          5, 6,  // out channel = 0, y = 1, x = 1
+          7, 8,  // out channel = 1, y = 0, x = 0
+          5, 6,  // out channel = 1, y = 0, x = 1
+          3, 4,  // out channel = 1, y = 1, x = 0
+          1, 2,  // out channel = 1, y = 1, x = 1
+      });
+  m.SetBias({3, -2});
+
+  // Invoke and verify output.
+  // output has dimension [1 * 1 * 2 * 2] as [batch, y, x, output_channel]
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 3);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            3);  // input, filter, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+TEST_F(NnApiSignedQuantizationTest,
+       Conv2dSignedPerChannelMapsToUnsignedOnSdk29) {
+  PerChannelQuantizedConvolutionOpModel m(
+      nnapi_mock_->GetNnApi(), ops::builtin::Register_CONVOLUTION_REF(),
+      {TensorType_INT8, {1, 2, 3, 2}, -63.5, 64, 0.5, -1},
+      {TensorType_INT8,
+       // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+       {2, 2, 2, 2},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{1, 2},
+       /*per_channel_quantization_offsets=*/{0, 0},
+       /*channel_index=*/0},
+      {TensorType_INT8, {}, -63.5, 64, 0.5, -1},
+      /*stride_width=*/1, /*stride_height=*/1);
+  m.SetInput({
+      // [1 * 2 * 3 * 2] as [batch, y, x, input_channel]
+      3, 2,    // batch = 0, y = 0, x = 0
+      1, -1,   // batch = 0, y = 0, x = 1
+      -2, -3,  // batch = 0, y = 0, x = 2
+      4, 3,    // batch = 0, y = 1, x = 0
+      2, -2,   // batch = 0, y = 1, x = 1
+      -3, -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetFilter(
+      // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+      {
+          1, 2,  // out channel = 0, y = 0, x = 0
+          3, 4,  // out channel = 0, y = 0, x = 1
+          3, 4,  // out channel = 0, y = 1, x = 0
+          5, 6,  // out channel = 0, y = 1, x = 1
+          7, 8,  // out channel = 1, y = 0, x = 0
+          5, 6,  // out channel = 1, y = 0, x = 1
+          3, 4,  // out channel = 1, y = 1, x = 0
+          1, 2,  // out channel = 1, y = 1, x = 1
+      });
+  m.SetBias({3, -2});
+
+  // Invoke and verify output.
+  // output has dimension [1 * 1 * 2 * 2] as [batch, y, x, output_channel]
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 4);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_SYMM_PER_CHANNEL),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            2);  // input, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_SYMM_PER_CHANNEL),
+            1);                                                   // filter
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+TEST_F(NnApiSignedQuantizationTest, Conv2dSignedPerChannelMapsToSignedOnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  PerChannelQuantizedConvolutionOpModel m(
+      nnapi_mock_->GetNnApi(), ops::builtin::Register_CONVOLUTION_REF(),
+      {TensorType_INT8, {1, 2, 3, 2}, -63.5, 64, 0.5, -1},
+      {TensorType_INT8,
+       // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+       {2, 2, 2, 2},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{1, 2},
+       /*per_channel_quantization_offsets=*/{0, 0},
+       /*channel_index=*/0},
+      {TensorType_INT8, {}, -63.5, 64, 0.5, -1},
+      /*stride_width=*/1, /*stride_height=*/1);
+  m.SetInput({
+      // [1 * 2 * 3 * 2] as [batch, y, x, input_channel]
+      3, 2,    // batch = 0, y = 0, x = 0
+      1, -1,   // batch = 0, y = 0, x = 1
+      -2, -3,  // batch = 0, y = 0, x = 2
+      4, 3,    // batch = 0, y = 1, x = 0
+      2, -2,   // batch = 0, y = 1, x = 1
+      -3, -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetFilter(
+      // [2 * 2 * 2 * 2] as [output_channel, y, x, input_channel]
+      {
+          1, 2,  // out channel = 0, y = 0, x = 0
+          3, 4,  // out channel = 0, y = 0, x = 1
+          3, 4,  // out channel = 0, y = 1, x = 0
+          5, 6,  // out channel = 0, y = 1, x = 1
+          7, 8,  // out channel = 1, y = 0, x = 0
+          5, 6,  // out channel = 1, y = 0, x = 1
+          3, 4,  // out channel = 1, y = 1, x = 0
+          1, 2,  // out channel = 1, y = 1, x = 1
+      });
+  m.SetBias({3, -2});
+
+  // Invoke and verify output.
+  // output has dimension [1 * 1 * 2 * 2] as [batch, y, x, output_channel]
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 4);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_SYMM_PER_CHANNEL),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_INT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_INT32), tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            2);  // input, output
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_SYMM_PER_CHANNEL),
+            1);                                                   // filter
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_INT32), 1);  // bias
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_INT32),
+            4);  //  padding, stride_width, stride_height, activation
+}
+
+class QuantizeOpModel : public SingleOpModelWithNNAPI {
+ public:
+  QuantizeOpModel(const NnApi* nnapi, const TensorData& input,
+                  const TensorData& output) {
+    SingleOpModelWithNNAPI::Init(nnapi);
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_QUANTIZE, BuiltinOptions_QuantizeOptions,
+                 CreateQuantizeOptions(builder_).Union());
+
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1,
+                     /* allow_fp32_relax_to_fp16 */ false,
+                     /*apply_delegate=*/false);
+    compilation_status_ = ApplyDelegate();
+  }
+
+  void SetInput(std::initializer_list data) {
+    PopulateTensor(input_, data);
+  }
+
+  template 
+  void SetInputAndQuantize(std::initializer_list data) {
+    QuantizeAndPopulate(input_, data);
+  }
+
+  template 
+  std::vector GetOutput() {
+    return ExtractVector(output_);
+  }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST_F(NnApiSignedQuantizationTest, QuantizeUint8MapsToUint8OnSdk29) {
+  // [-63.5, 64] -> scale=0.5 zero_point=127 for UINT8
+  QuantizeOpModel m(nnapi_mock_->GetNnApi(), {TensorType_FLOAT32, {2, 5}},
+                    {TensorType_UINT8, {2, 5}, 0, 0, 0.5, 127});
+
+  m.SetInput({-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            1);  // output
+}
+
+TEST_F(NnApiSignedQuantizationTest, QuantizeUint8MapsToUint8OnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  // [-63.5, 64] -> scale=0.5 zero_point=127 for UINT8
+  QuantizeOpModel m(nnapi_mock_->GetNnApi(), {TensorType_FLOAT32, {2, 5}},
+                    {TensorType_UINT8, {2, 5}, 0, 0, 0.5, 127});
+
+  m.SetInput({-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            1);  // output
+}
+
+// Quantize with Int8 output is only supported since SDK level 30.
+TEST_F(NnApiSignedQuantizationTest, QuantizeInt8MapsToInt8OnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  // [-63.5, 64] -> scale=0.5 zero_point=1 for INT8
+  QuantizeOpModel m(nnapi_mock_->GetNnApi(), {TensorType_FLOAT32, {2, 5}},
+                    {TensorType_INT8, {2, 5}, 0, 0, 0.5, -1});
+
+  m.SetInput({-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            1);  // output
+}
+
+class DequantizeOpModel : public SingleOpModelWithNNAPI {
+ public:
+  DequantizeOpModel(const NnApi* nnapi, TensorType type,
+                    std::initializer_list shape, float scale,
+                    int32_t zero_point, int version) {
+    SingleOpModelWithNNAPI::Init(nnapi);
+    const TensorData input_tensor_data = {type, shape, 0, 0, scale, zero_point};
+    input_ = AddInput(input_tensor_data);
+    output_ = AddOutput({TensorType_FLOAT32, shape});
+    SetBuiltinOp(BuiltinOperator_DEQUANTIZE, BuiltinOptions_DequantizeOptions,
+                 CreateDequantizeOptions(builder_).Union());
+
+    resolver_ = absl::make_unique(
+        BuiltinOperator_DEQUANTIZE, ops::builtin::Register_DEQUANTIZE(),
+        version);
+
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1,
+                     /* allow_fp32_relax_to_fp16 */ false,
+                     /*apply_delegate=*/false);
+    compilation_status_ = ApplyDelegate();
+  }
+
+  template 
+  void SetInput(std::initializer_list data) {
+    PopulateTensor(input_, data);
+  }
+
+  std::vector GetOutput() { return ExtractVector(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST_F(NnApiSignedQuantizationTest, DequantizeUint8MapsToUint8OnSdk29) {
+  // [-63.5, 64] -> scale=0.5 zero_point=127 for UINT8
+  DequantizeOpModel m(nnapi_mock_->GetNnApi(), TensorType_UINT8, {2, 5}, 0.5,
+                      127, 1);
+
+  m.SetInput({0, 1, 2, 3, 4, 251, 252, 253, 254, 255});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // output
+}
+
+TEST_F(NnApiSignedQuantizationTest, DequantizeUint8MapsToUint8OnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  // [-63.5, 64] -> scale=0.5 zero_point=127 for UINT8
+  DequantizeOpModel m(nnapi_mock_->GetNnApi(), TensorType_UINT8, {2, 5}, 0.5,
+                      127, 1);
+
+  m.SetInput({0, 1, 2, 3, 4, 251, 252, 253, 254, 255});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // output
+}
+
+// Dequantize with Int8 input is only supported for symmetric quantization on
+// SDK level 29
+TEST_F(NnApiSignedQuantizationTest,
+       DequantizeTestInt8SymmMapsToInt8SymmOnSdk29) {
+  // [-63.5, 64] -> scale=0.5, zero_point=0 for INT8
+  DequantizeOpModel m(nnapi_mock_->GetNnApi(), TensorType_INT8, {2, 5}, 0.5, 0,
+                      2);
+
+  m.SetInput({-128, -127, -126, -125, -124, 123, 124, 125, 126, 127});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_SYMM),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // output
+}
+
+// Dequantize with Int8 input is only supported since SDK level 30.
+TEST_F(NnApiSignedQuantizationTest, DequantizeTestInt8MapsToInt8OnSdk30) {
+  nnapi_mock_->SetAndroidSdkVersion(30);
+  // [-63.5, 64] -> scale=0.5, zero_point=1 for INT8
+  DequantizeOpModel m(nnapi_mock_->GetNnApi(), TensorType_INT8, {2, 5}, 0.5, -1,
+                      2);
+
+  m.SetInput({-128, -127, -126, -125, -124, 123, 124, 125, 126, 127});
+  m.Invoke();
+  EXPECT_EQ(m.GetCompilationStatus(), kTfLiteOk);
+
+  ASSERT_EQ(tensors_count->size(), 2);
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            tensors_count->end());
+  ASSERT_NE(tensors_count->find(ANEURALNETWORKS_TENSOR_FLOAT32),
+            tensors_count->end());
+
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED),
+            1);  // input
+  EXPECT_EQ(tensors_count->at(ANEURALNETWORKS_TENSOR_FLOAT32),
+            1);  // output
+}
+
+}  // namespace
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  ::tflite::LogToStderr();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
diff --git a/tensorflow/lite/nnapi/NeuralNetworksTypes.h b/tensorflow/lite/nnapi/NeuralNetworksTypes.h
index 3c30a0479fa..8415df58b8b 100644
--- a/tensorflow/lite/nnapi/NeuralNetworksTypes.h
+++ b/tensorflow/lite/nnapi/NeuralNetworksTypes.h
@@ -46,6 +46,7 @@ enum {
   ANEURALNETWORKS_TENSOR_QUANT16_SYMM = 7,
   ANEURALNETWORKS_TENSOR_QUANT8_SYMM_PER_CHANNEL = 11,
   ANEURALNETWORKS_TENSOR_QUANT8_SYMM = 13,
+  ANEURALNETWORKS_TENSOR_QUANT8_ASYMM_SIGNED = 14,
 };
 
 /**
diff --git a/tensorflow/lite/nnapi/nnapi_handler.h b/tensorflow/lite/nnapi/nnapi_handler.h
index 00c0b23e3cf..16e1e9fea10 100644
--- a/tensorflow/lite/nnapi/nnapi_handler.h
+++ b/tensorflow/lite/nnapi/nnapi_handler.h
@@ -118,6 +118,11 @@ class NnApiHandler {
            const ANeuralNetworksOperandType* type) { return Value; };
   }
 
+  void StubAddOperandWith(int(stub)(ANeuralNetworksModel* model,
+                                    const ANeuralNetworksOperandType* type)) {
+    nnapi_->ANeuralNetworksModel_addOperand = stub;
+  }
+
   template 
   void SetOperandValueReturns() {
     nnapi_->ANeuralNetworksModel_setOperandValue =
@@ -268,6 +273,23 @@ class NnApiHandler {
     };
   }
 
+  template 
+  void SetPriorityReturns() {
+    nnapi_->ANeuralNetworksCompilation_setPriority =
+        [](ANeuralNetworksCompilation* compilation, int priority) -> int {
+      return Value;
+    };
+  }
+
+  template 
+  void SetOperandSymmPerChannelQuantParamsReturns() {
+    nnapi_->ANeuralNetworksModel_setOperandSymmPerChannelQuantParams =
+        [](ANeuralNetworksModel* model, int32_t index,
+           const ANeuralNetworksSymmPerChannelQuantParams* channelQuant) {
+          return Value;
+        };
+  }
+
   /*
    * Sets the SDK Version in the nnapi structure.
    * If set_unsupported_ops_to_null is set to true, all the functions not

From 4fc5c50e15d48ae50e3d7c192a6efad19fafdb3d Mon Sep 17 00:00:00 2001
From: Lukas Geiger 
Date: Tue, 23 Jun 2020 16:00:35 +0200
Subject: [PATCH 0882/1390] Prefer generator expressions over list
 comprehensions

---
 tensorflow/python/eager/function.py                  |  4 ++--
 tensorflow/python/keras/engine/training_v1.py        |  4 ++--
 .../keras/layers/preprocessing/category_crossing.py  |  8 ++++----
 .../python/keras/layers/preprocessing/hashing.py     | 12 ++++++------
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index e2f5d86fbd2..beb58d8ff91 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -1192,7 +1192,7 @@ class _TapeGradientFunctions(object):
   def _wrap_backward_function(self, forward_graph, backward, outputs):
     """Create a backward function given `outputs` from the forward function."""
     capture_mapping = dict(
-        zip([ops.tensor_id(t) for t in forward_graph.outputs], outputs))
+        zip((ops.tensor_id(t) for t in forward_graph.outputs), outputs))
     remapped_captures = [
         capture_mapping.get(ops.tensor_id(capture), capture)
         for capture in backward.captured_inputs
@@ -1491,7 +1491,7 @@ class ConcreteFunction(object):
     self._ndarrays_list = (
         isinstance(structured_outputs, (list, tuple)) and
         structured_outputs and
-        all([isinstance(o, np_arrays.ndarray) for o in structured_outputs]))
+        all(isinstance(o, np_arrays.ndarray) for o in structured_outputs))
     self._ndarray_singleton = isinstance(structured_outputs, np_arrays.ndarray)
 
     # function_spec defines the structured signature.
diff --git a/tensorflow/python/keras/engine/training_v1.py b/tensorflow/python/keras/engine/training_v1.py
index c137c6e517a..c901b73c892 100644
--- a/tensorflow/python/keras/engine/training_v1.py
+++ b/tensorflow/python/keras/engine/training_v1.py
@@ -1492,8 +1492,8 @@ class Model(training_lib.Model):
   def _recompile_weights_loss_and_weighted_metrics(self):
     if not self._is_compiled:
       return False
-    recompile = any([e.sample_weights_mismatch()
-                     for e in self._training_endpoints])
+    recompile = any(e.sample_weights_mismatch()
+                    for e in self._training_endpoints)
 
     if recompile:
       self._compile_weights_loss_and_weighted_metrics()
diff --git a/tensorflow/python/keras/layers/preprocessing/category_crossing.py b/tensorflow/python/keras/layers/preprocessing/category_crossing.py
index 594b9741946..e949bd4c87c 100644
--- a/tensorflow/python/keras/layers/preprocessing/category_crossing.py
+++ b/tensorflow/python/keras/layers/preprocessing/category_crossing.py
@@ -188,15 +188,15 @@ class CategoryCrossing(Layer):
   def compute_output_signature(self, input_spec):
     input_shapes = [x.shape for x in input_spec]
     output_shape = self.compute_output_shape(input_shapes)
-    if any([
+    if any(
         isinstance(inp_spec, ragged_tensor.RaggedTensorSpec)
         for inp_spec in input_spec
-    ]):
+    ):
       return tensor_spec.TensorSpec(shape=output_shape, dtype=dtypes.string)
-    elif any([
+    elif any(
         isinstance(inp_spec, sparse_tensor.SparseTensorSpec)
         for inp_spec in input_spec
-    ]):
+    ):
       return sparse_tensor.SparseTensorSpec(
           shape=output_shape, dtype=dtypes.string)
     return tensor_spec.TensorSpec(shape=output_shape, dtype=dtypes.string)
diff --git a/tensorflow/python/keras/layers/preprocessing/hashing.py b/tensorflow/python/keras/layers/preprocessing/hashing.py
index faeeec63a86..89c3042ae24 100644
--- a/tensorflow/python/keras/layers/preprocessing/hashing.py
+++ b/tensorflow/python/keras/layers/preprocessing/hashing.py
@@ -158,10 +158,10 @@ class Hashing(Layer):
   def _preprocess_inputs(self, inputs):
     if isinstance(inputs, (tuple, list)):
       # If any of them is tensor or ndarray, then treat as list
-      if any([
+      if any(
           tensor_util.is_tensor(inp) or isinstance(inp, np.ndarray)
           for inp in inputs
-      ]):
+      ):
         return [self._preprocess_single_input(inp) for inp in inputs]
     return self._preprocess_single_input(inputs)
 
@@ -261,15 +261,15 @@ class Hashing(Layer):
         return tensor_spec.TensorSpec(shape=output_shape, dtype=output_dtype)
     input_shapes = [x.shape for x in input_spec]
     output_shape = self.compute_output_shape(input_shapes)
-    if any([
+    if any(
         isinstance(inp_spec, ragged_tensor.RaggedTensorSpec)
         for inp_spec in input_spec
-    ]):
+    ):
       return tensor_spec.TensorSpec(shape=output_shape, dtype=dtypes.int64)
-    elif any([
+    elif any(
         isinstance(inp_spec, sparse_tensor.SparseTensorSpec)
         for inp_spec in input_spec
-    ]):
+    ):
       return sparse_tensor.SparseTensorSpec(
           shape=output_shape, dtype=dtypes.int64)
     return tensor_spec.TensorSpec(shape=output_shape, dtype=dtypes.int64)

From ba915231715348096e18775581027d7e625d8485 Mon Sep 17 00:00:00 2001
From: Stephan Herhut 
Date: Tue, 23 Jun 2020 06:18:20 -0700
Subject: [PATCH 0883/1390] Add call-back to further configure the
 llvm::TargetMachine that was selected for PTX generation.

This is used in kernel generation to enable FMA operation fusion.

PiperOrigin-RevId: 317851871
Change-Id: Ifaaa1ba99e5aecb90f380d085583bb8f56805d9f
---
 .../compiler/mlir/tools/kernel_gen/cubin_creator.cc  | 12 +++++++++---
 .../service/gpu/llvm_gpu_backend/gpu_backend_lib.cc  | 12 +++++++++---
 .../service/gpu/llvm_gpu_backend/gpu_backend_lib.h   |  8 +++++---
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
index 79969a22572..85a53e042e1 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/cubin_creator.cc
@@ -268,6 +268,7 @@ StatusOr> tensorflow::kernel_gen::GenerateCubinForTfCode(
     options.tile_sizes = tile_sizes;
     options.unroll_factors = unroll_factors;
     options.collapse_parallel_loops = false;
+    options.use_approximations = true;
     TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerLHLOToGPU(module.get(), options));
   }
   TF_RETURN_IF_ERROR(xla::mlir_gpu::LowerKernelBodiesToNVVM(module.get()));
@@ -287,10 +288,15 @@ StatusOr> tensorflow::kernel_gen::GenerateCubinForTfCode(
   xla::HloModuleConfig config;
   config.set_debug_options(xla::GetDebugOptionsFromFlags());
 
+  auto enable_fusion = [](llvm::TargetMachine* target) {
+    target->Options.AllowFPOpFusion = llvm::FPOpFusion::FPOpFusionMode::Fast;
+  };
+
   TF_ASSIGN_OR_RETURN(std::string libdevice_dir, GetLibdeviceDir(config));
-  TF_ASSIGN_OR_RETURN(std::string ptx, xla::gpu::nvptx::CompileToPtx(
-                                           llvmModule.get(), compute_capability,
-                                           config, libdevice_dir));
+  TF_ASSIGN_OR_RETURN(
+      std::string ptx,
+      xla::gpu::nvptx::CompileToPtx(llvmModule.get(), compute_capability,
+                                    config, libdevice_dir, enable_fusion));
   VLOG(1) << ptx;
 
 #if GOOGLE_CUDA
diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
index 497dcda4361..d2126a8d17d 100644
--- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
+++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -492,9 +492,10 @@ void NVPTXBackendInit(const HloModuleConfig& hlo_module_config) {
 
 namespace nvptx {
 
-StatusOr CompileToPtx(llvm::Module* module, GpuVersion gpu_version,
-                              const HloModuleConfig& hlo_module_config,
-                              const string& libdevice_dir_path) {
+StatusOr CompileToPtx(
+    llvm::Module* module, GpuVersion gpu_version,
+    const HloModuleConfig& hlo_module_config, const string& libdevice_dir_path,
+    std::function configure_target) {
   static absl::once_flag backend_init_flag;
   absl::call_once(backend_init_flag, NVPTXBackendInit, hlo_module_config);
 
@@ -525,6 +526,11 @@ StatusOr CompileToPtx(llvm::Module* module, GpuVersion gpu_version,
     std::unique_ptr target_machine = NVPTXGetTargetMachine(
         default_target_triple, *compute_capability, hlo_module_config);
 
+    // Apply target machine configuration from call-back if available.
+    if (configure_target) {
+      configure_target(target_machine.get());
+    }
+
     // Link with libdevice, and optimize the LLVM module.
     TF_RETURN_IF_ERROR(LinkAndOptimizeModule(
         module, gpu_version, hlo_module_config, libdevice_dir_path,
diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h
index 526621de7a5..33ef9280c7a 100644
--- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h
+++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/strings/string_view.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Target/TargetMachine.h"
 #include "tensorflow/compiler/xla/service/gpu/gpu_types.h"
 #include "tensorflow/compiler/xla/service/hlo_module_config.h"
 #include "tensorflow/compiler/xla/statusor.h"
@@ -38,9 +39,10 @@ namespace nvptx {
 // The Compile.* interfaces each create their own llvm::LLVMContext objects for
 // thread safety, but note that LLVM's multithreaded support is very
 // preliminary; multithreaded use is not recommended at this time.
-StatusOr CompileToPtx(llvm::Module* module, GpuVersion gpu_version,
-                              const HloModuleConfig& hlo_module_config,
-                              const string& libdevice_dir_path);
+StatusOr CompileToPtx(
+    llvm::Module* module, GpuVersion gpu_version,
+    const HloModuleConfig& hlo_module_config, const string& libdevice_dir_path,
+    std::function configure_target = nullptr);
 }  // namespace nvptx
 
 namespace amdgpu {

From c67719c6c34207e0e8a1f5e6b429d20130cd71f9 Mon Sep 17 00:00:00 2001
From: Tres Popp 
Date: Tue, 23 Jun 2020 06:58:18 -0700
Subject: [PATCH 0884/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/f1c671925b1c

PiperOrigin-RevId: 317856949
Change-Id: I192ce19253524e7957d7fc0a83abaac4adb31772
---
 tensorflow/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 7c47628818b..d34e7d973d3 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "7a55d984971c11daa55e9423934f98bdc9c04f2f"
-    LLVM_SHA256 = "6ef6d1f92f51936ed3027433ea26875ac4d2ac8eed88b1fbce472018c4fb7720"
+    LLVM_COMMIT = "f1c671925b1c60ded3e4e7b3c6b1ec984b2d9b93"
+    LLVM_SHA256 = "57fc8f0ab46bdfdff52b03c2196d658c094bc4179cd1cf9495becf6a8466123a"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),

From 18dc2474023de36cb686f9900d5cac85763cbf6f Mon Sep 17 00:00:00 2001
From: Yash Katariya 
Date: Tue, 23 Jun 2020 08:30:44 -0700
Subject: [PATCH 0885/1390] Save output in lite notebooks to stop them from
 being tested.

PiperOrigin-RevId: 317870420
Change-Id: Ie69635ac083b7788e00acabf53ed3975f8b019f2
---
 .../post_training_float16_quant.ipynb         | 236 ++++++++++++----
 .../post_training_integer_quant.ipynb         | 178 +++++++++---
 .../performance/post_training_quant.ipynb     | 263 ++++++++++++++----
 3 files changed, 534 insertions(+), 143 deletions(-)

diff --git a/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
index ef08902865e..6015d3e1a65 100644
--- a/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
@@ -12,7 +12,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 1,
       "metadata": {
         "cellView": "form",
         "colab": {},
@@ -105,7 +105,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 2,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -124,13 +124,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 3,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "c6nb7OPlXs_3"
+        "id": "c6nb7OPlXs_3",
+        "outputId": "be7e4e14-cd67-4554-e928-ad803f36dad9"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "tf.float16"
+            ]
+          },
+          "execution_count": 3,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tf.float16"
       ]
@@ -147,13 +163,39 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 4,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 102
+        },
         "colab_type": "code",
-        "id": "hWSAjQWagIHl"
+        "id": "hWSAjQWagIHl",
+        "outputId": "9bf2b530-5a05-415f-f856-cab3642256e9"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
+            "11493376/11490434 [==============================] - 0s 0us/step\n",
+            "11501568/11490434 [==============================] - 0s 0us/step\n",
+            "1875/1875 [==============================] - 12s 6ms/step - loss: 0.2864 - accuracy: 0.9207 - val_loss: 0.1467 - val_accuracy: 0.9560\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "\u003ctensorflow.python.keras.callbacks.History at 0x7fcd75df46a0\u003e"
+            ]
+          },
+          "execution_count": 4,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Load MNIST dataset\n",
         "mnist = keras.datasets.mnist\n",
@@ -211,7 +253,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 5,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -235,7 +277,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 6,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -249,13 +291,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 7,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "Ie9pQaQrn5ue"
+        "id": "Ie9pQaQrn5ue",
+        "outputId": "5df7381a-78ee-4f3e-e1a9-0f3a028384cf"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "84452"
+            ]
+          },
+          "execution_count": 7,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tflite_model_file = tflite_models_dir/\"mnist_model.tflite\"\n",
         "tflite_model_file.write_bytes(tflite_model)"
@@ -273,7 +331,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 8,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -297,13 +355,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 9,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "yuNfl3CoHNK3"
+        "id": "yuNfl3CoHNK3",
+        "outputId": "839f02cd-0a8c-4551-aaa3-0c05c845ad2e"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "44272"
+            ]
+          },
+          "execution_count": 9,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tflite_fp16_model = converter.convert()\n",
         "tflite_model_fp16_file = tflite_models_dir/\"mnist_model_quant_f16.tflite\"\n",
@@ -322,13 +396,26 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 10,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 68
+        },
         "colab_type": "code",
-        "id": "JExfcfLDscu4"
+        "id": "JExfcfLDscu4",
+        "outputId": "6ca316c2-cb0e-40e9-ffb1-a8bcf267e101"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "total 128K\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828 44K Jun 23 06:04 mnist_model_quant_f16.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828 83K Jun 23 06:04 mnist_model.tflite\n"
+          ]
+        }
+      ],
       "source": [
         "!ls -lh {tflite_models_dir}"
       ]
@@ -365,7 +452,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 11,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -379,7 +466,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 12,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -403,7 +490,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 13,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -423,13 +510,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 14,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 281
+        },
         "colab_type": "code",
-        "id": "XZClM2vo3_bm"
+        "id": "XZClM2vo3_bm",
+        "outputId": "fec12377-9f68-45a7-b4a6-ad902d8db171"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAEICAYAAACQ6CLfAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAFxZJREFUeJzt3XtU1HXeB/D3cE0RVDSG4eKMPJBL\nIrI6ZqXhBTFrVwwpw5WEAGnLc9ZL2nbbI1arPPV4nix99jRR7aiFz7qmtIu6KhulVrJj4baYHiKI\nq6DCE4pyG7/PH51mI5nf4DAX9Pt+neM5zO/z/f2+H37ynt/M/GbmpxJCCBCRdDzc3QARuQfDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4yeF6enqgUqlQXV0NAMjOzsaGDRucPm9+fj5mzpzp9HluFgy/nYYN\nG2b55+HhgSFDhlhuv/vuu06fPzs7u1cPvr6+GDlypNPntUd+fj6effZZm+OmT5+OP/7xj07p4Ztv\nvum1v4YNGwaVSoXNmzc7Zb4bgZe7G7hRXbp0yfKzTqdDfn4+5syZY3V8T08PvLwct7vz8/ORn59v\nuZ2WloahQ4c6bPs/Zjab4enp6ZRtu0pERESv/7Ovv/4a48aNw8KFC93YlXvxyO8kzz//PB5++GEs\nXrwY/v7+2LFjB9LS0pCbm2sZc/jwYeh0Osvturo6JCcn49Zbb8XYsWOxdevWfs118eJF7NmzB+np\n6f0a/8O8L7zwAkaNGoWxY8di586dlnpaWhqWL1+OefPmwc/PD0eOHEFHRwdWr16N8PBwqNVqPPHE\nE+jo6LCsk5eXh+DgYISGhsJoNPaa76e/9/vvv4+4uDgEBAQgMjISBw8exG9/+1t8+umn+PWvf41h\nw4Zh5cqVAIBTp05hzpw5CAwMxM9+9jPs3r3bsp1z587hl7/8JQICAnDnnXeiqqqqX78/ABiNRsye\nPRvh4eH9XuemI2jAtFqtOHToUK9lzz33nPD29hYffPCBMJvN4vLly2LJkiVi3bp1ljGHDh0SWq1W\nCCFET0+PmDhxovj9738vOjs7RUVFhdBqteLw4cNCCCFKSkrEqFGj+pz/rbfeEpGRkf3u99ChQ8LT\n01OsWbNGdHR0iOLiYjFkyBBRUVEhhBBiyZIlYsSIEeKTTz4RZrNZdHR0iOXLl4sHHnhAtLS0iO++\n+07cd9994vnnnxdCCPGXv/xFBAcHi/LycnHp0iXx0EMPCQCiqqrKsr0ffu9jx46J4cOHi8OHDwuz\n2SxqamrE6dOnhRBCTJs2TbzzzjuWPtva2kRISIgwGo2iu7tbmEwmERgYaBmfkpIiUlNTRXt7uzh5\n8qQIDg4WM2bMsKw/b9488corr1zz+1+9elVotVqxffv2fu+zmxHD7wDWwj9r1qxey5TCf/ToUTF2\n7Nhe41944QWRnZ1tc/74+Hjx4osv9rvfQ4cOCW9vb9He3m5ZlpycLDZs2GDp89FHH7XUzGaz8PX1\nFdXV1ZZlH3/8seUO55FHHhHPPfecpVZeXm41/JmZmWLNmjV99vXT8O/YsUPMnDmz15jMzEzx0ksv\nia6uLuHp6Wm5wxJCiLVr1/YKvzV///vfhb+/f6/fX0Z8zu9E1/OQ8ttvv0VNTQ1GjBhhWWY2m22+\nel1VVYWjR49i27Zt19XbqFGjer1GoNVq0dDQYLn9497Pnj2Lzs5OTJw40bJM/OjzYA0NDZg2bVqv\nbVlTW1uLKVOm9KvHb7/9FseOHeu1T3p6epCRkYGmpiaYzeZefWq1WpSWltrcrtFoxEMPPeS010hu\nFAy/E6lUql63/fz8cPnyZcvts2fPWn4ODw9HVFQUvvrqq+uaY9u2bZgxY4Zi4Ppy4cIFXLlyBUOG\nDAEA1NTUQK/X99m7Wq2Gj48Pzpw5A7Vafc22NBoNamtrLbdramqszhseHo7Kyso+az/dX+Hh4UhI\nSMD+/fuvGdvd3Q0PDw/U1tYiMjLS5rw/aG9vx+7du1FUVGRz7M2OL/i5UFxcHIqKitDa2orGxka8\n9tprltpdd90FHx8fbNq0CR0dHTCbzfjyyy9x4sQJxW1u27YNGRkZ1yxPS0tDdna21fWuXr2K3Nxc\ndHV1oaSkBPv378eDDz7Y51hPT09kZ2dj5cqVOHfuHIQQqKurw8GDBwEAixYtwttvv43Tp0+jvb0d\n69evtzpvVlYW8vPz8eGHH+Lq1auoq6vDmTNnAHx/J/PNN99YxiYlJaG8vBzvvfceuru70d3djdLS\nUpw5cwbe3t544IEHsG7dOly5cgX/+te/sH37dsV9BQC7d+9GUFAQ7rnnHptjb3YMvwtlZGQgOjoa\nWq0W8+bNQ2pqqqXm5eWFffv2obS0FDqdDqNHj8Zjjz2GtrY2AEBJSUmvh78AcOTIETQ1NSElJeWa\nuWpra3s9FP+psLAw+Pn5QaPRID09Hfn5+YiKirI6ftOmTdBqtbjjjjswfPhwzJ07FxUVFQCA+fPn\nY/ny5ZgxYwZuu+02JCYmWt3O3XffjTfffBO/+c1vMHz4cMyaNcvyqGHlypUoKCjAiBEjsHr1agwf\nPhx/+9vfsGPHDmg0GgQHB+OZZ55BZ2cnAOAPf/gDWltboVarkZWVhUcffbTXXHPnzsXLL7/ca5nR\naMTSpUuveZQhI5UQ/DKPm01HRwd+/vOf48svv+zzvQWHDx9Gdna25R14JCc+578J3XLLLdf92gHJ\nhw/7iSTFh/1EkuKRn0hSLn3O76PyxS3wc+WURFLpQDu6RGe/xg4o/AcOHMCKFStgNpuRnZ2Np59+\nWnH8LfDDVFXCQKYkIgXHRXG/x9r9sN9sNmP58uXYv38/Tp06hYKCApw6dcrezRGRi9kd/tLSUkRG\nRiIiIgI+Pj5ITU1FYWGhI3sjIieyO/z19fW9PlQRFhaG+vr6a8YZDAbo9Xro9Xp0o3/PRYjI+ewO\nf19nCPt6y2ROTg5MJhNMJhO84WvvdETkYHaHPywsrNcnuerq6hASEuKQpojI+ewO/5QpU1BRUYGq\nqip0dXVh586dSEpKcmRvROREdp/q8/LywpYtW3DvvffCbDYjMzMT48ePd2RvRORELn17b4AqkOf5\niZzouChGm2jp11i+vZdIUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\nMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9I\nUgw/kaQYfiJJMfxEkmL4iSTF8BNJymsgK+t0Ovj7+8PT0xNeXl4wmUyO6ouInGxA4QeADz/8EKNH\nj3ZEL0TkQnzYTySpAYVfpVJh7ty5mDx5MgwGQ59jDAYD9Ho99Ho9utE5kOmIyIFUQghh78oNDQ0I\nCQlBc3MzEhMT8frrryM+Pt7q+ABVIKaqEuydjohsOC6K0SZa+jV2QEf+kJAQAEBQUBCSk5NRWlo6\nkM0RkQvZHf729nZcvHjR8vPBgwcRExPjsMaIyLnsfrW/qakJycnJAICenh786le/wrx58xzWGBE5\nl93hj4iIwMmTJx3ZCxG5EE/1EUmK4SeSFMNPJCmGn0hSDD+RpAb8wR5ZXFh2l9XamEe+Vlz3dLNa\nsd7V6a1YDy1Qrg+tu2S1drXslOK6JC8e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSfE8fz89\ntfY9q7UUv1bllf9jgJPPVC5X91y2Wtt8btYAJ79xlTZrrdb8Ng1XXNer+ISj2xl0eOQnkhTDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4iSQ1oCv2XK8b+Yo97Q9OtVo7H6t8HzryK+Vd3BqtUqz7xP6fYv3l\nmPet1hKHXFFct+jyMMX6L4Za/66AgboiuhTrxzv9FOszb+m2e+7IoscU67fl/MPubbuTy67YQ0Q3\nLoafSFIMP5GkGH4iSTH8RJJi+IkkxfATSYqf5+8nvz8fV6gNbNsBA1sdrwfPtFp7aZpOee6PlK85\n8PLMSDs66h+vK1cV637/bFSsj/p4t2J9go/16x0MrVa+FoIMbB75MzMzERQUhJiYGMuylpYWJCYm\nIioqComJiWhttfFlFkQ06NgMf0ZGBg4cONBrWV5eHhISElBRUYGEhATk5eU5rUEicg6b4Y+Pj0dg\nYGCvZYWFhUhPTwcApKenY+/evc7pjoicxq7n/E1NTdBoNAAAjUaD5uZmq2MNBgMMBgMAoBud9kxH\nRE7g9Ff7c3JyYDKZYDKZ4A1fZ09HRP1kV/jVajUaG79/JbaxsRFBQUEObYqInM+u8CclJcFoNAIA\njEYjFixY4NCmiMj5bD7nX7x4MUpKSnD+/HmEhYVh/fr1ePrpp7Fo0SK89dZbGDNmDHbt2uWKXsmK\nnrNNVmt+u63XAMBsY9t+f75gR0eO0ZR9l2J9vI/yn+9/tYyzWtO9843iuj2K1ZuDzfAXFBT0uby4\nuNjhzRCR6/DtvUSSYviJJMXwE0mK4SeSFMNPJCl+pJfcxksbrljf8uwWxbq3ylOxvmvzHKu1UY2f\nKq4rAx75iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ8Tw/uc3pVaGK9Sm+ypcuL+9Svvx44KnL\n192TTHjkJ5IUw08kKYafSFIMP5GkGH4iSTH8RJJi+IkkxfP85FSdv5hitfb5g/9tY23lKzw9vmKF\nYn3IJ6U2ti83HvmJJMXwE0mK4SeSFMNPJCmGn0hSDD+RpBh+IknxPD85Vc191o8vw1TK5/EXVyUq\n1oceOKlYF4pVsnnkz8zMRFBQEGJiYizLcnNzERoairi4OMTFxWHfvn1ObZKIHM9m+DMyMnDgwIFr\nlq9atQplZWUoKyvD/fff75TmiMh5bIY/Pj4egYGBruiFiFzI7hf8tmzZgtjYWGRmZqK1tdXqOIPB\nAL1eD71ej2502jsdETmYXeF//PHHUVlZibKyMmg0Gjz55JNWx+bk5MBkMsFkMsHbxgc1iMh17Aq/\nWq2Gp6cnPDw8sGzZMpSW8tNTRDcau8Lf2Nho+XnPnj29zgQQ0Y3B5nn+xYsXo6SkBOfPn0dYWBjW\nr1+PkpISlJWVQaVSQafT4Y033nBFrzQIefj7K9Yfueeo1Vrb1Q7FdZs3RCjWfTv/oVgnZTbDX1BQ\ncM2yrKwspzRDRK7Dt/cSSYrhJ5IUw08kKYafSFIMP5Gk+JFeGpCK3PGK9b+O/h+rtQUVKYrr+u7j\nqTxn4pGfSFIMP5GkGH4iSTH8RJJi+IkkxfATSYrhJ5IUz/OTou/S7lSs//Ph1xTrlT3dVmuX/jNM\ncV1fNCrWaWB45CeSFMNPJCmGn0hSDD+RpBh+Ikkx/ESSYviJJMXz/JLzCg1RrK/83f8q1n1Vyn9C\nqScfsVq7dT8/r+9OPPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8RJKyeZ6/trYWS5cuxdmzZ+Hh\n4YGcnBysWLECLS0tePjhh1FdXQ2dToc//elPGDlypCt6puug8lL+L5741zrF+kPDLijW370YpFhX\n/8768eWq4prkbDaP/F5eXti0aRO++uorfPbZZ9i6dStOnTqFvLw8JCQkoKKiAgkJCcjLy3NFv0Tk\nIDbDr9FoMGnSJACAv78/oqOjUV9fj8LCQqSnpwMA0tPTsXfvXud2SkQOdV3P+aurq/HFF19g6tSp\naGpqgkajAfD9HURzc7NTGiQi5+j3e/svXbqElJQUvPrqqwgICOj3BAaDAQaDAQDQjc7r75CInKJf\nR/7u7m6kpKRgyZIlWLhwIQBArVajsfH7L1hsbGxEUFDfL/zk5OTAZDLBZDLBG74OapuIBspm+IUQ\nyMrKQnR0NFavXm1ZnpSUBKPRCAAwGo1YsGCB87okIodTCSGE0oCjR4/innvuwYQJE+Dh8f19xYYN\nGzB16lQsWrQINTU1GDNmDHbt2oXAwEDFyQJUgZiqSnBc92STarLyJbSLPtg+oO3f/cxyxfqIbZ8O\naPt0fY6LYrSJln6Ntfmcf/r06bB2/1BcXHx9nRHRoMF3+BFJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\n8au7bwKet99mtZazs3BA2779beXz+Lrtnw1o++Q+PPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8\nRJLief6bwOknrH9l+vyhbQPadlhJl/IA5a+DoEGMR34iSTH8RJJi+IkkxfATSYrhJ5IUw08kKYaf\nSFI8z38D6Jh/h2K9eP4mhepQxzZDNw0e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSdk8z19b\nW4ulS5fi7Nmz8PDwQE5ODlasWIHc3Fy8+eabuPXWWwEAGzZswP333+/0hmXUMM1TsT7Gy/5z+e9e\nDFKse7cpf56fn+a/cdkMv5eXFzZt2oRJkybh4sWLmDx5MhITEwEAq1atwpo1a5zeJBE5ns3wazQa\naDQaAIC/vz+io6NRX1/v9MaIyLmu6zl/dXU1vvjiC0ydOhUAsGXLFsTGxiIzMxOtra19rmMwGKDX\n66HX69GNzoF3TEQO0e/wX7p0CSkpKXj11VcREBCAxx9/HJWVlSgrK4NGo8GTTz7Z53o5OTkwmUww\nmUzwhq/DGieigelX+Lu7u5GSkoIlS5Zg4cKFAAC1Wg1PT094eHhg2bJlKC0tdWqjRORYNsMvhEBW\nVhaio6OxevVqy/LGxkbLz3v27EFMTIxzOiQip7D5gt+xY8ewfft2TJgwAXFxcQC+P61XUFCAsrIy\nqFQq6HQ6vPHGG05vlq7fxgu3K9Y/vVenWBeNXzqwGxpMbIZ/+vTpEH18NzvP6RPd2PgOPyJJMfxE\nkmL4iSTF8BNJiuEnkhTDTyQplejrPJ6TBKgCMVWV4KrpiKRzXBSjTbT0ayyP/ESSYviJJMXwE0mK\n4SeSFMNPJCmGn0hSDD+RpFx6iW6fUR5o1VVZbp87d87y1d+DzWDtbbD2BbA3ezmyN5/q/h/PXfom\nn5/S6/UwmUzuml7RYO1tsPYFsDd7uas3PuwnkhTDTyQpz9zc3Fx3NjB58mR3Tq9osPY2WPsC2Ju9\n3NGbW5/zE5H78GE/kaQYfiJJuSX8Bw4cwLhx4xAZGYm8vDx3tGCVTqezXKNAr9e7tZfMzEwEBQX1\nuiBKS0sLEhMTERUVhcTERKvXSHRHb7m5uQgNDUVcXBzi4uKwb98+t/RWW1uLWbNmITo6GuPHj8fm\nzZsBuH/fWevLbftNuFhPT4+IiIgQlZWVorOzU8TGxory8nJXt2GVVqsV586dc3cbQgghPvroI3Hi\nxAkxfvx4y7K1a9eKjRs3CiGE2Lhxo3jqqacGTW/r1q0Tr7zyilv6+bGGhgZx4sQJIYQQbW1tIioq\nSpSXl7t931nry137zeVH/tLSUkRGRiIiIgI+Pj5ITU1FYWGhq9u4IcTHxyMwMLDXssLCQqSnpwMA\n0tPTsXfvXne01mdvg4VGo8GkSZMA9L6svLv3nbW+3MXl4a+vr0d4eLjldlhYmFt3wE+pVCrMnTsX\nkydPhsFgcHc712hqaoJGowHw/R9Tc3OzmzvqrT+XbXelH19WfjDtO3sud+9oLg+/6OPMokqlcnUb\nVh07dgyff/459u/fj61bt+Ljjz92d0s3jP5ett1VfnpZ+cHC3svdO5rLwx8WFoba2lrL7bq6OoSE\nhLi6Dat+6CUoKAjJycmD7tLjarXacoXkxsZGBAUFubmjfxtMl223dll5d++7wXS5e5eHf8qUKaio\nqEBVVRW6urqwc+dOJCUlubqNPrW3t+PixYuWnw8ePDjoLj2elJQEo9EIADAajViwYIGbO/q3wXLZ\ndmHlsvLu3nfW+nLbfnP5S4xCiKKiIhEVFSUiIiLESy+95I4W+lRZWSliY2NFbGysuP32293eW2pq\nqggODhZeXl4iNDRU5Ofni/Pnz4vZs2eLyMhIMXv2bHHhwoVB01taWpqIiYkREyZMEPPnzxcNDQ1u\n6e3IkSMCgJgwYYKYOHGimDhxoigqKnL7vrPWl7v2G9/eSyQpvsOPSFIMP5GkGH4iSTH8RJJi+Ikk\nxfATSYrhJ5LU/wOdAGX9nfSgHgAAAABJRU5ErkJggg==\n",
+            "text/plain": [
+              "\u003cFigure size 600x400 with 1 Axes\u003e"
+            ]
+          },
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "import matplotlib.pylab as plt\n",
         "\n",
@@ -442,7 +545,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 15,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -462,13 +565,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 16,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 281
+        },
         "colab_type": "code",
-        "id": "CIH7G_MwbY2x"
+        "id": "CIH7G_MwbY2x",
+        "outputId": "6a65e499-6618-4b3e-94f6-1d12af8fb251"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAEICAYAAACQ6CLfAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAFxZJREFUeJzt3XtU1HXeB/D3cE0RVDSG4eKMPJBL\nIrI6ZqXhBTFrVwwpw5WEAGnLc9ZL2nbbI1arPPV4nix99jRR7aiFz7qmtIu6KhulVrJj4baYHiKI\nq6DCE4pyG7/PH51mI5nf4DAX9Pt+neM5zO/z/f2+H37ynt/M/GbmpxJCCBCRdDzc3QARuQfDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4yeF6enqgUqlQXV0NAMjOzsaGDRucPm9+fj5mzpzp9HluFgy/nYYN\nG2b55+HhgSFDhlhuv/vuu06fPzs7u1cPvr6+GDlypNPntUd+fj6effZZm+OmT5+OP/7xj07p4Ztv\nvum1v4YNGwaVSoXNmzc7Zb4bgZe7G7hRXbp0yfKzTqdDfn4+5syZY3V8T08PvLwct7vz8/ORn59v\nuZ2WloahQ4c6bPs/Zjab4enp6ZRtu0pERESv/7Ovv/4a48aNw8KFC93YlXvxyO8kzz//PB5++GEs\nXrwY/v7+2LFjB9LS0pCbm2sZc/jwYeh0Osvturo6JCcn49Zbb8XYsWOxdevWfs118eJF7NmzB+np\n6f0a/8O8L7zwAkaNGoWxY8di586dlnpaWhqWL1+OefPmwc/PD0eOHEFHRwdWr16N8PBwqNVqPPHE\nE+jo6LCsk5eXh+DgYISGhsJoNPaa76e/9/vvv4+4uDgEBAQgMjISBw8exG9/+1t8+umn+PWvf41h\nw4Zh5cqVAIBTp05hzpw5CAwMxM9+9jPs3r3bsp1z587hl7/8JQICAnDnnXeiqqqqX78/ABiNRsye\nPRvh4eH9XuemI2jAtFqtOHToUK9lzz33nPD29hYffPCBMJvN4vLly2LJkiVi3bp1ljGHDh0SWq1W\nCCFET0+PmDhxovj9738vOjs7RUVFhdBqteLw4cNCCCFKSkrEqFGj+pz/rbfeEpGRkf3u99ChQ8LT\n01OsWbNGdHR0iOLiYjFkyBBRUVEhhBBiyZIlYsSIEeKTTz4RZrNZdHR0iOXLl4sHHnhAtLS0iO++\n+07cd9994vnnnxdCCPGXv/xFBAcHi/LycnHp0iXx0EMPCQCiqqrKsr0ffu9jx46J4cOHi8OHDwuz\n2SxqamrE6dOnhRBCTJs2TbzzzjuWPtva2kRISIgwGo2iu7tbmEwmERgYaBmfkpIiUlNTRXt7uzh5\n8qQIDg4WM2bMsKw/b9488corr1zz+1+9elVotVqxffv2fu+zmxHD7wDWwj9r1qxey5TCf/ToUTF2\n7Nhe41944QWRnZ1tc/74+Hjx4osv9rvfQ4cOCW9vb9He3m5ZlpycLDZs2GDp89FHH7XUzGaz8PX1\nFdXV1ZZlH3/8seUO55FHHhHPPfecpVZeXm41/JmZmWLNmjV99vXT8O/YsUPMnDmz15jMzEzx0ksv\nia6uLuHp6Wm5wxJCiLVr1/YKvzV///vfhb+/f6/fX0Z8zu9E1/OQ8ttvv0VNTQ1GjBhhWWY2m22+\nel1VVYWjR49i27Zt19XbqFGjer1GoNVq0dDQYLn9497Pnj2Lzs5OTJw40bJM/OjzYA0NDZg2bVqv\nbVlTW1uLKVOm9KvHb7/9FseOHeu1T3p6epCRkYGmpiaYzeZefWq1WpSWltrcrtFoxEMPPeS010hu\nFAy/E6lUql63/fz8cPnyZcvts2fPWn4ODw9HVFQUvvrqq+uaY9u2bZgxY4Zi4Ppy4cIFXLlyBUOG\nDAEA1NTUQK/X99m7Wq2Gj48Pzpw5A7Vafc22NBoNamtrLbdramqszhseHo7Kyso+az/dX+Hh4UhI\nSMD+/fuvGdvd3Q0PDw/U1tYiMjLS5rw/aG9vx+7du1FUVGRz7M2OL/i5UFxcHIqKitDa2orGxka8\n9tprltpdd90FHx8fbNq0CR0dHTCbzfjyyy9x4sQJxW1u27YNGRkZ1yxPS0tDdna21fWuXr2K3Nxc\ndHV1oaSkBPv378eDDz7Y51hPT09kZ2dj5cqVOHfuHIQQqKurw8GDBwEAixYtwttvv43Tp0+jvb0d\n69evtzpvVlYW8vPz8eGHH+Lq1auoq6vDmTNnAHx/J/PNN99YxiYlJaG8vBzvvfceuru70d3djdLS\nUpw5cwbe3t544IEHsG7dOly5cgX/+te/sH37dsV9BQC7d+9GUFAQ7rnnHptjb3YMvwtlZGQgOjoa\nWq0W8+bNQ2pqqqXm5eWFffv2obS0FDqdDqNHj8Zjjz2GtrY2AEBJSUmvh78AcOTIETQ1NSElJeWa\nuWpra3s9FP+psLAw+Pn5QaPRID09Hfn5+YiKirI6ftOmTdBqtbjjjjswfPhwzJ07FxUVFQCA+fPn\nY/ny5ZgxYwZuu+02JCYmWt3O3XffjTfffBO/+c1vMHz4cMyaNcvyqGHlypUoKCjAiBEjsHr1agwf\nPhx/+9vfsGPHDmg0GgQHB+OZZ55BZ2cnAOAPf/gDWltboVarkZWVhUcffbTXXHPnzsXLL7/ca5nR\naMTSpUuveZQhI5UQ/DKPm01HRwd+/vOf48svv+zzvQWHDx9Gdna25R14JCc+578J3XLLLdf92gHJ\nhw/7iSTFh/1EkuKRn0hSLn3O76PyxS3wc+WURFLpQDu6RGe/xg4o/AcOHMCKFStgNpuRnZ2Np59+\nWnH8LfDDVFXCQKYkIgXHRXG/x9r9sN9sNmP58uXYv38/Tp06hYKCApw6dcrezRGRi9kd/tLSUkRG\nRiIiIgI+Pj5ITU1FYWGhI3sjIieyO/z19fW9PlQRFhaG+vr6a8YZDAbo9Xro9Xp0o3/PRYjI+ewO\nf19nCPt6y2ROTg5MJhNMJhO84WvvdETkYHaHPywsrNcnuerq6hASEuKQpojI+ewO/5QpU1BRUYGq\nqip0dXVh586dSEpKcmRvROREdp/q8/LywpYtW3DvvffCbDYjMzMT48ePd2RvRORELn17b4AqkOf5\niZzouChGm2jp11i+vZdIUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\nMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9I\nUgw/kaQYfiJJMfxEkmL4iSTF8BNJymsgK+t0Ovj7+8PT0xNeXl4wmUyO6ouInGxA4QeADz/8EKNH\nj3ZEL0TkQnzYTySpAYVfpVJh7ty5mDx5MgwGQ59jDAYD9Ho99Ho9utE5kOmIyIFUQghh78oNDQ0I\nCQlBc3MzEhMT8frrryM+Pt7q+ABVIKaqEuydjohsOC6K0SZa+jV2QEf+kJAQAEBQUBCSk5NRWlo6\nkM0RkQvZHf729nZcvHjR8vPBgwcRExPjsMaIyLnsfrW/qakJycnJAICenh786le/wrx58xzWGBE5\nl93hj4iIwMmTJx3ZCxG5EE/1EUmK4SeSFMNPJCmGn0hSDD+RpAb8wR5ZXFh2l9XamEe+Vlz3dLNa\nsd7V6a1YDy1Qrg+tu2S1drXslOK6JC8e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSfE8fz89\ntfY9q7UUv1bllf9jgJPPVC5X91y2Wtt8btYAJ79xlTZrrdb8Ng1XXNer+ISj2xl0eOQnkhTDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4iSQ1oCv2XK8b+Yo97Q9OtVo7H6t8HzryK+Vd3BqtUqz7xP6fYv3l\nmPet1hKHXFFct+jyMMX6L4Za/66AgboiuhTrxzv9FOszb+m2e+7IoscU67fl/MPubbuTy67YQ0Q3\nLoafSFIMP5GkGH4iSTH8RJJi+IkkxfATSYqf5+8nvz8fV6gNbNsBA1sdrwfPtFp7aZpOee6PlK85\n8PLMSDs66h+vK1cV637/bFSsj/p4t2J9go/16x0MrVa+FoIMbB75MzMzERQUhJiYGMuylpYWJCYm\nIioqComJiWhttfFlFkQ06NgMf0ZGBg4cONBrWV5eHhISElBRUYGEhATk5eU5rUEicg6b4Y+Pj0dg\nYGCvZYWFhUhPTwcApKenY+/evc7pjoicxq7n/E1NTdBoNAAAjUaD5uZmq2MNBgMMBgMAoBud9kxH\nRE7g9Ff7c3JyYDKZYDKZ4A1fZ09HRP1kV/jVajUaG79/JbaxsRFBQUEObYqInM+u8CclJcFoNAIA\njEYjFixY4NCmiMj5bD7nX7x4MUpKSnD+/HmEhYVh/fr1ePrpp7Fo0SK89dZbGDNmDHbt2uWKXsmK\nnrNNVmt+u63XAMBsY9t+f75gR0eO0ZR9l2J9vI/yn+9/tYyzWtO9843iuj2K1ZuDzfAXFBT0uby4\nuNjhzRCR6/DtvUSSYviJJMXwE0mK4SeSFMNPJCl+pJfcxksbrljf8uwWxbq3ylOxvmvzHKu1UY2f\nKq4rAx75iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ8Tw/uc3pVaGK9Sm+ypcuL+9Svvx44KnL\n192TTHjkJ5IUw08kKYafSFIMP5GkGH4iSTH8RJJi+IkkxfP85FSdv5hitfb5g/9tY23lKzw9vmKF\nYn3IJ6U2ti83HvmJJMXwE0mK4SeSFMNPJCmGn0hSDD+RpBh+IknxPD85Vc191o8vw1TK5/EXVyUq\n1oceOKlYF4pVsnnkz8zMRFBQEGJiYizLcnNzERoairi4OMTFxWHfvn1ObZKIHM9m+DMyMnDgwIFr\nlq9atQplZWUoKyvD/fff75TmiMh5bIY/Pj4egYGBruiFiFzI7hf8tmzZgtjYWGRmZqK1tdXqOIPB\nAL1eD71ej2502jsdETmYXeF//PHHUVlZibKyMmg0Gjz55JNWx+bk5MBkMsFkMsHbxgc1iMh17Aq/\nWq2Gp6cnPDw8sGzZMpSW8tNTRDcau8Lf2Nho+XnPnj29zgQQ0Y3B5nn+xYsXo6SkBOfPn0dYWBjW\nr1+PkpISlJWVQaVSQafT4Y033nBFrzQIefj7K9Yfueeo1Vrb1Q7FdZs3RCjWfTv/oVgnZTbDX1BQ\ncM2yrKwspzRDRK7Dt/cSSYrhJ5IUw08kKYafSFIMP5Gk+JFeGpCK3PGK9b+O/h+rtQUVKYrr+u7j\nqTxn4pGfSFIMP5GkGH4iSTH8RJJi+IkkxfATSYrhJ5IUz/OTou/S7lSs//Ph1xTrlT3dVmuX/jNM\ncV1fNCrWaWB45CeSFMNPJCmGn0hSDD+RpBh+Ikkx/ESSYviJJMXz/JLzCg1RrK/83f8q1n1Vyn9C\nqScfsVq7dT8/r+9OPPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8RJKyeZ6/trYWS5cuxdmzZ+Hh\n4YGcnBysWLECLS0tePjhh1FdXQ2dToc//elPGDlypCt6puug8lL+L5741zrF+kPDLijW370YpFhX\n/8768eWq4prkbDaP/F5eXti0aRO++uorfPbZZ9i6dStOnTqFvLw8JCQkoKKiAgkJCcjLy3NFv0Tk\nIDbDr9FoMGnSJACAv78/oqOjUV9fj8LCQqSnpwMA0tPTsXfvXud2SkQOdV3P+aurq/HFF19g6tSp\naGpqgkajAfD9HURzc7NTGiQi5+j3e/svXbqElJQUvPrqqwgICOj3BAaDAQaDAQDQjc7r75CInKJf\nR/7u7m6kpKRgyZIlWLhwIQBArVajsfH7L1hsbGxEUFDfL/zk5OTAZDLBZDLBG74OapuIBspm+IUQ\nyMrKQnR0NFavXm1ZnpSUBKPRCAAwGo1YsGCB87okIodTCSGE0oCjR4/innvuwYQJE+Dh8f19xYYN\nGzB16lQsWrQINTU1GDNmDHbt2oXAwEDFyQJUgZiqSnBc92STarLyJbSLPtg+oO3f/cxyxfqIbZ8O\naPt0fY6LYrSJln6Ntfmcf/r06bB2/1BcXHx9nRHRoMF3+BFJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\n8au7bwKet99mtZazs3BA2779beXz+Lrtnw1o++Q+PPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8\nRJLief6bwOknrH9l+vyhbQPadlhJl/IA5a+DoEGMR34iSTH8RJJi+IkkxfATSYrhJ5IUw08kKYaf\nSFI8z38D6Jh/h2K9eP4mhepQxzZDNw0e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSdk8z19b\nW4ulS5fi7Nmz8PDwQE5ODlasWIHc3Fy8+eabuPXWWwEAGzZswP333+/0hmXUMM1TsT7Gy/5z+e9e\nDFKse7cpf56fn+a/cdkMv5eXFzZt2oRJkybh4sWLmDx5MhITEwEAq1atwpo1a5zeJBE5ns3wazQa\naDQaAIC/vz+io6NRX1/v9MaIyLmu6zl/dXU1vvjiC0ydOhUAsGXLFsTGxiIzMxOtra19rmMwGKDX\n66HX69GNzoF3TEQO0e/wX7p0CSkpKXj11VcREBCAxx9/HJWVlSgrK4NGo8GTTz7Z53o5OTkwmUww\nmUzwhq/DGieigelX+Lu7u5GSkoIlS5Zg4cKFAAC1Wg1PT094eHhg2bJlKC0tdWqjRORYNsMvhEBW\nVhaio6OxevVqy/LGxkbLz3v27EFMTIxzOiQip7D5gt+xY8ewfft2TJgwAXFxcQC+P61XUFCAsrIy\nqFQq6HQ6vPHGG05vlq7fxgu3K9Y/vVenWBeNXzqwGxpMbIZ/+vTpEH18NzvP6RPd2PgOPyJJMfxE\nkmL4iSTF8BNJiuEnkhTDTyQplejrPJ6TBKgCMVWV4KrpiKRzXBSjTbT0ayyP/ESSYviJJMXwE0mK\n4SeSFMNPJCmGn0hSDD+RpFx6iW6fUR5o1VVZbp87d87y1d+DzWDtbbD2BbA3ezmyN5/q/h/PXfom\nn5/S6/UwmUzuml7RYO1tsPYFsDd7uas3PuwnkhTDTyQpz9zc3Fx3NjB58mR3Tq9osPY2WPsC2Ju9\n3NGbW5/zE5H78GE/kaQYfiJJuSX8Bw4cwLhx4xAZGYm8vDx3tGCVTqezXKNAr9e7tZfMzEwEBQX1\nuiBKS0sLEhMTERUVhcTERKvXSHRHb7m5uQgNDUVcXBzi4uKwb98+t/RWW1uLWbNmITo6GuPHj8fm\nzZsBuH/fWevLbftNuFhPT4+IiIgQlZWVorOzU8TGxory8nJXt2GVVqsV586dc3cbQgghPvroI3Hi\nxAkxfvx4y7K1a9eKjRs3CiGE2Lhxo3jqqacGTW/r1q0Tr7zyilv6+bGGhgZx4sQJIYQQbW1tIioq\nSpSXl7t931nry137zeVH/tLSUkRGRiIiIgI+Pj5ITU1FYWGhq9u4IcTHxyMwMLDXssLCQqSnpwMA\n0tPTsXfvXne01mdvg4VGo8GkSZMA9L6svLv3nbW+3MXl4a+vr0d4eLjldlhYmFt3wE+pVCrMnTsX\nkydPhsFgcHc712hqaoJGowHw/R9Tc3OzmzvqrT+XbXelH19WfjDtO3sud+9oLg+/6OPMokqlcnUb\nVh07dgyff/459u/fj61bt+Ljjz92d0s3jP5ett1VfnpZ+cHC3svdO5rLwx8WFoba2lrL7bq6OoSE\nhLi6Dat+6CUoKAjJycmD7tLjarXacoXkxsZGBAUFubmjfxtMl223dll5d++7wXS5e5eHf8qUKaio\nqEBVVRW6urqwc+dOJCUlubqNPrW3t+PixYuWnw8ePDjoLj2elJQEo9EIADAajViwYIGbO/q3wXLZ\ndmHlsvLu3nfW+nLbfnP5S4xCiKKiIhEVFSUiIiLESy+95I4W+lRZWSliY2NFbGysuP32293eW2pq\nqggODhZeXl4iNDRU5Ofni/Pnz4vZs2eLyMhIMXv2bHHhwoVB01taWpqIiYkREyZMEPPnzxcNDQ1u\n6e3IkSMCgJgwYYKYOHGimDhxoigqKnL7vrPWl7v2G9/eSyQpvsOPSFIMP5GkGH4iSTH8RJJi+Ikk\nxfATSYrhJ5LU/wOdAGX9nfSgHgAAAABJRU5ErkJggg==\n",
+            "text/plain": [
+              "\u003cFigure size 600x400 with 1 Axes\u003e"
+            ]
+          },
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "plt.imshow(test_images[0])\n",
         "template = \"True:{true}, predicted:{predict}\"\n",
@@ -489,7 +608,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 17,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -531,13 +650,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 18,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "T5mWkSbMcU5z"
+        "id": "T5mWkSbMcU5z",
+        "outputId": "818e9142-70cf-420b-8e64-38c2ca11a370"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "0.956\n"
+          ]
+        }
+      ],
       "source": [
         "print(evaluate_model(interpreter))"
       ]
@@ -554,13 +684,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 19,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "-9cnwiPp6EGm"
+        "id": "-9cnwiPp6EGm",
+        "outputId": "53e00eac-51af-4030-be1a-3df986640f8d"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "0.956\n"
+          ]
+        }
+      ],
       "source": [
         "# NOTE: Colab runs on server CPUs. At the time of writing this, TensorFlow Lite\n",
         "# doesn't have super optimized server CPU kernels. For this reason this may be\n",
@@ -599,8 +740,7 @@
   "metadata": {
     "colab": {
       "collapsed_sections": [],
-      "name": "post_training-float16-quant.ipynb",
-      "private_outputs": true,
+      "name": "post_training_float16_quant.ipynb",
       "provenance": [],
       "toc_visible": true
     },
diff --git a/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
index ad461f56d6f..8397dbfa69f 100644
--- a/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
@@ -12,7 +12,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 17,
       "metadata": {
         "cellView": "form",
         "colab": {},
@@ -106,7 +106,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 18,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -135,13 +135,36 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 19,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 51
+        },
         "colab_type": "code",
-        "id": "eMsw_6HujaqM"
+        "id": "eMsw_6HujaqM",
+        "outputId": "5662a5f3-fc64-458f-958a-98f9c6348143"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "1875/1875 [==============================] - 2s 1ms/step - loss: 0.2782 - accuracy: 0.9221 - val_loss: 0.1230 - val_accuracy: 0.9664\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "\u003ctensorflow.python.keras.callbacks.History at 0x7f33f1817588\u003e"
+            ]
+          },
+          "execution_count": 19,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Load MNIST dataset\n",
         "mnist = keras.datasets.mnist\n",
@@ -199,7 +222,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 20,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -223,7 +246,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 21,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -237,13 +260,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 22,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "Ie9pQaQrn5ue"
+        "id": "Ie9pQaQrn5ue",
+        "outputId": "8580b835-61f0-42b3-a21e-b8d476042c11"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "84528"
+            ]
+          },
+          "execution_count": 22,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tflite_model_file = tflite_models_dir/\"mnist_model.tflite\"\n",
         "tflite_model_file.write_bytes(tflite_model)"
@@ -266,7 +305,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 23,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -289,7 +328,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 24,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -319,13 +358,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 25,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "yuNfl3CoHNK3"
+        "id": "yuNfl3CoHNK3",
+        "outputId": "79a19679-87a2-4dc6-eee4-b33f3e5c1c5d"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "24720"
+            ]
+          },
+          "execution_count": 25,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tflite_model_quant = converter.convert()\n",
         "tflite_model_quant_file = tflite_models_dir/\"mnist_model_quant.tflite\"\n",
@@ -344,13 +399,27 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 26,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 85
+        },
         "colab_type": "code",
-        "id": "JExfcfLDscu4"
+        "id": "JExfcfLDscu4",
+        "outputId": "58238f92-01b0-4faa-e293-35451d08dd7c"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "total 140K\n",
+            "-rw-rw-r-- 1 yashkatariya 10086651 25K Jun 23 06:06 mnist_model_quant_io.tflite\n",
+            "-rw-rw-r-- 1 yashkatariya 10086651 25K Jun 23 06:07 mnist_model_quant.tflite\n",
+            "-rw-rw-r-- 1 yashkatariya 10086651 83K Jun 23 06:06 mnist_model.tflite\n"
+          ]
+        }
+      ],
       "source": [
         "!ls -lh {tflite_models_dir}"
       ]
@@ -369,13 +438,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 27,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "kzjEjcDs3BHa"
+        "id": "kzjEjcDs3BHa",
+        "outputId": "8d7370ec-3f3f-41a2-8afb-4ecdd40e9efc"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "24784"
+            ]
+          },
+          "execution_count": 27,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
         "converter.inference_input_type = tf.uint8\n",
@@ -423,7 +508,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 28,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -437,7 +522,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 29,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -465,7 +550,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 30,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -484,13 +569,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 31,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 281
+        },
         "colab_type": "code",
-        "id": "XZClM2vo3_bm"
+        "id": "XZClM2vo3_bm",
+        "outputId": "3af2e31c-44c6-41f2-c51f-da9d7b71bdfb"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAEICAYAAACQ6CLfAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAFxZJREFUeJzt3XtU1HXeB/D3cE0RVDSG4eKMPJBL\nIrI6ZqXhBTFrVwwpw5WEAGnLc9ZL2nbbI1arPPV4nix99jRR7aiFz7qmtIu6KhulVrJj4baYHiKI\nq6DCE4pyG7/PH51mI5nf4DAX9Pt+neM5zO/z/f2+H37ynt/M/GbmpxJCCBCRdDzc3QARuQfDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4yeF6enqgUqlQXV0NAMjOzsaGDRucPm9+fj5mzpzp9HluFgy/nYYN\nG2b55+HhgSFDhlhuv/vuu06fPzs7u1cPvr6+GDlypNPntUd+fj6effZZm+OmT5+OP/7xj07p4Ztv\nvum1v4YNGwaVSoXNmzc7Zb4bgZe7G7hRXbp0yfKzTqdDfn4+5syZY3V8T08PvLwct7vz8/ORn59v\nuZ2WloahQ4c6bPs/Zjab4enp6ZRtu0pERESv/7Ovv/4a48aNw8KFC93YlXvxyO8kzz//PB5++GEs\nXrwY/v7+2LFjB9LS0pCbm2sZc/jwYeh0Osvturo6JCcn49Zbb8XYsWOxdevWfs118eJF7NmzB+np\n6f0a/8O8L7zwAkaNGoWxY8di586dlnpaWhqWL1+OefPmwc/PD0eOHEFHRwdWr16N8PBwqNVqPPHE\nE+jo6LCsk5eXh+DgYISGhsJoNPaa76e/9/vvv4+4uDgEBAQgMjISBw8exG9/+1t8+umn+PWvf41h\nw4Zh5cqVAIBTp05hzpw5CAwMxM9+9jPs3r3bsp1z587hl7/8JQICAnDnnXeiqqqqX78/ABiNRsye\nPRvh4eH9XuemI2jAtFqtOHToUK9lzz33nPD29hYffPCBMJvN4vLly2LJkiVi3bp1ljGHDh0SWq1W\nCCFET0+PmDhxovj9738vOjs7RUVFhdBqteLw4cNCCCFKSkrEqFGj+pz/rbfeEpGRkf3u99ChQ8LT\n01OsWbNGdHR0iOLiYjFkyBBRUVEhhBBiyZIlYsSIEeKTTz4RZrNZdHR0iOXLl4sHHnhAtLS0iO++\n+07cd9994vnnnxdCCPGXv/xFBAcHi/LycnHp0iXx0EMPCQCiqqrKsr0ffu9jx46J4cOHi8OHDwuz\n2SxqamrE6dOnhRBCTJs2TbzzzjuWPtva2kRISIgwGo2iu7tbmEwmERgYaBmfkpIiUlNTRXt7uzh5\n8qQIDg4WM2bMsKw/b9488corr1zz+1+9elVotVqxffv2fu+zmxHD7wDWwj9r1qxey5TCf/ToUTF2\n7Nhe41944QWRnZ1tc/74+Hjx4osv9rvfQ4cOCW9vb9He3m5ZlpycLDZs2GDp89FHH7XUzGaz8PX1\nFdXV1ZZlH3/8seUO55FHHhHPPfecpVZeXm41/JmZmWLNmjV99vXT8O/YsUPMnDmz15jMzEzx0ksv\nia6uLuHp6Wm5wxJCiLVr1/YKvzV///vfhb+/f6/fX0Z8zu9E1/OQ8ttvv0VNTQ1GjBhhWWY2m22+\nel1VVYWjR49i27Zt19XbqFGjer1GoNVq0dDQYLn9497Pnj2Lzs5OTJw40bJM/OjzYA0NDZg2bVqv\nbVlTW1uLKVOm9KvHb7/9FseOHeu1T3p6epCRkYGmpiaYzeZefWq1WpSWltrcrtFoxEMPPeS010hu\nFAy/E6lUql63/fz8cPnyZcvts2fPWn4ODw9HVFQUvvrqq+uaY9u2bZgxY4Zi4Ppy4cIFXLlyBUOG\nDAEA1NTUQK/X99m7Wq2Gj48Pzpw5A7Vafc22NBoNamtrLbdramqszhseHo7Kyso+az/dX+Hh4UhI\nSMD+/fuvGdvd3Q0PDw/U1tYiMjLS5rw/aG9vx+7du1FUVGRz7M2OL/i5UFxcHIqKitDa2orGxka8\n9tprltpdd90FHx8fbNq0CR0dHTCbzfjyyy9x4sQJxW1u27YNGRkZ1yxPS0tDdna21fWuXr2K3Nxc\ndHV1oaSkBPv378eDDz7Y51hPT09kZ2dj5cqVOHfuHIQQqKurw8GDBwEAixYtwttvv43Tp0+jvb0d\n69evtzpvVlYW8vPz8eGHH+Lq1auoq6vDmTNnAHx/J/PNN99YxiYlJaG8vBzvvfceuru70d3djdLS\nUpw5cwbe3t544IEHsG7dOly5cgX/+te/sH37dsV9BQC7d+9GUFAQ7rnnHptjb3YMvwtlZGQgOjoa\nWq0W8+bNQ2pqqqXm5eWFffv2obS0FDqdDqNHj8Zjjz2GtrY2AEBJSUmvh78AcOTIETQ1NSElJeWa\nuWpra3s9FP+psLAw+Pn5QaPRID09Hfn5+YiKirI6ftOmTdBqtbjjjjswfPhwzJ07FxUVFQCA+fPn\nY/ny5ZgxYwZuu+02JCYmWt3O3XffjTfffBO/+c1vMHz4cMyaNcvyqGHlypUoKCjAiBEjsHr1agwf\nPhx/+9vfsGPHDmg0GgQHB+OZZ55BZ2cnAOAPf/gDWltboVarkZWVhUcffbTXXHPnzsXLL7/ca5nR\naMTSpUuveZQhI5UQ/DKPm01HRwd+/vOf48svv+zzvQWHDx9Gdna25R14JCc+578J3XLLLdf92gHJ\nhw/7iSTFh/1EkuKRn0hSLn3O76PyxS3wc+WURFLpQDu6RGe/xg4o/AcOHMCKFStgNpuRnZ2Np59+\nWnH8LfDDVFXCQKYkIgXHRXG/x9r9sN9sNmP58uXYv38/Tp06hYKCApw6dcrezRGRi9kd/tLSUkRG\nRiIiIgI+Pj5ITU1FYWGhI3sjIieyO/z19fW9PlQRFhaG+vr6a8YZDAbo9Xro9Xp0o3/PRYjI+ewO\nf19nCPt6y2ROTg5MJhNMJhO84WvvdETkYHaHPywsrNcnuerq6hASEuKQpojI+ewO/5QpU1BRUYGq\nqip0dXVh586dSEpKcmRvROREdp/q8/LywpYtW3DvvffCbDYjMzMT48ePd2RvRORELn17b4AqkOf5\niZzouChGm2jp11i+vZdIUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\nMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9I\nUgw/kaQYfiJJMfxEkmL4iSTF8BNJymsgK+t0Ovj7+8PT0xNeXl4wmUyO6ouInGxA4QeADz/8EKNH\nj3ZEL0TkQnzYTySpAYVfpVJh7ty5mDx5MgwGQ59jDAYD9Ho99Ho9utE5kOmIyIFUQghh78oNDQ0I\nCQlBc3MzEhMT8frrryM+Pt7q+ABVIKaqEuydjohsOC6K0SZa+jV2QEf+kJAQAEBQUBCSk5NRWlo6\nkM0RkQvZHf729nZcvHjR8vPBgwcRExPjsMaIyLnsfrW/qakJycnJAICenh786le/wrx58xzWGBE5\nl93hj4iIwMmTJx3ZCxG5EE/1EUmK4SeSFMNPJCmGn0hSDD+RpAb8wR5ZXFh2l9XamEe+Vlz3dLNa\nsd7V6a1YDy1Qrg+tu2S1drXslOK6JC8e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSfE8fz89\ntfY9q7UUv1bllf9jgJPPVC5X91y2Wtt8btYAJ79xlTZrrdb8Ng1XXNer+ISj2xl0eOQnkhTDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4iSQ1oCv2XK8b+Yo97Q9OtVo7H6t8HzryK+Vd3BqtUqz7xP6fYv3l\nmPet1hKHXFFct+jyMMX6L4Za/66AgboiuhTrxzv9FOszb+m2e+7IoscU67fl/MPubbuTy67YQ0Q3\nLoafSFIMP5GkGH4iSTH8RJJi+IkkxfATSYqf5+8nvz8fV6gNbNsBA1sdrwfPtFp7aZpOee6PlK85\n8PLMSDs66h+vK1cV637/bFSsj/p4t2J9go/16x0MrVa+FoIMbB75MzMzERQUhJiYGMuylpYWJCYm\nIioqComJiWhttfFlFkQ06NgMf0ZGBg4cONBrWV5eHhISElBRUYGEhATk5eU5rUEicg6b4Y+Pj0dg\nYGCvZYWFhUhPTwcApKenY+/evc7pjoicxq7n/E1NTdBoNAAAjUaD5uZmq2MNBgMMBgMAoBud9kxH\nRE7g9Ff7c3JyYDKZYDKZ4A1fZ09HRP1kV/jVajUaG79/JbaxsRFBQUEObYqInM+u8CclJcFoNAIA\njEYjFixY4NCmiMj5bD7nX7x4MUpKSnD+/HmEhYVh/fr1ePrpp7Fo0SK89dZbGDNmDHbt2uWKXsmK\nnrNNVmt+u63XAMBsY9t+f75gR0eO0ZR9l2J9vI/yn+9/tYyzWtO9843iuj2K1ZuDzfAXFBT0uby4\nuNjhzRCR6/DtvUSSYviJJMXwE0mK4SeSFMNPJCl+pJfcxksbrljf8uwWxbq3ylOxvmvzHKu1UY2f\nKq4rAx75iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ8Tw/uc3pVaGK9Sm+ypcuL+9Svvx44KnL\n192TTHjkJ5IUw08kKYafSFIMP5GkGH4iSTH8RJJi+IkkxfP85FSdv5hitfb5g/9tY23lKzw9vmKF\nYn3IJ6U2ti83HvmJJMXwE0mK4SeSFMNPJCmGn0hSDD+RpBh+IknxPD85Vc191o8vw1TK5/EXVyUq\n1oceOKlYF4pVsnnkz8zMRFBQEGJiYizLcnNzERoairi4OMTFxWHfvn1ObZKIHM9m+DMyMnDgwIFr\nlq9atQplZWUoKyvD/fff75TmiMh5bIY/Pj4egYGBruiFiFzI7hf8tmzZgtjYWGRmZqK1tdXqOIPB\nAL1eD71ej2502jsdETmYXeF//PHHUVlZibKyMmg0Gjz55JNWx+bk5MBkMsFkMsHbxgc1iMh17Aq/\nWq2Gp6cnPDw8sGzZMpSW8tNTRDcau8Lf2Nho+XnPnj29zgQQ0Y3B5nn+xYsXo6SkBOfPn0dYWBjW\nr1+PkpISlJWVQaVSQafT4Y033nBFrzQIefj7K9Yfueeo1Vrb1Q7FdZs3RCjWfTv/oVgnZTbDX1BQ\ncM2yrKwspzRDRK7Dt/cSSYrhJ5IUw08kKYafSFIMP5Gk+JFeGpCK3PGK9b+O/h+rtQUVKYrr+u7j\nqTxn4pGfSFIMP5GkGH4iSTH8RJJi+IkkxfATSYrhJ5IUz/OTou/S7lSs//Ph1xTrlT3dVmuX/jNM\ncV1fNCrWaWB45CeSFMNPJCmGn0hSDD+RpBh+Ikkx/ESSYviJJMXz/JLzCg1RrK/83f8q1n1Vyn9C\nqScfsVq7dT8/r+9OPPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8RJKyeZ6/trYWS5cuxdmzZ+Hh\n4YGcnBysWLECLS0tePjhh1FdXQ2dToc//elPGDlypCt6puug8lL+L5741zrF+kPDLijW370YpFhX\n/8768eWq4prkbDaP/F5eXti0aRO++uorfPbZZ9i6dStOnTqFvLw8JCQkoKKiAgkJCcjLy3NFv0Tk\nIDbDr9FoMGnSJACAv78/oqOjUV9fj8LCQqSnpwMA0tPTsXfvXud2SkQOdV3P+aurq/HFF19g6tSp\naGpqgkajAfD9HURzc7NTGiQi5+j3e/svXbqElJQUvPrqqwgICOj3BAaDAQaDAQDQjc7r75CInKJf\nR/7u7m6kpKRgyZIlWLhwIQBArVajsfH7L1hsbGxEUFDfL/zk5OTAZDLBZDLBG74OapuIBspm+IUQ\nyMrKQnR0NFavXm1ZnpSUBKPRCAAwGo1YsGCB87okIodTCSGE0oCjR4/innvuwYQJE+Dh8f19xYYN\nGzB16lQsWrQINTU1GDNmDHbt2oXAwEDFyQJUgZiqSnBc92STarLyJbSLPtg+oO3f/cxyxfqIbZ8O\naPt0fY6LYrSJln6Ntfmcf/r06bB2/1BcXHx9nRHRoMF3+BFJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\n8au7bwKet99mtZazs3BA2779beXz+Lrtnw1o++Q+PPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8\nRJLief6bwOknrH9l+vyhbQPadlhJl/IA5a+DoEGMR34iSTH8RJJi+IkkxfATSYrhJ5IUw08kKYaf\nSFI8z38D6Jh/h2K9eP4mhepQxzZDNw0e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSdk8z19b\nW4ulS5fi7Nmz8PDwQE5ODlasWIHc3Fy8+eabuPXWWwEAGzZswP333+/0hmXUMM1TsT7Gy/5z+e9e\nDFKse7cpf56fn+a/cdkMv5eXFzZt2oRJkybh4sWLmDx5MhITEwEAq1atwpo1a5zeJBE5ns3wazQa\naDQaAIC/vz+io6NRX1/v9MaIyLmu6zl/dXU1vvjiC0ydOhUAsGXLFsTGxiIzMxOtra19rmMwGKDX\n66HX69GNzoF3TEQO0e/wX7p0CSkpKXj11VcREBCAxx9/HJWVlSgrK4NGo8GTTz7Z53o5OTkwmUww\nmUzwhq/DGieigelX+Lu7u5GSkoIlS5Zg4cKFAAC1Wg1PT094eHhg2bJlKC0tdWqjRORYNsMvhEBW\nVhaio6OxevVqy/LGxkbLz3v27EFMTIxzOiQip7D5gt+xY8ewfft2TJgwAXFxcQC+P61XUFCAsrIy\nqFQq6HQ6vPHGG05vlq7fxgu3K9Y/vVenWBeNXzqwGxpMbIZ/+vTpEH18NzvP6RPd2PgOPyJJMfxE\nkmL4iSTF8BNJiuEnkhTDTyQplejrPJ6TBKgCMVWV4KrpiKRzXBSjTbT0ayyP/ESSYviJJMXwE0mK\n4SeSFMNPJCmGn0hSDD+RpFx6iW6fUR5o1VVZbp87d87y1d+DzWDtbbD2BbA3ezmyN5/q/h/PXfom\nn5/S6/UwmUzuml7RYO1tsPYFsDd7uas3PuwnkhTDTyQpz9zc3Fx3NjB58mR3Tq9osPY2WPsC2Ju9\n3NGbW5/zE5H78GE/kaQYfiJJuSX8Bw4cwLhx4xAZGYm8vDx3tGCVTqezXKNAr9e7tZfMzEwEBQX1\nuiBKS0sLEhMTERUVhcTERKvXSHRHb7m5uQgNDUVcXBzi4uKwb98+t/RWW1uLWbNmITo6GuPHj8fm\nzZsBuH/fWevLbftNuFhPT4+IiIgQlZWVorOzU8TGxory8nJXt2GVVqsV586dc3cbQgghPvroI3Hi\nxAkxfvx4y7K1a9eKjRs3CiGE2Lhxo3jqqacGTW/r1q0Tr7zyilv6+bGGhgZx4sQJIYQQbW1tIioq\nSpSXl7t931nry137zeVH/tLSUkRGRiIiIgI+Pj5ITU1FYWGhq9u4IcTHxyMwMLDXssLCQqSnpwMA\n0tPTsXfvXne01mdvg4VGo8GkSZMA9L6svLv3nbW+3MXl4a+vr0d4eLjldlhYmFt3wE+pVCrMnTsX\nkydPhsFgcHc712hqaoJGowHw/R9Tc3OzmzvqrT+XbXelH19WfjDtO3sud+9oLg+/6OPMokqlcnUb\nVh07dgyff/459u/fj61bt+Ljjz92d0s3jP5ett1VfnpZ+cHC3svdO5rLwx8WFoba2lrL7bq6OoSE\nhLi6Dat+6CUoKAjJycmD7tLjarXacoXkxsZGBAUFubmjfxtMl223dll5d++7wXS5e5eHf8qUKaio\nqEBVVRW6urqwc+dOJCUlubqNPrW3t+PixYuWnw8ePDjoLj2elJQEo9EIADAajViwYIGbO/q3wXLZ\ndmHlsvLu3nfW+nLbfnP5S4xCiKKiIhEVFSUiIiLESy+95I4W+lRZWSliY2NFbGysuP32293eW2pq\nqggODhZeXl4iNDRU5Ofni/Pnz4vZs2eLyMhIMXv2bHHhwoVB01taWpqIiYkREyZMEPPnzxcNDQ1u\n6e3IkSMCgJgwYYKYOHGimDhxoigqKnL7vrPWl7v2G9/eSyQpvsOPSFIMP5GkGH4iSTH8RJJi+Ikk\nxfATSYrhJ5LU/wOdAGX9nfSgHgAAAABJRU5ErkJggg==\n",
+            "text/plain": [
+              "\u003cFigure size 600x400 with 1 Axes\u003e"
+            ]
+          },
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "import matplotlib.pylab as plt\n",
         "\n",
@@ -513,7 +614,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 33,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -530,7 +631,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": null,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -557,7 +658,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": null,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -599,7 +700,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": null,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -622,7 +723,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": null,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -653,7 +754,6 @@
     "colab": {
       "collapsed_sections": [],
       "name": "post_training_integer_quant.ipynb",
-      "private_outputs": true,
       "provenance": [],
       "toc_visible": true
     },
diff --git a/tensorflow/lite/g3doc/performance/post_training_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
index 201ccf5bdc3..5341fe5e4fb 100644
--- a/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
@@ -12,7 +12,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 1,
       "metadata": {
         "cellView": "form",
         "colab": {},
@@ -126,7 +126,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 2,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -155,13 +155,36 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 3,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 51
+        },
         "colab_type": "code",
-        "id": "hWSAjQWagIHl"
+        "id": "hWSAjQWagIHl",
+        "outputId": "961899f8-1597-4417-b21d-cae94a330ecc"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "1875/1875 [==============================] - 10s 5ms/step - loss: 0.2787 - accuracy: 0.9203 - val_loss: 0.1323 - val_accuracy: 0.9624\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "\u003ctensorflow.python.keras.callbacks.History at 0x7f6443480e80\u003e"
+            ]
+          },
+          "execution_count": 3,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Load MNIST dataset\n",
         "mnist = keras.datasets.mnist\n",
@@ -200,8 +223,7 @@
         "id": "5NMaNZQCkW9X"
       },
       "source": [
-        "For the example, since you trained the model for just a single epoch, so it only trains to ~96% accuracy.\n",
-        "\n"
+        "For the example, since you trained the model for just a single epoch, so it only trains to ~96% accuracy.\n"
       ]
     },
     {
@@ -220,7 +242,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 4,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -244,7 +266,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 5,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -258,13 +280,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 6,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "Ie9pQaQrn5ue"
+        "id": "Ie9pQaQrn5ue",
+        "outputId": "046db0bc-1745-4e94-9f21-f7e91bdaebda"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "84452"
+            ]
+          },
+          "execution_count": 6,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "tflite_model_file = tflite_models_dir/\"mnist_model.tflite\"\n",
         "tflite_model_file.write_bytes(tflite_model)"
@@ -282,13 +320,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 7,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "g8PUvLWDlmmz"
+        "id": "g8PUvLWDlmmz",
+        "outputId": "d79b45d3-babf-4890-8036-de2f497da88a"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "23840"
+            ]
+          },
+          "execution_count": 7,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
         "tflite_quant_model = converter.convert()\n",
@@ -308,13 +362,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 8,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 119
+        },
         "colab_type": "code",
-        "id": "JExfcfLDscu4"
+        "id": "JExfcfLDscu4",
+        "outputId": "d1fda4c2-343e-40fb-f90f-b6bde00c523e"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "total 214M\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  44K Jun 23 06:04 mnist_model_quant_f16.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  24K Jun 23 06:12 mnist_model_quant.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  83K Jun 23 06:12 mnist_model.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  44M Jun 23 06:10 resnet_v2_101_quantized.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828 171M Jun 23 06:09 resnet_v2_101.tflite\n"
+          ]
+        }
+      ],
       "source": [
         "!ls -lh {tflite_models_dir}"
       ]
@@ -329,8 +399,7 @@
         "## Run the TFLite models\n",
         "\n",
         "Run the TensorFlow Lite model using the Python TensorFlow Lite\n",
-        "Interpreter.\n",
-        "\n"
+        "Interpreter.\n"
       ]
     },
     {
@@ -345,7 +414,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 9,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -359,7 +428,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 10,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -383,7 +452,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 11,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -403,13 +472,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 12,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 281
+        },
         "colab_type": "code",
-        "id": "XZClM2vo3_bm"
+        "id": "XZClM2vo3_bm",
+        "outputId": "0fa4155b-01f8-4fea-f586-d9044d73572e"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAEICAYAAACQ6CLfAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAFxZJREFUeJzt3XtU1HXeB/D3cE0RVDSG4eKMPJBL\nIrI6ZqXhBTFrVwwpw5WEAGnLc9ZL2nbbI1arPPV4nix99jRR7aiFz7qmtIu6KhulVrJj4baYHiKI\nq6DCE4pyG7/PH51mI5nf4DAX9Pt+neM5zO/z/f2+H37ynt/M/GbmpxJCCBCRdDzc3QARuQfDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4yeF6enqgUqlQXV0NAMjOzsaGDRucPm9+fj5mzpzp9HluFgy/nYYN\nG2b55+HhgSFDhlhuv/vuu06fPzs7u1cPvr6+GDlypNPntUd+fj6effZZm+OmT5+OP/7xj07p4Ztv\nvum1v4YNGwaVSoXNmzc7Zb4bgZe7G7hRXbp0yfKzTqdDfn4+5syZY3V8T08PvLwct7vz8/ORn59v\nuZ2WloahQ4c6bPs/Zjab4enp6ZRtu0pERESv/7Ovv/4a48aNw8KFC93YlXvxyO8kzz//PB5++GEs\nXrwY/v7+2LFjB9LS0pCbm2sZc/jwYeh0Osvturo6JCcn49Zbb8XYsWOxdevWfs118eJF7NmzB+np\n6f0a/8O8L7zwAkaNGoWxY8di586dlnpaWhqWL1+OefPmwc/PD0eOHEFHRwdWr16N8PBwqNVqPPHE\nE+jo6LCsk5eXh+DgYISGhsJoNPaa76e/9/vvv4+4uDgEBAQgMjISBw8exG9/+1t8+umn+PWvf41h\nw4Zh5cqVAIBTp05hzpw5CAwMxM9+9jPs3r3bsp1z587hl7/8JQICAnDnnXeiqqqqX78/ABiNRsye\nPRvh4eH9XuemI2jAtFqtOHToUK9lzz33nPD29hYffPCBMJvN4vLly2LJkiVi3bp1ljGHDh0SWq1W\nCCFET0+PmDhxovj9738vOjs7RUVFhdBqteLw4cNCCCFKSkrEqFGj+pz/rbfeEpGRkf3u99ChQ8LT\n01OsWbNGdHR0iOLiYjFkyBBRUVEhhBBiyZIlYsSIEeKTTz4RZrNZdHR0iOXLl4sHHnhAtLS0iO++\n+07cd9994vnnnxdCCPGXv/xFBAcHi/LycnHp0iXx0EMPCQCiqqrKsr0ffu9jx46J4cOHi8OHDwuz\n2SxqamrE6dOnhRBCTJs2TbzzzjuWPtva2kRISIgwGo2iu7tbmEwmERgYaBmfkpIiUlNTRXt7uzh5\n8qQIDg4WM2bMsKw/b9488corr1zz+1+9elVotVqxffv2fu+zmxHD7wDWwj9r1qxey5TCf/ToUTF2\n7Nhe41944QWRnZ1tc/74+Hjx4osv9rvfQ4cOCW9vb9He3m5ZlpycLDZs2GDp89FHH7XUzGaz8PX1\nFdXV1ZZlH3/8seUO55FHHhHPPfecpVZeXm41/JmZmWLNmjV99vXT8O/YsUPMnDmz15jMzEzx0ksv\nia6uLuHp6Wm5wxJCiLVr1/YKvzV///vfhb+/f6/fX0Z8zu9E1/OQ8ttvv0VNTQ1GjBhhWWY2m22+\nel1VVYWjR49i27Zt19XbqFGjer1GoNVq0dDQYLn9497Pnj2Lzs5OTJw40bJM/OjzYA0NDZg2bVqv\nbVlTW1uLKVOm9KvHb7/9FseOHeu1T3p6epCRkYGmpiaYzeZefWq1WpSWltrcrtFoxEMPPeS010hu\nFAy/E6lUql63/fz8cPnyZcvts2fPWn4ODw9HVFQUvvrqq+uaY9u2bZgxY4Zi4Ppy4cIFXLlyBUOG\nDAEA1NTUQK/X99m7Wq2Gj48Pzpw5A7Vafc22NBoNamtrLbdramqszhseHo7Kyso+az/dX+Hh4UhI\nSMD+/fuvGdvd3Q0PDw/U1tYiMjLS5rw/aG9vx+7du1FUVGRz7M2OL/i5UFxcHIqKitDa2orGxka8\n9tprltpdd90FHx8fbNq0CR0dHTCbzfjyyy9x4sQJxW1u27YNGRkZ1yxPS0tDdna21fWuXr2K3Nxc\ndHV1oaSkBPv378eDDz7Y51hPT09kZ2dj5cqVOHfuHIQQqKurw8GDBwEAixYtwttvv43Tp0+jvb0d\n69evtzpvVlYW8vPz8eGHH+Lq1auoq6vDmTNnAHx/J/PNN99YxiYlJaG8vBzvvfceuru70d3djdLS\nUpw5cwbe3t544IEHsG7dOly5cgX/+te/sH37dsV9BQC7d+9GUFAQ7rnnHptjb3YMvwtlZGQgOjoa\nWq0W8+bNQ2pqqqXm5eWFffv2obS0FDqdDqNHj8Zjjz2GtrY2AEBJSUmvh78AcOTIETQ1NSElJeWa\nuWpra3s9FP+psLAw+Pn5QaPRID09Hfn5+YiKirI6ftOmTdBqtbjjjjswfPhwzJ07FxUVFQCA+fPn\nY/ny5ZgxYwZuu+02JCYmWt3O3XffjTfffBO/+c1vMHz4cMyaNcvyqGHlypUoKCjAiBEjsHr1agwf\nPhx/+9vfsGPHDmg0GgQHB+OZZ55BZ2cnAOAPf/gDWltboVarkZWVhUcffbTXXHPnzsXLL7/ca5nR\naMTSpUuveZQhI5UQ/DKPm01HRwd+/vOf48svv+zzvQWHDx9Gdna25R14JCc+578J3XLLLdf92gHJ\nhw/7iSTFh/1EkuKRn0hSLn3O76PyxS3wc+WURFLpQDu6RGe/xg4o/AcOHMCKFStgNpuRnZ2Np59+\nWnH8LfDDVFXCQKYkIgXHRXG/x9r9sN9sNmP58uXYv38/Tp06hYKCApw6dcrezRGRi9kd/tLSUkRG\nRiIiIgI+Pj5ITU1FYWGhI3sjIieyO/z19fW9PlQRFhaG+vr6a8YZDAbo9Xro9Xp0o3/PRYjI+ewO\nf19nCPt6y2ROTg5MJhNMJhO84WvvdETkYHaHPywsrNcnuerq6hASEuKQpojI+ewO/5QpU1BRUYGq\nqip0dXVh586dSEpKcmRvROREdp/q8/LywpYtW3DvvffCbDYjMzMT48ePd2RvRORELn17b4AqkOf5\niZzouChGm2jp11i+vZdIUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\nMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJMfxEkmL4iSTF8BNJiuEnkhTDTyQphp9I\nUgw/kaQYfiJJMfxEkmL4iSTF8BNJymsgK+t0Ovj7+8PT0xNeXl4wmUyO6ouInGxA4QeADz/8EKNH\nj3ZEL0TkQnzYTySpAYVfpVJh7ty5mDx5MgwGQ59jDAYD9Ho99Ho9utE5kOmIyIFUQghh78oNDQ0I\nCQlBc3MzEhMT8frrryM+Pt7q+ABVIKaqEuydjohsOC6K0SZa+jV2QEf+kJAQAEBQUBCSk5NRWlo6\nkM0RkQvZHf729nZcvHjR8vPBgwcRExPjsMaIyLnsfrW/qakJycnJAICenh786le/wrx58xzWGBE5\nl93hj4iIwMmTJx3ZCxG5EE/1EUmK4SeSFMNPJCmGn0hSDD+RpAb8wR5ZXFh2l9XamEe+Vlz3dLNa\nsd7V6a1YDy1Qrg+tu2S1drXslOK6JC8e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSfE8fz89\ntfY9q7UUv1bllf9jgJPPVC5X91y2Wtt8btYAJ79xlTZrrdb8Ng1XXNer+ISj2xl0eOQnkhTDTyQp\nhp9IUgw/kaQYfiJJMfxEkmL4iSQ1oCv2XK8b+Yo97Q9OtVo7H6t8HzryK+Vd3BqtUqz7xP6fYv3l\nmPet1hKHXFFct+jyMMX6L4Za/66AgboiuhTrxzv9FOszb+m2e+7IoscU67fl/MPubbuTy67YQ0Q3\nLoafSFIMP5GkGH4iSTH8RJJi+IkkxfATSYqf5+8nvz8fV6gNbNsBA1sdrwfPtFp7aZpOee6PlK85\n8PLMSDs66h+vK1cV637/bFSsj/p4t2J9go/16x0MrVa+FoIMbB75MzMzERQUhJiYGMuylpYWJCYm\nIioqComJiWhttfFlFkQ06NgMf0ZGBg4cONBrWV5eHhISElBRUYGEhATk5eU5rUEicg6b4Y+Pj0dg\nYGCvZYWFhUhPTwcApKenY+/evc7pjoicxq7n/E1NTdBoNAAAjUaD5uZmq2MNBgMMBgMAoBud9kxH\nRE7g9Ff7c3JyYDKZYDKZ4A1fZ09HRP1kV/jVajUaG79/JbaxsRFBQUEObYqInM+u8CclJcFoNAIA\njEYjFixY4NCmiMj5bD7nX7x4MUpKSnD+/HmEhYVh/fr1ePrpp7Fo0SK89dZbGDNmDHbt2uWKXsmK\nnrNNVmt+u63XAMBsY9t+f75gR0eO0ZR9l2J9vI/yn+9/tYyzWtO9843iuj2K1ZuDzfAXFBT0uby4\nuNjhzRCR6/DtvUSSYviJJMXwE0mK4SeSFMNPJCl+pJfcxksbrljf8uwWxbq3ylOxvmvzHKu1UY2f\nKq4rAx75iSTF8BNJiuEnkhTDTyQphp9IUgw/kaQYfiJJ8Tw/uc3pVaGK9Sm+ypcuL+9Svvx44KnL\n192TTHjkJ5IUw08kKYafSFIMP5GkGH4iSTH8RJJi+IkkxfP85FSdv5hitfb5g/9tY23lKzw9vmKF\nYn3IJ6U2ti83HvmJJMXwE0mK4SeSFMNPJCmGn0hSDD+RpBh+IknxPD85Vc191o8vw1TK5/EXVyUq\n1oceOKlYF4pVsnnkz8zMRFBQEGJiYizLcnNzERoairi4OMTFxWHfvn1ObZKIHM9m+DMyMnDgwIFr\nlq9atQplZWUoKyvD/fff75TmiMh5bIY/Pj4egYGBruiFiFzI7hf8tmzZgtjYWGRmZqK1tdXqOIPB\nAL1eD71ej2502jsdETmYXeF//PHHUVlZibKyMmg0Gjz55JNWx+bk5MBkMsFkMsHbxgc1iMh17Aq/\nWq2Gp6cnPDw8sGzZMpSW8tNTRDcau8Lf2Nho+XnPnj29zgQQ0Y3B5nn+xYsXo6SkBOfPn0dYWBjW\nr1+PkpISlJWVQaVSQafT4Y033nBFrzQIefj7K9Yfueeo1Vrb1Q7FdZs3RCjWfTv/oVgnZTbDX1BQ\ncM2yrKwspzRDRK7Dt/cSSYrhJ5IUw08kKYafSFIMP5Gk+JFeGpCK3PGK9b+O/h+rtQUVKYrr+u7j\nqTxn4pGfSFIMP5GkGH4iSTH8RJJi+IkkxfATSYrhJ5IUz/OTou/S7lSs//Ph1xTrlT3dVmuX/jNM\ncV1fNCrWaWB45CeSFMNPJCmGn0hSDD+RpBh+Ikkx/ESSYviJJMXz/JLzCg1RrK/83f8q1n1Vyn9C\nqScfsVq7dT8/r+9OPPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8RJKyeZ6/trYWS5cuxdmzZ+Hh\n4YGcnBysWLECLS0tePjhh1FdXQ2dToc//elPGDlypCt6puug8lL+L5741zrF+kPDLijW370YpFhX\n/8768eWq4prkbDaP/F5eXti0aRO++uorfPbZZ9i6dStOnTqFvLw8JCQkoKKiAgkJCcjLy3NFv0Tk\nIDbDr9FoMGnSJACAv78/oqOjUV9fj8LCQqSnpwMA0tPTsXfvXud2SkQOdV3P+aurq/HFF19g6tSp\naGpqgkajAfD9HURzc7NTGiQi5+j3e/svXbqElJQUvPrqqwgICOj3BAaDAQaDAQDQjc7r75CInKJf\nR/7u7m6kpKRgyZIlWLhwIQBArVajsfH7L1hsbGxEUFDfL/zk5OTAZDLBZDLBG74OapuIBspm+IUQ\nyMrKQnR0NFavXm1ZnpSUBKPRCAAwGo1YsGCB87okIodTCSGE0oCjR4/innvuwYQJE+Dh8f19xYYN\nGzB16lQsWrQINTU1GDNmDHbt2oXAwEDFyQJUgZiqSnBc92STarLyJbSLPtg+oO3f/cxyxfqIbZ8O\naPt0fY6LYrSJln6Ntfmcf/r06bB2/1BcXHx9nRHRoMF3+BFJiuEnkhTDTyQphp9IUgw/kaQYfiJJ\n8au7bwKet99mtZazs3BA2779beXz+Lrtnw1o++Q+PPITSYrhJ5IUw08kKYafSFIMP5GkGH4iSTH8\nRJLief6bwOknrH9l+vyhbQPadlhJl/IA5a+DoEGMR34iSTH8RJJi+IkkxfATSYrhJ5IUw08kKYaf\nSFI8z38D6Jh/h2K9eP4mhepQxzZDNw0e+YkkxfATSYrhJ5IUw08kKYafSFIMP5GkGH4iSdk8z19b\nW4ulS5fi7Nmz8PDwQE5ODlasWIHc3Fy8+eabuPXWWwEAGzZswP333+/0hmXUMM1TsT7Gy/5z+e9e\nDFKse7cpf56fn+a/cdkMv5eXFzZt2oRJkybh4sWLmDx5MhITEwEAq1atwpo1a5zeJBE5ns3wazQa\naDQaAIC/vz+io6NRX1/v9MaIyLmu6zl/dXU1vvjiC0ydOhUAsGXLFsTGxiIzMxOtra19rmMwGKDX\n66HX69GNzoF3TEQO0e/wX7p0CSkpKXj11VcREBCAxx9/HJWVlSgrK4NGo8GTTz7Z53o5OTkwmUww\nmUzwhq/DGieigelX+Lu7u5GSkoIlS5Zg4cKFAAC1Wg1PT094eHhg2bJlKC0tdWqjRORYNsMvhEBW\nVhaio6OxevVqy/LGxkbLz3v27EFMTIxzOiQip7D5gt+xY8ewfft2TJgwAXFxcQC+P61XUFCAsrIy\nqFQq6HQ6vPHGG05vlq7fxgu3K9Y/vVenWBeNXzqwGxpMbIZ/+vTpEH18NzvP6RPd2PgOPyJJMfxE\nkmL4iSTF8BNJiuEnkhTDTyQplejrPJ6TBKgCMVWV4KrpiKRzXBSjTbT0ayyP/ESSYviJJMXwE0mK\n4SeSFMNPJCmGn0hSDD+RpFx6iW6fUR5o1VVZbp87d87y1d+DzWDtbbD2BbA3ezmyN5/q/h/PXfom\nn5/S6/UwmUzuml7RYO1tsPYFsDd7uas3PuwnkhTDTyQpz9zc3Fx3NjB58mR3Tq9osPY2WPsC2Ju9\n3NGbW5/zE5H78GE/kaQYfiJJuSX8Bw4cwLhx4xAZGYm8vDx3tGCVTqezXKNAr9e7tZfMzEwEBQX1\nuiBKS0sLEhMTERUVhcTERKvXSHRHb7m5uQgNDUVcXBzi4uKwb98+t/RWW1uLWbNmITo6GuPHj8fm\nzZsBuH/fWevLbftNuFhPT4+IiIgQlZWVorOzU8TGxory8nJXt2GVVqsV586dc3cbQgghPvroI3Hi\nxAkxfvx4y7K1a9eKjRs3CiGE2Lhxo3jqqacGTW/r1q0Tr7zyilv6+bGGhgZx4sQJIYQQbW1tIioq\nSpSXl7t931nry137zeVH/tLSUkRGRiIiIgI+Pj5ITU1FYWGhq9u4IcTHxyMwMLDXssLCQqSnpwMA\n0tPTsXfvXne01mdvg4VGo8GkSZMA9L6svLv3nbW+3MXl4a+vr0d4eLjldlhYmFt3wE+pVCrMnTsX\nkydPhsFgcHc712hqaoJGowHw/R9Tc3OzmzvqrT+XbXelH19WfjDtO3sud+9oLg+/6OPMokqlcnUb\nVh07dgyff/459u/fj61bt+Ljjz92d0s3jP5ett1VfnpZ+cHC3svdO5rLwx8WFoba2lrL7bq6OoSE\nhLi6Dat+6CUoKAjJycmD7tLjarXacoXkxsZGBAUFubmjfxtMl223dll5d++7wXS5e5eHf8qUKaio\nqEBVVRW6urqwc+dOJCUlubqNPrW3t+PixYuWnw8ePDjoLj2elJQEo9EIADAajViwYIGbO/q3wXLZ\ndmHlsvLu3nfW+nLbfnP5S4xCiKKiIhEVFSUiIiLESy+95I4W+lRZWSliY2NFbGysuP32293eW2pq\nqggODhZeXl4iNDRU5Ofni/Pnz4vZs2eLyMhIMXv2bHHhwoVB01taWpqIiYkREyZMEPPnzxcNDQ1u\n6e3IkSMCgJgwYYKYOHGimDhxoigqKnL7vrPWl7v2G9/eSyQpvsOPSFIMP5GkGH4iSTH8RJJi+Ikk\nxfATSYrhJ5LU/wOdAGX9nfSgHgAAAABJRU5ErkJggg==\n",
+            "text/plain": [
+              "\u003cFigure size 600x400 with 1 Axes\u003e"
+            ]
+          },
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "import matplotlib.pylab as plt\n",
         "\n",
@@ -432,7 +517,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 13,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -474,13 +559,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 14,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "DqXBnDfJ7qxL"
+        "id": "DqXBnDfJ7qxL",
+        "outputId": "78f393f8-c4a5-41e0-abe4-ab6a5c394e51"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "0.9624\n"
+          ]
+        }
+      ],
       "source": [
         "print(evaluate_model(interpreter))"
       ]
@@ -497,13 +593,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 15,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "-9cnwiPp6EGm"
+        "id": "-9cnwiPp6EGm",
+        "outputId": "d82552d7-8a2c-49dc-a19a-56010a013102"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "0.9626\n"
+          ]
+        }
+      ],
       "source": [
         "print(evaluate_model(interpreter_quant))"
       ]
@@ -515,7 +622,6 @@
         "id": "L7lfxkor8pgv"
       },
       "source": [
-        "\n",
         "In this example, the compressed model has no difference in the accuracy."
       ]
     },
@@ -537,7 +643,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 16,
       "metadata": {
         "colab": {},
         "colab_type": "code",
@@ -557,13 +663,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 17,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "LwnV4KxwVEoG"
+        "id": "LwnV4KxwVEoG",
+        "outputId": "7d50f90d-6104-43a3-863c-28db9465d483"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "178509092"
+            ]
+          },
+          "execution_count": 17,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Convert to TF Lite without quantization\n",
         "resnet_tflite_file = tflite_models_dir/\"resnet_v2_101.tflite\"\n",
@@ -572,13 +694,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 18,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 34
+        },
         "colab_type": "code",
-        "id": "2qkZD0VoVExe"
+        "id": "2qkZD0VoVExe",
+        "outputId": "76a47590-fa91-49b9-f568-4e00b46c9537"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "45182656"
+            ]
+          },
+          "execution_count": 18,
+          "metadata": {
+            "tags": []
+          },
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Convert to TF Lite with quantization\n",
         "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
@@ -588,13 +726,28 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 0,
+      "execution_count": 19,
       "metadata": {
-        "colab": {},
+        "colab": {
+          "height": 102
+        },
         "colab_type": "code",
-        "id": "vhOjeg1x9Knp"
+        "id": "vhOjeg1x9Knp",
+        "outputId": "c643a660-f815-49f0-ac4b-ac48af3c1203"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  44K Jun 23 06:04 /tmp/mnist_tflite_models/mnist_model_quant_f16.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  24K Jun 23 06:12 /tmp/mnist_tflite_models/mnist_model_quant.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  83K Jun 23 06:12 /tmp/mnist_tflite_models/mnist_model.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828  44M Jun 23 06:13 /tmp/mnist_tflite_models/resnet_v2_101_quantized.tflite\n",
+            "-rw-rw-r-- 1 colaboratory-playground 50844828 171M Jun 23 06:12 /tmp/mnist_tflite_models/resnet_v2_101.tflite\n"
+          ]
+        }
+      ],
       "source": [
         "!ls -lh {tflite_models_dir}/*.tflite"
       ]
@@ -606,7 +759,6 @@
         "id": "qqHLaqFMCjRZ"
       },
       "source": [
-        "\n",
         "The model size reduces from 171 MB to 43 MB.\n",
         "The accuracy of this model on imagenet can be evaluated using the scripts provided for [TFLite accuracy measurement](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/accuracy/ilsvrc).\n",
         "\n",
@@ -618,7 +770,6 @@
     "colab": {
       "collapsed_sections": [],
       "name": "post_training_quant.ipynb",
-      "private_outputs": true,
       "provenance": [],
       "toc_visible": true
     },

From 89851a6a7725d92735e9817ee6a1f551dd7492ab Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 23 Jun 2020 08:42:05 -0700
Subject: [PATCH 0886/1390] Replace deprecated thread annotations macros.

PiperOrigin-RevId: 317872415
Change-Id: I2e3fcd5cebc7a2b4d63c1322a2386484b336479c
---
 tensorflow/compiler/xla/pjrt/distributed/service.h | 10 +++++-----
 tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc  |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/xla/pjrt/distributed/service.h b/tensorflow/compiler/xla/pjrt/distributed/service.h
index 725a76791ce..9ecbdb3cc7c 100644
--- a/tensorflow/compiler/xla/pjrt/distributed/service.h
+++ b/tensorflow/compiler/xla/pjrt/distributed/service.h
@@ -54,15 +54,15 @@ class DistributedRuntimeServiceImpl final
 
   absl::Mutex mu_;
   enum class State { kInitializing, kRunning };
-  State state_ GUARDED_BY(mu_) = State::kInitializing;
+  State state_ ABSL_GUARDED_BY(mu_) = State::kInitializing;
 
-  std::vector local_topologies_ GUARDED_BY(mu_);
-  GlobalTopologyProto topology_ GUARDED_BY(mu_);
+  std::vector local_topologies_ ABSL_GUARDED_BY(mu_);
+  GlobalTopologyProto topology_ ABSL_GUARDED_BY(mu_);
   struct Node {
     bool present = false;
   };
-  int num_nodes_present_ GUARDED_BY(mu_) = 0;
-  std::vector nodes_ GUARDED_BY(mu_);
+  int num_nodes_present_ ABSL_GUARDED_BY(mu_) = 0;
+  std::vector nodes_ ABSL_GUARDED_BY(mu_);
 
   KeyValueStore key_value_store_;
 };
diff --git a/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc b/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
index de760af8fd9..edffaf6c877 100644
--- a/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
+++ b/tensorflow/compiler/xla/pjrt/nvidia_gpu_device.cc
@@ -169,7 +169,7 @@ class NcclIdStore {
   const std::shared_ptr client_;
 
   absl::Mutex mu_;
-  absl::flat_hash_map cache_ GUARDED_BY(mu_);
+  absl::flat_hash_map cache_ ABSL_GUARDED_BY(mu_);
 };
 
 StatusOr NcclIdStore::GetNcclUniqueId(const NcclCliqueKey& key) {

From f5a9d24c847ffcc7ae09e850aad39e2cb55ae4f3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 23 Jun 2020 08:57:59 -0700
Subject: [PATCH 0887/1390] Added shape inference for tf_device.LaunchOp.

PiperOrigin-RevId: 317875384
Change-Id: Idb070c9e92d07ee19cd8ed26c1beec3de86f43df
---
 .../mlir/tensorflow/tests/shape_inference.mlir      | 13 +++++++++++++
 .../mlir/tensorflow/transforms/shape_inference.cc   |  4 ++++
 2 files changed, 17 insertions(+)

diff --git a/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir b/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir
index 1599d53ed15..1af4ba6b3dc 100644
--- a/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir
+++ b/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir
@@ -433,4 +433,17 @@ func @multiple_blocks_one_return(%arg0: tensor) -> tensor<*xf32> {
     // CHECK: return %[[CAST_RESULT_0]], %[[CAST_RESULT_1]], %[[ADDI]]
     return %27, %28, %2 : tensor<*xui8>, tensor<*xi8>, tensor<*xi8>
   }
+
+  // CHECK-LABEL: infer_device_launch
+  func @infer_device_launch(%arg0: tensor<1x8x2xi32>) -> (tensor<*xf32>, tensor<*xf32>) {
+    %0 = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor
+    %1 = "tf_device.launch"() ({
+      %2 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x8x2xi32>) -> tensor<1x8x2xf32>
+      tf_device.return %2 : tensor<1x8x2xf32>
+    // CHECK: () -> tensor<1x8x2xf32>
+    }) {device = "/device:CPU:0"} : () -> tensor<*xf32>
+    // CHECK: (tensor, tensor<1x8x2xf32>) -> (tensor<1x8x1xf32>, tensor<1x8x1xf32>)
+    %3:2 = "tf.Split"(%0, %1) {device = ""} : (tensor, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>)
+    return %3#0, %3#1 : tensor<*xf32>, tensor<*xf32>
+  }
 }
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc
index 7e4baadc397..33ccf5caff2 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc
@@ -215,6 +215,10 @@ bool InferShapeForNonTFDialectOperation(Operation* op, Dialect* tf_dialect) {
     return InferShapeForPassThroughOps(
         tensor_cast.getOperation()->getOperands(), op, tf_dialect);
   }
+  if (auto launch_op = dyn_cast(op)) {
+    return InferShapeForPassThroughOps(
+        launch_op.GetBody().getTerminator()->getOperands(), op, tf_dialect);
+  }
   return false;
 }
 

From a9e2bceddf4588f1ef0bae6fb4054111ff3d2225 Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 23 Jun 2020 09:51:29 -0700
Subject: [PATCH 0888/1390] ConvConstants converted to new style. Added memory
 types for buffer.

PiperOrigin-RevId: 317885492
Change-Id: I12dfede368bdf157f746f7823d7d47fd3807ffa1
---
 tensorflow/lite/delegates/gpu/cl/arguments.cc |  16 +-
 tensorflow/lite/delegates/gpu/cl/buffer.cc    |  14 ++
 tensorflow/lite/delegates/gpu/cl/buffer.h     |   3 +
 tensorflow/lite/delegates/gpu/cl/gpu_object.h |   3 +
 .../gpu/cl/kernels/conv_constants.cc          | 143 ++++++++++--------
 .../delegates/gpu/cl/kernels/conv_constants.h |  30 ++--
 .../lite/delegates/gpu/cl/linear_storage.cc   |   1 +
 .../lite/delegates/gpu/cl/linear_storage.h    |   1 +
 8 files changed, 132 insertions(+), 79 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/arguments.cc b/tensorflow/lite/delegates/gpu/cl/arguments.cc
index d26c869957c..1fd58ef2454 100644
--- a/tensorflow/lite/delegates/gpu/cl/arguments.cc
+++ b/tensorflow/lite/delegates/gpu/cl/arguments.cc
@@ -457,8 +457,20 @@ std::string Arguments::GetListOfArgs() {
   for (auto& t : buffers_) {
     const std::string type_name =
         t.second.data_type == DataType::FLOAT32 ? "float" : "half";
-    AppendArgument(absl::StrCat("__global ", type_name, t.second.element_size,
-                                "* ", t.first),
+    std::string memory_type;
+    switch (t.second.memory_type) {
+      case MemoryType::GLOBAL:
+        memory_type = "__global";
+        break;
+      case MemoryType::CONSTANT:
+        memory_type = "__constant";
+        break;
+      case MemoryType::LOCAL:
+        memory_type = "__local";
+        break;
+    }
+    AppendArgument(absl::StrCat(memory_type, " ", type_name,
+                                t.second.element_size, "* ", t.first),
                    &result);
   }
   for (auto& t : image_buffers_) {
diff --git a/tensorflow/lite/delegates/gpu/cl/buffer.cc b/tensorflow/lite/delegates/gpu/cl/buffer.cc
index 68e85593e5d..436d8751e18 100644
--- a/tensorflow/lite/delegates/gpu/cl/buffer.cc
+++ b/tensorflow/lite/delegates/gpu/cl/buffer.cc
@@ -50,6 +50,7 @@ GPUResources BufferDescriptor::GetGPUResources(AccessType access_type) const {
   desc.data_type = element_type;
   desc.access_type = access_type;
   desc.element_size = element_size;
+  desc.memory_type = memory_type;
   resources.buffers.push_back({"buffer", desc});
   return resources;
 }
@@ -59,6 +60,8 @@ absl::Status BufferDescriptor::PerformSelector(
     const std::vector& template_args, std::string* result) const {
   if (selector == "Read") {
     return PerformReadSelector(args, result);
+  } else if (selector == "GetPtr") {
+    return PerformGetPtrSelector(args, result);
   } else {
     return absl::NotFoundError(absl::StrCat(
         "BufferDescriptor don't have selector with name - ", selector));
@@ -76,6 +79,17 @@ absl::Status BufferDescriptor::PerformReadSelector(
   return absl::OkStatus();
 }
 
+absl::Status BufferDescriptor::PerformGetPtrSelector(
+    const std::vector& args, std::string* result) const {
+  if (!args.empty()) {
+    return absl::NotFoundError(
+        absl::StrCat("BufferDescriptor GetPtr require zero arguments, but ",
+                     args.size(), " was passed"));
+  }
+  *result = "buffer";
+  return absl::OkStatus();
+}
+
 Buffer::Buffer(cl_mem buffer, size_t size_in_bytes)
     : buffer_(buffer), size_(size_in_bytes) {}
 
diff --git a/tensorflow/lite/delegates/gpu/cl/buffer.h b/tensorflow/lite/delegates/gpu/cl/buffer.h
index 771aae9e002..0d1072040c1 100644
--- a/tensorflow/lite/delegates/gpu/cl/buffer.h
+++ b/tensorflow/lite/delegates/gpu/cl/buffer.h
@@ -32,6 +32,7 @@ namespace cl {
 struct BufferDescriptor : public GPUObjectDescriptor {
   DataType element_type;  // FLOAT32 or FLOAT16
   int element_size;
+  MemoryType memory_type = MemoryType::GLOBAL;
 
   absl::Status PerformSelector(const std::string& selector,
                                const std::vector& args,
@@ -41,6 +42,8 @@ struct BufferDescriptor : public GPUObjectDescriptor {
   GPUResources GetGPUResources(AccessType access_type) const override;
   absl::Status PerformReadSelector(const std::vector& args,
                                    std::string* result) const;
+  absl::Status PerformGetPtrSelector(const std::vector& args,
+                                     std::string* result) const;
 };
 
 // Buffer represent linear GPU data storage with arbitrary data format.
diff --git a/tensorflow/lite/delegates/gpu/cl/gpu_object.h b/tensorflow/lite/delegates/gpu/cl/gpu_object.h
index faf18b539e2..711c4726bc2 100644
--- a/tensorflow/lite/delegates/gpu/cl/gpu_object.h
+++ b/tensorflow/lite/delegates/gpu/cl/gpu_object.h
@@ -54,10 +54,13 @@ struct GPUImageBufferDescriptor {
   cl_mem memory;
 };
 
+enum class MemoryType { GLOBAL, CONSTANT, LOCAL };
+
 struct GPUBufferDescriptor {
   DataType data_type;
   AccessType access_type;
   int element_size;
+  MemoryType memory_type = MemoryType::GLOBAL;
   cl_mem memory;
 };
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.cc b/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.cc
index d4dc206ffce..e6fc5da36a2 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.cc
@@ -27,16 +27,29 @@ namespace gpu {
 namespace cl {
 namespace {
 
-std::string GenerateConvolutionConstantCode(
-    const OperationDef& op_def, const int2& kernel_size, int src_channels,
-    int dst_channels, bool stride_correction, const CLDevice& device,
-    const std::vector& linked_operations) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+std::string GenerateConvolutionConstantCode(const OperationDef& op_def,
+                                            const int2& kernel_size,
+                                            int src_channels, int dst_channels,
+                                            bool stride_correction,
+                                            const CLDevice& device,
+                                            Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  src_desc->SetTextureAddressMode(GetFastestZeroMode(device));
+  if (op_def.IsBatchSupported()) {
+    src_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    dst_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
+  args->AddInt("stride_x");
+  args->AddInt("stride_y");
+  args->AddInt("padding_x");
+  args->AddInt("padding_y");
+  args->AddInt("dilation_x");
+  args->AddInt("dilation_y");
 
   std::string c = GetCommonDefines(op_def.precision);
 
@@ -89,33 +102,24 @@ std::string GenerateConvolutionConstantCode(
   const std::string postfixes[] = {".x", ".xy", ".xyz", ""};
 
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ) + ",\n";
-  c += "    __constant FLT4* filters,  \n";
-  c += "    __constant FLT4* biases";
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int2 stride,               \n";
-  c += "    int2 padding,              \n";
-  c += "    int2 dilation,             \n";
-  c += "    int4 src_size,             \n";
-  c += "    int4 dst_size              \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
   c += "  int Y = get_global_id(1);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y) return;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height()) "
+       "return;\n";
   if (stride_correction) {
     c += "  int start_x = " +
-         GetXStrideCorrected("X", "src_size.w", "stride.x", "padding.x") +
+         GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
+                             "args.padding_x") +
          ";\n";
   } else {
-    c += "  int start_x = X * stride.x + padding.x;\n";
+    c += "  int start_x = X * args.stride_x + args.padding_x;\n";
   }
-  c += "  int start_y = Y * stride.y + padding.y;\n";
+  c += "  int start_y = Y * args.stride_y + args.padding_y;\n";
   c += "  ACCUM_FLT4 r[" + kOutZ + "];\n";
   c += "  for (int i = 0; i < " + kOutZ + "; ++i) {\n";
   c += "    r[i] = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);\n";
   c += "  }\n";
-  const auto address_mode = GetFastestZeroMode(device);
   int filters_counter = 0;
   for (int s = 0; s < src_depth; ++s) {
     const int ch_count = std::min(4, src_channels - s * 4);
@@ -124,27 +128,29 @@ std::string GenerateConvolutionConstantCode(
     const std::string s_type = absl::StrCat("FLT", s_count);
     const std::string s_postfix = postfixes[ch_count - 1];
     for (int ky = 0; ky < kernel_size.y; ++ky) {
-      std::string s_y = absl::StrCat("(start_y + ", ky, " * dilation.y)");
+      std::string s_y = absl::StrCat("(start_y + ", ky, " * args.dilation_y)");
       if (manual_clamp) {
         c += "  {\n";
-        c += "  bool y_out = " + s_y + " < 0 || " + s_y + " >= src_size.y;\n";
+        c += "  bool y_out = " + s_y + " < 0 || " + s_y +
+             " >= args.src_tensor.Height();\n";
       }
       for (int kx = 0; kx < kernel_size.x; ++kx) {
         c += "  {\n";
-        std::string s_x = absl::StrCat("(start_x + ", kx, " * dilation.x)");
+        std::string s_x =
+            absl::StrCat("(start_x + ", kx, " * args.dilation_x)");
         if (manual_clamp) {
-          c += "    bool x_out = " + s_x + "< 0 || " + s_x + ">= src_size.x;\n";
+          c += "    bool x_out = " + s_x + "< 0 || " + s_x +
+               ">= args.src_tensor.Width();\n";
           c += "    " + s_type + " src = x_out || y_out ?";
-          c += "(" + s_type + ")(0.0) : ";
-          c += src_tensor.ReadWHS(s_x, s_y, std::to_string(s)) + s_postfix +
-               ";\n";
+          c += "(" + s_type + ")(0.0) : args.src_tensor.Read(" + s_x + ", " +
+               s_y + ", " + std::to_string(s) + ")" + s_postfix + ";\n";
         } else {
-          c += "    " + s_type + " src = " +
-               src_tensor.ReadWHS(s_x, s_y, std::to_string(s), address_mode) +
-               s_postfix + ";\n";
+          c += "    " + s_type + " src = args.src_tensor.Read(" + s_x + ", " +
+               s_y + ", " + std::to_string(s) + ")" + s_postfix + ";\n";
         }
         for (int d = 0; d < out_z; ++d) {
-          c += "    " + s_conv + "(r[" + std::to_string(d) + "], src, filters,";
+          c += "    " + s_conv + "(r[" + std::to_string(d) +
+               "], src, args.weigths.GetPtr(),";
           c += " " + std::to_string(filters_counter) + ");\n";
           filters_counter += ch_count;
         }
@@ -158,10 +164,9 @@ std::string GenerateConvolutionConstantCode(
   for (int i = 0; i < out_z; ++i) {
     std::string s_i = std::to_string(i);
     c += "  {\n";
-    c += "    FLT4 res = TO_FLT4(r[" + s_i + "]) + biases[" + s_i + "];\n";
-    const LinkingContext context{"res", "X", "Y", s_i};
-    c += PostProcess(linked_operations, context);
-    c += "  " + dst_tensor.WriteWHS("res", "X", "Y", s_i);
+    c += "    FLT4 res = TO_FLT4(r[" + s_i + "]) + args.biases.Read(" + s_i +
+         ");\n";
+    c += "  args.dst_tensor.Write(res, X, Y, " + s_i + ");\n";
     c += "  }\n";
   }
   c += "}\n";
@@ -191,8 +196,6 @@ int GetOptimalMaxConstantSize(const DeviceInfo& info) {
 
 ConvConstants::ConvConstants(ConvConstants&& kernel)
     : GPUOperation(std::move(kernel)),
-      weights_(std::move(kernel.weights_)),
-      biases_(std::move(kernel.biases_)),
       kernel_size_(kernel.kernel_size_),
       stride_(kernel.stride_),
       padding_(kernel.padding_),
@@ -204,8 +207,6 @@ ConvConstants::ConvConstants(ConvConstants&& kernel)
 
 ConvConstants& ConvConstants::operator=(ConvConstants&& kernel) {
   if (this != &kernel) {
-    weights_ = std::move(kernel.weights_);
-    biases_ = std::move(kernel.biases_);
     std::swap(kernel_size_, kernel.kernel_size_);
     std::swap(stride_, kernel.stride_);
     std::swap(padding_, kernel.padding_);
@@ -222,9 +223,15 @@ ConvConstants& ConvConstants::operator=(ConvConstants&& kernel) {
 absl::Status ConvConstants::Compile(const CreationContext& creation_context) {
   const bool stride_correction =
       definition_.IsBatchSupported() && stride_.x != 1;
-  const auto code = GenerateConvolutionConstantCode(
+  std::string code = GenerateConvolutionConstantCode(
       definition_, kernel_size_, src_channels_, dst_channels_,
-      stride_correction, *creation_context.device, linked_operations_);
+      stride_correction, *creation_context.device, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   std::vector options;
   if (definition_.precision == CalculationsPrecision::F16 &&
       creation_context.device->IsAdreno3xx()) {
@@ -241,20 +248,16 @@ absl::Status ConvConstants::Compile(const CreationContext& creation_context) {
 }
 
 absl::Status ConvConstants::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(weights_.GetMemoryPtr()));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(biases_.GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(stride_));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(padding_.x * src_[0]->Batch(), padding_.y)));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(dilation_.x * src_[0]->Batch(), dilation_.y)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHSB()));
-  return absl::OkStatus();
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(args_.SetInt("stride_x", stride_.x));
+  RETURN_IF_ERROR(args_.SetInt("stride_y", stride_.y));
+  RETURN_IF_ERROR(args_.SetInt("padding_x", padding_.x * src_[0]->Batch()));
+  RETURN_IF_ERROR(args_.SetInt("padding_y", padding_.y));
+  RETURN_IF_ERROR(args_.SetInt("dilation_x", dilation_.x * src_[0]->Batch()));
+  RETURN_IF_ERROR(args_.SetInt("dilation_y", dilation_.y));
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 ConvConstants::GetGridSize() const {
@@ -304,12 +307,18 @@ absl::Status CreateConvConstants(const CreationContext& creation_context,
   *result = ConvConstants(definition, attr);
   RETURN_IF_ERROR(
       result->UploadWeights(attr.weights, creation_context.context));
-  LinearStorageCreateInfo create_info;
-  create_info.storage_type = LinearStorageType::BUFFER;
-  create_info.data_type = definition.GetDataType();
-  create_info.aligned_size = attr.weights.shape.o;
-  RETURN_IF_ERROR(CreateLinearStorage(
-      create_info, attr.bias, creation_context.context, &result->biases_));
+
+  TensorLinearDescriptor desc;
+  desc.storage_type = LinearStorageType::BUFFER;
+  desc.element_type = definition.GetDataType();
+  desc.memory_type = MemoryType::CONSTANT;
+
+  LinearStorage lt;
+  RETURN_IF_ERROR(
+      CreateLinearStorage(desc, attr.bias, creation_context.context, <));
+  result->args_.AddObject("biases", AccessType::READ,
+                          absl::make_unique(std::move(lt)),
+                          absl::make_unique(desc));
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.h b/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.h
index 8d80d48314d..b9cc52f7e94 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/conv_constants.h
@@ -71,9 +71,6 @@ class ConvConstants : public GPUOperation {
   absl::Status BindArguments();
   int3 GetGridSize() const;
 
-  Buffer weights_;
-  LinearStorage biases_;
-
   int2 kernel_size_;
   int2 stride_;
   int2 padding_;
@@ -92,21 +89,34 @@ absl::Status ConvConstants::UploadWeights(
   const int kernel_x = weights.shape.w;
   const int kernel_y = weights.shape.h;
 
-  const int float_size =
-      definition_.precision == CalculationsPrecision::F32 ? 4 : 2;
+  const bool f32_weights = definition_.precision == CalculationsPrecision::F32;
+
+  BufferDescriptor desc;
+  desc.element_type = f32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+  desc.element_size = 4;
+  desc.memory_type = MemoryType::CONSTANT;
+
+  const int float_size = f32_weights ? 4 : 2;
   const int float_count = src_channels_ * dst_depth * 4 * kernel_x * kernel_y;
 
-  if (definition_.GetDataType() == DataType::FLOAT32) {
+  Buffer weights_buffer;
+  if (f32_weights) {
     std::vector gpu_data(float_count / 4);
     RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
-    return CreateReadOnlyBuffer(float_size * float_count, gpu_data.data(),
-                                context, &weights_);
+    RETURN_IF_ERROR(CreateReadOnlyBuffer(
+        float_size * float_count, gpu_data.data(), context, &weights_buffer));
   } else {
     std::vector gpu_data(float_count / 4);
     RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
-    return CreateReadOnlyBuffer(float_size * float_count, gpu_data.data(),
-                                context, &weights_);
+    RETURN_IF_ERROR(CreateReadOnlyBuffer(
+        float_size * float_count, gpu_data.data(), context, &weights_buffer));
   }
+
+  args_.AddObject("weigths", AccessType::READ,
+                  absl::make_unique(std::move(weights_buffer)),
+                  absl::make_unique(desc));
+
+  return absl::OkStatus();
 }
 
 template 
diff --git a/tensorflow/lite/delegates/gpu/cl/linear_storage.cc b/tensorflow/lite/delegates/gpu/cl/linear_storage.cc
index 47504a34c2b..ee0ea3efbec 100644
--- a/tensorflow/lite/delegates/gpu/cl/linear_storage.cc
+++ b/tensorflow/lite/delegates/gpu/cl/linear_storage.cc
@@ -32,6 +32,7 @@ GPUResources TensorLinearDescriptor::GetGPUResources(
     desc.data_type = element_type;
     desc.access_type = access_type;
     desc.element_size = 4;
+    desc.memory_type = memory_type;
     resources.buffers.push_back({"buffer", desc});
   } else {
     GPUImage2DDescriptor desc;
diff --git a/tensorflow/lite/delegates/gpu/cl/linear_storage.h b/tensorflow/lite/delegates/gpu/cl/linear_storage.h
index 2c0770ef3dc..14c8460bf80 100644
--- a/tensorflow/lite/delegates/gpu/cl/linear_storage.h
+++ b/tensorflow/lite/delegates/gpu/cl/linear_storage.h
@@ -40,6 +40,7 @@ enum class LinearStorageType { BUFFER, TEXTURE_2D };
 struct TensorLinearDescriptor : public GPUObjectDescriptor {
   LinearStorageType storage_type;
   DataType element_type;  // FLOAT32 or FLOAT16
+  MemoryType memory_type = MemoryType::GLOBAL;  // applicable for BUFFER
 
   absl::Status PerformSelector(const std::string& selector,
                                const std::vector& args,

From 6262e61826764c5236c1e74596583973c7248f3c Mon Sep 17 00:00:00 2001
From: Akshay Modi 
Date: Tue, 23 Jun 2020 09:51:51 -0700
Subject: [PATCH 0889/1390] Add more np array interop tests.

PiperOrigin-RevId: 317885566
Change-Id: I26dc3643547f0bd834b8e913e3c6eda7847f1aea
---
 .../python/ops/numpy_ops/np_interop_test.py   | 120 ++++++++++++++++--
 1 file changed, 106 insertions(+), 14 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_interop_test.py b/tensorflow/python/ops/numpy_ops/np_interop_test.py
index 17a3bf81a01..33abb58f260 100644
--- a/tensorflow/python/ops/numpy_ops/np_interop_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_interop_test.py
@@ -93,9 +93,8 @@ class InteropTest(tf.test.TestCase):
 
     dx, dy = t.gradient([xx, yy], [x, y])
 
-    # # TODO(nareshmodi): Figure out a way to rewrap ndarray as tensors.
-    # self.assertIsInstance(dx, np.ndarray)
-    # self.assertIsInstance(dy, np.ndarray)
+    self.assertIsInstance(dx, np.ndarray)
+    self.assertIsInstance(dy, np.ndarray)
     self.assertAllClose(dx, 2.0)
     self.assertAllClose(dy, 3.0)
 
@@ -181,19 +180,17 @@ class InteropTest(tf.test.TestCase):
 
     multiplier = np.asarray(5.)
 
-    with strategy.scope():
+    @tf.function
+    def run():
+      ctx = tf.distribute.get_replica_context()
+      val = np.asarray(ctx.replica_id_in_sync_group)
+      return val * multiplier
 
-      @tf.function
-      def run():
-        ctx = tf.distribute.get_replica_context()
-        val = np.asarray(ctx.replica_id_in_sync_group)
-        return val * multiplier
+    distributed_values = strategy.run(run)
+    reduced = strategy.reduce(
+        tf.distribute.ReduceOp.SUM, distributed_values, axis=None)
 
-      distributed_values = strategy.run(run)
-      reduced = strategy.reduce(
-          tf.distribute.ReduceOp.SUM, distributed_values, axis=None)
-
-    values = distributed_values.values
+    values = strategy.experimental_local_results(distributed_values)
 
     # Note that this should match the number of virtual CPUs.
     self.assertLen(values, 3)
@@ -208,6 +205,101 @@ class InteropTest(tf.test.TestCase):
     # self.assertIsInstance(reduced, np.ndarray)
     self.assertAllClose(reduced, 15)
 
+  def testPyFuncInterop(self):
+    def py_func_fn(a, b):
+      return a + b
+
+    @tf.function
+    def fn(a, b):
+      result = tf.py_function(py_func_fn, [a, b], a.dtype)
+      return np.asarray(result)
+
+    a = np.asarray(1.)
+    b = np.asarray(2.)
+
+    result = fn(a, b)
+    self.assertIsInstance(result, np.ndarray)
+    self.assertAllClose(result, 3.)
+
+  def testDatasetInterop(self):
+    values = [1, 2, 3, 4, 5, 6]
+    values_as_array = np.asarray(values)
+
+    # Tensor dataset
+    dataset = tf.data.Dataset.from_tensors(values_as_array)
+
+    for value, value_from_dataset in zip([values_as_array], dataset):
+      self.assertIsInstance(value_from_dataset, np.ndarray)
+      self.assertAllEqual(value_from_dataset, value)
+
+    # Tensor slice dataset
+    dataset = tf.data.Dataset.from_tensor_slices(values_as_array)
+
+    for value, value_from_dataset in zip(values, dataset):
+      self.assertIsInstance(value_from_dataset, np.ndarray)
+      self.assertAllEqual(value_from_dataset, value)
+
+    # # TODO(nareshmodi): as_numpy_iterator() doesn't work.
+    # items = list(dataset.as_numpy_iterator())
+
+    # Map over a dataset.
+    dataset = dataset.map(lambda x: np.add(x, 1))
+
+    for value, value_from_dataset in zip(values, dataset):
+      self.assertIsInstance(value_from_dataset, np.ndarray)
+      self.assertAllEqual(value_from_dataset, value + 1)
+
+    # Batch a dataset.
+    dataset = tf.data.Dataset.from_tensor_slices(values_as_array).batch(2)
+
+    for value, value_from_dataset in zip([[1, 2], [3, 4], [5, 6]], dataset):
+      self.assertIsInstance(value_from_dataset, np.ndarray)
+      self.assertAllEqual(value_from_dataset, value)
+
+  def testKerasInterop(self):
+    # Return an ndarray from the model.
+    inputs = tf.keras.layers.Input(shape=(10,))
+    output_layer = tf.keras.layers.Lambda(np.square)(inputs)
+    model = tf.keras.Model([inputs], output_layer)
+
+    values = onp.arange(10, dtype=onp.float32)
+    values_as_array = np.asarray(values)
+
+    result = model(values)
+    self.assertIsInstance(result, np.ndarray)
+    self.assertAllClose(result, onp.square(values))
+
+    result = model(values_as_array)
+    self.assertIsInstance(result, np.ndarray)
+    self.assertAllClose(result, onp.square(values))
+
+  def testPForInterop(self):
+    def outer_product(a):
+      return np.tensordot(a, a, 0)
+
+    batch_size = 100
+    a = np.ones((batch_size, 32, 32))
+    c = tf.vectorized_map(outer_product, a)
+
+    # # TODO(nareshmodi): vectorized_map doesn't rewrap tensors in ndarray.
+    # self.assertIsInstance(c, np.ndarray)
+    self.assertEqual(c.shape, (batch_size, 32, 32, 32, 32))
+
+  def testJacobian(self):
+    with tf.GradientTape() as g:
+      x = np.asarray([1., 2.])
+      y = np.asarray([3., 4.])
+      g.watch(x)
+      g.watch(y)
+      z = x * x * y
+
+    jacobian = g.jacobian(z, [x, y])
+    answer = [tf.linalg.diag(2 * x * y), tf.linalg.diag(x * x)]
+
+    self.assertIsInstance(jacobian[0], np.ndarray)
+    self.assertIsInstance(jacobian[1], np.ndarray)
+    self.assertAllClose(jacobian, answer)
+
 
 class FunctionTest(InteropTest):
 

From d34bb0c10657b29322b231b08a818738d4eb010f Mon Sep 17 00:00:00 2001
From: Raman Sarokin 
Date: Tue, 23 Jun 2020 09:52:04 -0700
Subject: [PATCH 0890/1390] DepthwiseConvolution converted to new style.
 DepthwiseConvolution3D merged into DepthwiseConvolution.

PiperOrigin-RevId: 317885608
Change-Id: Ifead6b0998c75d70f2be18ba2d32936e9e8ecdf0
---
 .../lite/delegates/gpu/cl/kernels/BUILD       |  23 --
 .../gpu/cl/kernels/depthwise_conv.cc          | 328 ++++++++++-------
 .../delegates/gpu/cl/kernels/depthwise_conv.h | 152 +++++++-
 .../gpu/cl/kernels/depthwise_conv_3d.cc       | 338 ------------------
 .../gpu/cl/kernels/depthwise_conv_3d.h        | 170 ---------
 5 files changed, 337 insertions(+), 674 deletions(-)
 delete mode 100644 tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.cc
 delete mode 100644 tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.h

diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/BUILD b/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
index 24a9a962296..21fb65e8909 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/BUILD
@@ -594,29 +594,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "depthwise_conv_3d",
-    srcs = ["depthwise_conv_3d.cc"],
-    hdrs = ["depthwise_conv_3d.h"],
-    deps = [
-        ":gpu_operation",
-        ":util",
-        ":work_group_picking",
-        "//tensorflow/lite/delegates/gpu/cl:buffer",
-        "//tensorflow/lite/delegates/gpu/cl:cl_device",
-        "//tensorflow/lite/delegates/gpu/cl:linear_storage",
-        "//tensorflow/lite/delegates/gpu/cl:tensor",
-        "//tensorflow/lite/delegates/gpu/cl:texture2d",
-        "//tensorflow/lite/delegates/gpu/cl:util",
-        "//tensorflow/lite/delegates/gpu/common:data_type",
-        "//tensorflow/lite/delegates/gpu/common:operations",
-        "//tensorflow/lite/delegates/gpu/common:shape",
-        "//tensorflow/lite/delegates/gpu/common:status",
-        "//tensorflow/lite/delegates/gpu/common:tensor",
-        "//tensorflow/lite/delegates/gpu/common:types",
-    ],
-)
-
 cc_test(
     name = "depthwise_conv_test",
     srcs = ["depthwise_conv_test.cc"],
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.cc b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.cc
index 4c5e20abde3..de1a04befa8 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.cc
@@ -34,57 +34,71 @@ bool IsSpecializedCase(int channel_multiplier) {
          channel_multiplier == 4;
 }
 
-std::string GetSrcValue(const TensorCodeGenerator& src_tensor,
-                        int channel_multiplier,
-                        TextureAddressMode address_mode) {
+std::string GetSrcValue(int channel_multiplier, const std::string coords) {
   std::string c;
   if (channel_multiplier == 1) {
-    c += "      FLT4 src_final =" +
-         src_tensor.ReadWHS("x_c", "y_c", "Z", address_mode) + ";\n";
+    c += "      FLT4 src_final = args.src_tensor.Read(" + coords + ", S);\n";
   } else if (channel_multiplier == 2) {
-    c += "      int z_layer = Z / 2;\n";
-    c += "      FLT4 src =" +
-         src_tensor.ReadWHS("x_c", "y_c", "z_layer", address_mode) + ";\n";
-    c += "      FLT2 t0 = Z % 2 == 0 ? src.xy : src.zw;\n";
+    c += "      int s_layer = S / 2;\n";
+    c += "      FLT4 src = args.src_tensor.Read(" + coords + ", s_layer);\n";
+    c += "      FLT2 t0 = S % 2 == 0 ? src.xy : src.zw;\n";
     c += "      FLT4 src_final = (FLT4)(t0.x, t0.x, t0.y, t0.y);\n";
   } else if (channel_multiplier == 4) {
-    c += "      int z_layer = Z / 4;\n";
-    c += "      FLT4 src =" +
-         src_tensor.ReadWHS("x_c", "y_c", "z_layer", address_mode) + ";\n";
+    c += "      int s_layer = S / 4;\n";
+    c += "      FLT4 src = args.src_tensor.Read(" + coords + ", s_layer);\n";
     c += "      FLT t0 = src.x;\n";
-    c += "      int reminder = Z % 4;\n";
+    c += "      int reminder = S % 4;\n";
     c += "      if (reminder == 1) t0 = src.y;\n";
     c += "      if (reminder == 2) t0 = src.z;\n";
     c += "      if (reminder == 3) t0 = src.w;\n";
     c += "      FLT4 src_final = (FLT4)(t0, t0, t0, t0);\n";
   } else {
-    c += "      int z_layer = Z / channel_multiplier;\n";
-    c += "      FLT4 src =" +
-         src_tensor.ReadWHS("x_c", "y_c", "z_layer", address_mode) + ";\n";
-    c += "      int z_offset = (Z % channel_multiplier) * 4;\n";
+    c += "      int s_layer = S / args.ch_multiplier;\n";
+    c += "      FLT4 src = args.src_tensor.Read(" + coords + ", s_layer);\n";
+    c += "      int s_offset = (S % args.ch_multiplier) * 4;\n";
     c += "      FLT4 src_final;\n";
     c += "      FLT temp_arr[4] = {src.x, src.y, src.z, src.w};\n";
-    c += "      src_final.x = temp_arr[(z_offset + 0) / channel_multiplier];\n";
-    c += "      src_final.y = temp_arr[(z_offset + 1) / channel_multiplier];\n";
-    c += "      src_final.z = temp_arr[(z_offset + 2) / channel_multiplier];\n";
-    c += "      src_final.w = temp_arr[(z_offset + 3) / channel_multiplier];\n";
+    c += "      src_final.x = temp_arr[(s_offset + 0) / args.ch_multiplier];\n";
+    c += "      src_final.y = temp_arr[(s_offset + 1) / args.ch_multiplier];\n";
+    c += "      src_final.z = temp_arr[(s_offset + 2) / args.ch_multiplier];\n";
+    c += "      src_final.w = temp_arr[(s_offset + 3) / args.ch_multiplier];\n";
   }
 
   return c;
 }
 
 std::string GenerateDepthwiseConvolutionCode(
-    const OperationDef& op_def, bool stride_correction,
-    const LinearStorage& biases, int channel_multiplier,
-    bool weights_are_buffer,
-    const std::vector& linked_operations,
-    const CLDevice& device) {
-  TensorCodeGenerator src_tensor(
-      "src_data", WHSPoint{"src_size.x", "src_size.y", "src_size.z"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data", WHSPoint{"dst_size.x", "dst_size.y", "dst_size.z"},
-      op_def.dst_tensors[0]);
+    const OperationDef& op_def, bool stride_correction, int channel_multiplier,
+    bool weights_are_buffer, const CLDevice& device, Arguments* args) {
+  auto src_desc = absl::make_unique(op_def.src_tensors[0]);
+  src_desc->SetTextureAddressMode(GetFastestZeroMode(device));
+  if (op_def.IsBatchSupported()) {
+    src_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("src_tensor", AccessType::READ, std::move(src_desc));
+  auto dst_desc = absl::make_unique(op_def.dst_tensors[0]);
+  if (op_def.IsBatchSupported()) {
+    dst_desc->SetStateVar("BatchedWidth", "true");
+  }
+  args->AddObjectRef("dst_tensor", AccessType::WRITE, std::move(dst_desc));
+  args->AddInt("kernel_size_x");
+  args->AddInt("stride_x");
+  args->AddInt("padding_x");
+  args->AddInt("dilation_x");
+  args->AddInt("kernel_size_y");
+  args->AddInt("stride_y");
+  args->AddInt("padding_y");
+  args->AddInt("dilation_y");
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    args->AddInt("kernel_size_z");
+    args->AddInt("stride_z");
+    args->AddInt("padding_z");
+    args->AddInt("dilation_z");
+  }
+  if (!IsSpecializedCase(channel_multiplier)) {
+    args->AddInt("ch_multiplier");
+  }
+
   const auto src_tensor_type = op_def.src_tensors[0].storage_type;
 
   std::string c = GetCommonDefines(op_def.precision);
@@ -93,86 +107,110 @@ std::string GenerateDepthwiseConvolutionCode(
                             src_tensor_type == TensorStorageType::IMAGE_BUFFER;
 
   c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ) + ",\n";
-  if (weights_are_buffer) {
-    c += "    __global FLT4* filters,  \n";
-  } else {
-    c += "    __read_only image2d_t filters,  \n";
-  }
-  c += biases.GetDeclaration();
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int2 kernel_size,                \n";
-  c += "    int2 stride,                     \n";
-  c += "    int2 padding,                    \n";
-  c += "    int2 dilation,                   \n";
-  if (!IsSpecializedCase(channel_multiplier)) {
-    c += "    int channel_multiplier,            \n";
-  }
-  c += "    int4 src_size,                   \n";
-  c += "    int4 dst_size                    \n";
-  c += ") {\n";
+  c += "$0) {\n";
   c += "  int X = get_global_id(0);\n";
   c += "  int Y = get_global_id(1);\n";
-  c += "  int Z = get_global_id(2);\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_2 = get_global_id(2);\n";
+    c += "  int S = linear_id_2 / args.dst_tensor.Depth();\n";
+    c += "  int Z = linear_id_2 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int S = get_global_id(2);\n";
+  }
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
   c += "  ACCUM_FLT4 r = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);\n";
   if (stride_correction) {
     c += "  int x_offseted = " +
-         GetXStrideCorrected("X", "src_size.w", "stride.x", "padding.x") +
+         GetXStrideCorrected("X", "args.src_tensor.Batch()", "args.stride_x",
+                             "args.padding_x") +
          ";\n";
   } else {
-    c += "  int x_offseted = X * stride.x + padding.x;\n";
+    c += "  int x_offseted = X * args.stride_x + args.padding_x;\n";
+  }
+  c += "  int y_offseted = Y * args.stride_y + args.padding_y;\n";
+  std::string weights_offset = "args.kernel_size_x * args.kernel_size_y";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int z_offseted = Z * args.stride_z + args.padding_z;\n";
+    weights_offset += " * args.kernel_size_z";
   }
-  c += "  int y_offseted = Y * stride.y + padding.y;\n";
   if (weights_are_buffer) {
-    c += "  int fx_c = Z * kernel_size.x * kernel_size.y;\n";
+    c += "  int fx_c = S * " + weights_offset + ";\n";
   } else {
     c += "  int fx_c = 0;\n";
   }
 
+  std::string flat_coords = "x_c, y_c";
   if (manual_clamp) {
-    c += "  for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-    c += "    int y_c = y_offseted + ky * dilation.y;\n";
-    c += "    bool outside_y = y_c < 0 || y_c >= src_size.y;\n";
-    c += "    for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
-    c += "      int x_c = x_offseted + kx * dilation.x;\n";
-    c += "      bool outside_x = x_c < 0 || x_c >= src_size.x;\n";
-    c += "      if (!outside_x && !outside_y) {\n";
-    if (weights_are_buffer) {
-      c += "        FLT4 f = filters[fx_c];\n";
-    } else {
-      c += "        FLT4 f = READ_IMAGE(filters, smp_none, (int2)(fx_c, Z));\n";
+    std::string check = "!outside_x && !outside_y";
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      check += " && !outside_z";
+      flat_coords += ", z_c";
+      c += "  for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+      c += "    int z_c = z_offseted + kz * args.dilation_z;\n";
+      c += "    bool outside_z = z_c < 0 || z_c >= args.src_tensor.Depth();\n";
     }
-    c += GetSrcValue(src_tensor, channel_multiplier,
-                     TextureAddressMode::DONT_CARE);
+    c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
+    c += "    int y_c = y_offseted + ky * args.dilation_y;\n";
+    c += "    bool outside_y = y_c < 0 || y_c >= args.src_tensor.Height();\n";
+    c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
+    c += "      int x_c = x_offseted + kx * args.dilation_x;\n";
+    c += "      bool outside_x = x_c < 0 || x_c >= args.src_tensor.Width();\n";
+    c += "      if (" + check + ") {\n";
+    if (weights_are_buffer) {
+      c += "        FLT4 f = args.weights.Read(fx_c);\n";
+    } else {
+      c += "        FLT4 f = args.weights.Read(fx_c, S);\n";
+    }
+    c += GetSrcValue(channel_multiplier, flat_coords);
     c += "        r += TO_ACCUM_TYPE(src_final * f);\n";
     c += "      };\n";
     c += "      fx_c++;\n";
     c += "    }\n";
     c += "  }\n";
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      c += "  }\n";
+    }
   } else {  // Texture types with ZERO clamping
-    c += "  for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-    c += "    int y_c = y_offseted + ky * dilation.y;\n";
-    c += "    for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
-    c += "      int x_c = x_offseted + kx * dilation.x;\n";
-    const auto access_mode = GetFastestZeroMode(device);
-    c += GetSrcValue(src_tensor, channel_multiplier, access_mode);
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      flat_coords += ", z_c";
+      c += "  for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+      c += "    int z_c = z_offseted + kz * args.dilation_z;\n";
+      if (src_tensor_type !=
+          TensorStorageType::TEXTURE_3D) {  // Only TEXTURE_3D supports clamping
+                                            // in DEPTH dimension
+        c += "    if (z_c < 0 || z_c >= args.src_tensor.Depth()) {\n";
+        c += "      fx_c += args.kernel_size_y * args.kernel_size_x;\n";
+        c += "      continue;\n";
+        c += "    }\n";
+      }
+    }
+    c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
+    c += "    int y_c = y_offseted + ky * args.dilation_y;\n";
+    c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
+    c += "      int x_c = x_offseted + kx * args.dilation_x;\n";
+    c += GetSrcValue(channel_multiplier, flat_coords);
     if (weights_are_buffer) {
-      c += "      FLT4 f = filters[fx_c];\n";
+      c += "      FLT4 f = args.weights.Read(fx_c);\n";
     } else {
-      c += "      FLT4 f = READ_IMAGE(filters, smp_none, (int2)(fx_c, Z));\n";
+      c += "      FLT4 f = args.weights.Read(fx_c, S);\n";
     }
     c += "      fx_c++;\n";
     c += "      r += TO_ACCUM_TYPE(src_final * f);\n";
     c += "    }\n";
     c += "  }\n";
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      c += "  }\n";
+    }
+  }
+  c += "  FLT4 res0 = TO_FLT4(r) + args.biases.Read(S);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  args.dst_tensor.Write(res0, X, Y, Z, S);\n";
+  } else {
+    c += "  args.dst_tensor.Write(res0, X, Y, S);\n";
   }
-  c += "  FLT4 bias_val = " + biases.ReadLinearFLT4("Z") + ";\n";
-  c += "  FLT4 res0 = TO_FLT4(r) + bias_val;\n";
-  const LinkingContext context{"res0", "X", "Y", "Z"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHS("res0", "X", "Y", "Z") + "\n";
   c += "}\n";
 
   return c;
@@ -184,20 +222,30 @@ DepthwiseConvolution::DepthwiseConvolution(
     const DepthwiseConvolution2DAttributes& attr, bool weights_are_buffer)
     : GPUOperation(definition),
       weights_are_buffer_(weights_are_buffer),
-      kernel_size_(attr.weights.shape.w, attr.weights.shape.h),
-      stride_(attr.strides.w, attr.strides.h),
-      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h),
-      dilation_(attr.dilations.w, attr.dilations.h),
+      kernel_size_(attr.weights.shape.w, attr.weights.shape.h, 0, 0),
+      stride_(attr.strides.w, attr.strides.h, 0, 0),
+      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h, 0, 0),
+      dilation_(attr.dilations.w, attr.dilations.h, 0, 0),
+      channel_multiplier_(attr.weights.shape.o),
+      work_group_size_(8, 8, 1) {}
+
+DepthwiseConvolution::DepthwiseConvolution(
+    const OperationDef& definition,
+    const DepthwiseConvolution3DAttributes& attr, bool weights_are_buffer)
+    : GPUOperation(definition),
+      weights_are_buffer_(weights_are_buffer),
+      kernel_size_(attr.weights.shape.w, attr.weights.shape.h,
+                   attr.weights.shape.d, 0),
+      stride_(attr.strides.w, attr.strides.h, attr.strides.d, 0),
+      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h,
+               -attr.padding.prepended.d, 0),
+      dilation_(attr.dilations.w, attr.dilations.h, attr.dilations.d, 0),
       channel_multiplier_(attr.weights.shape.o),
       work_group_size_(8, 8, 1) {}
 
 DepthwiseConvolution::DepthwiseConvolution(DepthwiseConvolution&& operation)
     : GPUOperation(std::move(operation)),
       weights_are_buffer_(operation.weights_are_buffer_),
-      weights_tex2d_(std::move(operation.weights_tex2d_)),
-      weights_buf_(std::move(operation.weights_buf_)),
-      weights_(operation.weights_),
-      biases_(std::move(operation.biases_)),
       kernel_size_(operation.kernel_size_),
       stride_(operation.stride_),
       padding_(operation.padding_),
@@ -210,10 +258,6 @@ DepthwiseConvolution& DepthwiseConvolution::operator=(
     DepthwiseConvolution&& operation) {
   if (this != &operation) {
     std::swap(weights_are_buffer_, operation.weights_are_buffer_);
-    weights_tex2d_ = std::move(operation.weights_tex2d_);
-    weights_buf_ = std::move(operation.weights_buf_);
-    std::swap(weights_, operation.weights_);
-    biases_ = std::move(operation.biases_);
     std::swap(kernel_size_, operation.kernel_size_);
     std::swap(stride_, operation.stride_);
     std::swap(padding_, operation.padding_);
@@ -230,39 +274,48 @@ absl::Status DepthwiseConvolution::Compile(
     const CreationContext& creation_context) {
   const bool stride_correction =
       definition_.IsBatchSupported() && stride_.x != 1;
-  const auto code = GenerateDepthwiseConvolutionCode(
-      definition_, stride_correction, biases_, channel_multiplier_,
-      weights_are_buffer_, linked_operations_, *creation_context.device);
+  std::string code = GenerateDepthwiseConvolutionCode(
+      definition_, stride_correction, channel_multiplier_, weights_are_buffer_,
+      *creation_context.device, &args_);
+  std::string element_wise_code;
+  RETURN_IF_ERROR(
+      MergeOperations(linked_operations_, &args_, &element_wise_code));
+  RETURN_IF_ERROR(args_.TransformToCLCode(creation_context.device->GetInfo(),
+                                          {{"dst_tensor", element_wise_code}},
+                                          &code));
   return creation_context.cache->GetOrCreateCLKernel(
       code, "main_function", *creation_context.context,
       *creation_context.device, &kernel_);
 }
 
 absl::Status DepthwiseConvolution::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(weights_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(biases_.GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(kernel_size_));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(stride_));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(padding_.x * src_[0]->Batch(), padding_.y)));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int2(dilation_.x * src_[0]->Batch(), dilation_.y)));
-  if (!IsSpecializedCase(channel_multiplier_)) {
-    RETURN_IF_ERROR(kernel_.SetBytesAuto(int32_t(channel_multiplier_)));
+  RETURN_IF_ERROR(args_.SetObjectRef("src_tensor", src_[0]));
+  RETURN_IF_ERROR(args_.SetObjectRef("dst_tensor", dst_[0]));
+  RETURN_IF_ERROR(args_.SetInt("kernel_size_x", kernel_size_.x));
+  RETURN_IF_ERROR(args_.SetInt("stride_x", stride_.x));
+  RETURN_IF_ERROR(args_.SetInt("padding_x", padding_.x * src_[0]->Batch()));
+  RETURN_IF_ERROR(args_.SetInt("dilation_x", dilation_.x * src_[0]->Batch()));
+  RETURN_IF_ERROR(args_.SetInt("kernel_size_y", kernel_size_.y));
+  RETURN_IF_ERROR(args_.SetInt("stride_y", stride_.y));
+  RETURN_IF_ERROR(args_.SetInt("padding_y", padding_.y));
+  RETURN_IF_ERROR(args_.SetInt("dilation_y", dilation_.y));
+  if (definition_.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    RETURN_IF_ERROR(args_.SetInt("kernel_size_z", kernel_size_.z));
+    RETURN_IF_ERROR(args_.SetInt("stride_z", stride_.z));
+    RETURN_IF_ERROR(args_.SetInt("padding_z", padding_.z));
+    RETURN_IF_ERROR(args_.SetInt("dilation_z", dilation_.z));
   }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHSB()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHSB()));
-  return absl::OkStatus();
+  if (!IsSpecializedCase(channel_multiplier_)) {
+    RETURN_IF_ERROR(args_.SetInt("ch_multiplier", channel_multiplier_));
+  }
+  RETURN_IF_ERROR(SetArguments(linked_operations_, &args_));
+  return args_.Bind(kernel_.kernel());
 }
 
 int3 DepthwiseConvolution::GetGridSize() const {
   const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
   const int grid_y = dst_[0]->Height();
-  const int grid_z = dst_[0]->Slices();
+  const int grid_z = dst_[0]->Slices() * dst_[0]->Depth();
   return int3(grid_x, grid_y, grid_z);
 }
 
@@ -284,14 +337,41 @@ absl::Status CreateDepthwiseConvolution(
   *result = DepthwiseConvolution(definition, attr, weights_are_buffer);
   RETURN_IF_ERROR(
       result->UploadWeights(attr.weights, creation_context.context));
-  LinearStorageCreateInfo create_info;
-  create_info.storage_type = weights_are_buffer ? LinearStorageType::BUFFER
-                                                : LinearStorageType::TEXTURE_2D;
-  create_info.data_type = definition.GetDataType();
-  create_info.name = "biases";
-  create_info.aligned_size = attr.weights.shape.o * attr.weights.shape.i;
-  RETURN_IF_ERROR(CreateLinearStorage(
-      create_info, attr.bias, creation_context.context, &result->biases_));
+
+  TensorLinearDescriptor desc;
+  desc.storage_type = weights_are_buffer ? LinearStorageType::BUFFER
+                                         : LinearStorageType::TEXTURE_2D;
+  desc.element_type = definition.GetDataType();
+
+  LinearStorage lt;
+  RETURN_IF_ERROR(
+      CreateLinearStorage(desc, attr.bias, creation_context.context, <));
+  result->args_.AddObject("biases", AccessType::READ,
+                          absl::make_unique(std::move(lt)),
+                          absl::make_unique(desc));
+  return absl::OkStatus();
+}
+
+absl::Status CreateDepthwiseConvolution(
+    const CreationContext& creation_context, const OperationDef& definition,
+    const DepthwiseConvolution3DAttributes& attr,
+    DepthwiseConvolution* result) {
+  bool weights_are_buffer = creation_context.device->IsMali();
+  *result = DepthwiseConvolution(definition, attr, weights_are_buffer);
+  RETURN_IF_ERROR(
+      result->UploadWeights(attr.weights, creation_context.context));
+
+  TensorLinearDescriptor desc;
+  desc.storage_type = weights_are_buffer ? LinearStorageType::BUFFER
+                                         : LinearStorageType::TEXTURE_2D;
+  desc.element_type = definition.GetDataType();
+
+  LinearStorage lt;
+  RETURN_IF_ERROR(
+      CreateLinearStorage(desc, attr.bias, creation_context.context, <));
+  result->args_.AddObject("biases", AccessType::READ,
+                          absl::make_unique(std::move(lt)),
+                          absl::make_unique(desc));
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.h b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.h
index 7655f2abae0..30cd3d06a5a 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.h
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv.h
@@ -54,9 +54,17 @@ class DepthwiseConvolution : public GPUOperation {
       const CreationContext& creation_context, const OperationDef& definition,
       const DepthwiseConvolution2DAttributes& attr,
       DepthwiseConvolution* result);
+  friend absl::Status CreateDepthwiseConvolution(
+      const CreationContext& creation_context, const OperationDef& definition,
+      const DepthwiseConvolution3DAttributes& attr,
+      DepthwiseConvolution* result);
   DepthwiseConvolution(const OperationDef& definition,
                        const DepthwiseConvolution2DAttributes& attr,
                        bool weights_are_buffer);
+  DepthwiseConvolution(const OperationDef& definition,
+                       const DepthwiseConvolution3DAttributes& attr,
+                       bool weights_are_buffer);
+
   template 
   absl::Status UploadWeights(const tflite::gpu::Tensor& weights,
                              CLContext* context);
@@ -65,20 +73,23 @@ class DepthwiseConvolution : public GPUOperation {
   void RearrangeWeightsData(const tflite::gpu::Tensor& weights,
                             absl::Span dst);
 
+  template 
+  absl::Status UploadWeights(const tflite::gpu::Tensor& weights,
+                             CLContext* context);
+
+  template 
+  void RearrangeWeightsData(const tflite::gpu::Tensor& weights,
+                            absl::Span dst);
+
   absl::Status BindArguments();
   int3 GetGridSize() const;
 
   bool weights_are_buffer_;
-  Texture2D weights_tex2d_;
-  Buffer weights_buf_;
-  cl_mem weights_;
 
-  LinearStorage biases_;
-
-  int2 kernel_size_;
-  int2 stride_;
-  int2 padding_;
-  int2 dilation_;
+  int4 kernel_size_;
+  int4 stride_;
+  int4 padding_;
+  int4 dilation_;
   int channel_multiplier_;
 
   CLKernel kernel_;
@@ -89,26 +100,28 @@ template 
 absl::Status DepthwiseConvolution::UploadWeights(
     const tflite::gpu::Tensor& weights, CLContext* context) {
   const int dst_channels = weights.shape.i * weights.shape.o;
-  const int dst_depth = DivideRoundUp(dst_channels, 4);
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
   const int kernel_x = weights.shape.w;
   const int kernel_y = weights.shape.h;
 
-  const int elements_count = kernel_x * kernel_y * dst_depth;
+  const int elements_count = kernel_x * kernel_y * dst_slices;
 
   const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
   const int float4_size = fp32_weights ? 16 : 8;
 
+  Texture2D weights_tex2d;
+  Buffer weights_buf;
   if (fp32_weights) {
     std::vector gpu_data(elements_count);
     RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
     if (weights_are_buffer_) {
       RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
                                            gpu_data.data(), context,
-                                           &weights_buf_));
+                                           &weights_buf));
     } else {
       RETURN_IF_ERROR(CreateTexture2DRGBA(
-          definition_.GetDataType(), kernel_x * kernel_y, dst_depth,
-          gpu_data.data(), context, &weights_tex2d_));
+          definition_.GetDataType(), kernel_x * kernel_y, dst_slices,
+          gpu_data.data(), context, &weights_tex2d));
     }
   } else {
     std::vector gpu_data(elements_count);
@@ -116,18 +129,27 @@ absl::Status DepthwiseConvolution::UploadWeights(
     if (weights_are_buffer_) {
       RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
                                            gpu_data.data(), context,
-                                           &weights_buf_));
+                                           &weights_buf));
     } else {
       RETURN_IF_ERROR(CreateTexture2DRGBA(
-          definition_.GetDataType(), kernel_x * kernel_y, dst_depth,
-          gpu_data.data(), context, &weights_tex2d_));
+          definition_.GetDataType(), kernel_x * kernel_y, dst_slices,
+          gpu_data.data(), context, &weights_tex2d));
     }
   }
 
   if (weights_are_buffer_) {
-    weights_ = weights_buf_.GetMemoryPtr();
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    args_.AddObject("weights", AccessType::READ,
+                    absl::make_unique(std::move(weights_buf)),
+                    absl::make_unique(desc));
   } else {
-    weights_ = weights_tex2d_.GetMemoryPtr();
+    Texture2DDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    args_.AddObject("weights", AccessType::READ,
+                    absl::make_unique(std::move(weights_tex2d)),
+                    absl::make_unique(desc));
   }
 
   return absl::OkStatus();
@@ -162,6 +184,98 @@ void DepthwiseConvolution::RearrangeWeightsData(
   }
 }
 
+template 
+absl::Status DepthwiseConvolution::UploadWeights(
+    const tflite::gpu::Tensor& weights, CLContext* context) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+  const int kernel_z = weights.shape.d;
+
+  const int elements_count = kernel_x * kernel_y * kernel_z * dst_slices;
+
+  const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int float4_size = fp32_weights ? 16 : 8;
+
+  Texture2D weights_tex2d;
+  Buffer weights_buf;
+  if (fp32_weights) {
+    std::vector gpu_data(elements_count);
+    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
+    if (weights_are_buffer_) {
+      RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
+                                           gpu_data.data(), context,
+                                           &weights_buf));
+    } else {
+      RETURN_IF_ERROR(CreateTexture2DRGBA(
+          definition_.GetDataType(), kernel_x * kernel_y * kernel_z, dst_slices,
+          gpu_data.data(), context, &weights_tex2d));
+    }
+  } else {
+    std::vector gpu_data(elements_count);
+    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
+    if (weights_are_buffer_) {
+      RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
+                                           gpu_data.data(), context,
+                                           &weights_buf));
+    } else {
+      RETURN_IF_ERROR(CreateTexture2DRGBA(
+          definition_.GetDataType(), kernel_x * kernel_y * kernel_z, dst_slices,
+          gpu_data.data(), context, &weights_tex2d));
+    }
+  }
+
+  if (weights_are_buffer_) {
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    args_.AddObject("weights", AccessType::READ,
+                    absl::make_unique(std::move(weights_buf)),
+                    absl::make_unique(desc));
+  } else {
+    Texture2DDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    args_.AddObject("weights", AccessType::READ,
+                    absl::make_unique(std::move(weights_tex2d)),
+                    absl::make_unique(desc));
+  }
+
+  return absl::OkStatus();
+}
+
+template 
+void DepthwiseConvolution::RearrangeWeightsData(
+    const tflite::gpu::Tensor& weights, absl::Span dst) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+  const int kernel_z = weights.shape.d;
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int z = 0; z < kernel_z; ++z) {
+      for (int y = 0; y < kernel_y; ++y) {
+        for (int x = 0; x < kernel_x; ++x) {
+          T filter_val;
+          for (int i = 0; i < 4; ++i) {
+            const int d_ch = d * 4 + i;
+            if (d_ch < dst_channels) {
+              const int f_index = weights.shape.LinearIndex(
+                  {d_ch % weights.shape.o, y, x, z, d_ch / weights.shape.o});
+              filter_val[i] = weights.data[f_index];
+            } else {
+              filter_val[i] = 0.0f;
+            }
+          }
+          dst[counter++] = filter_val;
+        }
+      }
+    }
+  }
+}
+
 absl::Status CreateDepthwiseConvolution(
     const CreationContext& creation_context, const OperationDef& definition,
     const DepthwiseConvolution2DAttributes& attr, DepthwiseConvolution* result);
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.cc b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.cc
deleted file mode 100644
index f9926a9f466..00000000000
--- a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.cc
+++ /dev/null
@@ -1,338 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.h"
-
-#include 
-#include 
-#include 
-
-#include "tensorflow/lite/delegates/gpu/cl/cl_device.h"
-#include "tensorflow/lite/delegates/gpu/cl/kernels/util.h"
-#include "tensorflow/lite/delegates/gpu/cl/kernels/work_group_picking.h"
-
-namespace tflite {
-namespace gpu {
-namespace cl {
-namespace {
-
-bool IsSpecializedCase(int channel_multiplier) {
-  return channel_multiplier == 1 || channel_multiplier == 2 ||
-         channel_multiplier == 4;
-}
-
-std::string GetSrcValue(const TensorCodeGenerator& src_tensor,
-                        int channel_multiplier,
-                        TextureAddressMode address_mode) {
-  std::string c;
-  if (channel_multiplier == 1) {
-    c += "        FLT4 src_final =" +
-         src_tensor.ReadWHDS("x_c", "y_c", "z_c", "S", address_mode) + ";\n";
-  } else if (channel_multiplier == 2) {
-    c += "        int z_layer = S / 2;\n";
-    c += "        FLT4 src =" +
-         src_tensor.ReadWHDS("x_c", "y_c", "z_c", "z_layer", address_mode) +
-         ";\n";
-    c += "        FLT2 t0 = S % 2 == 0 ? src.xy : src.zw;\n";
-    c += "        FLT4 src_final = (FLT4)(t0.x, t0.x, t0.y, t0.y);\n";
-  } else if (channel_multiplier == 4) {
-    c += "        int z_layer = S / 4;\n";
-    c += "        FLT4 src =" +
-         src_tensor.ReadWHDS("x_c", "y_c", "z_c", "z_layer", address_mode) +
-         ";\n";
-    c += "        FLT t0 = src.x;\n";
-    c += "        int reminder = S % 4;\n";
-    c += "        if (reminder == 1) t0 = src.y;\n";
-    c += "        if (reminder == 2) t0 = src.z;\n";
-    c += "        if (reminder == 3) t0 = src.w;\n";
-    c += "        FLT4 src_final = (FLT4)(t0, t0, t0, t0);\n";
-  } else {
-    c += "        int z_layer = S / channel_multiplier;\n";
-    c += "        FLT4 src =" +
-         src_tensor.ReadWHDS("x_c", "y_c", "z_c", "z_layer", address_mode) +
-         ";\n";
-    c += "        int z_offset = (S % channel_multiplier) * 4;\n";
-    c += "        FLT4 src_final;\n";
-    c += "        FLT temp_arr[4] = {src.x, src.y, src.z, src.w};\n";
-    c += "        src_final.x = temp_arr[(z_offset + 0) / "
-         "channel_multiplier];\n";
-    c += "        src_final.y = temp_arr[(z_offset + 1) / "
-         "channel_multiplier];\n";
-    c += "        src_final.z = temp_arr[(z_offset + 2) / "
-         "channel_multiplier];\n";
-    c += "        src_final.w = temp_arr[(z_offset + 3) / "
-         "channel_multiplier];\n";
-  }
-
-  return c;
-}
-
-std::string GenerateDepthwiseConvolution3DCode(
-    const OperationDef& op_def, bool stride_correction,
-    const LinearStorage& biases, int channel_multiplier,
-    bool weights_are_buffer,
-    const std::vector& linked_operations,
-    const CLDevice& device) {
-  TensorCodeGenerator src_tensor(
-      "src_data",
-      WHDSPoint{"src_size.x", "src_size.y", "src_size.z", "src_size.w"},
-      op_def.src_tensors[0]);
-  TensorCodeGenerator dst_tensor(
-      "dst_data",
-      WHDSPoint{"dst_size.x", "dst_size.y", "dst_size.z", "dst_size.w"},
-      op_def.dst_tensors[0]);
-  const auto src_tensor_type = op_def.src_tensors[0].storage_type;
-
-  std::string c = GetCommonDefines(op_def.precision);
-
-  const bool manual_clamp = src_tensor_type == TensorStorageType::BUFFER ||
-                            src_tensor_type == TensorStorageType::IMAGE_BUFFER;
-
-  c += "__kernel void main_function(\n";
-  c += src_tensor.GetDeclaration(AccessType::READ) + ",\n";
-  if (weights_are_buffer) {
-    c += "    __global FLT4* filters,  \n";
-  } else {
-    c += "    __read_only image2d_t filters,  \n";
-  }
-  c += biases.GetDeclaration();
-  c += GetArgsDeclaration(linked_operations);
-  c += dst_tensor.GetDeclaration(AccessType::WRITE) + ",\n";
-  c += "    int4 kernel_size,                \n";
-  c += "    int4 stride,                     \n";
-  c += "    int4 padding,                    \n";
-  c += "    int4 dilation,                   \n";
-  if (!IsSpecializedCase(channel_multiplier)) {
-    c += "    int channel_multiplier,            \n";
-  }
-  if (op_def.IsBatchSupported()) {
-    c += "    int batch_size,          \n";
-  }
-  c += "    int4 src_size,                   \n";
-  c += "    int4 dst_size                    \n";
-  c += ") {\n";
-  c += "  int X = get_global_id(0);\n";
-  c += "  int Y = get_global_id(1);\n";
-  c += "  int linear_id_z = get_global_id(2);\n";
-  c += "  int S = linear_id_z % dst_size.w;\n";
-  c += "  int Z = linear_id_z / dst_size.w;\n";
-  c += "  if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
-  c += "  ACCUM_FLT4 r = (ACCUM_FLT4)(0.0f, 0.0f, 0.0f, 0.0f);\n";
-  if (stride_correction) {
-    c += "  int x_offseted = " +
-         GetXStrideCorrected("X", "batch_size", "stride.x", "padding.x") +
-         ";\n";
-  } else {
-    c += "  int x_offseted = X * stride.x + padding.x;\n";
-  }
-  c += "  int y_offseted = Y * stride.y + padding.y;\n";
-  c += "  int z_offseted = Z * stride.z + padding.z;\n";
-  if (weights_are_buffer) {
-    c += "  int fx_c = S * kernel_size.x * kernel_size.y * kernel_size.z;\n";
-  } else {
-    c += "  int fx_c = 0;\n";
-  }
-
-  if (manual_clamp) {
-    c += "  for (int kz = 0; kz < kernel_size.z; ++kz) {\n";
-    c += "    int z_c = z_offseted + kz * dilation.z;\n";
-    c += "    bool outside_z = z_c < 0 || z_c >= src_size.z;\n";
-    c += "    for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-    c += "      int y_c = y_offseted + ky * dilation.y;\n";
-    c += "      bool outside_y = y_c < 0 || y_c >= src_size.y;\n";
-    c += "      for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
-    c += "        int x_c = x_offseted + kx * dilation.x;\n";
-    c += "        bool outside_x = x_c < 0 || x_c >= src_size.x;\n";
-    c += "        if (!outside_x && !outside_y && !outside_z) {\n";
-    if (weights_are_buffer) {
-      c += "          FLT4 f = filters[fx_c];\n";
-    } else {
-      c += "          FLT4 f = READ_IMAGE(filters, smp_none, (int2)(fx_c, "
-           "S));\n";
-    }
-    c += GetSrcValue(src_tensor, channel_multiplier,
-                     TextureAddressMode::DONT_CARE);
-    c += "          r += TO_ACCUM_TYPE(src_final * f);\n";
-    c += "        };\n";
-    c += "        fx_c++;\n";
-    c += "      }\n";
-    c += "    }\n";
-    c += "  }\n";
-  } else {  // Texture types with ZERO clamping
-    c += "  for (int kz = 0; kz < kernel_size.z; ++kz) {\n";
-    c += "    int z_c = z_offseted + kz * dilation.z;\n";
-    if (src_tensor_type !=
-        TensorStorageType::TEXTURE_3D) {  // Only TEXTURE_3D supports clamping
-                                          // in DEPTH dimension
-      c += "    if (z_c < 0 || z_c >= src_size.z) {\n";
-      c += "      fx_c += kernel_size.y * kernel_size.x;\n";
-      c += "      continue;\n";
-      c += "    }\n";
-    }
-    c += "    for (int ky = 0; ky < kernel_size.y; ++ky) {\n";
-    c += "      int y_c = y_offseted + ky * dilation.y;\n";
-    c += "      for (int kx = 0; kx < kernel_size.x; ++kx) {\n";
-    c += "        int x_c = x_offseted + kx * dilation.x;\n";
-    const auto access_mode = GetFastestZeroMode(device);
-    c += GetSrcValue(src_tensor, channel_multiplier, access_mode);
-    if (weights_are_buffer) {
-      c += "        FLT4 f = filters[fx_c];\n";
-    } else {
-      c += "        FLT4 f = READ_IMAGE(filters, smp_none, (int2)(fx_c, S));\n";
-    }
-    c += "        fx_c++;\n";
-    c += "        r += TO_ACCUM_TYPE(src_final * f);\n";
-    c += "      }\n";
-    c += "    }\n";
-    c += "  }\n";
-  }
-  c += "  FLT4 bias_val = " + biases.ReadLinearFLT4("S") + ";\n";
-  c += "  FLT4 res0 = TO_FLT4(r) + bias_val;\n";
-  const LinkingContext context{"res0", "X", "Y", "S"};
-  c += PostProcess(linked_operations, context);
-  c += "  " + dst_tensor.WriteWHDS("res0", "X", "Y", "Z", "S") + "\n";
-  c += "}\n";
-  return c;
-}
-}  // namespace
-
-DepthwiseConvolution3D::DepthwiseConvolution3D(
-    const OperationDef& definition,
-    const DepthwiseConvolution3DAttributes& attr, const CLDevice& device)
-    : GPUOperation(definition),
-      weights_are_buffer_(device.IsMali()),
-      kernel_size_(attr.weights.shape.w, attr.weights.shape.h,
-                   attr.weights.shape.d),
-      stride_(attr.strides.w, attr.strides.h, attr.strides.d),
-      padding_(-attr.padding.prepended.w, -attr.padding.prepended.h,
-               -attr.padding.prepended.d),
-      dilation_(attr.dilations.w, attr.dilations.h, attr.dilations.d),
-      channel_multiplier_(attr.weights.shape.o),
-      work_group_size_(8, 8, 1) {}
-
-DepthwiseConvolution3D::DepthwiseConvolution3D(
-    DepthwiseConvolution3D&& operation)
-    : GPUOperation(std::move(operation)),
-      weights_tex2d_(std::move(operation.weights_tex2d_)),
-      weights_buf_(std::move(operation.weights_buf_)),
-      weights_are_buffer_(operation.weights_are_buffer_),
-      biases_(std::move(operation.biases_)),
-      kernel_size_(operation.kernel_size_),
-      stride_(operation.stride_),
-      padding_(operation.padding_),
-      dilation_(operation.dilation_),
-      channel_multiplier_(operation.channel_multiplier_),
-      kernel_(std::move(operation.kernel_)),
-      work_group_size_(operation.work_group_size_) {}
-
-DepthwiseConvolution3D& DepthwiseConvolution3D::operator=(
-    DepthwiseConvolution3D&& operation) {
-  if (this != &operation) {
-    weights_tex2d_ = std::move(operation.weights_tex2d_);
-    weights_buf_ = std::move(operation.weights_buf_);
-    std::swap(weights_are_buffer_, operation.weights_are_buffer_);
-    biases_ = std::move(operation.biases_);
-    std::swap(kernel_size_, operation.kernel_size_);
-    std::swap(stride_, operation.stride_);
-    std::swap(padding_, operation.padding_);
-    std::swap(dilation_, operation.dilation_);
-    std::swap(channel_multiplier_, operation.channel_multiplier_);
-    kernel_ = std::move(operation.kernel_);
-    std::swap(work_group_size_, operation.work_group_size_);
-    GPUOperation::operator=(std::move(operation));
-  }
-  return *this;
-}
-
-absl::Status DepthwiseConvolution3D::Compile(
-    const CreationContext& creation_context) {
-  const bool stride_correction =
-      definition_.IsBatchSupported() && stride_.x != 1;
-  const auto code = GenerateDepthwiseConvolution3DCode(
-      definition_, stride_correction, biases_, channel_multiplier_,
-      weights_are_buffer_, linked_operations_, *creation_context.device);
-  return creation_context.cache->GetOrCreateCLKernel(
-      code, "main_function", *creation_context.context,
-      *creation_context.device, &kernel_);
-}
-
-absl::Status DepthwiseConvolution3D::BindArguments() {
-  kernel_.ResetBindingCounter();
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[0]->GetMemoryPtr()));
-  if (weights_are_buffer_) {
-    RETURN_IF_ERROR(kernel_.SetMemoryAuto(weights_buf_.GetMemoryPtr()));
-  } else {
-    RETURN_IF_ERROR(kernel_.SetMemoryAuto(weights_tex2d_.GetMemoryPtr()));
-  }
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(biases_.GetMemoryPtr()));
-  RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
-  RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(
-      int4(kernel_size_.x, kernel_size_.y, kernel_size_.z, 1)));
-  RETURN_IF_ERROR(
-      kernel_.SetBytesAuto(int4(stride_.x, stride_.y, stride_.z, 1)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(
-      int4(padding_.x * src_[0]->Batch(), padding_.y, padding_.z, 1)));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(
-      int4(dilation_.x * src_[0]->Batch(), dilation_.y, dilation_.z, 1)));
-  if (!IsSpecializedCase(channel_multiplier_)) {
-    RETURN_IF_ERROR(kernel_.SetBytesAuto(int32_t(channel_multiplier_)));
-  }
-  if (definition_.IsBatchSupported()) {
-    RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->Batch()));
-  }
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHDS()));
-  RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHDS()));
-  return absl::OkStatus();
-}
-
-int3 DepthwiseConvolution3D::GetGridSize() const {
-  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
-  const int grid_y = dst_[0]->Height();
-  const int grid_z = dst_[0]->Slices() * dst_[0]->Depth();
-  return int3(grid_x, grid_y, grid_z);
-}
-
-absl::Status DepthwiseConvolution3D::Tune(const TuningParameters& params) {
-  RETURN_IF_ERROR(BindArguments());
-  return GetBestWorkGroup(params, kernel_, GetGridSize(), &work_group_size_);
-}
-
-absl::Status DepthwiseConvolution3D::AddToQueue(CLCommandQueue* queue) {
-  RETURN_IF_ERROR(BindArguments());
-  return queue->DispatchImplicit(kernel_, GetGridSize(), work_group_size_);
-}
-
-absl::Status CreateDepthwiseConvolution3D(
-    const CreationContext& creation_context, const OperationDef& definition,
-    const DepthwiseConvolution3DAttributes& attr,
-    DepthwiseConvolution3D* result) {
-  *result = DepthwiseConvolution3D(definition, attr, *creation_context.device);
-  RETURN_IF_ERROR(
-      result->UploadWeights(attr.weights, creation_context.context));
-  LinearStorageCreateInfo create_info;
-  create_info.storage_type =
-      DeduceLinearStorageType(definition.GetPrimaryStorageType());
-  create_info.data_type = definition.GetDataType();
-  create_info.name = "biases";
-  create_info.aligned_size = attr.weights.shape.o * attr.weights.shape.i;
-  RETURN_IF_ERROR(CreateLinearStorage(
-      create_info, attr.bias, creation_context.context, &result->biases_));
-  return absl::OkStatus();
-}
-
-}  // namespace cl
-}  // namespace gpu
-}  // namespace tflite
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.h b/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.h
deleted file mode 100644
index 3c87ba5832c..00000000000
--- a/tensorflow/lite/delegates/gpu/cl/kernels/depthwise_conv_3d.h
+++ /dev/null
@@ -1,170 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_LITE_DELEGATES_GPU_CL_KERNELS_DEPTHWISE_CONV_3D_H_
-#define TENSORFLOW_LITE_DELEGATES_GPU_CL_KERNELS_DEPTHWISE_CONV_3D_H_
-
-#include 
-
-#include "tensorflow/lite/delegates/gpu/cl/buffer.h"
-#include "tensorflow/lite/delegates/gpu/cl/kernels/gpu_operation.h"
-#include "tensorflow/lite/delegates/gpu/cl/linear_storage.h"
-#include "tensorflow/lite/delegates/gpu/cl/tensor.h"
-#include "tensorflow/lite/delegates/gpu/cl/texture2d.h"
-#include "tensorflow/lite/delegates/gpu/cl/util.h"
-#include "tensorflow/lite/delegates/gpu/common/data_type.h"
-#include "tensorflow/lite/delegates/gpu/common/operations.h"
-#include "tensorflow/lite/delegates/gpu/common/shape.h"
-#include "tensorflow/lite/delegates/gpu/common/status.h"
-#include "tensorflow/lite/delegates/gpu/common/tensor.h"
-#include "tensorflow/lite/delegates/gpu/common/types.h"
-
-namespace tflite {
-namespace gpu {
-namespace cl {
-
-class DepthwiseConvolution3D : public GPUOperation {
- public:
-  DepthwiseConvolution3D() = default;
-  absl::Status AddToQueue(CLCommandQueue* queue) override;
-  absl::Status Tune(const TuningParameters& params) override;
-
-  absl::Status Compile(const CreationContext& creation_context) override;
-
-  // Move only
-  DepthwiseConvolution3D(DepthwiseConvolution3D&& operation);
-  DepthwiseConvolution3D& operator=(DepthwiseConvolution3D&& operation);
-  DepthwiseConvolution3D(const DepthwiseConvolution3D&) = delete;
-  DepthwiseConvolution3D& operator=(const DepthwiseConvolution3D&) = delete;
-
- private:
-  friend absl::Status CreateDepthwiseConvolution3D(
-      const CreationContext& creation_context, const OperationDef& definition,
-      const DepthwiseConvolution3DAttributes& attr,
-      DepthwiseConvolution3D* result);
-  DepthwiseConvolution3D(const OperationDef& definition,
-                         const DepthwiseConvolution3DAttributes& attr,
-                         const CLDevice& device);
-  template 
-  absl::Status UploadWeights(const tflite::gpu::Tensor& weights,
-                             CLContext* context);
-
-  template 
-  void RearrangeWeightsData(const tflite::gpu::Tensor& weights,
-                            absl::Span dst);
-
-  absl::Status BindArguments();
-  int3 GetGridSize() const;
-
-  Texture2D weights_tex2d_;
-  Buffer weights_buf_;
-  bool weights_are_buffer_;
-
-  LinearStorage biases_;
-
-  int3 kernel_size_;
-  int3 stride_;
-  int3 padding_;
-  int3 dilation_;
-  int channel_multiplier_;
-
-  CLKernel kernel_;
-  int3 work_group_size_;
-};
-
-template 
-absl::Status DepthwiseConvolution3D::UploadWeights(
-    const tflite::gpu::Tensor& weights, CLContext* context) {
-  const int dst_channels = weights.shape.i * weights.shape.o;
-  const int dst_slices = DivideRoundUp(dst_channels, 4);
-  const int kernel_x = weights.shape.w;
-  const int kernel_y = weights.shape.h;
-  const int kernel_z = weights.shape.d;
-
-  const int elements_count = kernel_x * kernel_y * kernel_z * dst_slices;
-  const bool f32_weights = definition_.precision == CalculationsPrecision::F32;
-
-  const int float4_size = f32_weights ? 16 : 8;
-
-  if (f32_weights) {
-    std::vector gpu_data(elements_count);
-    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
-    if (weights_are_buffer_) {
-      RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
-                                           gpu_data.data(), context,
-                                           &weights_buf_));
-    } else {
-      RETURN_IF_ERROR(CreateTexture2DRGBA(
-          definition_.GetDataType(), kernel_x * kernel_y * kernel_z, dst_slices,
-          gpu_data.data(), context, &weights_tex2d_));
-    }
-  } else {
-    std::vector gpu_data(elements_count);
-    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data));
-    if (weights_are_buffer_) {
-      RETURN_IF_ERROR(CreateReadOnlyBuffer(float4_size * elements_count,
-                                           gpu_data.data(), context,
-                                           &weights_buf_));
-    } else {
-      RETURN_IF_ERROR(CreateTexture2DRGBA(
-          definition_.GetDataType(), kernel_x * kernel_y * kernel_z, dst_slices,
-          gpu_data.data(), context, &weights_tex2d_));
-    }
-  }
-  return absl::OkStatus();
-}
-
-template 
-void DepthwiseConvolution3D::RearrangeWeightsData(
-    const tflite::gpu::Tensor& weights, absl::Span dst) {
-  const int dst_channels = weights.shape.i * weights.shape.o;
-  const int dst_slices = DivideRoundUp(dst_channels, 4);
-  const int kernel_x = weights.shape.w;
-  const int kernel_y = weights.shape.h;
-  const int kernel_z = weights.shape.d;
-
-  int counter = 0;
-  for (int d = 0; d < dst_slices; ++d) {
-    for (int z = 0; z < kernel_z; ++z) {
-      for (int y = 0; y < kernel_y; ++y) {
-        for (int x = 0; x < kernel_x; ++x) {
-          T filter_val;
-          for (int i = 0; i < 4; ++i) {
-            const int d_ch = d * 4 + i;
-            if (d_ch < dst_channels) {
-              const int f_index = weights.shape.LinearIndex(
-                  {d_ch % weights.shape.o, y, x, z, d_ch / weights.shape.o});
-              filter_val[i] = weights.data[f_index];
-            } else {
-              filter_val[i] = 0.0f;
-            }
-          }
-          dst[counter++] = filter_val;
-        }
-      }
-    }
-  }
-}
-
-absl::Status CreateDepthwiseConvolution3D(
-    const CreationContext& creation_context, const OperationDef& definition,
-    const DepthwiseConvolution3DAttributes& attr,
-    DepthwiseConvolution3D* result);
-
-}  // namespace cl
-}  // namespace gpu
-}  // namespace tflite
-
-#endif  // TENSORFLOW_LITE_DELEGATES_GPU_CL_KERNELS_DEPTHWISE_CONV_3D_H_

From bd6f60bf01875b4151ebf7bb37a7172ffa9ec426 Mon Sep 17 00:00:00 2001
From: Robert David 
Date: Tue, 23 Jun 2020 09:55:28 -0700
Subject: [PATCH 0891/1390] Rename few LSTM symbols to a name that is
 consistent among different versions.

Also remove an unused std::vector variable.

PiperOrigin-RevId: 317886335
Change-Id: Ib1e131c42d86b983630a4f8130d9950a45119b4a
---
 tensorflow/lite/kernels/lstm_eval.cc | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/lite/kernels/lstm_eval.cc b/tensorflow/lite/kernels/lstm_eval.cc
index 5b4e8a8d479..ca8344d863b 100644
--- a/tensorflow/lite/kernels/lstm_eval.cc
+++ b/tensorflow/lite/kernels/lstm_eval.cc
@@ -490,7 +490,7 @@ inline void LstmStepFloat(
 // Temporary pre-allocated storage for quantized values:
 //   quantized_input_ptr (same size as input_ptr)
 //   quantized_output_state_ptr (same size as output_state_ptr)
-//   quantized_cell_state_ptr (same size as cell_state_ptr)
+//   quantized_output_scratch (same size as cell_state_ptr)
 // Temporary pre-allocated storage for recovered values:
 //   recovered_cell_weights (same size as cell_to_*_weights)
 //
@@ -540,7 +540,7 @@ inline void LstmStepHybrid(
     float* scratch2, float* scratch3, float* scaling_factors,
     float* scaling_factors_scratch, float* recovered_cell_weights,
     int8_t* quantized_input_ptr, int8_t* quantized_aux_input_ptr,
-    int8_t* quantized_output_state_ptr, int8_t* quantized_cell_state_ptr,
+    int8_t* quantized_output_state_ptr, int8_t* quantized_output_scratch,
     float* output_state_ptr, float* cell_state_ptr, int32_t* accum_scratch_ptr,
     float* output_ptr, int32_t* zero_points, int32_t* row_sums,
     int row_sums_size, bool* compute_row_sums, bool asymmetric_quantize_inputs,
@@ -882,10 +882,10 @@ inline void LstmStepHybrid(
     if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) {
       // Save quantization and matmul computation for all zero input.
       tensor_utils::BatchQuantizeFloats(
-          output_gate_scratch, n_batch, n_cell, quantized_cell_state_ptr,
+          output_gate_scratch, n_batch, n_cell, quantized_output_scratch,
           scaling_factors, zero_points, asymmetric_quantize_inputs);
       tensor_utils::MatrixBatchVectorMultiplyAccumulate(
-          projection_weights_ptr, n_output, n_cell, quantized_cell_state_ptr,
+          projection_weights_ptr, n_output, n_cell, quantized_output_scratch,
           projection_weights_scale, scaling_factors, n_batch, output_state_ptr,
           /*per_channel_scale=*/nullptr,
           asymmetric_quantize_inputs ? zero_points : nullptr, accum_scratch_ptr,
@@ -996,7 +996,7 @@ inline void LstmStepHybrid(
 //   output_state_ptr - size 'n_batch * n_output'
 //   cell_state_ptr   - size 'n_batch * n_cell'
 //   output_ptr       - size 'n_batch * n_output'
-inline void LstmStepInteger(
+inline void LstmStepInteger8x8_16(
     const int8_t* input_ptr, const int8_t* input_to_input_weight_ptr,
     int32_t effective_input_to_input_scale_a,
     int32_t effective_input_to_input_scale_b,
@@ -1060,7 +1060,7 @@ inline void LstmStepInteger(
     int32_t output_state_zp, int16_t* cell_state_ptr, int8_t* output_ptr,
     int16_t* scratch0, int16_t* scratch1, int16_t* scratch2, int16_t* scratch3,
     int8_t* scratch4, int32_t* scratch5, CpuBackendContext* context) {
-  ruy::profiler::ScopeLabel label("LstmStepInteger");
+  ruy::profiler::ScopeLabel label("LstmStepInteger8x8_16");
   // Make named scratch buffers for the different gates.
   int16_t* input_gate_scratch = scratch0;
   int16_t* forget_gate_scratch = scratch1;
@@ -1336,7 +1336,7 @@ inline void LstmStepInteger(
 //   cell_state_ptr   - size 'n_batch * n_cell'
 //   output_ptr       - size 'n_batch * n_output'
 // TODO(b/148688698): Move zero point calculation into Prepare().
-void LstmStepInteger(
+inline void LstmStepInteger8x8_8(
     const int8_t* input_ptr, int32_t input_zp,
     const int8_t* input_to_input_weight_ptr,
     int32_t effective_input_to_input_scale_a,
@@ -1391,6 +1391,7 @@ void LstmStepInteger(
     int8_t* scratch0, int8_t* scratch1, int16_t* scratch2, int16_t* scratch3,
     int16_t* scratch4, int16_t* scratch5, int16_t* scratch6,
     int16_t* scratch7) {
+  ruy::profiler::ScopeLabel label("LstmStepInteger8x8_8");
   // Make named scratch buffers for the different gates.
   int16_t* input_gate_scratch = scratch5;
   int16_t* forget_gate_scratch = scratch2;
@@ -1426,7 +1427,7 @@ void LstmStepInteger(
   tensor_utils::ApplySigmoidFloat(forget_gate_scratch, n_batch, n_cell,
                                   forget_gate_scratch);
 
-  // Update gate.
+  // Cell gate.
   std::fill_n(scratch0, n_batch * n_cell, 0);
   std::fill_n(scratch1, n_batch * n_cell, 0);
   tensor_utils::MatrixBatchVectorMultiply(
@@ -1444,13 +1445,13 @@ void LstmStepInteger(
       intermediate_scale_a[4], intermediate_scale_b[4], intermediate_scale_a[5],
       intermediate_scale_b[5], n_batch, n_cell, cell_gate_scratch);
 
-  // Update gate layer norm.
+  // Cell gate layer norm.
   tensor_utils::ApplyLayerNormFloat(
       cell_gate_scratch, layer_norm_cell_weight_ptr, layer_norm_cell_scale_a,
       layer_norm_cell_scale_b, cell_gate_bias_ptr, n_batch, n_cell,
       cell_gate_scratch);
 
-  // Update gate tanh.
+  // Cell gate tanh.
   tensor_utils::ApplyTanhFloat(cell_gate_scratch, n_batch, n_cell, -12,
                                cell_gate_scratch);
 
@@ -1505,7 +1506,6 @@ void LstmStepInteger(
   tensor_utils::ApplyTanhFloat(cell_state_ptr, n_batch, n_cell, -15,
                                forget_gate_scratch);
 
-  std::vector hidden(n_batch * n_cell);
   tensor_utils::CwiseMul(output_gate_scratch, forget_gate_scratch, n_batch,
                          n_cell, 15 + 15 - 15, cell_gate_scratch);
 
@@ -2004,7 +2004,7 @@ TfLiteStatus EvalInteger8x8_16(
     const int t_rel = t;
     int8_t* output_ptr = GetTensorData(output) + t_rel * output_step;
     const int8_t* input_ptr = GetTensorData(input) + t_rel * input_step;
-    LstmStepInteger(
+    LstmStepInteger8x8_16(
         input_ptr, GetTensorData(input_to_input_weights),
         integer_lstm_param->effective_input_to_input_scale_a,
         integer_lstm_param->effective_input_to_input_scale_b,
@@ -2140,7 +2140,7 @@ TfLiteStatus EvalInteger8x8_8(
     int8_t* output_ptr = GetTensorData(output) + t_rel * output_step;
     // Input can be int8 asymmetric or int16 symmetric.
     const int8_t* input_ptr = GetTensorData(input) + t_rel * input_step;
-    lstm_eval::LstmStepInteger(
+    lstm_eval::LstmStepInteger8x8_8(
         input_ptr, input_zp,
 
         GetTensorData(input_to_input_weights),

From b07691301fa26602c35531aa6349e95136c9ef3d Mon Sep 17 00:00:00 2001
From: Smit Hinsu 
Date: Tue, 23 Jun 2020 10:05:49 -0700
Subject: [PATCH 0892/1390] MultiProcessRunner: Add
 UnexpectedSubprocessExitError to be raised if the exit code from a subprocess
 in unexpected. This results in subprocess having segfault failing the test,
 which would not have before this change.

PiperOrigin-RevId: 317888643
Change-Id: Id1c372ec695e8478bbb9c0888516ca9af4bb9799
---
 .../python/distribute/multi_process_runner.py | 38 ++-----------------
 .../distribute/multi_process_runner_test.py   | 35 -----------------
 2 files changed, 3 insertions(+), 70 deletions(-)

diff --git a/tensorflow/python/distribute/multi_process_runner.py b/tensorflow/python/distribute/multi_process_runner.py
index cb460c8fff5..db31b9c4dd4 100644
--- a/tensorflow/python/distribute/multi_process_runner.py
+++ b/tensorflow/python/distribute/multi_process_runner.py
@@ -475,28 +475,14 @@ class MultiProcessRunner(object):
     process_statuses = self._queue_to_list(self._process_status_queue)
     if not self._all_forced_terminated and len(
         process_statuses) != self._outstanding_subprocess_count:
-      raise UnexpectedSubprocessExitError(
-          'Missing status(es) from %d subprocess(es). See logs for details.' %
-          (self._outstanding_subprocess_count - len(process_statuses)),
-          self._get_mpr_result(process_statuses))
+      raise RuntimeError(
+          'missing statuses from %d subproceses.' %
+          (self._outstanding_subprocess_count - len(process_statuses)))
     for process_status in process_statuses:
       assert isinstance(process_status, _ProcessStatusInfo)
       if not process_status.is_successful:
         six.reraise(*process_status.exc_info)
 
-    # Checking all the processes that are expected to exit properly.
-    for (task_type, task_id), p in self._processes.items():
-      if self._dependence_on_chief and task_type != 'chief':
-        # If _dependence_on_chief, other processes may have been
-        # forced-terminated, which is expected.
-        continue
-      # Successfully exiting process has exit code 0.
-      if p.exitcode > 0:
-        raise UnexpectedSubprocessExitError(
-            'Subprocess %s-%d exited with exit code %d. See logs for details.' %
-            (task_type, task_id, p.exitcode),
-            self._get_mpr_result(process_statuses))
-
     logging.info('Joining log reading threads.')
     for thread in self._reading_threads:
       thread.join()
@@ -532,8 +518,6 @@ class MultiProcessRunner(object):
     for (task_type, task_id), p in self._processes.items():
       try:
         os.kill(p.pid, sig)
-        logging.info('%s-%d terminated with signal %r.', task_type, task_id,
-                     sig)
       except ProcessLookupError:
         logging.info('Attempting to kill %s-%d but it does not exist.',
                      task_type, task_id)
@@ -686,9 +670,6 @@ class _ProcFunc(object):
       self._resources.process_status_queue.put(info)
       self._close_streaming()
 
-    # Exit with code 0 as it's considered successful exit at this point.
-    sys.exit(0)
-
 
 class SubprocessTimeoutError(RuntimeError):
   """An error that indicates there is at least one subprocess timing out.
@@ -703,19 +684,6 @@ class SubprocessTimeoutError(RuntimeError):
     self.mpr_result = mpr_result
 
 
-class UnexpectedSubprocessExitError(RuntimeError):
-  """An error indicating there is at least one subprocess with unexpected exit.
-
-  When this is raised, a `MultiProcessRunnerResult` object can be retrieved by
-  `UnexpectedSubprocessExitError`'s mpr_result attribute. See
-  `MultiProcessRunner.join()` for more information.
-  """
-
-  def __init__(self, msg, mpr_result):
-    super(UnexpectedSubprocessExitError, self).__init__(msg)
-    self.mpr_result = mpr_result
-
-
 def _set_tf_config(task_type, task_id, cluster_spec, rpc_layer=None):
   """Set TF_CONFIG environment variable."""
   tf_config_dict = {
diff --git a/tensorflow/python/distribute/multi_process_runner_test.py b/tensorflow/python/distribute/multi_process_runner_test.py
index 529d7fd91a5..d6e04010e34 100644
--- a/tensorflow/python/distribute/multi_process_runner_test.py
+++ b/tensorflow/python/distribute/multi_process_runner_test.py
@@ -18,7 +18,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import ctypes
 import json
 import os
 import threading
@@ -299,40 +298,6 @@ class MultiProcessRunnerTest(test.TestCase):
     self.assertTrue(
         any('something printed' in line for line in list_to_assert))
 
-  def test_seg_fault_raises_error(self):
-
-    def proc_func_expected_to_seg_fault():
-      ctypes.string_at(0)  # Intentionally made seg fault.
-
-    with self.assertRaises(
-        multi_process_runner.UnexpectedSubprocessExitError) as cm:
-      multi_process_runner.run(
-          proc_func_expected_to_seg_fault,
-          multi_worker_test_base.create_cluster_spec(num_workers=1),
-          list_stdout=True)
-    self.assertIn('Missing status(es) from 1 subprocess(es).',
-                  str(cm.exception))
-    list_to_assert = cm.exception.mpr_result.stdout
-    self.assertTrue(any('SIGSEGV' in line for line in list_to_assert))
-
-  def test_seg_fault_in_chief_raises_error(self):
-
-    def proc_func_expected_to_seg_fault():
-      if multi_worker_test_base.get_task_type() == 'worker':
-        time.sleep(10000)
-      ctypes.string_at(0)  # Intentionally made seg fault.
-
-    with self.assertRaises(
-        multi_process_runner.UnexpectedSubprocessExitError) as cm:
-      multi_process_runner.run(
-          proc_func_expected_to_seg_fault,
-          multi_worker_test_base.create_cluster_spec(
-              has_chief=True, num_workers=1),
-          list_stdout=True)
-    self.assertIn('Subprocess chief-0 exited with exit code',
-                  str(cm.exception))
-    list_to_assert = cm.exception.mpr_result.stdout
-    self.assertTrue(any('SIGSEGV' in line for line in list_to_assert))
 
 if __name__ == '__main__':
   multi_process_runner.test_main()

From e4d6335bcb7a73cd8967c2c12339380aa1ae284f Mon Sep 17 00:00:00 2001
From: Hye Soo Yang 
Date: Tue, 23 Jun 2020 10:21:31 -0700
Subject: [PATCH 0893/1390] Refactor DecodeImageOp for the purpose of removing
 redundant data parsing and format checks from python wrapper and having them
 take place only in kernels. Remove security concerns. This change:

- Creates new op kernel (`DecodeImageV2Op`) that can decode all four image formats (jpg, png, gif, bmp). `DecodeImage` is the op name. `DecodeBmpOp` is moved into `DecodeImageV2Op`. (Now we have `gen_image_ops.decode_image` as opposed to previous `decode_image` which was a pure python implementation.)
- Updates GIF decoder to take in `expand_animation` flag for decoding just one frame.
- Removes data parsing and format checking logic from python layer entirely.
- Updates magic bytes for detecting image formats.
- Replicates portions of `convert_image_dtype` functionality in kernel (for optionally converting uint8/uint16 -> float32).

PiperOrigin-RevId: 317891936
Change-Id: I84f18e053f6dad845d9f2a61e1119f4de131c85d
---
 .../base_api/api_def_DecodeImage.pbtxt        |  51 ++
 .../python_api/api_def_DecodeImage.pbtxt      |   4 +
 tensorflow/core/kernels/decode_image_op.cc    | 459 +++++++++++++++++-
 tensorflow/core/lib/gif/gif_io.cc             |  13 +-
 tensorflow/core/lib/gif/gif_io.h              |   2 +-
 tensorflow/core/ops/image_ops.cc              |  45 ++
 tensorflow/core/ops/image_ops_test.cc         |  33 +-
 tensorflow/python/ops/image_ops_impl.py       |  17 +
 tensorflow/python/ops/image_ops_test.py       | 201 ++++----
 .../api/golden/v1/tensorflow.raw_ops.pbtxt    |   4 +
 .../api/golden/v2/tensorflow.raw_ops.pbtxt    |   4 +
 11 files changed, 738 insertions(+), 95 deletions(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_DecodeImage.pbtxt
 create mode 100644 tensorflow/core/api_def/python_api/api_def_DecodeImage.pbtxt

diff --git a/tensorflow/core/api_def/base_api/api_def_DecodeImage.pbtxt b/tensorflow/core/api_def/base_api/api_def_DecodeImage.pbtxt
new file mode 100644
index 00000000000..c534425eb24
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_DecodeImage.pbtxt
@@ -0,0 +1,51 @@
+op {
+  graph_op_name: "DecodeImage"
+  in_arg {
+    name: "contents"
+    description: <
A legal label file is the plain text file whose contents are split into lines, and each line
-   * is an individual value. The file should be in assets of the context.
-   *
-   * @param context The context holds assets.
-   * @param filePath The path of the label file, relative with assets directory.
-   * @return a list of labels.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadLabels(@NonNull Context context, @NonNull String filePath)
-      throws IOException {
-    return loadLabels(context, filePath, Charset.defaultCharset());
-  }
-
-  /**
-   * Loads labels from the label file into a list of strings.
-   *
-   * 
A legal label file is the plain text file whose contents are split into lines, and each line
-   * is an individual value. The empty lines will be ignored. The file should be in assets of the
-   * context.
-   *
-   * @param context The context holds assets.
-   * @param filePath The path of the label file, relative with assets directory.
-   * @param cs {@code Charset} to use when decoding content of label file.
-   * @return a list of labels.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadLabels(
-      @NonNull Context context, @NonNull String filePath, Charset cs) throws IOException {
-    SupportPreconditions.checkNotNull(context, "Context cannot be null.");
-    SupportPreconditions.checkNotNull(filePath, "File path cannot be null.");
-    try (InputStream inputStream = context.getAssets().open(filePath)) {
-      return loadLabels(inputStream, cs);
-    }
-  }
-
-  /**
-   * Loads labels from an input stream of an opened label file. See details for label files in
-   * {@link FileUtil#loadLabels(Context, String)}.
-   *
-   * @param inputStream the input stream of an opened label file.
-   * @return a list of labels.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadLabels(@NonNull InputStream inputStream) throws IOException {
-    return loadLabels(inputStream, Charset.defaultCharset());
-  }
-
-  /**
-   * Loads labels from an input stream of an opened label file. See details for label files in
-   * {@link FileUtil#loadLabels(Context, String)}.
-   *
-   * @param inputStream the input stream of an opened label file.
-   * @param cs {@code Charset} to use when decoding content of label file.
-   * @return a list of labels.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadLabels(@NonNull InputStream inputStream, Charset cs)
-      throws IOException {
-    List labels = new ArrayList<>();
-    try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, cs))) {
-      String line;
-      while ((line = reader.readLine()) != null) {
-        if (line.trim().length() > 0) {
-          labels.add(line);
-        }
-      }
-      return labels;
-    }
-  }
-
-  /**
-   * Loads a vocabulary file (a single-column text file) into a list of strings.
-   *
-   * 
A vocabulary file is a single-column plain text file whose contents are split into lines,
-   * and each line is an individual value. The file should be in assets of the context.
-   *
-   * @param context The context holds assets.
-   * @param filePath The path of the vocabulary file, relative with assets directory.
-   * @return a list of vocabulary words.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadSingleColumnTextFile(
-      @NonNull Context context, @NonNull String filePath, Charset cs) throws IOException {
-    return loadLabels(context, filePath, cs);
-  }
-
-  /**
-   * Loads vocabulary from an input stream of an opened vocabulary file (which is a single-column
-   * text file). See details for vocabulary files in {@link FileUtil#loadVocabularyFile(Context,
-   * String)}.
-   *
-   * @param inputStream the input stream of an opened vocabulary file.
-   * @return a list of vocabulary words.
-   * @throws IOException if error occurs to open or read the file.
-   */
-  @NonNull
-  public static List loadSingleColumnTextFile(@NonNull InputStream inputStream, Charset cs)
-      throws IOException {
-    return loadLabels(inputStream, cs);
-  }
-
-  /**
-   * Loads a file from the asset folder through memory mapping.
-   *
-   * @param context Application context to access assets.
-   * @param filePath Asset path of the file.
-   * @return the loaded memory mapped file.
-   * @throws IOException if an I/O error occurs when loading the tflite model.
-   */
-  @NonNull
-  public static MappedByteBuffer loadMappedFile(@NonNull Context context, @NonNull String filePath)
-      throws IOException {
-    SupportPreconditions.checkNotNull(context, "Context should not be null.");
-    SupportPreconditions.checkNotNull(filePath, "File path cannot be null.");
-    try (AssetFileDescriptor fileDescriptor = context.getAssets().openFd(filePath);
-    FileInputStream inputStream = new FileInputStream(fileDescriptor.getFileDescriptor())) {
-      FileChannel fileChannel = inputStream.getChannel();
-      long startOffset = fileDescriptor.getStartOffset();
-      long declaredLength = fileDescriptor.getDeclaredLength();
-      return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
-    }
-  }
-
-  /**
-   * Loads a binary file from the asset folder.
-   *
-   * @param context Application context to access assets.
-   * @param filePath Asset path of the file.
-   * @return the byte array for the binary file.
-   * @throws IOException if an I/O error occurs when loading file.
-   */
-  @NonNull
-  public static byte[] loadByteFromFile(@NonNull Context context, @NonNull String filePath)
-      throws IOException {
-    ByteBuffer buffer = loadMappedFile(context, filePath);
-    byte[] byteArray = new byte[buffer.remaining()];
-    buffer.get(byteArray);
-    return byteArray;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Operator.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Operator.java
deleted file mode 100644
index 38dfe8818cb..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Operator.java
+++ /dev/null
@@ -1,31 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-/**
- * The common interface for classes that carries an "apply" method, which converts T to another one.
- * @param  The class which Operator handles.
- */
-public interface Operator {
-
-  /**
-   * Applies an operation on a T object, returning a T object.
-   *
-   * 
Note: The returned object could probably be the same one with given input, and given input
-   * could probably be changed.
-   */
-  T apply(T x);
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Processor.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Processor.java
deleted file mode 100644
index 07d7e2bda43..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/Processor.java
+++ /dev/null
@@ -1,23 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-/**
- * Processes T object with prepared {@link Operator}.
- */
-public interface Processor {
-  T process(T input);
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SequentialProcessor.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SequentialProcessor.java
deleted file mode 100644
index ff0c6406f03..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SequentialProcessor.java
+++ /dev/null
@@ -1,82 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import org.checkerframework.checker.nullness.qual.NonNull;
-
-/**
- * A processor base class that chains a serial of {@link Operator} and executes them.
- *
- * 
Typically, users could use its subclasses, e.g. {@link
- * org.tensorflow.lite.support.image.ImageProcessor} rather than directly use this one.
- *
- * @param  The type that the Operator is handling.
- */
-public class SequentialProcessor implements Processor {
-
-  /** List of operators added to this {@link SequentialProcessor}. */
-  protected final List> operatorList;
-  /**
-   * The {@link Map} between the operator name and the corresponding op indexes in {@code
-   * operatorList}. An operator may be added multiple times into this {@link SequentialProcessor}.
-   */
-  protected final Map> operatorIndex;
-
-  protected SequentialProcessor(Builder builder) {
-    operatorList = builder.operatorList;
-    operatorIndex = Collections.unmodifiableMap(builder.operatorIndex);
-  }
-
-  @Override
-  public T process(T x) {
-    for (Operator op : operatorList) {
-      x = op.apply(x);
-    }
-    return x;
-  }
-
-  /** The inner builder class to build a Sequential Processor. */
-  protected static class Builder {
-
-    private final List> operatorList;
-    private final Map> operatorIndex;
-
-    protected Builder() {
-      operatorList = new ArrayList<>();
-      operatorIndex = new HashMap<>();
-    }
-
-    public Builder add(@NonNull Operator op) {
-      SupportPreconditions.checkNotNull(op, "Adding null Op is illegal.");
-      operatorList.add(op);
-      String operatorName = op.getClass().getName();
-      if (!operatorIndex.containsKey(operatorName)) {
-        operatorIndex.put(operatorName, new ArrayList());
-      }
-      operatorIndex.get(operatorName).add(operatorList.size() - 1);
-      return this;
-    }
-
-    public SequentialProcessor build() {
-      return new SequentialProcessor(this);
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SupportPreconditions.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SupportPreconditions.java
deleted file mode 100644
index 8620e13eec7..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/SupportPreconditions.java
+++ /dev/null
@@ -1,184 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-import org.checkerframework.checker.nullness.qual.Nullable;
-
-/** Static error checking util methods. */
-public final class SupportPreconditions {
-  /**
-   * Ensures that an object reference passed as a parameter to the calling method is not null.
-   *
-   * @param reference an object reference
-   * @return the non-null reference that was validated
-   * @throws NullPointerException if {@code reference} is null
-   */
-  public static  T checkNotNull(T reference) {
-    if (reference == null) {
-      throw new NullPointerException("The object reference is null.");
-    }
-    return reference;
-  }
-
-  /**
-   * Ensures that an object reference passed as a parameter to the calling method is not null.
-   *
-   * @param reference an object reference
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @return the non-null reference that was validated
-   * @throws NullPointerException if {@code reference} is null
-   */
-  public static  T checkNotNull(T reference, @Nullable Object errorMessage) {
-    if (reference == null) {
-      throw new NullPointerException(String.valueOf(errorMessage));
-    }
-    return reference;
-  }
-
-  /**
-   * Ensures that the given String is not empty and not null.
-   *
-   * @param string the String to test
-   * @return the non-null non-empty String that was validated
-   * @throws IllegalArgumentException if {@code string} is null or empty
-   */
-  public static String checkNotEmpty(String string) {
-    if (string == null || string.length() == 0) {
-      throw new IllegalArgumentException("Given String is empty or null.");
-    }
-    return string;
-  }
-
-  /**
-   * Ensures that the given String is not empty and not null.
-   *
-   * @param string the String to test
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @return the non-null non-empty String that was validated
-   * @throws IllegalArgumentException if {@code string} is null or empty
-   */
-  public static String checkNotEmpty(String string, Object errorMessage) {
-    if (string == null || string.length() == 0) {
-      throw new IllegalArgumentException(String.valueOf(errorMessage));
-    }
-    return string;
-  }
-
-  /**
-   * Ensures the truth of an expression involving one or more parameters to the calling method.
-   *
-   * @param expression a boolean expression.
-   * @throws IllegalArgumentException if {@code expression} is false.
-   */
-  public static void checkArgument(boolean expression) {
-    if (!expression) {
-      throw new IllegalArgumentException();
-    }
-  }
-
-  /**
-   * Ensures the truth of an expression involving one or more parameters to the calling method.
-   *
-   * @param expression a boolean expression.
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}.
-   * @throws IllegalArgumentException if {@code expression} is false.
-   */
-  public static void checkArgument(boolean expression, @Nullable Object errorMessage) {
-    if (!expression) {
-      throw new IllegalArgumentException(String.valueOf(errorMessage));
-    }
-  }
-
-  /**
-   * Ensures that {@code index} specifies a valid element in an array, list or string of size
-   * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive.
-   *
-   * @param index a user-supplied index identifying an element of an array, list or string
-   * @param size the size of that array, list or string
-   * @return the value of {@code index}
-   * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size}
-   * @throws IllegalArgumentException if {@code size} is negative
-   */
-  public static int checkElementIndex(int index, int size) {
-    return checkElementIndex(index, size, "index");
-  }
-
-  /**
-   * Ensures that {@code index} specifies a valid element in an array, list or string of size
-   * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive.
-   *
-   * @param index a user-supplied index identifying an element of an array, list or string
-   * @param size the size of that array, list or string
-   * @param desc the text to use to describe this index in an error message
-   * @return the value of {@code index}
-   * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size}
-   * @throws IllegalArgumentException if {@code size} is negative
-   */
-  public static int checkElementIndex(int index, int size, @Nullable String desc) {
-    // Carefully optimized for execution by hotspot (explanatory comment above)
-    if (index < 0 || index >= size) {
-      throw new IndexOutOfBoundsException(badElementIndex(index, size, desc));
-    }
-    return index;
-  }
-
-  /**
-   * Ensures the truth of an expression involving the state of the calling instance, but not
-   * involving any parameters to the calling method.
-   *
-   * @param expression a boolean expression
-   * @throws IllegalStateException if {@code expression} is false
-   * @see Verify#verify Verify.verify()
-   */
-  public static void checkState(boolean expression) {
-    if (!expression) {
-      throw new IllegalStateException();
-    }
-  }
-
-  /**
-   * Ensures the truth of an expression involving the state of the calling instance, but not
-   * involving any parameters to the calling method.
-   *
-   * @param expression a boolean expression
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @throws IllegalStateException if {@code expression} is false
-   * @see Verify#verify Verify.verify()
-   */
-  public static void checkState(boolean expression, @Nullable Object errorMessage) {
-    if (!expression) {
-      throw new IllegalStateException(String.valueOf(errorMessage));
-    }
-  }
-
-  private static String badElementIndex(int index, int size, @Nullable String desc) {
-    if (index < 0) {
-      return String.format("%s (%s) must not be negative", desc, index);
-    } else if (size < 0) {
-      throw new IllegalArgumentException("negative size: " + size);
-    } else { // index >= size
-      return String.format("%s (%s) must be less than size (%s)", desc, index, size);
-    }
-  }
-
-  private SupportPreconditions() {
-    throw new AssertionError("SupportPreconditions is Uninstantiable.");
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorOperator.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorOperator.java
deleted file mode 100644
index d1b7021df25..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorOperator.java
+++ /dev/null
@@ -1,27 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Applies some operation on TensorBuffers.
- */
-public interface TensorOperator extends Operator {
-  /** @see Operator#apply(Object) . */
-  @Override
-  TensorBuffer apply(TensorBuffer input);
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorProcessor.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorProcessor.java
deleted file mode 100644
index 31531b2eb6a..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/TensorProcessor.java
+++ /dev/null
@@ -1,68 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common;
-
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * TensorProcessor is a helper class for preprocessing and postprocessing tensors. It could
- * transform a {@link TensorBuffer} to another by executing a chain of {@link TensorOperator}.
- *
- * 
Example Usage:
- *
- * 
- *   TensorProcessor processor = new TensorProcessor.Builder().add(new NormalizeOp(1, 2)).build();
- *   TensorBuffer anotherTensorBuffer = processor.process(tensorBuffer);
- * 

- *
- * @see TensorProcessor.Builder to build a {@link TensorProcessor} instance.
- * @see TensorProcessor#process(TensorBuffer) to apply the processor on a {@link TensorBuffer}.
- */
-public class TensorProcessor extends SequentialProcessor {
-  private TensorProcessor(Builder builder) {
-    super(builder);
-  }
-
-  /** The Builder to create an {@link TensorProcessor}, which could be executed later. */
-  public static class Builder extends SequentialProcessor.Builder {
-
-    /**
-     * Creates a Builder to build {@link TensorProcessor}.
-     *
-     * @see #add(TensorOperator) to add an Op.
-     * @see #build() to complete the building process and get a built Processor.
-     */
-    public Builder() {
-      super();
-    }
-
-    /**
-     * Adds an {@link TensorOperator} into the Operator chain.
-     *
-     * @param op the Operator instance to be executed then.
-     */
-    public TensorProcessor.Builder add(TensorOperator op) {
-      super.add(op);
-      return this;
-    }
-
-    /** Completes the building process and gets the {@link TensorProcessor} instance. */
-    @Override
-    public TensorProcessor build() {
-      return new TensorProcessor(this);
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/CastOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/CastOp.java
deleted file mode 100644
index 3355b185655..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/CastOp.java
+++ /dev/null
@@ -1,55 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common.ops;
-
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/** Casts a {@link TensorBuffer} to a specified data type. */
-public class CastOp implements TensorOperator {
-
-  private final DataType destinationType;
-
-  /**
-   * Constructs a CastOp.
-   *
-   * 
Note: For only converting type for a certain {@link TensorBuffer} on-the-fly rather than in
-   * a processor, please directly use {@link TensorBuffer#createFrom(TensorBuffer, DataType)}.
-   *
-   * 
When this Op is executed, if the original {@link TensorBuffer} is already in {@code
-   * destinationType}, the original buffer will be directly returned.
-   *
-   * @param destinationType: The type of the casted {@link TensorBuffer}.
-   * @throws IllegalArgumentException if {@code destinationType} is neither {@link DataType#UINT8}
-   * nor {@link DataType#FLOAT32}.
-   */
-  public CastOp(DataType destinationType) {
-    SupportPreconditions.checkArgument(
-        destinationType == DataType.UINT8 || destinationType == DataType.FLOAT32,
-        "Destination type " + destinationType + " is not supported.");
-    this.destinationType = destinationType;
-  }
-
-  @Override
-  public TensorBuffer apply(TensorBuffer input) {
-    if (input.getDataType() == destinationType) {
-      return input;
-    }
-    return TensorBuffer.createFrom(input, destinationType);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/DequantizeOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/DequantizeOp.java
deleted file mode 100644
index 1881747870b..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/DequantizeOp.java
+++ /dev/null
@@ -1,40 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common.ops;
-
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Dequantizes a {@link TensorBuffer} with given {@code zeroPoint} and {@code scale}.
- *
- * 
Note: The data type of output tensor is always {@code FLOAT32} except when the DequantizeOp is
- * created effectively as an identity Op such as setting {@code zeroPoint} to 0 and {@code scale} to
- * 1 (in this case, the output tensor is the same instance as input).
- *
- * 
If both {@code zeroPoint} and {@code scale} are 0, the {@link DequantizeOp} will be bypassed,
- * which is equivalent to setting {@code zeroPoint} to 0 and {@code scale} to 1. This can be useful
- * when passing in the quantization parameters that are extracted directly from the TFLite model
- * flatbuffer. If the tensor is not quantized, both {@code zeroPoint} and {@code scale} will be read
- * as 0.
- */
-public class DequantizeOp extends NormalizeOp implements TensorOperator {
-
-  public DequantizeOp(float zeroPoint, float scale) {
-    // Quantization: f = (q - z) * s
-    super(zeroPoint, 1 / scale);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java
deleted file mode 100644
index 8ac57eed286..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/NormalizeOp.java
+++ /dev/null
@@ -1,160 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common.ops;
-
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-import org.tensorflow.lite.support.tensorbuffer.TensorBufferFloat;
-
-/**
- * Normalizes a {@link TensorBuffer} with given mean and stddev: output = (input - mean) / stddev.
- */
-public class NormalizeOp implements TensorOperator {
-
-  // mean.length should always be equal to stddev.length and always >= 1.
-  private final float[] mean;
-  private final float[] stddev;
-  private final int numChannels;
-  private final boolean isIdentityOp;
-
-  /**
-   * Initializes a NormalizeOp. When being called, it creates a new {@link TensorBuffer}, which
-   * satisfies:
-   *
-   * 
-   *   output = (input - mean) / stddev
-   * 

-   *
-   * 
In the following two cases, reset {@code mean} to 0 and {@code stddev} to 1 to bypass the
-   * normalization. 

-   * 1. Both {@code mean} and {code stddev} are 0. 

-   * 2. {@code mean} is 0 and {stddev} is Infinity.
-   *
-   * 
Note: If {@code mean} is set to 0 and {@code stddev} is set to 1, no computation will
-   * happen, and original input will be directly returned in execution.
-   *
-   * 
Note: The returned {@link TensorBuffer} is always a {@link DataType#FLOAT32} tensor at
-   * present, except that the input is a {@link DataType#UINT8} tensor, {@code mean} is set to 0 and
-   * {@code stddev} is set to 1.
-   *
-   * @param mean the mean value to be subtracted first.
-   * @param stddev the standard deviation value to divide then.
-   * @throws IllegalArgumentException if {@code stddev} is zero.
-   */
-  public NormalizeOp(float mean, float stddev) {
-    // Make exceptions to the cases that
-    // 1. Both mean and stddev are 0.0f. This may happen when reading the normalization parameters
-    // from a tensor which does not have the values populated in the metadata. The same situation
-    // may also happen to the quantization parameters.
-    // 2. mean is 0.0f and stddev is Infinity. This may happen when reading the quantization
-    // parameters from a tensor which does not have the values populated in the metadata, and then
-    // passing the parameters into the DequantizeOp.
-    // Bypass both of the two cases, by reseting stddev to 1.0f.
-    if (mean == 0.0f && (stddev == 0.0f || Float.isInfinite(stddev))) {
-      stddev = 1.0f;
-    }
-
-    SupportPreconditions.checkArgument(stddev != 0.0f, "Stddev cannot be zero.");
-    boolean meansIsZeroAndDevsIs1 = false;
-    if (mean == 0.0f && stddev == 1.0f) {
-      meansIsZeroAndDevsIs1 = true;
-    }
-
-    this.isIdentityOp = meansIsZeroAndDevsIs1;
-    this.mean = new float[] {mean};
-    this.stddev = new float[] {stddev};
-    this.numChannels = 1;
-  }
-
-  /**
-   * Initializes a NormalizeOp. When being called, it creates a new {@link TensorBuffer}, which
-   * satisfies:
-   *
-   * 
-   *   // Pseudo code. [...][i] means a certain element whose channel id is i.
-   *   output[...][i] = (input[...][i] - mean[i]) / stddev[i]
-   * 

-   *
-   * 
Note: If all values in {@code mean} are set to 0 and all {@code stddev} are set to 1, no
-   * computation will happen, and original input will be directly returned in execution.
-   *
-   * 
Note: The returned {@link TensorBuffer} is always a {@link DataType#FLOAT32} tensor at
-   * present, except that the input is a {@link DataType#UINT8} tensor, all {@code mean} are set to
-   * 0 and all {@code stddev} are set to 1.
-   *
-   * @param mean the mean values to be subtracted first for each channel.
-   * @param stddev the standard deviation values to divide then for each channel.
-   * @throws IllegalArgumentException if any {@code stddev} is zero, or {@code mean} has different
-   *     number of elements with {@code stddev}, or any of them is empty.
-   */
-  public NormalizeOp(@NonNull float[] mean, @NonNull float[] stddev) {
-    SupportPreconditions.checkNotNull(mean, "Mean cannot be null");
-    SupportPreconditions.checkNotNull(stddev, "Stddev cannot be null");
-    SupportPreconditions.checkArgument(
-        mean.length == stddev.length,
-        "Per channel normalization requires same number of means and stddevs");
-    SupportPreconditions.checkArgument(mean.length > 0, "Means and stddevs are empty.");
-    this.mean = mean.clone();
-    this.stddev = stddev.clone();
-    boolean allMeansAreZeroAndAllDevsAre1 = true;
-    this.numChannels = mean.length;
-    for (int i = 0; i < numChannels; i++) {
-      SupportPreconditions.checkArgument(this.stddev[i] != 0, "Stddev cannot be zero.");
-      if (this.stddev[i] != 1 || this.mean[i] != 0) {
-        allMeansAreZeroAndAllDevsAre1 = false;
-      }
-    }
-    this.isIdentityOp = allMeansAreZeroAndAllDevsAre1;
-  }
-
-  /**
-   * Applies the defined normalization on given tensor and returns the result.
-   *
-   * 
Note: {@code input} is possibly the same instance with the output.
-   *
-   * @param input input tensor. It may be the same instance with the output.
-   * @return output tensor.
-   */
-  @Override
-  @NonNull
-  public TensorBuffer apply(@NonNull TensorBuffer input) {
-    if (isIdentityOp) {
-      return input;
-    }
-    int[] shape = input.getShape();
-    SupportPreconditions.checkArgument(
-        numChannels == 1 || (shape.length != 0 && shape[shape.length - 1] == numChannels),
-        "Number of means (stddevs) is not same with number of channels (size of last axis).");
-    // TODO(136750944): Eliminate the array copy here.
-    float[] values = input.getFloatArray();
-    int j = 0;
-    for (int i = 0; i < values.length; i++) {
-      values[i] = (values[i] - mean[j]) / stddev[j];
-      j = (j + 1) % numChannels;
-    }
-    TensorBuffer output;
-    if (input.isDynamic()) {
-      output = TensorBufferFloat.createDynamic(DataType.FLOAT32);
-    } else {
-      output = TensorBufferFloat.createFixedSize(shape, DataType.FLOAT32);
-    }
-    output.loadArray(values, shape);
-    return output;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/QuantizeOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/QuantizeOp.java
deleted file mode 100644
index 8b3e82aee13..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/common/ops/QuantizeOp.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.common.ops;
-
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Quantizes a {@link TensorBuffer} with given {@code zeroPoint} and {@code scale}.
- *
- * 
Note: {@link QuantizeOp} does not cast output to UINT8, but only performs the quantization
- * math on top of input. The data type of output tensor is always {@code FLOAT32} except that the Op
- * is effectively an identity Op (in this case, the output tensor is the same instance as the
- * input). To connect with quantized model, a {@link CastOp} is probably needed.
- *
- * 
If both {@code zeroPoint} and {@code scale} are 0, the {@link QuantizeOp} will be bypassed,
- * which is equivalent to setting {@code zeroPoint} to 0 and {@code scale} to 1. This can be useful
- * when passing in the quantization parameters that are extracted directly from the TFLite model
- * flatbuffer. If the tensor is not quantized, both {@code zeroPoint} and {@code scale} will be read
- * as 0.
- */
-public class QuantizeOp extends NormalizeOp implements TensorOperator {
-
-  public QuantizeOp(float zeroPoint, float scale) {
-    // Quantization: f = (q - z) * s, i.e. q = f / s + z = (f - (-z * s)) / s
-    super(-zeroPoint * scale, scale);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/BoundingBoxUtil.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/BoundingBoxUtil.java
deleted file mode 100644
index 30f562063f3..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/BoundingBoxUtil.java
+++ /dev/null
@@ -1,202 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image;
-
-import static org.tensorflow.lite.support.common.SupportPreconditions.checkArgument;
-
-import android.graphics.RectF;
-import java.nio.ByteBuffer;
-import java.nio.FloatBuffer;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Helper class for converting values that represents bounding boxes into rectangles.
- *
- * 
The class provides a static function to create bounding boxes as {@link RectF} from different
- * types of configurations.
- *
- * 
Generally, a bounding box could be represented by 4 float values, but the values could be
- * interpreted in many ways. We now support 3 {@link Type} of configurations, and the order of
- * elements in each type is configurable as well.
- */
-public final class BoundingBoxUtil {
-
-  /** Denotes how a bounding box is represented. */
-  public enum Type {
-    /**
-     * Represents the bounding box by using the combination of boundaries, {left, top, right,
-     * bottom}. The default order is {left, top, right, bottom}. Other orders can be indicated by an
-     * index array.
-     */
-    BOUNDARIES,
-    /**
-     * Represents the bounding box by using the upper_left corner, width and height. The default
-     * order is {upper_left_x, upper_left_y, width, height}. Other orders can be indicated by an
-     * index array.
-     */
-    UPPER_LEFT,
-    /**
-     * Represents the bounding box by using the center of the box, width and height. The default
-     * order is {center_x, center_y, width, height}. Other orders can be indicated by an index
-     * array.
-     */
-    CENTER,
-  }
-
-  /** Denotes if the coordinates are actual pixels or relative ratios. */
-  public enum CoordinateType {
-    /** The coordinates are relative ratios in range [0, 1]. */
-    RATIO,
-    /** The coordinates are actual pixel values. */
-    PIXEL
-  }
-
-  /**
-   * Creates a list of bounding boxes from a {@link TensorBuffer} which represents bounding boxes.
-   *
-   * @param tensor holds the data representing some boxes.
-   * @param valueIndex denotes the order of the elements defined in each bounding box type. An empty
-   *     index array represent the default order of each bounding box type. For example, to denote
-   *     the default order of BOUNDARIES, {left, top, right, bottom}, the index should be {0, 1, 2,
-   *     3}. To denote the order {left, right, top, bottom}, the order should be {0, 2, 1, 3}.
-   *     
The index array can be applied to all bounding box types to adjust the order of their
-   *     corresponding underlying elements.
-   * @param boundingBoxAxis specifies the index of the dimension that represents bounding box. The
-   *     size of that dimension is required to be 4. Index here starts from 0. For example, if the
-   *     tensor has shape 4x10, the axis for bounding boxes is likely to be 0. For shape 10x4, the
-   *     axis is likely to be 1 (or -1, equivalently).
-   * @param type defines how values should be converted into boxes. See {@link Type}
-   * @param coordinateType defines how values are interpreted to coordinates. See {@link
-   *     CoordinateType}
-   * @param height the height of the image which the boxes belong to. Only has effects when {@code
-   *     coordinateType} is {@link CoordinateType#RATIO}
-   * @param width the width of the image which the boxes belong to. Only has effects when {@code
-   *     coordinateType} is {@link CoordinateType#RATIO}
-   * @return A list of bounding boxes that the {@code tensor} represents. All dimensions except
-   *     {@code boundingBoxAxis} will be collapsed with order kept. For example, given {@code
-   *     tensor} with shape {1, 4, 10, 2} and {@code boundingBoxAxis = 1}, The result will be a list
-   *     of 20 bounding boxes.
-   * @throws IllegalArgumentException if size of bounding box dimension (set by {@code
-   *     boundingBoxAxis}) is not 4.
-   * @throws IllegalArgumentException if {@code boundingBoxAxis} is not in {@code (-(D+1), D)} where
-   *     {@code D} is the number of dimensions of the {@code tensor}.
-   * @throws IllegalArgumentException if {@code tensor} has data type other than {@link
-   *     DataType#FLOAT32}.
-   */
-  public static List convert(
-      TensorBuffer tensor,
-      int[] valueIndex,
-      int boundingBoxAxis,
-      Type type,
-      CoordinateType coordinateType,
-      int height,
-      int width) {
-    int[] shape = tensor.getShape();
-    checkArgument(
-        boundingBoxAxis >= -shape.length && boundingBoxAxis < shape.length,
-        String.format(
-            "Axis %d is not in range (-(D+1), D), where D is the number of dimensions of input"
-                + " tensor (shape=%s)",
-            boundingBoxAxis, Arrays.toString(shape)));
-    if (boundingBoxAxis < 0) {
-      boundingBoxAxis = shape.length + boundingBoxAxis;
-    }
-    checkArgument(
-        shape[boundingBoxAxis] == 4,
-        String.format(
-            "Size of bounding box dimension %d is not 4. Got %d in shape %s",
-            boundingBoxAxis, shape[boundingBoxAxis], Arrays.toString(shape)));
-    checkArgument(
-        valueIndex.length == 4,
-        String.format(
-            "Bounding box index array length %d is not 4. Got index array %s",
-            valueIndex.length, Arrays.toString(valueIndex)));
-    checkArgument(
-        tensor.getDataType() == DataType.FLOAT32,
-        "Bounding Boxes only create from FLOAT32 buffers. Got: " + tensor.getDataType().name());
-    List boundingBoxList = new ArrayList<>();
-    // Collapse dimensions to {a, 4, b}. So each bounding box could be represent as (i, j), and its
-    // four values are (i, k, j), where 0 <= k < 4. We can compute the 4 flattened index by
-    // i * 4b + k * b + j.
-    int a = 1;
-    for (int i = 0; i < boundingBoxAxis; i++) {
-      a *= shape[i];
-    }
-    int b = 1;
-    for (int i = boundingBoxAxis + 1; i < shape.length; i++) {
-      b *= shape[i];
-    }
-    float[] values = new float[4];
-    ByteBuffer byteBuffer = tensor.getBuffer();
-    byteBuffer.rewind();
-    FloatBuffer floatBuffer = byteBuffer.asFloatBuffer();
-    for (int i = 0; i < a; i++) {
-      for (int j = 0; j < b; j++) {
-        for (int k = 0; k < 4; k++) {
-          values[k] = floatBuffer.get((i * 4 + k) * b + j);
-        }
-        boundingBoxList.add(
-            convertOneBoundingBox(values, valueIndex, type, coordinateType, height, width));
-      }
-    }
-    byteBuffer.rewind();
-    return boundingBoxList;
-  }
-
-  private static RectF convertOneBoundingBox(
-      float[] values,
-      int[] valueIndex,
-      Type type,
-      CoordinateType coordinateType,
-      int height,
-      int width) {
-    float[] orderedValues = new float[4];
-    for (int i = 0; i < 4; i++) {
-      orderedValues[i] = values[valueIndex[i]];
-    }
-    return convertOneBoundingBox(orderedValues, type, coordinateType, height, width);
-  }
-
-  private static RectF convertOneBoundingBox(
-      float[] values, Type type, CoordinateType coordinateType, int height, int width) {
-    switch (type) {
-      case BOUNDARIES:
-        return convertFromBoundaries(values, coordinateType, height, width);
-      case UPPER_LEFT:
-      case CENTER:
-        // TODO(b/150824448): convertFrom{UpperLeft, Center}
-        throw new IllegalArgumentException("BoundingBox.Type " + type + " is not yet supported.");
-    }
-    throw new IllegalArgumentException("Cannot recognize BoundingBox.Type " + type);
-  }
-
-  private static RectF convertFromBoundaries(
-      float[] values, CoordinateType coordinateType, int height, int width) {
-    if (coordinateType == CoordinateType.RATIO) {
-      return new RectF(
-          values[0] * width, values[1] * height, values[2] * width, values[3] * height);
-    } else {
-      return new RectF(values[0], values[1], values[2], values[3]);
-    }
-  }
-
-  // Private constructor to prevent initialization.
-  private BoundingBoxUtil() {}
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageConversions.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageConversions.java
deleted file mode 100644
index b2b7a339a75..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageConversions.java
+++ /dev/null
@@ -1,108 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image;
-
-import android.graphics.Bitmap;
-import android.graphics.Color;
-import java.util.Arrays;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Implements some stateless image conversion methods.
- *
- * This class is an internal helper for {@link org.tensorflow.lite.support.image}.
- */
-class ImageConversions {
-
-  /**
-   * Converts an Image in a TensorBuffer to a Bitmap, whose memory is already allocated.
-   *
-   * 
Notice: We only support ARGB_8888 at this point.
-   *
-   * @param buffer The TensorBuffer object representing the image. It should be an UInt8 buffer with
-   *     3 dimensions: width, height, channel. Size of each dimension should be positive and the
-   *     size of channels should be 3 (representing R, G, B). An optional 4th dimension "batch" is
-   *     acceptable, and dimensions look like: batch, width, height, channel. In this case, size of
-   *     batches should be 1.
-   * @param bitmap The destination of the conversion. Needs to be created in advance, needs to be
-   *     mutable, and needs to have the same width and height with the buffer.
-   * @throws IllegalArgumentException 1) if the {@code buffer} is not uint8 (e.g. a float buffer),
-   *     or has an invalid shape. 2) if the {@code bitmap} is not mutable. 3) if the {@code bitmap}
-   *     has different height or width with the buffer.
-   */
-  static void convertTensorBufferToBitmap(TensorBuffer buffer, Bitmap bitmap) {
-    if (buffer.getDataType() != DataType.UINT8) {
-      // We will add support to FLOAT format conversion in the future, as it may need other configs.
-      throw new UnsupportedOperationException(
-          String.format(
-              "Converting TensorBuffer of type %s to ARGB_8888 Bitmap is not supported yet.",
-              buffer.getDataType()));
-    }
-    int[] shape = buffer.getShape();
-    TensorImage.checkImageTensorShape(shape);
-    int h = shape[shape.length - 3];
-    int w = shape[shape.length - 2];
-    if (bitmap.getWidth() != w || bitmap.getHeight() != h) {
-      throw new IllegalArgumentException(String.format(
-          "Given bitmap has different width or height %s with the expected ones %s.",
-          Arrays.toString(new int[]{bitmap.getWidth(), bitmap.getHeight()}),
-          Arrays.toString(new int[]{w, h})));
-    }
-    if (!bitmap.isMutable()) {
-      throw new IllegalArgumentException("Given bitmap is not mutable");
-    }
-    // TODO(b/138904567): Find a way to avoid creating multiple intermediate buffers every time.
-    int[] intValues = new int[w * h];
-    int[] rgbValues = buffer.getIntArray();
-    for (int i = 0, j = 0; i < intValues.length; i++) {
-      int r = rgbValues[j++];
-      int g = rgbValues[j++];
-      int b = rgbValues[j++];
-      intValues[i] = Color.rgb(r, g, b);
-    }
-    bitmap.setPixels(intValues, 0, w, 0, 0, w, h);
-  }
-
-  /**
-   * Converts an Image in a Bitmap to a TensorBuffer (3D Tensor: Width-Height-Channel) whose memory
-   * is already allocated, or could be dynamically allocated.
-   *
-   * @param bitmap The Bitmap object representing the image. Currently we only support ARGB_8888
-   * config.
-   * @param buffer The destination of the conversion. Needs to be created in advance. If it's
-   * fixed-size, its flat size should be w*h*3.
-   * @throws IllegalArgumentException if the buffer is fixed-size, but the size doesn't match.
-   */
-  static void convertBitmapToTensorBuffer(Bitmap bitmap, TensorBuffer buffer) {
-    int w = bitmap.getWidth();
-    int h = bitmap.getHeight();
-    int[] intValues = new int[w * h];
-    bitmap.getPixels(intValues, 0, w, 0, 0, w, h);
-    // TODO(b/138904567): Find a way to avoid creating multiple intermediate buffers every time.
-    int[] rgbValues = new int[w * h * 3];
-    for (int i = 0, j = 0; i < intValues.length; i++) {
-      rgbValues[j++] = ((intValues[i] >> 16) & 0xFF);
-      rgbValues[j++] = ((intValues[i] >> 8) & 0xFF);
-      rgbValues[j++] = (intValues[i] & 0xFF);
-    }
-    int[] shape = new int[] {h, w, 3};
-    buffer.loadArray(rgbValues, shape);
-  }
-
-  // Hide the constructor as the class is static.
-  private ImageConversions() {}
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageOperator.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageOperator.java
deleted file mode 100644
index 1e546634e90..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageOperator.java
+++ /dev/null
@@ -1,43 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image;
-
-import android.graphics.PointF;
-import org.tensorflow.lite.support.common.Operator;
-
-/** Operates a TensorImage object. Used in ImageProcessor. */
-public interface ImageOperator extends Operator {
-  /** @see org.tensorflow.lite.support.common.Operator#apply(java.lang.Object) */
-  @Override
-  TensorImage apply(TensorImage image);
-
-  /** Computes the width of the expected output image when input image size is given. */
-  int getOutputImageWidth(int inputImageHeight, int inputImageWidth);
-
-  /** Computes the height of the expected output image when input image size is given. */
-  int getOutputImageHeight(int inputImageHeight, int inputImageWidth);
-
-  /**
-   * Transforms a point from coordinates system of the result image back to the one of the input
-   * image.
-   *
-   * @param point the point from the result coordinates system.
-   * @param inputImageHeight the height of input image.
-   * @param inputImageWidth the width of input image.
-   * @return the point with the coordinates from the coordinates system of the input image.
-   */
-  PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth);
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageProcessor.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageProcessor.java
deleted file mode 100644
index e1ef1309bbe..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ImageProcessor.java
+++ /dev/null
@@ -1,198 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image;
-
-import android.graphics.PointF;
-import android.graphics.RectF;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.ListIterator;
-import org.tensorflow.lite.support.common.Operator;
-import org.tensorflow.lite.support.common.SequentialProcessor;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.image.ops.Rot90Op;
-import org.tensorflow.lite.support.image.ops.TensorOperatorWrapper;
-
-/**
- * ImageProcessor is a helper class for preprocessing and postprocessing {@link TensorImage}. It
- * could transform a {@link TensorImage} to another by executing a chain of {@link ImageOperator}.
- *
- * 
Example Usage:
- *
- * 
- *   ImageProcessor processor = new ImageProcessor.Builder()
- *       .add(new ResizeOp(224, 224, ResizeMethod.NEAREST_NEIGHBOR)
- *       .add(new Rot90Op())
- *       .add(new NormalizeOp(127.5f, 127.5f))
- *       .build();
- *   TensorImage anotherTensorImage = processor.process(tensorImage);
- * 

- *
- * 
WARNING: Instances of an {@code ImageProcessor} are not thread-safe with {@link
- * #updateNumberOfRotations}. Updating the number of rotations and then processing images (using
- * {@link #process}) must be protected from concurrent access. It is recommended to create separate
- * {@code ImageProcessor} instances for each thread. If multiple threads access a {@code
- * ImageProcessor} concurrently, it must be synchronized externally.
- *
- * @see ImageProcessor.Builder to build a {@link ImageProcessor} instance
- * @see ImageProcessor#process(TensorImage) to apply the processor on a {@link TensorImage}
- */
-public class ImageProcessor extends SequentialProcessor {
-  private ImageProcessor(Builder builder) {
-    super(builder);
-  }
-
-  /**
-   * Transforms a point from coordinates system of the result image back to the one of the input
-   * image.
-   *
-   * @param point the point from the result coordinates system.
-   * @param inputImageHeight the height of input image.
-   * @param inputImageWidth the width of input image.
-   * @return the point with the coordinates from the coordinates system of the input image.
-   */
-  public PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth) {
-    List widths = new ArrayList<>();
-    List heights = new ArrayList<>();
-    int currentWidth = inputImageWidth;
-    int currentHeight = inputImageHeight;
-    for (Operator op : operatorList) {
-      widths.add(currentWidth);
-      heights.add(currentHeight);
-      ImageOperator imageOperator = (ImageOperator) op;
-      int newHeight = imageOperator.getOutputImageHeight(currentHeight, currentWidth);
-      int newWidth = imageOperator.getOutputImageWidth(currentHeight, currentWidth);
-      currentHeight = newHeight;
-      currentWidth = newWidth;
-    }
-    ListIterator> opIterator = operatorList.listIterator(operatorList.size());
-    ListIterator widthIterator = widths.listIterator(widths.size());
-    ListIterator heightIterator = heights.listIterator(heights.size());
-    while (opIterator.hasPrevious()) {
-      ImageOperator imageOperator = (ImageOperator) opIterator.previous();
-      int height = heightIterator.previous();
-      int width = widthIterator.previous();
-      point = imageOperator.inverseTransform(point, height, width);
-    }
-    return point;
-  }
-
-  /**
-   * Transforms a rectangle from coordinates system of the result image back to the one of the input
-   * image.
-   *
-   * @param rect the rectangle from the result coordinates system.
-   * @param inputImageHeight the height of input image.
-   * @param inputImageWidth the width of input image.
-   * @return the rectangle with the coordinates from the coordinates system of the input image.
-   */
-  public RectF inverseTransform(RectF rect, int inputImageHeight, int inputImageWidth) {
-    // when rotation is involved, corner order may change - top left changes to bottom right, .etc
-    PointF p1 =
-        inverseTransform(new PointF(rect.left, rect.top), inputImageHeight, inputImageWidth);
-    PointF p2 =
-        inverseTransform(new PointF(rect.right, rect.bottom), inputImageHeight, inputImageWidth);
-    return new RectF(
-        Math.min(p1.x, p2.x), Math.min(p1.y, p2.y), Math.max(p1.x, p2.x), Math.max(p1.y, p2.y));
-  }
-
-  /**
-   * The Builder to create an ImageProcessor, which could be executed later.
-   *
-   * @see #add(TensorOperator) to add a general TensorOperator
-   * @see #add(ImageOperator) to add an ImageOperator
-   * @see #build() complete the building process and get a built Processor
-   */
-  public static class Builder extends SequentialProcessor.Builder {
-    public Builder() {
-      super();
-    }
-
-    /**
-     * Adds an {@link ImageOperator} into the Operator chain.
-     *
-     * @param op the Operator instance to be executed then
-     */
-    public Builder add(ImageOperator op) {
-      super.add(op);
-      return this;
-    }
-
-    /**
-     * Adds a {@link TensorOperator} into the Operator chain. In execution, the processor calls
-     * {@link TensorImage#getTensorBuffer()} to transform the {@link TensorImage} by transforming
-     * the underlying {@link org.tensorflow.lite.support.tensorbuffer.TensorBuffer}.
-     *
-     * @param op the Operator instance to be executed then
-     */
-    public Builder add(TensorOperator op) {
-      return add(new TensorOperatorWrapper(op));
-    }
-
-    /** Completes the building process and gets the {@link ImageProcessor} instance. */
-    @Override
-    public ImageProcessor build() {
-      return new ImageProcessor(this);
-    }
-  }
-
-  /**
-   * Updates the number of rotations for the first {@link Rot90Op} in this {@link ImageProcessor}.
-   *
-   * 
WARNING:this method is not thread-safe. Updating the number of rotations and
-   * then processing images (using {@link #process}) must be protected from concurrent access with
-   * additional synchronization.
-   *
-   * @param k the number of rotations
-   * @throws IllegalStateException if {@link Rot90Op} has not been added to this {@link
-   *     ImageProcessor}
-   */
-  public void updateNumberOfRotations(int k) {
-    updateNumberOfRotations(k, /*occurrence=*/ 0);
-  }
-
-  /**
-   * Updates the number of rotations for the {@link Rot90Op} specified by {@code occurrence} in this
-   * {@link ImageProcessor}.
-   *
-   * 
WARNING:this method is not thread-safe. Updating the number of rotations and
-   * then processing images (using {@link #process}) must be protected from concurrent access with
-   * additional synchronization.
-   *
-   * @param k the number of rotations
-   * @param occurrence the index of perticular {@link Rot90Op} in this {@link ImageProcessor}. For
-   *     example, if the second {@link Rot90Op} needs to be updated, {@code occurrence} should be
-   *     set to 1.
-   * @throws IndexOutOfBoundsException if {@code occurrence} is negative or is not less than the
-   *     number of {@link Rot90Op} in this {@link ImageProcessor}
-   * @throws IllegalStateException if {@link Rot90Op} has not been added to this {@link
-   *     ImageProcessor}
-   */
-  public synchronized void updateNumberOfRotations(int k, int occurrence) {
-    SupportPreconditions.checkState(
-        operatorIndex.containsKey(Rot90Op.class.getName()),
-        "The Rot90Op has not been added to the ImageProcessor.");
-
-    List indexes = operatorIndex.get(Rot90Op.class.getName());
-    SupportPreconditions.checkElementIndex(occurrence, indexes.size(), "occurrence");
-
-    // The index of the Rot90Op to be replaced in operatorList.
-    int index = indexes.get(occurrence);
-    Rot90Op newRot = new Rot90Op(k);
-    operatorList.set(index, newRot);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/TensorImage.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/TensorImage.java
deleted file mode 100644
index bced23e6f67..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/TensorImage.java
+++ /dev/null
@@ -1,381 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image;
-
-import android.graphics.Bitmap;
-import android.graphics.Bitmap.Config;
-import java.nio.ByteBuffer;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * TensorImage is the wrapper class for Image object. When using image processing utils in
- * TFLite.support library, it's common to convert image objects in variant types to TensorImage at
- * first.
- *
- * 
At present, only RGB images are supported, and the A channel is always ignored.
- *
- * 
Details of data storage: a {@link TensorImage} object may have 2 potential sources of truth: a
- * {@link Bitmap} or a {@link TensorBuffer}. {@link TensorImage} maintains the state and only
- * convert one to the other when needed.
- *
- * 
IMPORTANT: The container doesn't own its data. Callers should not modify data objects those
- * are passed to {@link ImageContainer#set(Bitmap)} or {@link ImageContainer#set(TensorBuffer)}.
- *
- * 
IMPORTANT: All methods are not proved thread-safe.
- *
- * @see ImageProcessor which is often used for transforming a {@link TensorImage}.
- */
-// TODO(b/138906681): Support basic Image properties (ColorType, DataType)
-// TODO(b/138907116): Support loading images from TensorBuffer with properties.
-// TODO(b/138905544): Support directly loading RGBBytes, YUVBytes and other types if necessary.
-public class TensorImage {
-
-  private final ImageContainer container;
-
-  /**
-   * Initialize a TensorImage object.
-   *
-   * Note: The data type of this TensorImage is UINT8, which means it could naturally accept Bitmaps
-   * whose pixel value range is [0, 255]. However, any image with float value pixels will not be
-   * loaded correctly. In those cases, please use {@link TensorImage(DataType)}.
-   */
-  public TensorImage() {
-    this(DataType.UINT8);
-  }
-
-  /**
-   * Initializes a TensorImage object with data type specified.
-   *
-   * 
Note: The shape of a TensorImage is not fixed. It is determined when {@code load} methods
-   * called, and could be change later.
-   *
-   * @param dataType the expected internal data type of underlying tensor. The type is always fixed
-   *     during the lifetime of the {@link TensorImage}. To convert the data type, use {@link
-   *     TensorImage#createFrom(TensorImage, DataType)} to create a copy and convert data type at
-   *     the same time.
-   * @throws IllegalArgumentException if {@code dataType} is neither {@link DataType#UINT8} nor
-   *     {@link DataType#FLOAT32}.
-   */
-  public TensorImage(DataType dataType) {
-    SupportPreconditions.checkArgument(
-        dataType == DataType.UINT8 || dataType == DataType.FLOAT32,
-        "Illegal data type for TensorImage: Only FLOAT32 and UINT8 are accepted");
-    container = new ImageContainer(dataType);
-  }
-
-  /**
-   * Initializes a {@link TensorImage} object with a {@link Bitmap}.
-   *
-   * @see TensorImage#load(Bitmap) for reusing the object when it's expensive to create objects
-   *     frequently, because every call of {@code fromBitmap} creates a new {@link TensorImage}.
-   */
-  public static TensorImage fromBitmap(Bitmap bitmap) {
-    TensorImage image = new TensorImage();
-    image.load(bitmap);
-    return image;
-  }
-
-  /**
-   * Creates a deep-copy of a given {@link TensorImage} and converts internal tensor data type.
-   *
-   * 
If the given {@code dataType} is different with {@code src.getDataType()}, an implicit data
-   * conversion will be applied. Converting data from {@link DataType#FLOAT32} to {@link
-   * DataType#UINT8} may involve default float->int conversion and value clamping, because {@link
-   * DataType#UINT8} stores value from 0 to 255 (inclusively).
-   *
-   * @param src the TensorImage to copy from.
-   * @param dataType the expected data type of newly created {@link TensorImage}.
-   * @return a TensorImage whose data is copied from {@code src} and data type is {@code dataType}.
-   */
-  @NonNull
-  public static TensorImage createFrom(@NonNull TensorImage src, DataType dataType) {
-    TensorImage dst = new TensorImage(dataType);
-    if (src.container.isBufferUpdated) {
-      dst.container.set(TensorBuffer.createFrom(src.getTensorBuffer(), dataType));
-    } else if (src.container.isBitmapUpdated) {
-      Bitmap srcBitmap = src.getBitmap();
-      dst.container.set(srcBitmap.copy(srcBitmap.getConfig(), srcBitmap.isMutable()));
-    }
-    return dst;
-  }
-
-  /**
-   * Loads a Bitmap image object into TensorImage.
-   *
-   * Important: When loading a bitmap, DO NOT MODIFY the bitmap from the caller side anymore. The
-   * {@code TensorImage} object will rely on the bitmap. It will probably modify the bitmap as well.
-   * In this method, we perform a zero-copy approach for that bitmap, by simply holding its
-   * reference. Use {@code bitmap.copy(bitmap.getConfig(), true)} to create a copy if necessary.
-   *
-   * Note: To get the best performance, please load images in the same shape to avoid memory
-   * re-allocation.
-   *
-   * @throws IllegalArgumentException if {@code bitmap} is not in ARGB_8888.
-   */
-  public void load(@NonNull Bitmap bitmap) {
-    SupportPreconditions.checkNotNull(bitmap, "Cannot load null bitmap.");
-    SupportPreconditions.checkArgument(
-        bitmap.getConfig().equals(Config.ARGB_8888), "Only supports loading ARGB_8888 bitmaps.");
-    container.set(bitmap);
-  }
-
-  /**
-   * Loads a float array as RGB pixels into TensorImage, representing the pixels inside.
-   *
-   * 
Note: If the TensorImage has data type {@link DataType#UINT8}, numeric casting and clamping
-   * will be applied.
-   *
-   * @param pixels The RGB pixels representing the image.
-   * @param shape The shape of the image, should either in form (h, w, 3), or in form (1, h, w, 3).
-   */
-  public void load(@NonNull float[] pixels, @NonNull int[] shape) {
-    checkImageTensorShape(shape);
-    TensorBuffer buffer = TensorBuffer.createDynamic(getDataType());
-    buffer.loadArray(pixels, shape);
-    load(buffer);
-  }
-
-  /**
-   * Loads an uint8 array as RGB pixels into TensorImage, representing the pixels inside.
-   *
-   * 
Note: If the TensorImage has data type {@link DataType#UINT8}, all pixel values will clamp
-   * into [0, 255].
-   *
-   * @param pixels The RGB pixels representing the image.
-   * @param shape The shape of the image, should either in form (h, w, 3), or in form (1, h, w, 3).
-   */
-  public void load(@NonNull int[] pixels, @NonNull int[] shape) {
-    checkImageTensorShape(shape);
-    TensorBuffer buffer = TensorBuffer.createDynamic(getDataType());
-    buffer.loadArray(pixels, shape);
-    load(buffer);
-  }
-
-  /**
-   * Loads a TensorBuffer containing pixel values. The color layout should be RGB.
-   *
-   * @param buffer The TensorBuffer to load. Its shape should be either (h, w, 3) or (1, h, w, 3).
-   */
-  public void load(TensorBuffer buffer) {
-    checkImageTensorShape(buffer.getShape());
-    container.set(buffer);
-  }
-
-  /**
-   * Returns a bitmap representation of this TensorImage.
-   *
-   * 
Important: It's only a reference. DO NOT MODIFY. We don't create a copy here for performance
-   * concern, but if modification is necessary, please make a copy.
-   *
-   * @return a reference to a Bitmap in ARGB_8888 config. "A" channel is always opaque.
-   * @throws IllegalStateException if the TensorImage never loads data, or if the TensorImage is
-   *     holding a float-value image in {@code TensorBuffer}.
-   */
-  @NonNull
-  public Bitmap getBitmap() {
-    return container.getBitmap();
-  }
-
-  /**
-   * Returns a ByteBuffer representation of this TensorImage.
-   *
-   * 
Important: It's only a reference. DO NOT MODIFY. We don't create a copy here for performance
-   * concern, but if modification is necessary, please make a copy.
-   *
-   * 
It's essentially a short cut for {@code getTensorBuffer().getBuffer()}.
-   *
-   * @return a reference to a ByteBuffer which holds the image data.
-   * @throws IllegalStateException if the TensorImage never loads data.
-   */
-  @NonNull
-  public ByteBuffer getBuffer() {
-    return container.getTensorBuffer().getBuffer();
-  }
-
-  /**
-   * Returns a ByteBuffer representation of this TensorImage.
-   *
-   * 
Important: It's only a reference. DO NOT MODIFY. We don't create a copy here for performance
-   * concern, but if modification is necessary, please make a copy.
-   *
-   * @return a reference to a TensorBuffer which holds the image data.
-   * @throws IllegalStateException if the TensorImage never loads data.
-   */
-  @NonNull
-  public TensorBuffer getTensorBuffer() {
-    return container.getTensorBuffer();
-  }
-
-  /**
-   * Gets the current data type.
-   *
-   * @return a data type. Currently only UINT8 and FLOAT32 are possible.
-   */
-  public DataType getDataType() {
-    return container.getDataType();
-  }
-
-  /**
-   * Gets the image width.
-   *
-   * @throws IllegalStateException if the TensorImage never loads data.
-   * @throws IllegalArgumentException if the container data is corrupted.
-   */
-  public int getWidth() {
-    return container.getWidth();
-  }
-
-  /**
-   * Gets the image height.
-   *
-   * @throws IllegalStateException if the TensorImage never loads data.
-   * @throws IllegalArgumentException if the container data is corrupted.
-   */
-  public int getHeight() {
-    return container.getHeight();
-  }
-
-  // Requires tensor shape [h, w, 3] or [1, h, w, 3].
-  static void checkImageTensorShape(int[] shape) {
-    SupportPreconditions.checkArgument(
-        (shape.length == 3 || (shape.length == 4 && shape[0] == 1))
-            && shape[shape.length - 3] > 0
-            && shape[shape.length - 2] > 0
-            && shape[shape.length - 1] == 3,
-        "Only supports image shape in (h, w, c) or (1, h, w, c), and channels representing R, G, B"
-            + " in order.");
-  }
-
-  // Handles RGB image data storage strategy of TensorBuffer.
-  private static class ImageContainer {
-
-    private TensorBuffer bufferImage;
-    private boolean isBufferUpdated;
-    private Bitmap bitmapImage;
-    private boolean isBitmapUpdated;
-
-    private final DataType dataType;
-
-    private static final int ARGB_8888_ELEMENT_BYTES = 4;
-
-    ImageContainer(DataType dataType) {
-      this.dataType = dataType;
-    }
-
-    // Internal method to set the image source-of-truth with a bitmap. The bitmap has to be
-    // ARGB_8888.
-    void set(Bitmap bitmap) {
-      bitmapImage = bitmap;
-      isBufferUpdated = false;
-      isBitmapUpdated = true;
-    }
-
-    // Internal method to set the image source-of-truth with a TensorBuffer.
-    void set(TensorBuffer buffer) {
-      bufferImage = buffer;
-      isBitmapUpdated = false;
-      isBufferUpdated = true;
-    }
-
-    int getWidth() {
-      SupportPreconditions.checkState(
-          isBitmapUpdated || isBufferUpdated,
-          "Both buffer and bitmap data are obsolete. Forgot to call TensorImage#load?");
-      if (isBitmapUpdated) {
-        return bitmapImage.getWidth();
-      }
-      return getBufferDimensionSize(-2);
-    }
-
-    int getHeight() {
-      SupportPreconditions.checkState(
-          isBitmapUpdated || isBufferUpdated,
-          "Both buffer and bitmap data are obsolete. Forgot to call TensorImage#load?");
-      if (isBitmapUpdated) {
-        return bitmapImage.getHeight();
-      }
-      return getBufferDimensionSize(-3);
-    }
-
-    // Internal helper method to get the size of one dimension in the shape of the `bufferImage`.
-    // Requires `isBufferUpdated` is true.
-    // Throws `IllegalArgumentException` if data is corrupted.
-    private int getBufferDimensionSize(int dim) {
-      int[] shape = bufferImage.getShape();
-      // The defensive check is needed because bufferImage might be invalidly changed by user
-      // (a.k.a internal data is corrupted)
-      TensorImage.checkImageTensorShape(shape);
-      dim = dim % shape.length;
-      if (dim < 0) {
-        dim += shape.length;
-      }
-      return shape[dim];
-    }
-
-    public DataType getDataType() {
-      return dataType;
-    }
-
-    // Internal method to update the internal Bitmap data by TensorBuffer data.
-    @NonNull
-    Bitmap getBitmap() {
-      if (isBitmapUpdated) {
-        return bitmapImage;
-      }
-      if (!isBufferUpdated) {
-        throw new IllegalStateException(
-            "Both buffer and bitmap data are obsolete. Forgot to call TensorImage#load?");
-      }
-      if (bufferImage.getDataType() != DataType.UINT8) {
-        throw new IllegalStateException(
-            "TensorImage is holding a float-value image which is not able to convert a Bitmap.");
-      }
-      int requiredAllocation = bufferImage.getFlatSize() * ARGB_8888_ELEMENT_BYTES;
-      // Create a new bitmap and reallocate memory for it.
-      if (bitmapImage == null || bitmapImage.getAllocationByteCount() < requiredAllocation) {
-        int[] shape = bufferImage.getShape();
-        int h = shape[shape.length - 3];
-        int w = shape[shape.length - 2];
-        bitmapImage = Bitmap.createBitmap(w, h, Config.ARGB_8888);
-      }
-      ImageConversions.convertTensorBufferToBitmap(bufferImage, bitmapImage);
-      isBitmapUpdated = true;
-      return bitmapImage;
-    }
-
-    // Internal method to update the internal TensorBuffer data by Bitmap data.
-    @NonNull
-    TensorBuffer getTensorBuffer() {
-      if (isBufferUpdated) {
-        return bufferImage;
-      }
-      SupportPreconditions.checkArgument(
-          isBitmapUpdated,
-          "Both buffer and bitmap data are obsolete. Forgot to call TensorImage#load?");
-      int requiredFlatSize = bitmapImage.getWidth() * bitmapImage.getHeight() * 3;
-      if (bufferImage == null
-          || (!bufferImage.isDynamic() && bufferImage.getFlatSize() != requiredFlatSize)) {
-        bufferImage = TensorBuffer.createDynamic(dataType);
-      }
-      ImageConversions.convertBitmapToTensorBuffer(bitmapImage, bufferImage);
-      isBufferUpdated = true;
-      return bufferImage;
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeOp.java
deleted file mode 100644
index 35606dd66d3..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeOp.java
+++ /dev/null
@@ -1,89 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image.ops;
-
-import android.graphics.Bitmap;
-import android.graphics.PointF;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.image.ImageOperator;
-import org.tensorflow.lite.support.image.TensorImage;
-
-/**
- * As a computation unit for processing images, it can resize an image to user-specified size.
- *
- * 
It interpolates pixels when image is stretched, and discards pixels when image is compressed.
- *
- * @see ResizeWithCropOrPadOp for resizing without content distortion.
- */
-public class ResizeOp implements ImageOperator {
-
-  /** Algorithms for resizing. */
-  public enum ResizeMethod {
-    BILINEAR,
-    NEAREST_NEIGHBOR
-  }
-
-  private final int targetHeight;
-  private final int targetWidth;
-  private final boolean useBilinear;
-
-  /**
-   * Creates a ResizeOp which can resize images to specified size in specified method.
-   *
-   * @param targetHeight: The expected height of resized image.
-   * @param targetWidth: The expected width of resized image.
-   * @param resizeMethod: The algorithm to use for resizing. Options: {@link ResizeMethod}
-   */
-  public ResizeOp(int targetHeight, int targetWidth, ResizeMethod resizeMethod) {
-    this.targetHeight = targetHeight;
-    this.targetWidth = targetWidth;
-    useBilinear = (resizeMethod == ResizeMethod.BILINEAR);
-  }
-
-  /**
-   * Applies the defined resizing on given image and returns the result.
-   *
-   * 
Note: the content of input {@code image} will change, and {@code image} is the same instance
-   * with the output.
-   *
-   * @param image input image.
-   * @return output image.
-   */
-  @Override
-  @NonNull
-  public TensorImage apply(@NonNull TensorImage image) {
-    Bitmap scaled =
-        Bitmap.createScaledBitmap(image.getBitmap(), targetWidth, targetHeight, useBilinear);
-    image.load(scaled);
-    return image;
-  }
-
-  @Override
-  public int getOutputImageHeight(int inputImageHeight, int inputImageWidth) {
-    return targetHeight;
-  }
-
-  @Override
-  public int getOutputImageWidth(int inputImageHeight, int inputImageWidth) {
-    return targetWidth;
-  }
-
-  @Override
-  public PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth) {
-    return new PointF(
-        point.x * inputImageWidth / targetWidth, point.y * inputImageHeight / targetHeight);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeWithCropOrPadOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeWithCropOrPadOp.java
deleted file mode 100644
index 404429efa06..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/ResizeWithCropOrPadOp.java
+++ /dev/null
@@ -1,125 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image.ops;
-
-import android.graphics.Bitmap;
-import android.graphics.Bitmap.Config;
-import android.graphics.Canvas;
-import android.graphics.PointF;
-import android.graphics.Rect;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.image.ImageOperator;
-import org.tensorflow.lite.support.image.TensorImage;
-
-/**
- * As a computation unit for processing images, it could resize image to predefined size.
- *
- * 
It will not stretch or compress the content of image. However, to fit the new size, it crops
- * or pads pixels. When it crops image, it performs a center-crop; when it pads pixels, it performs
- * a zero-padding.
- *
- * @see ResizeOp for reszing images while stretching / compressing the content.
- */
-public class ResizeWithCropOrPadOp implements ImageOperator {
-  private final int targetHeight;
-  private final int targetWidth;
-  private final Bitmap output;
-
-  /**
-   * Creates a ResizeWithCropOrPadOp which could crop/pad images to specified size. It adopts
-   * center-crop and zero-padding.
-   *
-   * @param targetHeight: The expected height of cropped/padded image.
-   * @param targetWidth: The expected width of cropped/padded image.
-   */
-  public ResizeWithCropOrPadOp(int targetHeight, int targetWidth) {
-    this.targetHeight = targetHeight;
-    this.targetWidth = targetWidth;
-    output = Bitmap.createBitmap(this.targetWidth, this.targetHeight, Config.ARGB_8888);
-  }
-
-  /**
-   * Applies the defined resizing with cropping or/and padding on given image and returns the
-   * result.
-   *
-   * 
Note: the content of input {@code image} will change, and {@code image} is the same instance
-   * with the output.
-   *
-   * @param image input image.
-   * @return output image.
-   */
-  @Override
-  @NonNull
-  public TensorImage apply(@NonNull TensorImage image) {
-    Bitmap input = image.getBitmap();
-    int srcL;
-    int srcR;
-    int srcT;
-    int srcB;
-    int dstL;
-    int dstR;
-    int dstT;
-    int dstB;
-    int w = input.getWidth();
-    int h = input.getHeight();
-    if (targetWidth > w) { // padding
-      srcL = 0;
-      srcR = w;
-      dstL = (targetWidth - w) / 2;
-      dstR = dstL + w;
-    } else { // cropping
-      dstL = 0;
-      dstR = targetWidth;
-      srcL = (w - targetWidth) / 2;
-      srcR = srcL + targetWidth;
-    }
-    if (targetHeight > h) { // padding
-      srcT = 0;
-      srcB = h;
-      dstT = (targetHeight - h) / 2;
-      dstB = dstT + h;
-    } else { // cropping
-      dstT = 0;
-      dstB = targetHeight;
-      srcT = (h - targetHeight) / 2;
-      srcB = srcT + targetHeight;
-    }
-    Rect src = new Rect(srcL, srcT, srcR, srcB);
-    Rect dst = new Rect(dstL, dstT, dstR, dstB);
-    new Canvas(output).drawBitmap(input, src, dst, null);
-    image.load(output);
-    return image;
-  }
-
-  @Override
-  public int getOutputImageHeight(int inputImageHeight, int inputImageWidth) {
-    return targetHeight;
-  }
-
-  @Override
-  public int getOutputImageWidth(int inputImageHeight, int inputImageWidth) {
-    return targetWidth;
-  }
-
-  @Override
-  public PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth) {
-    return transformImpl(point, targetHeight, targetWidth, inputImageHeight, inputImageWidth);
-  }
-
-  private static PointF transformImpl(PointF point, int srcH, int srcW, int dstH, int dstW) {
-    return new PointF(point.x + (dstW - srcW) / 2, point.y + (dstH - srcH) / 2);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/Rot90Op.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/Rot90Op.java
deleted file mode 100644
index 2fa2293763c..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/Rot90Op.java
+++ /dev/null
@@ -1,103 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image.ops;
-
-import android.graphics.Bitmap;
-import android.graphics.Matrix;
-import android.graphics.PointF;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.image.ImageOperator;
-import org.tensorflow.lite.support.image.TensorImage;
-
-/** Rotates image counter-clockwise. */
-public class Rot90Op implements ImageOperator {
-
-  private final int numRotation;
-
-  /** Creates a Rot90 Op which will rotate image by 90 degree counter-clockwise. */
-  public Rot90Op() {
-    this(1);
-  }
-
-  /**
-   * Creates a Rot90 Op which will rotate image by 90 degree for {@code k} times counter-clockwise.
-   *
-   * @param k: The number of times the image is rotated by 90 degrees. If it's positive, the image
-   *     will be rotated counter-clockwise. If it's negative, the op will rotate image clockwise.
-   */
-  public Rot90Op(int k) {
-    numRotation = k % 4;
-  }
-
-  /**
-   * Applies the defined rotation on given image and returns the result.
-   *
-   * 
Note: the content of input {@code image} will change, and {@code image} is the same instance
-   * with the output.
-   *
-   * @param image input image.
-   * @return output image.
-   */
-  @NonNull
-  @Override
-  public TensorImage apply(@NonNull TensorImage image) {
-    Bitmap input = image.getBitmap();
-    if (numRotation == 0) {
-      return image;
-    }
-    int w = input.getWidth();
-    int h = input.getHeight();
-    Matrix matrix = new Matrix();
-    matrix.postTranslate(w * 0.5f, h * 0.5f);
-    matrix.postRotate(-90 * numRotation);
-    int newW = (numRotation % 2 == 0) ? w : h;
-    int newH = (numRotation % 2 == 0) ? h : w;
-    matrix.postTranslate(newW * 0.5f, newH * 0.5f);
-    Bitmap output = Bitmap.createBitmap(input, 0, 0, w, h, matrix, false);
-    image.load(output);
-    return image;
-  }
-
-  @Override
-  public int getOutputImageHeight(int inputImageHeight, int inputImageWidth) {
-    return (numRotation % 2 == 0) ? inputImageHeight : inputImageWidth;
-  }
-
-  @Override
-  public int getOutputImageWidth(int inputImageHeight, int inputImageWidth) {
-    return (numRotation % 2 == 0) ? inputImageWidth : inputImageHeight;
-  }
-
-  @Override
-  public PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth) {
-    int inverseNumRotation = (4 - numRotation) % 4;
-    int height = getOutputImageHeight(inputImageHeight, inputImageWidth);
-    int width = getOutputImageWidth(inputImageHeight, inputImageWidth);
-    return transformImpl(point, height, width, inverseNumRotation);
-  }
-
-  private static PointF transformImpl(PointF point, int height, int width, int numRotation) {
-    if (numRotation == 0) {
-      return point;
-    } else if (numRotation == 1) {
-      return new PointF(point.y, width - point.x);
-    } else if (numRotation == 2) {
-      return new PointF(width - point.x, height - point.y);
-    } else { // numRotation == 3
-      return new PointF(height - point.y, point.x);
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/TensorOperatorWrapper.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/TensorOperatorWrapper.java
deleted file mode 100644
index 75ccdac9b83..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/image/ops/TensorOperatorWrapper.java
+++ /dev/null
@@ -1,70 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.image.ops;
-
-import android.graphics.PointF;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.common.TensorOperator;
-import org.tensorflow.lite.support.image.ImageOperator;
-import org.tensorflow.lite.support.image.TensorImage;
-
-/**
- * The adapter that makes a TensorOperator able to run with TensorImage.
- *
- * @see org.tensorflow.lite.support.common.TensorOperator
- * @see org.tensorflow.lite.support.image.TensorImage
- */
-public class TensorOperatorWrapper implements ImageOperator {
-
-  private final TensorOperator tensorOp;
-
-  /**
-   * Wraps a {@link TensorOperator} object as an {@link ImageOperator}, so that the {@link
-   * TensorOperator} could handle {@link TensorImage} objects by handling its underlying {@link
-   * org.tensorflow.lite.support.tensorbuffer.TensorBuffer}.
-   *
-   * 
Requirement: The {@code op} should not change coordinate system when applied on an image.
-   *
-   * @param op The created operator.
-   */
-  public TensorOperatorWrapper(TensorOperator op) {
-    tensorOp = op;
-  }
-
-  @Override
-  @NonNull
-  public TensorImage apply(@NonNull TensorImage image) {
-    SupportPreconditions.checkNotNull(image, "Op cannot apply on null image.");
-    image.load(tensorOp.apply(image.getTensorBuffer()));
-    return image;
-  }
-
-  @Override
-  public int getOutputImageHeight(int inputImageHeight, int inputImageWidth) {
-    return inputImageHeight;
-  }
-
-  @Override
-  public int getOutputImageWidth(int inputImageHeight, int inputImageWidth) {
-    return inputImageWidth;
-  }
-
-  @Override
-  public PointF inverseTransform(PointF point, int inputImageHeight, int inputImageWidth) {
-    return point;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/Category.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/Category.java
deleted file mode 100644
index ea369c3ac12..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/Category.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.label;
-
-import java.util.Objects;
-
-/**
- * Category is a util class, contains a label and a float value. Typically it's used as result of
- * classification tasks.
- */
-public final class Category {
-  private final String label;
-  private final float score;
-
-  /** Constructs a Category. */
-  public Category(String label, float score) {
-    this.label = label;
-    this.score = score;
-  }
-
-  /** Gets the reference of category's label. */
-  public String getLabel() {
-    return label;
-  }
-
-  /** Gets the score of the category. */
-  public float getScore() {
-    return score;
-  }
-
-  @Override
-  public boolean equals(Object o) {
-    if (o instanceof Category) {
-      Category other = (Category) o;
-      return (other.getLabel().equals(this.label) && other.getScore() == this.score);
-    }
-    return false;
-  }
-
-  @Override
-  public int hashCode() {
-    return Objects.hash(label, score);
-  }
-
-  @Override
-  public String toString() {
-    return "";
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/LabelUtil.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/LabelUtil.java
deleted file mode 100644
index 840ed5fb77d..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/LabelUtil.java
+++ /dev/null
@@ -1,64 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.label;
-
-import android.util.Log;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/** Label operation utils. */
-public class LabelUtil {
-  /**
-   * Maps an int value tensor to a list of string labels. It takes an array of strings as the
-   * dictionary. Example: if the given tensor is [3, 1, 0], and given labels is ["background",
-   * "apple", "banana", "cherry", "date"], the result will be ["date", "banana", "apple"].
-   *
-   * @param tensorBuffer: A tensor with index values. The values should be non-negative integers,
-   *     and each value {@code x} will be converted to {@code labels[x + offset]}. If the tensor is
-   *     given as a float {@link TensorBuffer}, values will be cast to integers. All values that are
-   *     out of bound will map to empty string.
-   * @param labels: A list of strings, used as a dictionary to look up. The index of the array
-   *     element will be used as the key. To get better performance, use an object that implements
-   *     RandomAccess, such as {@link ArrayList}.
-   * @param offset: The offset value when look up int values in the {@code labels}.
-   * @return the mapped strings. The length of the list is {@link TensorBuffer#getFlatSize}.
-   * @throws IllegalArgumentException if {@code tensorBuffer} or {@code labels} is null.
-   */
-  public static List mapValueToLabels(
-      @NonNull TensorBuffer tensorBuffer, @NonNull List labels, int offset) {
-    SupportPreconditions.checkNotNull(tensorBuffer, "Given tensor should not be null");
-    SupportPreconditions.checkNotNull(labels, "Given labels should not be null");
-    int[] values = tensorBuffer.getIntArray();
-    Log.d("values", Arrays.toString(values));
-    List result = new ArrayList<>();
-    for (int v : values) {
-      int index = v + offset;
-      if (index < 0 || index >= labels.size()) {
-        result.add("");
-      } else {
-        result.add(labels.get(index));
-      }
-    }
-    return result;
-  }
-
-  // Private constructor to prevent initialization.
-  private LabelUtil() {}
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/TensorLabel.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/TensorLabel.java
deleted file mode 100644
index 10763a1a065..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/TensorLabel.java
+++ /dev/null
@@ -1,224 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.label;
-
-import android.content.Context;
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * TensorLabel is an util wrapper for TensorBuffers with meaningful labels on an axis.
- *
- * 
For example, an image classification model may have an output tensor with shape as {1, 10},
- * where 1 is the batch size and 10 is the number of categories. In fact, on the 2nd axis, we could
- * label each sub-tensor with the name or description of each corresponding category. {@link
- * TensorLabel} could help converting the plain Tensor in {@link TensorBuffer} into a map from
- * predefined labels to sub-tensors. In this case, if provided 10 labels for the 2nd axis, {@link
- * TensorLabel} could convert the original {1, 10} Tensor to a 10 element map, each value of which
- * is Tensor in shape {} (scalar). Usage example:
- *
- * 
- *   TensorBuffer outputTensor = ...;
- *   {@literal List} labels = FileUtil.loadLabels(context, labelFilePath);
- *   // labels the first axis with size greater than one
- *   TensorLabel labeled = new TensorLabel(labels, outputTensor);
- *   // If each sub-tensor has effectively size 1, we can directly get a float value
- *   {@literal Map} probabilities = labeled.getMapWithFloatValue();
- *   // Or get sub-tensors, when each sub-tensor has elements more than 1
- *   {@literal Map} subTensors = labeled.getMapWithTensorBuffer();
- * 

- *
- * 
Note: currently we only support tensor-to-map conversion for the first label with size greater
- * than 1.
- *
- * @see org.tensorflow.lite.support.common.FileUtil#loadLabels(Context, String) to load labels from
- *     a label file (plain text file whose each line is a label) in assets simply.
- */
-public class TensorLabel {
-  private final Map> axisLabels;
-  private final TensorBuffer tensorBuffer;
-  private final int[] shape;
-
-  /**
-   * Creates a TensorLabel object which is able to label on the axes of multi-dimensional tensors.
-   *
-   * @param axisLabels A map, whose key is axis id (starting from 0) and value is corresponding
-   *     labels. Note: The size of labels should be same with the size of the tensor on that axis.
-   * @param tensorBuffer The TensorBuffer to be labeled.
-   * @throws NullPointerException if {@code axisLabels} or {@code tensorBuffer} is null, or any
-   *     value in {@code axisLabels} is null.
-   * @throws IllegalArgumentException if any key in {@code axisLabels} is out of range (compared to
-   *     the shape of {@code tensorBuffer}, or any value (labels) has different size with the {@code
-   *     tensorBuffer} on the given dimension.
-   */
-  public TensorLabel(
-      @NonNull Map> axisLabels, @NonNull TensorBuffer tensorBuffer) {
-    SupportPreconditions.checkNotNull(axisLabels, "Axis labels cannot be null.");
-    SupportPreconditions.checkNotNull(tensorBuffer, "Tensor Buffer cannot be null.");
-    this.axisLabels = axisLabels;
-    this.tensorBuffer = tensorBuffer;
-    this.shape = tensorBuffer.getShape();
-    for (Map.Entry> entry : axisLabels.entrySet()) {
-      int axis = entry.getKey();
-      SupportPreconditions.checkArgument(
-          axis >= 0 && axis < shape.length, "Invalid axis id: " + axis);
-      SupportPreconditions.checkNotNull(entry.getValue(), "Label list is null on axis " + axis);
-      SupportPreconditions.checkArgument(
-          shape[axis] == entry.getValue().size(),
-          "Label number " + entry.getValue().size() + " mismatch the shape on axis " + axis);
-    }
-  }
-
-  /**
-   * Creates a TensorLabel object which is able to label on one axis of multi-dimensional tensors.
-   *
-   * 
Note: The labels are applied on the first axis whose size is larger than 1. For example, if
-   * the shape of the tensor is [1, 10, 3], the labels will be applied on axis 1 (id starting from
-   * 0), and size of {@code axisLabels} should be 10 as well.
-   *
-   * @param axisLabels A list of labels, whose size should be same with the size of the tensor on
-   *     the to-be-labeled axis.
-   * @param tensorBuffer The TensorBuffer to be labeled.
-   */
-  public TensorLabel(@NonNull List axisLabels, @NonNull TensorBuffer tensorBuffer) {
-    this(makeMap(getFirstAxisWithSizeGreaterThanOne(tensorBuffer), axisLabels), tensorBuffer);
-  }
-
-  /**
-   * Gets the map with a pair of the label and the corresponding TensorBuffer. Only allow the
-   * mapping on the first axis with size greater than 1 currently.
-   */
-  @NonNull
-  public Map getMapWithTensorBuffer() {
-    int labeledAxis = getFirstAxisWithSizeGreaterThanOne(tensorBuffer);
-
-    Map labelToTensorMap = new LinkedHashMap<>();
-    SupportPreconditions.checkArgument(
-        axisLabels.containsKey(labeledAxis),
-        "get a  map requires the labels are set on the first non-1 axis.");
-    List labels = axisLabels.get(labeledAxis);
-
-    DataType dataType = tensorBuffer.getDataType();
-    int typeSize = tensorBuffer.getTypeSize();
-    int flatSize = tensorBuffer.getFlatSize();
-
-    // Gets the underlying bytes that could be used to generate the sub-array later.
-    ByteBuffer byteBuffer = tensorBuffer.getBuffer();
-    byteBuffer.rewind();
-
-    // Note: computation below is only correct when labeledAxis is the first axis with size greater
-    // than 1.
-    int subArrayLength = flatSize / shape[labeledAxis] * typeSize;
-    int i = 0;
-    SupportPreconditions.checkNotNull(labels, "Label list should never be null");
-    for (String label : labels) {
-      // Gets the corresponding TensorBuffer.
-      byteBuffer.position(i * subArrayLength);
-      ByteBuffer subBuffer = byteBuffer.slice();
-      // ByteBuffer.slice doesn't keep order. Modify it to align with the original one.
-      subBuffer.order(byteBuffer.order()).limit(subArrayLength);
-      TensorBuffer labelBuffer = TensorBuffer.createDynamic(dataType);
-      labelBuffer.loadBuffer(subBuffer, Arrays.copyOfRange(shape, labeledAxis + 1, shape.length));
-      labelToTensorMap.put(label, labelBuffer);
-      i += 1;
-    }
-    return labelToTensorMap;
-  }
-
-  /**
-   * Gets a map that maps label to float. Only allow the mapping on the first axis with size greater
-   * than 1, and the axis should be effectively the last axis (which means every sub tensor
-   * specified by this axis should have a flat size of 1).
-   *
-   * 
{@link TensorLabel#getCategoryList()} is an alternative API to get the result.
-   *
-   * @throws IllegalStateException if size of a sub tensor on each label is not 1.
-   */
-  @NonNull
-  public Map getMapWithFloatValue() {
-    int labeledAxis = getFirstAxisWithSizeGreaterThanOne(tensorBuffer);
-    SupportPreconditions.checkState(
-        labeledAxis == shape.length - 1,
-        "get a  map is only valid when the only labeled axis is the last one.");
-    List labels = axisLabels.get(labeledAxis);
-    float[] data = tensorBuffer.getFloatArray();
-    SupportPreconditions.checkState(labels.size() == data.length);
-    Map result = new LinkedHashMap<>();
-    int i = 0;
-    for (String label : labels) {
-      result.put(label, data[i]);
-      i += 1;
-    }
-    return result;
-  }
-
-  /**
-   * Gets a list of {@link Category} from the {@link TensorLabel} object.
-   *
-   * 
The axis of label should be effectively the last axis (which means every sub tensor
-   * specified by this axis should have a flat size of 1), so that each labelled sub tensor could be
-   * converted into a float value score. Example: A {@link TensorLabel} with shape {@code {2, 5, 3}}
-   * and axis 2 is valid. If axis is 1 or 0, it cannot be converted into a {@link Category}.
-   *
-   * 
{@link TensorLabel#getMapWithFloatValue()} is an alternative but returns a {@link Map} as
-   * the result.
-   *
-   * @throws IllegalStateException if size of a sub tensor on each label is not 1.
-   */
-  @NonNull
-  public List getCategoryList() {
-    int labeledAxis = getFirstAxisWithSizeGreaterThanOne(tensorBuffer);
-    SupportPreconditions.checkState(
-        labeledAxis == shape.length - 1,
-        "get a Category list is only valid when the only labeled axis is the last one.");
-    List labels = axisLabels.get(labeledAxis);
-    float[] data = tensorBuffer.getFloatArray();
-    SupportPreconditions.checkState(labels.size() == data.length);
-    List result = new ArrayList<>();
-    int i = 0;
-    for (String label : labels) {
-      result.add(new Category(label, data[i]));
-      i += 1;
-    }
-    return result;
-  }
-
-  private static int getFirstAxisWithSizeGreaterThanOne(@NonNull TensorBuffer tensorBuffer) {
-    int[] shape = tensorBuffer.getShape();
-    for (int i = 0; i < shape.length; i++) {
-      if (shape[i] > 1) {
-        return i;
-      }
-    }
-    throw new IllegalArgumentException(
-        "Cannot find an axis to label. A valid axis to label should have size larger than 1.");
-  }
-
-  // Helper function to wrap the List to a one-entry map.
-  private static Map> makeMap(int axis, List labels) {
-    Map> map = new LinkedHashMap<>();
-    map.put(axis, labels);
-    return map;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/ops/LabelAxisOp.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/ops/LabelAxisOp.java
deleted file mode 100644
index c2de8c0baad..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/label/ops/LabelAxisOp.java
+++ /dev/null
@@ -1,74 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.label.ops;
-
-import android.content.Context;
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.support.common.FileUtil;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-import org.tensorflow.lite.support.label.TensorLabel;
-import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
-
-/**
- * Labels TensorBuffer with axisLabels for outputs.
- *
- * 
Apply on a {@code TensorBuffer} to get a {@code TensorLabel} that could output a Map, which is
- * a pair of the label name and the corresponding TensorBuffer value.
- */
-public class LabelAxisOp {
-  // Axis and its corresponding label names.
-  private final Map> axisLabels;
-
-  protected LabelAxisOp(Builder builder) {
-    axisLabels = builder.axisLabels;
-  }
-
-  public TensorLabel apply(@NonNull TensorBuffer buffer) {
-    SupportPreconditions.checkNotNull(buffer, "Tensor buffer cannot be null.");
-    return new TensorLabel(axisLabels, buffer);
-  }
-
-  /** The inner builder class to build a LabelTensor Operator. */
-  public static class Builder {
-    private final Map> axisLabels;
-
-    protected Builder() {
-      axisLabels = new HashMap<>();
-    }
-
-    public Builder addAxisLabel(@NonNull Context context, int axis, @NonNull String filePath)
-        throws IOException {
-      SupportPreconditions.checkNotNull(context, "Context cannot be null.");
-      SupportPreconditions.checkNotNull(filePath, "File path cannot be null.");
-      List labels = FileUtil.loadLabels(context, filePath);
-      axisLabels.put(axis, labels);
-      return this;
-    }
-
-    public Builder addAxisLabel(int axis, @NonNull List labels) {
-      axisLabels.put(axis, labels);
-      return this;
-    }
-
-    public LabelAxisOp build() {
-      return new LabelAxisOp(this);
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/GpuDelegateProxy.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/GpuDelegateProxy.java
deleted file mode 100644
index 9cfcf923ded..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/GpuDelegateProxy.java
+++ /dev/null
@@ -1,69 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.model;
-
-import android.util.Log;
-import java.io.Closeable;
-import java.io.IOException;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.tensorflow.lite.Delegate;
-
-/**
- * Helper class to create and call necessary methods of {@code GpuDelegate} which is not a strict
- * dependency.
- */
-class GpuDelegateProxy implements Delegate, Closeable {
-
-  private static final String TAG = "GpuDelegateProxy";
-
-  private final Delegate proxiedDelegate;
-  private final Closeable proxiedCloseable;
-
-  @Nullable
-  public static GpuDelegateProxy maybeNewInstance() {
-    try {
-      Class clazz = Class.forName("org.tensorflow.lite.gpu.GpuDelegate");
-      Object instance = clazz.getDeclaredConstructor().newInstance();
-      return new GpuDelegateProxy(instance);
-    } catch (ReflectiveOperationException e) {
-      Log.e(TAG, "Failed to create the GpuDelegate dynamically.", e);
-      return null;
-    }
-  }
-
-  /** Calls {@code close()} method of the delegate. */
-  @Override
-  public void close() {
-    try {
-      proxiedCloseable.close();
-    } catch (IOException e) {
-      // Should not trigger, because GpuDelegate#close never throws. The catch is required because
-      // of Closeable#close.
-      Log.e(TAG, "Failed to close the GpuDelegate.", e);
-    }
-  }
-
-  /** Calls {@code getNativeHandle()} method of the delegate. */
-  @Override
-  public long getNativeHandle() {
-    return proxiedDelegate.getNativeHandle();
-  }
-
-  private GpuDelegateProxy(Object instance) {
-    this.proxiedCloseable = (Closeable) instance;
-    this.proxiedDelegate = (Delegate) instance;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/Model.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/Model.java
deleted file mode 100644
index 8062d68d7b9..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/model/Model.java
+++ /dev/null
@@ -1,285 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.model;
-
-import android.content.Context;
-import java.io.IOException;
-import java.nio.MappedByteBuffer;
-import java.util.Map;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.tensorflow.lite.Interpreter;
-import org.tensorflow.lite.Tensor;
-import org.tensorflow.lite.support.common.FileUtil;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-
-/**
- * The wrapper class for a TFLite model and a TFLite interpreter.
- *
- * 
Note: A {@link Model} can only holds 1 TFLite model at a time, and always holds a TFLite
- * interpreter instance to run it.
- */
-public class Model {
-
-  /** The runtime device type used for executing classification. */
-  public enum Device {
-    CPU,
-    NNAPI,
-    GPU
-  }
-
-  /**
-   * Options for running the model. Configurable parameters includes:
-   *
-   * 
    
-   *   
  • {@code device} {@link Builder#setDevice(Device)} specifies the hardware to run the model.
-   *       The default value is {@link Device#CPU}.
-   *   
  • {@code numThreads} {@link Builder#setNumThreads(int)} specifies the number of threads
-   *       used by TFLite inference. It's only effective when device is set to {@link Device#CPU}
-   *       and default value is 1.
-   * 

-   */
-  public static class Options {
-    private final Device device;
-    private final int numThreads;
-
-    /** Builder of {@link Options}. See its doc for details. */
-    public static class Builder {
-      private Device device = Device.CPU;
-      private int numThreads = 1;
-
-      public Builder setDevice(Device device) {
-        this.device = device;
-        return this;
-      }
-
-      public Builder setNumThreads(int numThreads) {
-        this.numThreads = numThreads;
-        return this;
-      }
-
-      public Options build() {
-        return new Options(this);
-      }
-    }
-
-    private Options(Builder builder) {
-      device = builder.device;
-      numThreads = builder.numThreads;
-    }
-  }
-
-  /** An instance of the driver class to run model inference with Tensorflow Lite. */
-  private final Interpreter interpreter;
-
-  /** Path to tflite model file in asset folder. */
-  private final String modelPath;
-
-  /** The memory-mapped model data. */
-  private final MappedByteBuffer byteModel;
-
-  private final GpuDelegateProxy gpuDelegateProxy;
-
-  /**
-   * Builder for {@link Model}.
-   *
-   * @deprecated Please use {@link Model#createModel(Context, String, Options)}.
-   */
-  @Deprecated
-  public static class Builder {
-    private Device device = Device.CPU;
-    private int numThreads = 1;
-    private final String modelPath;
-    private final MappedByteBuffer byteModel;
-
-    /**
-     * Creates a builder which loads tflite model from asset folder using memory-mapped files.
-     *
-     * @param context: Application context to access assets.
-     * @param modelPath: Asset path of the model (.tflite file).
-     * @throws IOException if an I/O error occurs when loading the tflite model.
-     */
-    @NonNull
-    public Builder(@NonNull Context context, @NonNull String modelPath) throws IOException {
-      this.modelPath = modelPath;
-      byteModel = FileUtil.loadMappedFile(context, modelPath);
-    }
-
-    /** Sets running device. By default, TFLite will run on CPU. */
-    @NonNull
-    public Builder setDevice(Device device) {
-      this.device = device;
-      return this;
-    }
-
-    /** Sets number of threads. By default it's 1. */
-    @NonNull
-    public Builder setNumThreads(int numThreads) {
-      this.numThreads = numThreads;
-      return this;
-    }
-
-    // Note: The implementation is copied from `Model#createModel`. As the builder is going to be
-    // deprecated, this function is also to be removed.
-    @NonNull
-    public Model build() {
-      Options options = new Options.Builder().setNumThreads(numThreads).setDevice(device).build();
-      return createModel(byteModel, modelPath, options);
-    }
-  }
-
-  /**
-   * Loads a model from assets and initialize TFLite interpreter.
-   *
-   * 
The default options are: (1) CPU device; (2) one thread.
-   *
-   * @param context The App Context.
-   * @param modelPath The path of the model file.
-   * @throws IOException if any exception occurs when open the model file.
-   */
-  public static Model createModel(@NonNull Context context, @NonNull String modelPath)
-      throws IOException {
-    return createModel(context, modelPath, new Options.Builder().build());
-  }
-
-  /**
-   * Loads a model from assets and initialize TFLite interpreter with given options.
-   *
-   * @see Options for details.
-   * @param context The App Context.
-   * @param modelPath The path of the model file.
-   * @param options The options for running the model.
-   * @throws IOException if any exception occurs when open the model file.
-   */
-  public static Model createModel(
-      @NonNull Context context, @NonNull String modelPath, @NonNull Options options)
-      throws IOException {
-    SupportPreconditions.checkNotEmpty(
-        modelPath, "Model path in the asset folder cannot be empty.");
-    MappedByteBuffer byteModel = FileUtil.loadMappedFile(context, modelPath);
-    return createModel(byteModel, modelPath, options);
-  }
-
-  /**
-   * Creates a model with loaded {@link MappedByteBuffer}.
-   *
-   * @see Options for details.
-   * @param byteModel The loaded TFLite model.
-   * @param modelPath The original path of the model. It can be fetched later by {@link
-   *     Model#getPath()}.
-   * @param options The options for running the model.
-   * @throws IllegalArgumentException if {@code options.device} is {@link Device#GPU} but
-   *     "tensorflow-lite-gpu" is not linked to the project.
-   */
-  public static Model createModel(
-      @NonNull MappedByteBuffer byteModel, @NonNull String modelPath, @NonNull Options options) {
-    Interpreter.Options interpreterOptions = new Interpreter.Options();
-    GpuDelegateProxy gpuDelegateProxy = null;
-    switch (options.device) {
-      case NNAPI:
-        interpreterOptions.setUseNNAPI(true);
-        break;
-      case GPU:
-        gpuDelegateProxy = GpuDelegateProxy.maybeNewInstance();
-        SupportPreconditions.checkArgument(
-            gpuDelegateProxy != null,
-            "Cannot inference with GPU. Did you add \"tensorflow-lite-gpu\" as dependency?");
-        interpreterOptions.addDelegate(gpuDelegateProxy);
-        break;
-      case CPU:
-        break;
-    }
-    interpreterOptions.setNumThreads(options.numThreads);
-    Interpreter interpreter = new Interpreter(byteModel, interpreterOptions);
-    return new Model(modelPath, byteModel, interpreter, gpuDelegateProxy);
-  }
-
-  /** Returns the memory-mapped model data. */
-  @NonNull
-  public MappedByteBuffer getData() {
-    return byteModel;
-  }
-
-  /** Returns the path of the model file stored in Assets. */
-  @NonNull
-  public String getPath() {
-    return modelPath;
-  }
-
-  /**
-   * Gets the Tensor associated with the provdied input index.
-   *
-   * @throws IllegalStateException if the interpreter is closed.
-   */
-  public Tensor getInputTensor(int inputIndex) {
-    return interpreter.getInputTensor(inputIndex);
-  }
-
-  /**
-   * Gets the Tensor associated with the provdied output index.
-   *
-   * @throws IllegalStateException if the interpreter is closed.
-   */
-  public Tensor getOutputTensor(int outputIndex) {
-    return interpreter.getOutputTensor(outputIndex);
-  }
-
-  /**
-   * Returns the output shape. Useful if output shape is only determined when graph is created.
-   *
-   * @throws IllegalStateException if the interpreter is closed.
-   */
-  public int[] getOutputTensorShape(int outputIndex) {
-    return interpreter.getOutputTensor(outputIndex).shape();
-  }
-
-  /**
-   * Runs model inference on multiple inputs, and returns multiple outputs.
-   *
-   * @param inputs an array of input data. The inputs should be in the same order as inputs of the
-   *     model. Each input can be an array or multidimensional array, or a {@link
-   *     java.nio.ByteBuffer} of primitive types including int, float, long, and byte. {@link
-   *     java.nio.ByteBuffer} is the preferred way to pass large input data, whereas string types
-   *     require using the (multi-dimensional) array input path. When {@link java.nio.ByteBuffer} is
-   *     used, its content should remain unchanged until model inference is done.
-   * @param outputs a map mapping output indices to multidimensional arrays of output data or {@link
-   *     java.nio.ByteBuffer}s of primitive types including int, float, long, and byte. It only
-   *     needs to keep entries for the outputs to be used.
-   */
-  public void run(@NonNull Object[] inputs, @NonNull Map outputs) {
-    interpreter.runForMultipleInputsOutputs(inputs, outputs);
-  }
-
-  public void close() {
-    if (interpreter != null) {
-      interpreter.close();
-    }
-    if (gpuDelegateProxy != null) {
-      gpuDelegateProxy.close();
-    }
-  }
-
-  private Model(
-      @NonNull String modelPath,
-      @NonNull MappedByteBuffer byteModel,
-      @NonNull Interpreter interpreter,
-      @Nullable GpuDelegateProxy gpuDelegateProxy) {
-    this.modelPath = modelPath;
-    this.byteModel = byteModel;
-    this.interpreter = interpreter;
-    this.gpuDelegateProxy = gpuDelegateProxy;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBuffer.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBuffer.java
deleted file mode 100644
index fa05be363a6..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBuffer.java
+++ /dev/null
@@ -1,412 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.tensorbuffer;
-
-import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
-import java.util.Arrays;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-
-/** Represents the data buffer for either a model's input or its output. */
-public abstract class TensorBuffer {
-  /** Where the data is stored. */
-  protected ByteBuffer buffer;
-
-  /** Shape of the tensor stored in this buffer. */
-  protected int[] shape;
-
-  /** Number of elements in the buffer. It will be changed to a proper value in the constructor. */
-  protected int flatSize = -1;
-
-  /**
-   * Indicator of whether this buffer is dynamic or fixed-size. Fixed-size buffers will have
-   * pre-allocated memory and fixed size. While the size of dynamic buffers can be changed.
-   */
-  protected final boolean isDynamic;
-
-  /**
-   * Creates a {@link TensorBuffer} with specified {@code shape} and {@link DataType}. Here are some
-   * examples:
-   *
-   * 
-   * Creating a float TensorBuffer with shape {2, 3}:
-   * int[] shape = new int[] {2, 3};
-   * TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.FLOAT32);
-   * 

-   *
-   * 
-   * Creating an uint8 TensorBuffer of a scalar:
-   * int[] shape = new int[] {};
-   * TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.UINT8);
-   * 

-   *
-   * 
-   * Creating an empty uint8 TensorBuffer:
-   * int[] shape = new int[] {0};
-   * TensorBuffer tensorBuffer = TensorBuffer.createFixedSize(shape, DataType.UINT8);
-   * 

-   *
-   * 
The size of a fixed-size TensorBuffer cannot be changed once it is created.
-   *
-   * @param shape The shape of the {@link TensorBuffer} to be created.
-   * @param dataType The dataType of the {@link TensorBuffer} to be created.
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if {@code shape} has non-positive elements.
-   */
-  @NonNull
-  public static TensorBuffer createFixedSize(@NonNull int[] shape, DataType dataType) {
-    switch (dataType) {
-      case FLOAT32:
-        return new TensorBufferFloat(shape);
-      case UINT8:
-        return new TensorBufferUint8(shape);
-      default:
-        throw new AssertionError("TensorBuffer does not support data type: " + dataType);
-    }
-  }
-
-  /**
-   * Creates an empty dynamic {@link TensorBuffer} with specified {@link DataType}. The shape of the
-   * created {@link TensorBuffer} is {0}.
-   *
-   * 
Dynamic TensorBuffers will reallocate memory when loading arrays or data buffers of
-   * different buffer sizes.
-   *
-   * @param dataType The dataType of the {@link TensorBuffer} to be created.
-   */
-  @NonNull
-  public static TensorBuffer createDynamic(DataType dataType) {
-    switch (dataType) {
-      case FLOAT32:
-        return new TensorBufferFloat();
-      case UINT8:
-        return new TensorBufferUint8();
-      default:
-        throw new AssertionError("TensorBuffer does not support data type: " + dataType);
-    }
-  }
-
-  /**
-   * Creates a {@link TensorBuffer} deep-copying data from another, with specified {@link DataType}.
-   *
-   * @param buffer the source {@link TensorBuffer} to copy from.
-   * @param dataType the expected {@link DataType} of newly created {@link TensorBuffer}.
-   * @throws NullPointerException if {@code buffer} is null.
-   */
-  @NonNull
-  public static TensorBuffer createFrom(@NonNull TensorBuffer buffer, DataType dataType) {
-    SupportPreconditions.checkNotNull(buffer, "Cannot create a buffer from null");
-    TensorBuffer result;
-    if (buffer.isDynamic()) {
-      result = createDynamic(dataType);
-    } else {
-      result = createFixedSize(buffer.shape, dataType);
-    }
-    // The only scenario we need float array is FLOAT32->FLOAT32, or we can always use INT as
-    // intermediate container.
-    // The assumption is not true when we support other data types.
-    if (buffer.getDataType() == DataType.FLOAT32 && dataType == DataType.FLOAT32) {
-      float[] data = buffer.getFloatArray();
-      result.loadArray(data, buffer.shape);
-    } else {
-      int[] data = buffer.getIntArray();
-      result.loadArray(data, buffer.shape);
-    }
-    return result;
-  }
-
-  /** Returns the data buffer. */
-  @NonNull
-  public ByteBuffer getBuffer() {
-    return buffer;
-  }
-
-  /** Gets the {@link TensorBuffer#flatSize} of the buffer. */
-  public int getFlatSize() {
-    return flatSize;
-  }
-
-  /** Gets the current shape. (returning a copy here to avoid unexpected modification.) */
-  @NonNull
-  public int[] getShape() {
-    return Arrays.copyOf(shape, shape.length);
-  }
-
-  /** Returns the data type of this buffer. */
-  public abstract DataType getDataType();
-
-  /**
-   * Returns a float array of the values stored in this buffer. If the buffer is of different types
-   * than float, the values will be converted into float. For example, values in {@link
-   * TensorBufferUint8} will be converted from uint8 to float.
-   */
-  @NonNull
-  public abstract float[] getFloatArray();
-
-  /**
-   * Returns a float value at a given index. If the buffer is of different types than float, the
-   * value will be converted into float. For example, when reading a value from {@link
-   * TensorBufferUint8}, the value will be first read out as uint8, and then will be converted from
-   * uint8 to float.
-   *
-   * 
-   * For example, a TensorBuffer with shape {2, 3} that represents the following array,
-   * [[0.0f, 1.0f, 2.0f], [3.0f, 4.0f, 5.0f]].
-   *
-   * The fourth element (whose value is 3.0f) in the TensorBuffer can be retrived by:
-   * float v = tensorBuffer.getFloatValue(3);
-   * 

-   *
-   * @param absIndex The absolute index of the value to be read.
-   */
-  public abstract float getFloatValue(int absIndex);
-
-  /**
-   * Returns an int array of the values stored in this buffer. If the buffer is of different type
-   * than int, the values will be converted into int, and loss of precision may apply. For example,
-   * getting an int array from a {@link TensorBufferFloat} with values {400.32f, 23.04f}, the output
-   * is {400, 23}.
-   */
-  @NonNull
-  public abstract int[] getIntArray();
-
-  /**
-   * Returns an int value at a given index. If the buffer is of different types than int, the value
-   * will be converted into int. For example, when reading a value from {@link TensorBufferFloat},
-   * the value will be first read out as float, and then will be converted from float to int. Loss
-   * of precision may apply.
-   *
-   * 
-   * For example, a TensorBuffer with shape {2, 3} that represents the following array,
-   * [[0.0f, 1.0f, 2.0f], [3.0f, 4.0f, 5.0f]].
-   *
-   * The fourth element (whose value is 3.0f) in the TensorBuffer can be retrived by:
-   * int v = tensorBuffer.getIntValue(3);
-   * Note that v is converted from 3.0f to 3 as a result of type conversion.
-   * 

-   *
-   * @param absIndex The absolute index of the value to be read.
-   */
-  public abstract int getIntValue(int absIndex);
-
-  /**
-   * Returns the number of bytes of a single element in the array. For example, a float buffer will
-   * return 4, and a byte buffer will return 1.
-   */
-  public abstract int getTypeSize();
-
-  /** Returns if the {@link TensorBuffer} is dynamic sized (could resize arbitrarily). */
-  public boolean isDynamic() {
-    return isDynamic;
-  }
-
-  /**
-   * Loads an int array into this buffer with specific shape. If the buffer is of different types
-   * than int, the values will be converted into the buffer's type before being loaded into the
-   * buffer, and loss of precision may apply. For example, loading an int array with values {400,
-   * -23} into a {@link TensorBufferUint8} , the values will be clamped to [0, 255] and then be
-   * casted to uint8 by {255, 0}.
-   *
-   * @param src The source array to be loaded.
-   * @param shape Shape of the tensor that {@code src} represents.
-   * @throws NullPointerException if {@code src} is null.
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if the size of the array to be loaded does not match the
-   *     specified shape.
-   */
-  public abstract void loadArray(@NonNull int[] src, @NonNull int[] shape);
-
-  /**
-   * Loads an int array into this buffer. If the buffer is of different types than int, the values
-   * will be converted into the buffer's type before being loaded into the buffer, and loss of
-   * precision may apply. For example, loading an int array with values {400, -23} into a {@link
-   * TensorBufferUint8} , the values will be clamped to [0, 255] and then be casted to uint8 by
-   * {255, 0}.
-   *
-   * 
Size of {@code src} should always match the flat size of this {@link TensorBuffer}, for both
-   * fixed-size and dynamic {@link TensorBuffer}.
-   *
-   * @param src The source array to be loaded.
-   */
-  public void loadArray(@NonNull int[] src) {
-    loadArray(src, shape);
-  }
-
-  /**
-   * Loads a float array into this buffer with specific shape. If the buffer is of different types
-   * than float, the values will be converted into the buffer's type before being loaded into the
-   * buffer, and loss of precision may apply. For example, loading a float array into a {@link
-   * TensorBufferUint8} with values {400.32f, -23.04f}, the values will be clamped to [0, 255] and
-   * then be casted to uint8 by {255, 0}.
-   *
-   * @param src The source array to be loaded.
-   * @param shape Shape of the tensor that {@code src} represents.
-   * @throws NullPointerException if {@code src} is null.
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if the size of the array to be loaded does not match the
-   *     specified shape.
-   */
-  public abstract void loadArray(@NonNull float[] src, @NonNull int[] shape);
-
-  /**
-   * Loads a float array into this buffer. If the buffer is of different types than float, the
-   * values will be converted into the buffer's type before being loaded into the buffer, and loss
-   * of precision may apply. For example, loading a float array into a {@link TensorBufferUint8}
-   * with values {400.32f, -23.04f}, the values will be clamped to [0, 255] and then be casted to
-   * uint8 by {255, 0}.
-   *
-   * 
Size of {@code src} should always match the flat size of this {@link TensorBuffer}, for both
-   * fixed-size and dynamic {@link TensorBuffer}.
-   *
-   * @param src The source array to be loaded.
-   */
-  public void loadArray(@NonNull float[] src) {
-    loadArray(src, shape);
-  }
-
-  /**
-   * Loads a byte buffer into this {@link TensorBuffer} with specific shape.
-   *
-   * 
Important: The loaded buffer is a reference. DO NOT MODIFY. We don't create a copy here for
-   * performance concern, but if modification is necessary, please make a copy.
-   *
-   * @param buffer The byte buffer to load.
-   * @throws NullPointerException if {@code buffer} is null.
-   * @throws IllegalArgumentException if the size of {@code buffer} and {@code typeSize} do not
-   *     match or the size of {@code buffer} and {@code flatSize} do not match.
-   */
-  public void loadBuffer(@NonNull ByteBuffer buffer, @NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(buffer, "Byte buffer cannot be null.");
-    int flatSize = computeFlatSize(shape);
-    SupportPreconditions.checkArgument(
-        (buffer.limit() == getTypeSize() * flatSize),
-        "The size of byte buffer and the shape do not match.");
-
-    if (!isDynamic) {
-      SupportPreconditions.checkArgument(
-          flatSize == this.flatSize,
-          "The size of byte buffer and the size of the tensor buffer do not match.");
-    } else {
-      this.flatSize = flatSize;
-    }
-
-    this.shape = shape.clone();
-    buffer.rewind();
-    this.buffer = buffer;
-  }
-
-  /**
-   * Loads a byte buffer into this {@link TensorBuffer}. Buffer size must match the flat size of
-   * this {@link TensorBuffer}.
-   *
-   * 
Important: The loaded buffer is a reference. DO NOT MODIFY. We don't create a copy here for
-   * performance concern, but if modification is necessary, please make a copy.
-   *
-   * @param buffer The byte buffer to load.
-   */
-  public void loadBuffer(@NonNull ByteBuffer buffer) {
-    loadBuffer(buffer, shape);
-  }
-
-  /**
-   * Constructs a fixed size {@link TensorBuffer} with specified {@code shape}.
-   *
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if {@code shape} has non-positive elements.
-   */
-  protected TensorBuffer(@NonNull int[] shape) {
-    isDynamic = false;
-    allocateMemory(shape);
-  }
-
-  /** Constructs a dynamic {@link TensorBuffer} which can be resized. */
-  protected TensorBuffer() {
-    isDynamic = true;
-    // Initialize the dynamic TensorBuffer with an empty ByteBuffer.
-    allocateMemory(new int[] {0});
-  }
-
-  /** Calculates number of elements in the buffer. */
-  protected static int computeFlatSize(@NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(shape, "Shape cannot be null.");
-    int prod = 1;
-    for (int s : shape) {
-      prod = prod * s;
-    }
-    return prod;
-  }
-
-  /**
-   * For dynamic buffer, resize the memory if needed. For fixed-size buffer, check if the {@code
-   * shape} of src fits the buffer size.
-   */
-  protected void resize(@NonNull int[] shape) {
-    if (isDynamic) {
-      allocateMemory(shape);
-    } else {
-      // Make sure the new shape fits the buffer size when TensorBuffer has fixed size.
-      SupportPreconditions.checkArgument(Arrays.equals(shape, this.shape));
-      this.shape = shape.clone();
-    }
-  }
-
-  /**
-   * Allocates buffer with corresponding size of the {@code shape}. If shape is an empty array, this
-   * {@link TensorBuffer} will be created as a scalar and its flatSize will be 1.
-   *
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if {@code shape} has negative elements.
-   */
-  private void allocateMemory(@NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(shape, "TensorBuffer shape cannot be null.");
-    SupportPreconditions.checkArgument(
-        isShapeValid(shape), "Values in TensorBuffer shape should be non-negative.");
-
-    // Check if the new shape is the same as current shape.
-    int newFlatSize = computeFlatSize(shape);
-    this.shape = shape.clone();
-    if (flatSize == newFlatSize) {
-      return;
-    }
-
-    // Update to the new shape.
-    flatSize = newFlatSize;
-    buffer = ByteBuffer.allocateDirect(flatSize * getTypeSize());
-    buffer.order(ByteOrder.nativeOrder());
-  }
-
-  /**
-   * Checks if {@code shape} meets one of following two requirements: 1. Elements in {@code shape}
-   * are all non-negative numbers. 2. {@code shape} is an empty array, which corresponds to scalar.
-   */
-  private static boolean isShapeValid(@NonNull int[] shape) {
-    if (shape.length == 0) {
-      // This shape refers to a scalar.
-      return true;
-    }
-
-    // This shape refers to a multidimensional array.
-    for (int s : shape) {
-      // All elements in shape should be non-negative.
-      if (s < 0) {
-        return false;
-      }
-    }
-    return true;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferFloat.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferFloat.java
deleted file mode 100644
index c5b46b19f29..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferFloat.java
+++ /dev/null
@@ -1,110 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.tensorbuffer;
-
-import java.nio.FloatBuffer;
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-
-/** Represents data buffer with float values. */
-public final class TensorBufferFloat extends TensorBuffer {
-  private static final DataType DATA_TYPE = DataType.FLOAT32;
-
-  /**
-   * Creates a {@link TensorBufferFloat} with specified {@code shape}.
-   *
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if {@code shape} has non-positive elements.
-   */
-  TensorBufferFloat(@NonNull int[] shape) {
-    super(shape);
-  }
-
-  TensorBufferFloat() {
-    super();
-  }
-
-  @Override
-  public DataType getDataType() {
-    return DATA_TYPE;
-  }
-
-  @Override
-  @NonNull
-  public float[] getFloatArray() {
-    buffer.rewind();
-    float[] arr = new float[flatSize];
-
-    FloatBuffer floatBuffer = buffer.asFloatBuffer();
-    floatBuffer.get(arr);
-    return arr;
-  }
-
-  @Override
-  public float getFloatValue(int absIndex) {
-    return buffer.getFloat(absIndex << 2);
-  }
-
-  @Override
-  @NonNull
-  public int[] getIntArray() {
-    buffer.rewind();
-    int[] arr = new int[flatSize];
-
-    for (int i = 0; i < flatSize; i++) {
-      arr[i] = (int) buffer.getFloat();
-    }
-    return arr;
-  }
-
-  @Override
-  public int getIntValue(int absIndex) {
-    return (int) buffer.getFloat(absIndex << 2);
-  }
-
-  @Override
-  public int getTypeSize() {
-    return DATA_TYPE.byteSize();
-  }
-
-  @Override
-  public void loadArray(@NonNull float[] src, @NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(src, "The array to be loaded cannot be null.");
-    SupportPreconditions.checkArgument(
-        src.length == computeFlatSize(shape),
-        "The size of the array to be loaded does not match the specified shape.");
-    resize(shape);
-    buffer.rewind();
-
-    FloatBuffer floatBuffer = buffer.asFloatBuffer();
-    floatBuffer.put(src);
-  }
-
-  @Override
-  public void loadArray(@NonNull int[] src, @NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(src, "The array to be loaded cannot be null.");
-    SupportPreconditions.checkArgument(
-        src.length == computeFlatSize(shape),
-        "The size of the array to be loaded does not match the specified shape.");
-    resize(shape);
-    buffer.rewind();
-
-    for (int a : src) {
-      buffer.putFloat((float) a);
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferUint8.java b/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferUint8.java
deleted file mode 100644
index dce63ef2d51..00000000000
--- a/tensorflow/lite/experimental/support/java/src/java/org/tensorflow/lite/support/tensorbuffer/TensorBufferUint8.java
+++ /dev/null
@@ -1,111 +0,0 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.tensorbuffer;
-
-import org.checkerframework.checker.nullness.qual.NonNull;
-import org.tensorflow.lite.DataType;
-import org.tensorflow.lite.support.common.SupportPreconditions;
-
-/** Represents data buffer with 8-bit unsigned integer values. */
-public final class TensorBufferUint8 extends TensorBuffer {
-  private static final DataType DATA_TYPE = DataType.UINT8;
-
-  /**
-   * Creates a {@link TensorBufferUint8} with specified {@code shape}.
-   *
-   * @throws NullPointerException if {@code shape} is null.
-   * @throws IllegalArgumentException if {@code shape} has non-positive elements.
-   */
-  TensorBufferUint8(@NonNull int[] shape) {
-    super(shape);
-  }
-
-  TensorBufferUint8() {
-    super();
-  }
-
-  @Override
-  public DataType getDataType() {
-    return DATA_TYPE;
-  }
-
-  @Override
-  @NonNull
-  public float[] getFloatArray() {
-    buffer.rewind();
-    float[] arr = new float[flatSize];
-
-    for (int i = 0; i < flatSize; i++) {
-      arr[i] = (float) (buffer.get() & 0xff);
-    }
-    return arr;
-  }
-
-  @Override
-  public float getFloatValue(int index) {
-    return (float) (buffer.get(index) & 0xff);
-  }
-
-  @Override
-  @NonNull
-  public int[] getIntArray() {
-    buffer.rewind();
-    int[] arr = new int[flatSize];
-
-    for (int i = 0; i < flatSize; i++) {
-      arr[i] = buffer.get() & 0xff;
-    }
-    return arr;
-  }
-
-  @Override
-  public int getIntValue(int index) {
-    return buffer.get(index) & 0xff;
-  }
-
-  @Override
-  public int getTypeSize() {
-    return DATA_TYPE.byteSize();
-  }
-
-  @Override
-  public void loadArray(@NonNull float[] src, @NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(src, "The array to be loaded cannot be null.");
-    SupportPreconditions.checkArgument(
-        src.length == computeFlatSize(shape),
-        "The size of the array to be loaded does not match the specified shape.");
-    resize(shape);
-    buffer.rewind();
-
-    for (float a : src) {
-      buffer.put((byte) Math.max(Math.min(a, 255.0), 0.0));
-    }
-  }
-
-  @Override
-  public void loadArray(@NonNull int[] src, @NonNull int[] shape) {
-    SupportPreconditions.checkNotNull(src, "The array to be loaded cannot be null.");
-    SupportPreconditions.checkArgument(
-        src.length == computeFlatSize(shape),
-        "The size of the array to be loaded does not match the specified shape.");
-    resize(shape);
-    buffer.rewind();
-
-    for (int a : src) {
-      buffer.put((byte) Math.max(Math.min(a, 255), 0));
-    }
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/BUILD b/tensorflow/lite/experimental/support/metadata/BUILD
deleted file mode 100644
index ba410d914c7..00000000000
--- a/tensorflow/lite/experimental/support/metadata/BUILD
+++ /dev/null
@@ -1,113 +0,0 @@
-load("//tensorflow:tensorflow.bzl", "py_test")
-load("@flatbuffers//:build_defs.bzl", "flatbuffer_android_library", "flatbuffer_cc_library", "flatbuffer_java_library", "flatbuffer_py_library")
-load("//tensorflow/lite/experimental/support/metadata:build_defs.bzl", "stamp_metadata_parser_version")
-
-package(
-    default_visibility = [
-        "//visibility:public",
-    ],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-exports_files(["metadata_schema.fbs"])
-
-flatbuffer_py_library(
-    name = "schema_py",
-    srcs = ["//tensorflow/lite/schema:schema.fbs"],
-)
-
-# Generic schema for inference on device.
-flatbuffer_android_library(
-    name = "schema_fbs_android",
-    srcs = ["//tensorflow/lite/schema:schema.fbs"],
-    custom_package = "org.tensorflow.lite.schema",
-)
-
-flatbuffer_java_library(
-    name = "schema_fbs_java",
-    srcs = ["//tensorflow/lite/schema:schema.fbs"],
-    custom_package = "org.tensorflow.lite.schema",
-)
-
-# Generic schema for model metadata.
-flatbuffer_cc_library(
-    name = "metadata_schema_cc",
-    srcs = ["metadata_schema.fbs"],
-)
-
-flatbuffer_py_library(
-    name = "metadata_schema_py",
-    srcs = ["metadata_schema.fbs"],
-)
-
-flatbuffer_java_library(
-    name = "metadata_schema_java",
-    srcs = ["metadata_schema.fbs"],
-    custom_package = "org.tensorflow.lite.support.metadata.schema",
-)
-
-flatbuffer_android_library(
-    name = "metadata_schema_fbs_android",
-    srcs = ["metadata_schema.fbs"],
-    custom_package = "org.tensorflow.lite.support.metadata.schema",
-)
-
-# TODO(b/157813075): move the metadata python library to metadata/python/ when migrating to the new repo.
-stamp_metadata_parser_version(
-    name = "metadata_parser_py",
-    srcs = ["metadata_parser.py.template"],
-    outs = ["metadata_parser.py"],
-)
-
-py_library(
-    name = "metadata",
-    srcs = [
-        "metadata.py",
-        ":metadata_parser_py",
-    ],
-    data = [
-        "//tensorflow/lite/experimental/support/metadata:metadata_schema.fbs",
-    ],
-    srcs_version = "PY2AND3",
-    visibility = ["//visibility:public"],
-    deps = [
-        ":metadata_schema_py",
-        ":schema_py",
-        "//tensorflow/lite/experimental/support/metadata/cc/python:_pywrap_metadata_version",
-        "//tensorflow/lite/experimental/support/metadata/flatbuffers_lib:_pywrap_flatbuffers",
-        "//tensorflow/python:platform",
-        "@flatbuffers//:runtime_py",
-    ],
-)
-
-py_test(
-    name = "metadata_test",
-    srcs = ["metadata_test.py"],
-    data = ["testdata/golden_json.json"],
-    python_version = "PY3",
-    srcs_version = "PY2AND3",
-    tags = [
-        "no_mac",  # TODO(b/148247402): flatbuffers import broken on Mac OS.
-    ],
-    deps = [
-        ":metadata",
-        ":metadata_schema_py",
-        ":schema_py",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:platform",
-        "//tensorflow/python:platform_test",
-        "@flatbuffers//:runtime_py",
-        "@six_archive//:six",
-    ],
-)
-
-py_test(
-    name = "metadata_parser_test",
-    srcs = ["metadata_parser_test.py"],
-    python_version = "PY3",
-    srcs_version = "PY2AND3",
-    deps = [
-        ":metadata",
-        "//tensorflow/python:client_testlib",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/README.md b/tensorflow/lite/experimental/support/metadata/README.md
deleted file mode 100644
index ff7d25f27cb..00000000000
--- a/tensorflow/lite/experimental/support/metadata/README.md
+++ /dev/null
@@ -1,15 +0,0 @@
-# TensorFlow Lite Metadata and Android wrapper code generator
-
-Note: Both TensorFlow Lite Metadata and the Android wrapper code generator are
-in experimental (beta) phase.
-
-TensorFlow Lite metadata provides a structured framework for storing metadata
-to convey information for both the developer that will utilitised the model and
-code generators which can create wrapper around the model. For information on
-how to populate model metadata, please refer to the [TensorFlow Lite Metadata 
-documentation](https://www.tensorflow.org/lite/convert/metadata).
-
-The first code generator which takes advantage of this metadata format is the
-TensorFlow Lite Android Code Generator. For more information on how to use this
-generator, please refer to the [TensorFlow Lite Android wrapper code generator
-documentation](https://www.tensorflow.org/lite/guide/codegen).
diff --git a/tensorflow/lite/experimental/support/metadata/build_defs.bzl b/tensorflow/lite/experimental/support/metadata/build_defs.bzl
deleted file mode 100644
index 3ea945770e0..00000000000
--- a/tensorflow/lite/experimental/support/metadata/build_defs.bzl
+++ /dev/null
@@ -1,43 +0,0 @@
-"""Build rules to generate metadata schema versions."""
-
-METADATA_SCHEMA_FILE = "//tensorflow/lite/experimental/support/metadata:metadata_schema.fbs"
-
-def stamp_metadata_parser_version(
-        name,
-        srcs,
-        outs):
-    """Stamps the latest metadata parser version into the srcs files.
-
-    Replaces all the occurrences of "{LATEST_METADATA_PARSER_VERSION}" in the
-    srcs files with the metadata schema version extracted from
-    METADATA_SCHEMA_FILE and then outputs the generated file into outs,
-    respectively. The number of srcs files needs to match the number of outs
-    files.
-
-    Args:
-        name: Rule name. (required)
-        srcs: List of source files. (required)
-        outs: List of output files. (required)
-    """
-    if len(srcs) != len(outs):
-        fail(("The number of srcs files (%d) does not match that of the outs" +
-              " files (%d).") %
-             (len(srcs), len(outs)))
-
-    for i in range(0, len(srcs)):
-        native.genrule(
-            name = "%s_file%d" % (name, i),
-            srcs = [srcs[i]],
-            outs = [outs[i]],
-            tools = [METADATA_SCHEMA_FILE],
-            # Gets the metadata schema version from the file, and stamps it
-            # into the srcs file.
-            cmd = "version=$$(sed -n -e '/Schema Semantic version/ s/.*\\: *//p' $(location %s));" %
-                  METADATA_SCHEMA_FILE +
-                  'sed "s/{LATEST_METADATA_PARSER_VERSION}/$$version/" $< > $@',
-        )
-
-    native.filegroup(
-        name = name,
-        srcs = outs,
-    )
diff --git a/tensorflow/lite/experimental/support/metadata/cc/BUILD b/tensorflow/lite/experimental/support/metadata/cc/BUILD
deleted file mode 100644
index 8febc7a2237..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/BUILD
+++ /dev/null
@@ -1,29 +0,0 @@
-load("//tensorflow/lite/experimental/support/metadata:build_defs.bzl", "stamp_metadata_parser_version")
-
-package(
-    default_visibility = ["//tensorflow/lite/experimental/support:users"],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-stamp_metadata_parser_version(
-    name = "metadata_parser_h",
-    srcs = ["metadata_parser.h.template"],
-    outs = ["metadata_parser.h"],
-)
-
-cc_library(
-    name = "metadata_version",
-    srcs = ["metadata_version.cc"],
-    hdrs = [
-        "metadata_version.h",
-        ":metadata_parser_h",
-    ],
-    deps = [
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/experimental/support/metadata:metadata_schema_cc",
-        "//tensorflow/lite/kernels/internal:compatibility",
-        "//tensorflow/lite/tools:logging",
-        "@com_google_absl//absl/strings",
-        "@flatbuffers",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template b/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
deleted file mode 100644
index dfb62d0de81..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h.template
+++ /dev/null
@@ -1,28 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
-#define TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
-
-namespace tflite {
-namespace metadata {
-
-// The version of the metadata parser that this metadata versioning library is
-// depending on.
-inline constexpr char kMatadataParserVersion[] = "{LATEST_METADATA_PARSER_VERSION}";
-
-}  // namespace metadata
-}  // namespace tflite
-
-#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_PARSER_H_
diff --git a/tensorflow/lite/experimental/support/metadata/cc/metadata_version.cc b/tensorflow/lite/experimental/support/metadata/cc/metadata_version.cc
deleted file mode 100644
index 971465f7747..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/metadata_version.cc
+++ /dev/null
@@ -1,214 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include "tensorflow/lite/experimental/support/metadata/cc/metadata_version.h"
-
-#include 
-#include 
-
-#include 
-#include 
-#include 
-#include 
-
-#include "absl/strings/str_join.h"
-#include "absl/strings/str_split.h"
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/c/common.h"
-#include "tensorflow/lite/experimental/support/metadata/metadata_schema_generated.h"
-#include "tensorflow/lite/kernels/internal/compatibility.h"
-#include "tensorflow/lite/tools/logging.h"
-
-namespace tflite {
-namespace metadata {
-namespace {
-
-// Members that are added to the metadata schema after the initial version
-// of 1.0.0.
-enum class SchemaMembers {
-  kAssociatedFileTypeVocabulary = 0,
-};
-
-// Helper class to compare semantic versions in terms of three integers, major,
-// minor, and patch.
-class Version {
- public:
-  explicit Version(int major, int minor = 0, int patch = 0)
-      : version_({major, minor, patch}) {}
-
-  explicit Version(const std::string& version) {
-    const std::vector vec = absl::StrSplit(version, '.');
-    // The version string should always be less than four numbers.
-    TFLITE_DCHECK(vec.size() <= kElementNumber && !vec.empty());
-    version_[0] = std::stoi(vec[0]);
-    version_[1] = vec.size() > 1 ? std::stoi(vec[1]) : 0;
-    version_[2] = vec.size() > 2 ? std::stoi(vec[2]) : 0;
-  }
-
-  // Compares two semantic version numbers.
-  //
-  // Example results when comparing two versions strings:
-  //   "1.9" precedes "1.14";
-  //   "1.14" precedes "1.14.1";
-  //   "1.14" and "1.14.0" are equal.
-  //
-  // Returns the value 0 if the two versions are equal; a value less than 0 if
-  // *this precedes v; a value greater than 0 if v precedes *this.
-  int Compare(const Version& v) {
-    for (int i = 0; i < kElementNumber; ++i) {
-      if (version_[i] != v.version_[i]) {
-        return version_[i] < v.version_[i] ? -1 : 1;
-      }
-    }
-    return 0;
-  }
-
-  // Converts version_ into a version string.
-  std::string ToString() { return absl::StrJoin(version_, "."); }
-
- private:
-  static constexpr int kElementNumber = 3;
-  std::array version_;
-};
-
-Version GetMemberVersion(SchemaMembers member) {
-  switch (member) {
-    case SchemaMembers::kAssociatedFileTypeVocabulary:
-      return Version(1, 0, 1);
-    default:
-      TFLITE_LOG(FATAL) << "Unsupported schema member: "
-                        << static_cast(member);
-  }
-}
-
-// Updates min_version if it precedes the new_version.
-inline void UpdateMinimumVersion(const Version& new_version,
-                                 Version* min_version) {
-  if (min_version->Compare(new_version) < 0) {
-    *min_version = new_version;
-  }
-}
-
-void UpdateMinimumVersionForAssociatedFile(
-    const tflite::AssociatedFile* associated_file, Version* min_version) {
-  if (associated_file == nullptr) return;
-
-  if (associated_file->type() == AssociatedFileType_VOCABULARY) {
-    UpdateMinimumVersion(
-        GetMemberVersion(SchemaMembers::kAssociatedFileTypeVocabulary),
-        min_version);
-  }
-}
-
-void UpdateMinimumVersionForAssociatedFileArray(
-    const flatbuffers::Vector>*
-        associated_files,
-    Version* min_version) {
-  if (associated_files == nullptr) return;
-
-  for (int i = 0; i < associated_files->size(); ++i) {
-    UpdateMinimumVersionForAssociatedFile(associated_files->Get(i),
-                                          min_version);
-  }
-}
-
-void UpdateMinimumVersionForTensorMetadata(
-    const tflite::TensorMetadata* tensor_metadata, Version* min_version) {
-  if (tensor_metadata == nullptr) return;
-
-  // Checks the associated_files field.
-  UpdateMinimumVersionForAssociatedFileArray(
-      tensor_metadata->associated_files(), min_version);
-}
-
-void UpdateMinimumVersionForTensorMetadataArray(
-    const flatbuffers::Vector>*
-        tensor_metadata_array,
-    Version* min_version) {
-  if (tensor_metadata_array == nullptr) return;
-
-  for (int i = 0; i < tensor_metadata_array->size(); ++i) {
-    UpdateMinimumVersionForTensorMetadata(tensor_metadata_array->Get(i),
-                                          min_version);
-  }
-}
-
-void UpdateMinimumVersionForSubGraphMetadata(
-    const tflite::SubGraphMetadata* subgraph_metadata, Version* min_version) {
-  if (subgraph_metadata == nullptr) return;
-
-  // Checks in the input/output metadata arrays.
-  UpdateMinimumVersionForTensorMetadataArray(
-      subgraph_metadata->input_tensor_metadata(), min_version);
-  UpdateMinimumVersionForTensorMetadataArray(
-      subgraph_metadata->output_tensor_metadata(), min_version);
-
-  // Checks the associated_files field.
-  UpdateMinimumVersionForAssociatedFileArray(
-      subgraph_metadata->associated_files(), min_version);
-}
-
-void UpdateMinimumVersionForModelMetadata(
-    const tflite::ModelMetadata& model_metadata, Version* min_version) {
-  // Checks the subgraph_metadata field.
-  if (model_metadata.subgraph_metadata() != nullptr) {
-    for (int i = 0; i < model_metadata.subgraph_metadata()->size(); ++i) {
-      UpdateMinimumVersionForSubGraphMetadata(
-          model_metadata.subgraph_metadata()->Get(i), min_version);
-    }
-  }
-
-  // Checks the associated_files field.
-  UpdateMinimumVersionForAssociatedFileArray(model_metadata.associated_files(),
-                                             min_version);
-}
-
-}  // namespace
-
-TfLiteStatus GetMinimumMetadataParserVersion(const uint8_t* buffer_data,
-                                             size_t buffer_size,
-                                             std::string* min_version_str) {
-  flatbuffers::Verifier verifier =
-      flatbuffers::Verifier(buffer_data, buffer_size);
-  if (!tflite::VerifyModelMetadataBuffer(verifier)) {
-    TFLITE_LOG(ERROR) << "The model metadata is not a valid FlatBuffer buffer.";
-    return kTfLiteError;
-  }
-
-  static constexpr char kDefaultVersion[] = "1.0.0";
-  Version min_version = Version(kDefaultVersion);
-
-  // Checks if any member declared after 1.0.0 (such as those in
-  // SchemaMembers) exists, and updates min_version accordingly. The minimum
-  // metadata parser version will be the largest version number of all fields
-  // that has been added to a metadata flatbuffer
-  const tflite::ModelMetadata* model_metadata = GetModelMetadata(buffer_data);
-
-  // All tables in the metadata schema should have their dedicated
-  // UpdateMinimumVersionFor**() methods, respectively. We'll gradually add
-  // these methods when new fields show up in later schema versions.
-  //
-  // UpdateMinimumVersionFor() takes a const pointer of Foo. The pointer
-  // can be a nullptr if Foo is not populated into the corresponding table of
-  // the Flatbuffer object. In this case, UpdateMinimumVersionFor() will be
-  // skipped. An exception is UpdateMinimumVersionForModelMetadata(), where
-  // ModelMetadata is the root table, and it won't be null.
-  UpdateMinimumVersionForModelMetadata(*model_metadata, &min_version);
-
-  *min_version_str = min_version.ToString();
-  return kTfLiteOk;
-}
-
-}  // namespace metadata
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/cc/metadata_version.h b/tensorflow/lite/experimental/support/metadata/cc/metadata_version.h
deleted file mode 100644
index c4127118bc7..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/metadata_version.h
+++ /dev/null
@@ -1,38 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_VERSION_H_
-#define TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_VERSION_H_
-
-#include 
-#include 
-
-#include 
-
-#include "tensorflow/lite/c/common.h"
-
-namespace tflite {
-namespace metadata {
-
-// Gets the minimum metadata parser version that can fully understand all fields
-// in a given metadata flatbuffer. TFLite Metadata follows Semantic Versioning
-// 2.0. Each release version has the form MAJOR.MINOR.PATCH.
-TfLiteStatus GetMinimumMetadataParserVersion(const uint8_t* buffer_data,
-                                             size_t buffer_size,
-                                             std::string* min_version);
-
-}  // namespace metadata
-}  // namespace tflite
-
-#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SUPPORT_METADATA_CC_METADATA_VERSION_H_
diff --git a/tensorflow/lite/experimental/support/metadata/cc/python/BUILD b/tensorflow/lite/experimental/support/metadata/cc/python/BUILD
deleted file mode 100644
index 4128f0ac9d1..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/python/BUILD
+++ /dev/null
@@ -1,22 +0,0 @@
-load("//tensorflow:tensorflow.bzl", "pybind_extension")
-
-package(
-    default_visibility = [
-        "//tensorflow/lite/experimental/support/metadata:__pkg__",
-    ],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-pybind_extension(
-    name = "_pywrap_metadata_version",
-    srcs = [
-        "metadata_version.cc",
-    ],
-    features = ["-use_header_modules"],
-    module_name = "_pywrap_metadata_version",
-    deps = [
-        "//tensorflow/lite/c:common",
-        "//tensorflow/lite/experimental/support/metadata/cc:metadata_version",
-        "@pybind11",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/cc/python/metadata_version.cc b/tensorflow/lite/experimental/support/metadata/cc/python/metadata_version.cc
deleted file mode 100644
index 7d1f9d1e122..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/python/metadata_version.cc
+++ /dev/null
@@ -1,55 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "tensorflow/lite/experimental/support/metadata/cc/metadata_version.h"
-
-#include "pybind11/pybind11.h"
-#include "tensorflow/lite/c/common.h"
-
-namespace tflite {
-namespace metadata {
-
-PYBIND11_MODULE(_pywrap_metadata_version, m) {
-  m.doc() = R"pbdoc(
-    _pywrap_metadata_version
-    A module that returns the minimum metadata parser version of a given
-    metadata flatbuffer.
-  )pbdoc";
-
-  // Using pybind11 type conversions to convert between Python and native
-  // C++ types. There are other options to provide access to native Python types
-  // in C++ and vice versa. See the pybind 11 instrcution [1] for more details.
-  // Type converstions is recommended by pybind11, though the main downside
-  // is that a copy of the data must be made on every Python to C++ transition:
-  // this is needed since the C++ and Python versions of the same type generally
-  // won’t have the same memory layout.
-  //
-  // [1]: https://pybind11.readthedocs.io/en/stable/advanced/cast/index.html
-  m.def("GetMinimumMetadataParserVersion",
-        [](const std::string& buffer_data) -> std::string {
-          std::string min_version;
-          if (GetMinimumMetadataParserVersion(
-                  reinterpret_cast(buffer_data.c_str()),
-                  buffer_data.length(), &min_version) != kTfLiteOk) {
-            pybind11::value_error(
-                "Error occurred when getting the minimum metadata parser "
-                "version of the metadata flatbuffer.");
-          }
-          return min_version;
-        });
-}
-
-}  // namespace metadata
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/BUILD b/tensorflow/lite/experimental/support/metadata/cc/test/BUILD
deleted file mode 100644
index f9d78567d70..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/test/BUILD
+++ /dev/null
@@ -1,24 +0,0 @@
-package(
-    default_visibility = ["//visibility:public"],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-cc_test(
-    name = "metadata_version_test",
-    srcs = ["metadata_version_test.cc"],
-    deps = [
-        "//tensorflow/lite/experimental/support/metadata:metadata_schema_cc",
-        "//tensorflow/lite/experimental/support/metadata/cc:metadata_version",
-        "@com_google_googletest//:gtest_main",
-        "@flatbuffers",
-    ],
-)
-
-cc_test(
-    name = "metadata_parser_test",
-    srcs = ["metadata_parser_test.cc"],
-    deps = [
-        "//tensorflow/lite/experimental/support/metadata/cc:metadata_version",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
deleted file mode 100644
index af7b8791fe8..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_parser_test.cc
+++ /dev/null
@@ -1,33 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include "tensorflow/lite/experimental/support/metadata/cc/metadata_parser.h"
-
-#include 
-#include 
-
-namespace tflite {
-namespace metadata {
-namespace {
-
-using ::testing::MatchesRegex;
-
-TEST(MetadataParserTest, MatadataParserVersionIsWellFormed) {
-  // Validates that the version is well-formed (x.y.z).
-  EXPECT_THAT(kMatadataParserVersion, MatchesRegex("[0-9]+\\.[0-9]+\\.[0-9]+"));
-}
-
-}  // namespace
-}  // namespace metadata
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc b/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc
deleted file mode 100644
index 03f4d3bf28b..00000000000
--- a/tensorflow/lite/experimental/support/metadata/cc/test/metadata_version_test.cc
+++ /dev/null
@@ -1,187 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-#include "tensorflow/lite/experimental/support/metadata/cc/metadata_version.h"
-
-#include 
-
-#include 
-#include 
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
-#include "tensorflow/lite/experimental/support/metadata/metadata_schema_generated.h"
-
-namespace tflite {
-namespace metadata {
-namespace {
-
-using ::testing::MatchesRegex;
-using ::testing::StrEq;
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionSucceedsWithValidMetadata) {
-  // Creates a dummy metadata flatbuffer for test.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  auto name = builder.CreateString("Foo");
-  ModelMetadataBuilder metadata_builder(builder);
-  metadata_builder.add_name(name);
-  auto metadata = metadata_builder.Finish();
-  FinishModelMetadataBuffer(builder, metadata);
-
-  // Gets the mimimum metadata parser version.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteOk);
-  // Validates that the version is well-formed (x.y.z).
-  EXPECT_THAT(min_version, MatchesRegex("[0-9]+\\.[0-9]+\\.[0-9]+"));
-}
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionFailsWithInvalidIdentifier) {
-  // Creates a dummy metadata flatbuffer without identifier.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  ModelMetadataBuilder metadata_builder(builder);
-  auto metadata = metadata_builder.Finish();
-  builder.Finish(metadata);
-
-  // Gets the mimimum metadata parser version and triggers error.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteError);
-  EXPECT_TRUE(min_version.empty());
-}
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionForModelMetadataVocabAssociatedFiles) {
-  // Creates a metadata flatbuffer with the field,
-  // ModelMetadata.associated_fiels, populated with the vocabulary file type.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  AssociatedFileBuilder associated_file_builder(builder);
-  associated_file_builder.add_type(tflite::AssociatedFileType_VOCABULARY);
-  auto associated_files =
-      builder.CreateVector(std::vector>{
-          associated_file_builder.Finish()});
-  ModelMetadataBuilder metadata_builder(builder);
-  metadata_builder.add_associated_files(associated_files);
-  FinishModelMetadataBuffer(builder, metadata_builder.Finish());
-
-  // Gets the mimimum metadata parser version.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteOk);
-  // Validates that the version is exactly 1.0.1.
-  EXPECT_THAT(min_version, StrEq("1.0.1"));
-}
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionForSubGraphMetadataVocabAssociatedFiles) {
-  // Creates a metadata flatbuffer with the field,
-  // SubGraphMetadata.associated_fiels, populated with the vocabulary file type.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  AssociatedFileBuilder associated_file_builder(builder);
-  associated_file_builder.add_type(tflite::AssociatedFileType_VOCABULARY);
-  auto associated_files =
-      builder.CreateVector(std::vector>{
-          associated_file_builder.Finish()});
-  SubGraphMetadataBuilder subgraph_builder(builder);
-  subgraph_builder.add_associated_files(associated_files);
-  auto subgraphs =
-      builder.CreateVector(std::vector>{
-          subgraph_builder.Finish()});
-  ModelMetadataBuilder metadata_builder(builder);
-  metadata_builder.add_subgraph_metadata(subgraphs);
-  FinishModelMetadataBuffer(builder, metadata_builder.Finish());
-
-  // Gets the mimimum metadata parser version.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteOk);
-  // Validates that the version is exactly 1.0.1.
-  EXPECT_THAT(min_version, StrEq("1.0.1"));
-}
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionForInputMetadataVocabAssociatedFiles) {
-  // Creates a metadata flatbuffer with the field,
-  // SubGraphMetadata.input_tensor_metadata.associated_fiels, populated with the
-  // vocabulary file type.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  AssociatedFileBuilder associated_file_builder(builder);
-  associated_file_builder.add_type(tflite::AssociatedFileType_VOCABULARY);
-  auto associated_files =
-      builder.CreateVector(std::vector>{
-          associated_file_builder.Finish()});
-  TensorMetadataBuilder tensor_builder(builder);
-  tensor_builder.add_associated_files(associated_files);
-  auto tensors =
-      builder.CreateVector(std::vector>{
-          tensor_builder.Finish()});
-  SubGraphMetadataBuilder subgraph_builder(builder);
-  subgraph_builder.add_input_tensor_metadata(tensors);
-  auto subgraphs =
-      builder.CreateVector(std::vector>{
-          subgraph_builder.Finish()});
-  ModelMetadataBuilder metadata_builder(builder);
-  metadata_builder.add_subgraph_metadata(subgraphs);
-  FinishModelMetadataBuffer(builder, metadata_builder.Finish());
-
-  // Gets the mimimum metadata parser version.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteOk);
-  // Validates that the version is exactly 1.0.1.
-  EXPECT_THAT(min_version, StrEq("1.0.1"));
-}
-
-TEST(MetadataVersionTest,
-     GetMinimumMetadataParserVersionForOutputMetadataVocabAssociatedFiles) {
-  // Creates a metadata flatbuffer with the field,
-  // SubGraphMetadata.output_tensor_metadata.associated_fiels, populated with
-  // the vocabulary file type.
-  flatbuffers::FlatBufferBuilder builder(1024);
-  AssociatedFileBuilder associated_file_builder(builder);
-  associated_file_builder.add_type(tflite::AssociatedFileType_VOCABULARY);
-  auto associated_files =
-      builder.CreateVector(std::vector>{
-          associated_file_builder.Finish()});
-  TensorMetadataBuilder tensor_builder(builder);
-  tensor_builder.add_associated_files(associated_files);
-  auto tensors =
-      builder.CreateVector(std::vector>{
-          tensor_builder.Finish()});
-  SubGraphMetadataBuilder subgraph_builder(builder);
-  subgraph_builder.add_output_tensor_metadata(tensors);
-  auto subgraphs =
-      builder.CreateVector(std::vector>{
-          subgraph_builder.Finish()});
-  ModelMetadataBuilder metadata_builder(builder);
-  metadata_builder.add_subgraph_metadata(subgraphs);
-  FinishModelMetadataBuffer(builder, metadata_builder.Finish());
-
-  // Gets the mimimum metadata parser version.
-  std::string min_version;
-  EXPECT_EQ(GetMinimumMetadataParserVersion(builder.GetBufferPointer(),
-                                            builder.GetSize(), &min_version),
-            kTfLiteOk);
-  // Validates that the version is exactly 1.0.1.
-  EXPECT_EQ(min_version, "1.0.1");
-}
-
-}  // namespace
-}  // namespace metadata
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/BUILD b/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/BUILD
deleted file mode 100644
index ca9a79b7451..00000000000
--- a/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/BUILD
+++ /dev/null
@@ -1,23 +0,0 @@
-load("//tensorflow:tensorflow.bzl", "pybind_extension")
-
-package(
-    default_visibility = [
-        "//visibility:public",
-    ],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-pybind_extension(
-    name = "_pywrap_flatbuffers",
-    srcs = [
-        "flatbuffers_lib.cc",
-    ],
-    features = ["-use_header_modules"],
-    module_name = "_pywrap_flatbuffers",
-    deps = [
-        "//tensorflow/python:pybind11_lib",
-        "//third_party/python_runtime:headers",
-        "@flatbuffers",
-        "@pybind11",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/flatbuffers_lib.cc b/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/flatbuffers_lib.cc
deleted file mode 100644
index 6185722504f..00000000000
--- a/tensorflow/lite/experimental/support/metadata/flatbuffers_lib/flatbuffers_lib.cc
+++ /dev/null
@@ -1,59 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
-#include "flatbuffers/idl.h"  // from @flatbuffers
-#include "pybind11/pybind11.h"
-#include "pybind11/pytypes.h"
-#include "pybind11/stl.h"
-
-namespace tflite {
-namespace support {
-
-PYBIND11_MODULE(_pywrap_flatbuffers, m) {
-  pybind11::class_(m, "IDLOptions")
-      .def(pybind11::init<>())
-      .def_readwrite("strict_json", &flatbuffers::IDLOptions::strict_json);
-  pybind11::class_(m, "Parser")
-      .def(pybind11::init())
-      .def("parse",
-           [](flatbuffers::Parser* self, const std::string& source) {
-             return self->Parse(source.c_str());
-           })
-      .def_readonly("builder", &flatbuffers::Parser::builder_)
-      .def_readonly("error", &flatbuffers::Parser::error_);
-  pybind11::class_(m, "FlatBufferBuilder")
-      .def("clear", &flatbuffers::FlatBufferBuilder::Clear)
-      .def("push_flat_buffer", [](flatbuffers::FlatBufferBuilder* self,
-                                  const std::string& contents) {
-        self->PushFlatBuffer(reinterpret_cast(contents.c_str()),
-                             contents.length());
-      });
-  m.def("generate_text_file", &flatbuffers::GenerateTextFile);
-  m.def(
-      "generate_text",
-      [](const flatbuffers::Parser& parser,
-         const std::string& buffer) -> std::string {
-        std::string text;
-        if (!flatbuffers::GenerateText(
-                parser, reinterpret_cast(buffer.c_str()), &text)) {
-          return "";
-        }
-        return text;
-      });
-}
-
-}  // namespace support
-}  // namespace tflite
diff --git a/tensorflow/lite/experimental/support/metadata/java/AndroidManifest.xml b/tensorflow/lite/experimental/support/metadata/java/AndroidManifest.xml
deleted file mode 100644
index b2e22628db6..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/AndroidManifest.xml
+++ /dev/null
@@ -1,6 +0,0 @@
-
-
-    
-
-
diff --git a/tensorflow/lite/experimental/support/metadata/java/BUILD b/tensorflow/lite/experimental/support/metadata/java/BUILD
deleted file mode 100644
index 00d10bcca56..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/BUILD
+++ /dev/null
@@ -1,40 +0,0 @@
-# Description:
-# TensorFlow Lite Support API in Java for metadata.
-
-load("@build_bazel_rules_android//android:rules.bzl", "android_library")
-load("//tensorflow/java:build_defs.bzl", "JAVACOPTS")
-
-package(
-    default_visibility = ["//visibility:public"],
-    licenses = ["notice"],  # Apache 2.0
-)
-
-METADATA_SRCS = glob(
-    ["src/java/org/tensorflow/lite/support/metadata/**/*.java"],
-)
-
-android_library(
-    name = "tensorflow-lite-support-metadata",
-    srcs = METADATA_SRCS,
-    manifest = "AndroidManifest.xml",
-    deps = [
-        "//tensorflow/lite/experimental/support/metadata:metadata_schema_fbs_android",
-        "//tensorflow/lite/experimental/support/metadata:schema_fbs_android",
-        "@org_checkerframework_qual",
-    ],
-)
-
-java_library(
-    name = "tensorflow-lite-support-metadata-lib",
-    srcs = METADATA_SRCS,
-    javacopts = JAVACOPTS,
-    resource_jars = [
-        "//tensorflow/lite/experimental/support/metadata:libmetadata_schema_java.jar",
-        "//tensorflow/lite/experimental/support/metadata:libschema_fbs_java.jar",
-    ],
-    deps = [
-        "//tensorflow/lite/experimental/support/metadata:metadata_schema_java",
-        "//tensorflow/lite/experimental/support/metadata:schema_fbs_java",
-        "@org_checkerframework_qual",
-    ],
-)
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/BoundedInputStream.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/BoundedInputStream.java
deleted file mode 100644
index 6c3d23270f3..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/BoundedInputStream.java
+++ /dev/null
@@ -1,116 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkElementIndex;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkNotNull;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.ByteBuffer;
-
-/**
- * An {@link InputStream} that wraps a section of a {@link SeekableByteChannelCompat}.
- *
- * 
WARNING: Similar as {@link InputStream}, instances of an {@link BoundedInputStream} are
- * not thread-safe. If multiple threads concurrently reading from the same {@link
- * BoundedInputStream}, it must be synchronized externally. Also, if multiple instances of {@link
- * BoundedInputStream} are created on the same {@link SeekableByteChannelCompat}, it must be
- * synchronized as well.
- */
-final class BoundedInputStream extends InputStream {
-  private final ByteBuffer singleByteBuffer = ByteBuffer.allocate(1);
-  private final long end; // The valid data for the stream is between [start, end).
-  private long position;
-  private final SeekableByteChannelCompat channel;
-
-  /**
-   * Creates a {@link BoundedInputStream} with a {@link SeekableByteChannelCompat}.
-   *
-   * @param channel the {@link SeekableByteChannelCompat} that backs up this {@link
-   *     BoundedInputStream}
-   * @param start the starting position of this {@link BoundedInputStream} in the given {@link
-   *     SeekableByteChannelCompat}
-   * @param remaining the length of this {@link BoundedInputStream}
-   * @throws IllegalArgumentException if {@code start} or {@code remaining} is negative
-   */
-  BoundedInputStream(SeekableByteChannelCompat channel, long start, long remaining) {
-    checkArgument(
-        remaining >= 0 && start >= 0,
-        String.format("Invalid length of stream at offset=%d, length=%d", start, remaining));
-
-    end = start + remaining;
-    this.channel = channel;
-    position = start;
-  }
-
-  @Override
-  public int available() throws IOException {
-    return (int) (Math.min(end, channel.size()) - position);
-  }
-
-  @Override
-  public int read() throws IOException {
-    if (position >= end) {
-      return -1;
-    }
-
-    singleByteBuffer.rewind();
-    int count = read(position, singleByteBuffer);
-    if (count < 0) {
-      return count;
-    }
-
-    position++;
-    return singleByteBuffer.get() & 0xff;
-  }
-
-  @Override
-  public int read(byte[] b, int off, int len) throws IOException {
-    checkNotNull(b);
-    checkElementIndex(off, b.length, "The start offset");
-    checkElementIndex(len, b.length - off + 1, "The maximumn number of bytes to read");
-
-    if (len == 0) {
-      return 0;
-    }
-
-    if (len > end - position) {
-      if (position >= end) {
-        return -1;
-      }
-      len = (int) (end - position);
-    }
-
-    ByteBuffer buf = ByteBuffer.wrap(b, off, len);
-    int count = read(position, buf);
-    if (count > 0) {
-      position += count;
-    }
-    return count;
-  }
-
-  private int read(long position, ByteBuffer buf) throws IOException {
-    int count;
-    synchronized (channel) {
-      channel.position(position);
-      count = channel.read(buf);
-    }
-    buf.flip();
-    return count;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ByteBufferChannel.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ByteBufferChannel.java
deleted file mode 100644
index e5d54a415ed..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ByteBufferChannel.java
+++ /dev/null
@@ -1,130 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static java.lang.Math.min;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkNotNull;
-
-import java.nio.ByteBuffer;
-import java.nio.channels.NonWritableChannelException;
-
-/** Implements the {@link SeekableByteChannelCompat} on top of {@link ByteBuffer}. */
-final class ByteBufferChannel implements SeekableByteChannelCompat {
-
-  /** The ByteBuffer that holds the data. */
-  private final ByteBuffer buffer;
-
-  /**
-   * Creates a {@link ByteBufferChannel} that wraps a {@link ByteBuffer}.
-   *
-   * @param buffer the {@link ByteBuffer} that backs this {@link ByteBufferChannel}
-   * @throws NullPointerException if {@code buffer} is null
-   */
-  public ByteBufferChannel(ByteBuffer buffer) {
-    checkNotNull(buffer, "The ByteBuffer cannot be null.");
-    this.buffer = buffer;
-  }
-
-  @Override
-  public void close() {}
-
-  @Override
-  public boolean isOpen() {
-    return true;
-  }
-
-  @Override
-  public long position() {
-    return buffer.position();
-  }
-
-  /**
-   * Sets this channel's position.
-   *
-   * @param newPosition the new position, a non-negative integer counting the number of bytes from
-   *     the beginning of the entity
-   * @return this channel
-   * @throws IllegalArgumentException if the new position is negative, or greater than the size of
-   *     the underlying {@link ByteBuffer}, or greater than Integer.MAX_VALUE
-   */
-  @Override
-  public synchronized ByteBufferChannel position(long newPosition) {
-    checkArgument(
-        (newPosition >= 0 && newPosition <= Integer.MAX_VALUE),
-        "The new position should be non-negative and be less than Integer.MAX_VALUE.");
-    buffer.position((int) newPosition);
-    return this;
-  }
-
-  /**
-   * {@inheritDoc}
-   *
-   * 
Bytes are read starting at this channel's current position, and then the position is updated
-   * with the number of bytes actually read. Otherwise this method behaves exactly as specified in
-   * the {@link ReadableByteChannel} interface.
-   */
-  @Override
-  public synchronized int read(ByteBuffer dst) {
-    if (buffer.remaining() == 0) {
-      return -1;
-    }
-
-    int count = min(dst.remaining(), buffer.remaining());
-    if (count > 0) {
-      ByteBuffer tempBuffer = buffer.slice();
-      tempBuffer.order(buffer.order()).limit(count);
-      dst.put(tempBuffer);
-      buffer.position(buffer.position() + count);
-    }
-    return count;
-  }
-
-  @Override
-  public long size() {
-    return buffer.limit();
-  }
-
-  @Override
-  public synchronized ByteBufferChannel truncate(long size) {
-    checkArgument(
-        (size >= 0 && size <= Integer.MAX_VALUE),
-        "The new size should be non-negative and be less than Integer.MAX_VALUE.");
-
-    if (size < buffer.limit()) {
-      buffer.limit((int) size);
-      if (buffer.position() > size) {
-        buffer.position((int) size);
-      }
-    }
-    return this;
-  }
-
-  @Override
-  public synchronized int write(ByteBuffer src) {
-    if (buffer.isReadOnly()) {
-      throw new NonWritableChannelException();
-    }
-
-    int count = min(src.remaining(), buffer.remaining());
-    if (count > 0) {
-      ByteBuffer tempBuffer = src.slice();
-      tempBuffer.order(buffer.order()).limit(count);
-      buffer.put(tempBuffer);
-    }
-    return count;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java
deleted file mode 100644
index 9bf5ae93138..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataExtractor.java
+++ /dev/null
@@ -1,368 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.ByteBuffer;
-import java.util.zip.ZipException;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.tensorflow.lite.schema.Tensor;
-import org.tensorflow.lite.support.metadata.schema.ModelMetadata;
-import org.tensorflow.lite.support.metadata.schema.TensorMetadata;
-
-/**
- * Loads metadata from TFLite Model FlatBuffer.
- *
- * 
TFLite Model FlatBuffer can be generated using the TFLite
- * Model schema file.
- *
- * 
Some models contain a TFLite Metadata Flatbuffer, which records more information about what
- * the model does and how to interprete the model. TFLite Metadata Flatbuffer can be generated using
- * the TFLite
- * Metadata schema file.
- *
- * 
It is allowed to pass in a model FlatBuffer without TFLite metadata. However, invoking methods
- * that read from TFLite metadata will cause runtime errors.
- *
- * 
Similarly, it is allowed to pass in a model FlatBuffer without associated files. However,
- * invoking methods that read the associated files will cause runtime errors.
- *
- * 
Though TFLite model FlatBuffer supports multiple subgraphs, TFLite Interpreter only supports a
- * single subgraph so far. See the instruction
- * of how to specify subgraph during convertion for more information. Therefore, {@link
- * MetadataExtractor} omits subgraph index as an input in its methods.
- */
-public class MetadataExtractor {
-
-  /** The helper class to load metadata from TFLite model FlatBuffer. */
-  private final ModelInfo modelInfo;
-
-  /** The helper class to load metadata from TFLite metadata FlatBuffer. */
-  @Nullable private final ModelMetadataInfo metadataInfo;
-
-  /** The handler to load associated files through zip. */
-  @Nullable private final ZipFile zipFile;
-
-  /**
-   * Creates a {@link MetadataExtractor} with TFLite model FlatBuffer.
-   *
-   * @param buffer the TFLite model FlatBuffer
-   * @throws IllegalArgumentException if the number of input or output tensors in the model does not
-   *     match that in the metadata
-   * @throws IOException if an error occurs while reading the model as a Zip file
-   */
-  public MetadataExtractor(ByteBuffer buffer) throws IOException {
-    modelInfo = new ModelInfo(buffer);
-    ByteBuffer metadataBuffer = modelInfo.getMetadataBuffer();
-    if (metadataBuffer != null) {
-      metadataInfo = new ModelMetadataInfo(metadataBuffer);
-
-      // Prints warning message if the minimum parser version is not satisfied.
-      if (!isMinimumParserVersionSatisfied()) {
-        System.err.printf(
-            " Some fields in the metadata belong to a future schema. The minimum parser"
-                + " version required is %s, but the version of the current metadata parser is %s",
-            metadataInfo.getMininumParserVersion(), MetadataParser.VERSION);
-      }
-
-      checkArgument(
-          modelInfo.getInputTensorCount() == metadataInfo.getInputTensorCount(),
-          String.format(
-              "The number of input tensors in the model is %d. The number of input tensors that"
-                  + " recorded in the metadata is %d. These two values does not match.",
-              modelInfo.getInputTensorCount(), metadataInfo.getInputTensorCount()));
-      checkArgument(
-          modelInfo.getOutputTensorCount() == metadataInfo.getOutputTensorCount(),
-          String.format(
-              "The number of output tensors in the model is %d. The number of output tensors that"
-                  + " recorded in the metadata is %d. These two values does not match.",
-              modelInfo.getOutputTensorCount(), metadataInfo.getOutputTensorCount()));
-    } else {
-      // It is allowed to pass in a model FlatBuffer without TFLite metadata. However, invoking
-      // methods that read from TFLite metadata will cause runtime errors.
-      metadataInfo = null;
-    }
-
-    zipFile = createZipFile(buffer);
-  }
-
-  /**
-   * Quantization parameters that corresponds to the table, {@code QuantizationParameters}, in the
-   * TFLite
-   * Model schema file.
-   *
-   * 
Since per-channel quantization does not apply to input and output tensors, {@code scale} and
-   * {@code zero_point} are both single values instead of arrays.
-   *
-   * 
For tensor that are not quantized, the values of scale and zero_point are both 0.
-   *
-   * 
Given a quantized value q, the corresponding float value f should be: 

-   * f = scale * (q - zero_point) 

-   */
-  public static class QuantizationParams {
-    /** The scale value used in quantization. */
-    private final float scale;
-    /** The zero point value used in quantization. */
-    private final int zeroPoint;
-
-    /**
-     * Creates a {@link QuantizationParams} with {@code scale} and {@code zero_point}.
-     *
-     * @param scale The scale value used in quantization.
-     * @param zeroPoint The zero point value used in quantization.
-     */
-    public QuantizationParams(final float scale, final int zeroPoint) {
-      this.scale = scale;
-      this.zeroPoint = zeroPoint;
-    }
-
-    /** Returns the scale value. */
-    public float getScale() {
-      return scale;
-    }
-
-    /** Returns the zero point value. */
-    public int getZeroPoint() {
-      return zeroPoint;
-    }
-  }
-
-  /** Returns {@code true} if the model has metadata. Otherwise, returns {@code false}. */
-  public boolean hasMetadata() {
-    return metadataInfo != null;
-  }
-
-  /**
-   * Gets the packed associated file with the specified {@code fileName}.
-   *
-   * @param fileName the name of the associated file
-   * @return the raw input stream containing specified file
-   * @throws IllegalStateException if the model is not a zip file
-   * @throws IllegalArgumentException if the specified file does not exist in the model
-   */
-  public InputStream getAssociatedFile(String fileName) {
-    assertZipFile();
-    return zipFile.getRawInputStream(fileName);
-  }
-
-  /** Gets the count of input tensors in the model. */
-  public int getInputTensorCount() {
-    return modelInfo.getInputTensorCount();
-  }
-
-  /**
-   * Gets the metadata for the input tensor specified by {@code inputIndex}.
-   *
-   * @param inputIndex the index of the desired input tensor
-   * @throws IllegalStateException if this model does not contain model metadata
-   */
-  @Nullable
-  public TensorMetadata getInputTensorMetadata(int inputIndex) {
-    assertMetadataInfo();
-    return metadataInfo.getInputTensorMetadata(inputIndex);
-  }
-
-  /**
-   * Gets the quantization parameters for the input tensor specified by {@code inputIndex}.
-   *
-   * @param inputIndex the index of the desired input tensor
-   */
-  public QuantizationParams getInputTensorQuantizationParams(int inputIndex) {
-    Tensor tensor = modelInfo.getInputTensor(inputIndex);
-    return modelInfo.getQuantizationParams(tensor);
-  }
-
-  /**
-   * Gets the shape of the input tensor with {@code inputIndex}.
-   *
-   * @param inputIndex the index of the desired input tensor
-   */
-  public int[] getInputTensorShape(int inputIndex) {
-    return modelInfo.getInputTensorShape(inputIndex);
-  }
-
-  /**
-   * Gets the {@link TensorType} of the input tensor with {@code inputIndex}.
-   *
-   * @param inputIndex the index of the desired input tensor
-   */
-  public byte getInputTensorType(int inputIndex) {
-    return modelInfo.getInputTensorType(inputIndex);
-  }
-
-  /**
-   * Gets the root handler for the model metadata.
-   *
-   * @throws IllegalStateException if this model does not contain model metadata
-   */
-  public ModelMetadata getModelMetadata() {
-    assertMetadataInfo();
-    return metadataInfo.getModelMetadata();
-  }
-
-  /** Gets the count of output tensors in the model. */
-  public int getOutputTensorCount() {
-    return modelInfo.getOutputTensorCount();
-  }
-
-  /**
-   * Gets the metadata for the output tensor specified by {@code outputIndex}.
-   *
-   * @param outputIndex the index of the desired output tensor
-   * @throws IllegalStateException if this model does not contain model metadata
-   */
-  @Nullable
-  public TensorMetadata getOutputTensorMetadata(int outputIndex) {
-    assertMetadataInfo();
-    return metadataInfo.getOutputTensorMetadata(outputIndex);
-  }
-
-  /**
-   * Gets the quantization parameters for the output tensor specified by {@code outputIndex}.
-   *
-   * @param outputIndex the index of the desired output tensor
-   */
-  public QuantizationParams getOutputTensorQuantizationParams(int outputIndex) {
-    Tensor tensor = modelInfo.getOutputTensor(outputIndex);
-    return modelInfo.getQuantizationParams(tensor);
-  }
-
-  /**
-   * Gets the shape of the output tensor with {@code outputIndex}.
-   *
-   * @param outputIndex the index of the desired output tensor
-   */
-  public int[] getOutputTensorShape(int outputIndex) {
-    return modelInfo.getOutputTensorShape(outputIndex);
-  }
-
-  /**
-   * Gets the {@link TensorType} of the output tensor with {@code outputIndex}.
-   *
-   * @param outputIndex the index of the desired output tensor
-   */
-  public byte getOutputTensorType(int outputIndex) {
-    return modelInfo.getOutputTensorType(outputIndex);
-  }
-
-  /**
-   * Returns {@code true} if the minimum parser version required by the given metadata flatbuffer
-   * precedes or equals to the version of the metadata parser that this MetadataExtractor library is
-   * relying on. All fields in the metadata can be parsed correctly with this metadata extractor
-   * library in this case. Otherwise, it returns {@code false}.
-   *
-   * 
For example, assume the underlying metadata parser version is {@code 1.14.1},
-   *
-   * 
    
-   *   
  • it returns {@code true}, if the required minimum parser version is the same or older,
-   *       such as {@code 1.14.1} or {@code 1.14.0}. Null version precedes all numeric versions,
-   *       because some metadata flatbuffers are generated before the first versioned release; 
    
-   *   
  • it returns {@code false}, if the required minimum parser version is newer, such as {@code
-   *       1.14.2}.
-   * 

-   */
-  public final boolean isMinimumParserVersionSatisfied() {
-    String minVersion = metadataInfo.getMininumParserVersion();
-    if (minVersion == null) {
-      return true;
-    }
-    return compareVersions(minVersion, MetadataParser.VERSION) <= 0;
-  }
-
-  /**
-   * Asserts if {@link #metadataInfo} is not initialized. Some models may not have metadata and this
-   * is allowed. However, invoking methods that reads the metadata is not allowed.
-   *
-   * @throws IllegalStateException if this model does not contain model metadata
-   */
-  private void assertMetadataInfo() {
-    if (metadataInfo == null) {
-      throw new IllegalStateException("This model does not contain model metadata.");
-    }
-  }
-
-  /**
-   * Asserts if {@link #zipFile} is not initialized. Some models may not have associated files, thus
-   * are not Zip files. This is allowed. However, invoking methods that reads those associated files
-   * is not allowed.
-   *
-   * @throws IllegalStateException if this model is not a Zip file
-   */
-  private void assertZipFile() {
-    if (zipFile == null) {
-      throw new IllegalStateException(
-          "This model does not contain associated files, and is not a Zip file.");
-    }
-  }
-
-  /**
-   * Creates a Zip file handler to read the associated files. If the model is not a zip file, i.e.
-   * it does not have associated files, return a null handler.
-   *
-   * @param buffer the TFLite model FlatBuffer
-   * @throws IOException if an error occurs while reading the model as a Zip file
-   */
-  @Nullable
-  private static ZipFile createZipFile(ByteBuffer buffer) throws IOException {
-    try {
-      // Creates the handler to hold the associated files through the Zip.
-      ByteBufferChannel byteBufferChannel = new ByteBufferChannel(buffer);
-      return ZipFile.createFrom(byteBufferChannel);
-    } catch (ZipException e) {
-      // Some models may not have associate files. Therefore, Those models are not zip files.
-      // However, invoking methods that read associated files later will lead into errors.
-      return null;
-    }
-  }
-
-  /**
-   * Compares two semantic version numbers.
-   *
-   * 
Examples of comparing two versions: 

-   * {@code 1.9} precedes {@code 1.14}; 

-   * {@code 1.14} precedes {@code 1.14.1}; 

-   * {@code 1.14} and {@code 1.14.0} are euqal;
-   *
-   * @return the value {@code 0} if the two versions are equal; a value less than {@code 0} if
-   *     {@code version1} precedes {@code version2}; a value greater than {@code 0} if {@code
-   *     version2} precedes {@code version1}.
-   */
-  private static int compareVersions(String version1, String version2) {
-    // Using String.split instead of the recommanded Guava Splitter because we've been avoiding
-    // depending on other third party libraries in this project.
-    String[] levels1 = version1.split("\\.", 0);
-    String[] levels2 = version2.split("\\.", 0);
-
-    int length = Math.max(levels1.length, levels2.length);
-    for (int i = 0; i < length; i++) {
-      Integer v1 = i < levels1.length ? Integer.parseInt(levels1[i]) : 0;
-      Integer v2 = i < levels2.length ? Integer.parseInt(levels2[i]) : 0;
-      int compare = v1.compareTo(v2);
-      if (compare != 0) {
-        return compare;
-      }
-    }
-
-    return 0;
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
deleted file mode 100644
index 195a330462b..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/MetadataParser.java
+++ /dev/null
@@ -1,27 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-/** Information about the metadata parser that this metadata extractor library is depending on. */
-public final class MetadataParser {
-  /**
-   * The version of the metadata parser that this metadata extractor library is depending on. The
-   * value should match the value of "Schema Semantic version" in metadata_schema.fbs.
-   */
-  public static final String VERSION = "1.0.1";
-
-  private MetadataParser() {}
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelInfo.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelInfo.java
deleted file mode 100644
index 309a3dbe774..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelInfo.java
+++ /dev/null
@@ -1,266 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkNotNull;
-
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.tensorflow.lite.schema.Buffer;
-import org.tensorflow.lite.schema.Metadata;
-import org.tensorflow.lite.schema.Model;
-import org.tensorflow.lite.schema.QuantizationParameters;
-import org.tensorflow.lite.schema.SubGraph;
-import org.tensorflow.lite.schema.Tensor;
-import org.tensorflow.lite.schema.TensorType;
-import org.tensorflow.lite.support.metadata.MetadataExtractor.QuantizationParams;
-
-/** Extracts model information out of TFLite model FLatBuffer. */
-final class ModelInfo {
-  /** The model that is loaded from TFLite model FlatBuffer. */
-  private final Model model;
-
-  /** A list of input tensors. */
-  private final List inputTensors;
-
-  /** A list of output tensors. */
-  private final List outputTensors;
-
-  /** Identifier of the TFLite model metadata in the Metadata array. */
-  static final String METADATA_FIELD_NAME = "TFLITE_METADATA";
-
-  /**
-   * Creates a {@link ModelInfo} with the model FlatBuffer, {@code buffer}.
-   *
-   * 
Though TFLite model FlatBuffer supports multiple subgraphs, TFLite Interpreter only supports
-   * single subgraph so far. See the instruction
-   * of how to specify subgraph during convertion for more information. Therefore, all methods
-   * in {@link ModelInfo} retrieves metadata of the first subgrpah as default.
-   *
-   * @param buffer the TFLite model FlatBuffer
-   * @throws NullPointerException if {@code buffer} is null
-   * @throws IllegalArgumentException if the model does not contain any subgraph, or the model does
-   *     not contain the expected identifier
-   */
-  ModelInfo(ByteBuffer buffer) {
-    assertTFLiteModel(buffer);
-
-    model = Model.getRootAsModel(buffer);
-    checkArgument(model.subgraphsLength() > 0, "The model does not contain any subgraph.");
-
-    inputTensors = getInputTensors(model);
-    outputTensors = getOutputTensors(model);
-  }
-
-  /**
-   * Gets the input tensor with {@code inputIndex}.
-   *
-   * @param inputIndex The index of the desired input tensor.
-   * @throws IllegalArgumentException if the inputIndex specified is invalid.
-   */
-  @Nullable
-  Tensor getInputTensor(int inputIndex) {
-    checkArgument(
-        inputIndex >= 0 && inputIndex < inputTensors.size(),
-        "The inputIndex specified is invalid.");
-    return inputTensors.get(inputIndex);
-  }
-
-  int getInputTensorCount() {
-    return inputTensors.size();
-  }
-
-  /**
-   * Gets shape of the input tensor with {@code inputIndex}.
-   *
-   * @param inputIndex The index of the desired intput tensor.
-   */
-  int[] getInputTensorShape(int inputIndex) {
-    Tensor tensor = getInputTensor(inputIndex);
-    return getShape(tensor);
-  }
-
-  /**
-   * Gets the {@link TensorType} in byte of the input tensor with {@code inputIndex}.
-   *
-   * @param inputIndex The index of the desired intput tensor.
-   */
-  byte getInputTensorType(int inputIndex) {
-    return getInputTensor(inputIndex).type();
-  }
-
-  /** Gets the metadata FlatBuffer from the model FlatBuffer. */
-  @Nullable
-  ByteBuffer getMetadataBuffer() {
-    // Some models may not have metadata, and this is allowed.
-    if (model.metadataLength() == 0) {
-      return null;
-    }
-
-    for (int i = 0; i < model.metadataLength(); i++) {
-      Metadata meta = model.metadata(i);
-      if (METADATA_FIELD_NAME.equals(meta.name())) {
-        long bufferIndex = meta.buffer();
-        Buffer metadataBuf = model.buffers((int) bufferIndex);
-        return metadataBuf.dataAsByteBuffer();
-      }
-    }
-    return null;
-  }
-
-  /**
-   * Gets the output tensor with {@code outputIndex}.
-   *
-   * @param outputIndex The index of the desired outtput tensor.
-   * @throws IllegalArgumentException if the outputIndex specified is invalid.
-   */
-  @Nullable
-  Tensor getOutputTensor(int outputIndex) {
-    checkArgument(
-        outputIndex >= 0 && outputIndex < outputTensors.size(),
-        "The outputIndex specified is invalid.");
-    return outputTensors.get(outputIndex);
-  }
-
-  int getOutputTensorCount() {
-    return outputTensors.size();
-  }
-
-  /**
-   * Gets shape of the output tensor with {@code outputIndex}.
-   *
-   * @param outputIndex The index of the desired outtput tensor.
-   */
-  int[] getOutputTensorShape(int outputIndex) {
-    Tensor tensor = getOutputTensor(outputIndex);
-    return getShape(tensor);
-  }
-
-  /**
-   * Gets the {@link TensorType} in byte of the output tensor {@code outputIndex}.
-   *
-   * @param outputIndex The index of the desired outtput tensor.
-   */
-  byte getOutputTensorType(int outputIndex) {
-    return getOutputTensor(outputIndex).type();
-  }
-
-  /**
-   * Gets the quantization parameters of a tensor.
-   *
-   * 
Only quantized tensors have valid {@code QuantizationParameters}. For tensor that are not
-   * quantized, the values of scale and zero_point are both 0.
-   *
-   * @param tensor The tensor whoes quantization parameters is desired.
-   * @throws NullPointerException if the tensor is null.
-   * @throws IllegalArgumentException if {@code scale} and {@code zeroPoint} of the tensor's {@link
-   *     QuantizationParameters} are not single values.
-   */
-  QuantizationParams getQuantizationParams(Tensor tensor) {
-    checkNotNull(tensor, "Tensor cannot be null.");
-
-    float scale;
-    int zeroPoint;
-    QuantizationParameters quantization = tensor.quantization();
-
-    // Tensors that are not quantized do not have quantization parameters, which can be null when
-    // being extracted from the flatbuffer.
-    if (quantization == null) {
-      scale = 0.0f;
-      zeroPoint = 0;
-      return new QuantizationParams(scale, zeroPoint);
-    }
-
-    // Tensors that are not quantized do not have quantization parameters.
-    // quantization.scaleLength() and quantization.zeroPointLength() may both return 0.
-    checkArgument(
-        quantization.scaleLength() <= 1,
-        "Input and output tensors do not support per-channel quantization.");
-    checkArgument(
-        quantization.zeroPointLength() <= 1,
-        "Input and output tensors do not support per-channel quantization.");
-
-    // For tensors that are not quantized, quantization.scale(0) and quantization.zeroPoint(0) will
-    // both be the default value in flatbuffer, 0. This behavior is consistent with the TFlite C++
-    // runtime.
-    scale = quantization.scale(0);
-    // zeroPoint is a long value in the schema, but an integer in the C++ runtime. Here we keep it
-    // consistent with the C++ runtime.
-    zeroPoint = (int) quantization.zeroPoint(0);
-
-    return new QuantizationParams(scale, zeroPoint);
-  }
-
-  /**
-   * Verifies if the buffer is a valid TFLite model.
-   *
-   * @param buffer the TFLite model flatbuffer
-   * @throws NullPointerException if {@code buffer} is null.
-   * @throws IllegalArgumentException if {@code buffer} does not contain the expected identifier
-   */
-  private static void assertTFLiteModel(ByteBuffer buffer) {
-    checkNotNull(buffer, "Model flatbuffer cannot be null.");
-    checkArgument(
-        Model.ModelBufferHasIdentifier(buffer),
-        "The identifier of the model is invalid. The buffer may not be a valid TFLite model"
-            + " flatbuffer.");
-  }
-
-  /**
-   * Gets the shape of a tensor.
-   *
-   * @param tensor The tensor whoes shape is desired.
-   * @throws NullPointerException if the tensor is null.
-   */
-  private static int[] getShape(Tensor tensor) {
-    checkNotNull(tensor, "Tensor cannot be null.");
-    int shapeDim = tensor.shapeLength();
-    int[] tensorShape = new int[shapeDim];
-    for (int i = 0; i < shapeDim; i++) {
-      tensorShape[i] = tensor.shape(i);
-    }
-    return tensorShape;
-  }
-
-  /** Gets input tensors from a model. */
-  private static List getInputTensors(Model model) {
-    // TFLite only support one subgraph currently.
-    SubGraph subgraph = model.subgraphs(0);
-    int tensorNum = subgraph.inputsLength();
-    ArrayList inputTensors = new ArrayList<>(tensorNum);
-    for (int i = 0; i < tensorNum; i++) {
-      inputTensors.add(subgraph.tensors(subgraph.inputs(i)));
-    }
-    return Collections.unmodifiableList(inputTensors);
-  }
-
-  /** Gets output tensors from a model. */
-  private static List getOutputTensors(Model model) {
-    // TFLite only support one subgraph currently.
-    SubGraph subgraph = model.subgraphs(0);
-    int tensorNum = subgraph.outputsLength();
-    ArrayList outputTensors = new ArrayList<>(tensorNum);
-    for (int i = 0; i < tensorNum; i++) {
-      outputTensors.add(subgraph.tensors(subgraph.outputs(i)));
-    }
-    return Collections.unmodifiableList(outputTensors);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelMetadataInfo.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelMetadataInfo.java
deleted file mode 100644
index 751ed500dc2..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ModelMetadataInfo.java
+++ /dev/null
@@ -1,153 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkNotNull;
-
-import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import org.checkerframework.checker.nullness.qual.Nullable;
-import org.tensorflow.lite.support.metadata.schema.ModelMetadata;
-import org.tensorflow.lite.support.metadata.schema.SubGraphMetadata;
-import org.tensorflow.lite.support.metadata.schema.TensorMetadata;
-
-/** Extracts model metadata information out of TFLite metadata FlatBuffer. */
-final class ModelMetadataInfo {
-  /** The root handler for the model metadata. */
-  private final ModelMetadata modelMetadata;
-
-  /** Metadata array of input tensors. */
-  private final List inputsMetadata;
-
-  /** Metadata array of output tensors. */
-  private final List outputsMetadata;
-
-  /** The minimum parser version required to fully understand the metadata flatbuffer. */
-  private final String /* @Nullable */ minVersion;
-
-  /**
-   * Creates a {@link ModelMetadataInfo} with the metadata FlatBuffer, {@code buffer}.
-   *
-   * @param buffer the TFLite metadata FlatBuffer
-   * @throws NullPointerException if {@code buffer} is null
-   * @throws IllegalArgumentException if {@code buffer} does not contain any subgraph metadata, or
-   *     it does not contain the expected identifier
-   */
-  ModelMetadataInfo(ByteBuffer buffer) {
-    assertTFLiteMetadata(buffer);
-
-    modelMetadata = ModelMetadata.getRootAsModelMetadata(buffer);
-    checkArgument(
-        modelMetadata.subgraphMetadataLength() > 0,
-        "The metadata flatbuffer does not contain any subgraph metadata.");
-
-    inputsMetadata = getInputsMetadata(modelMetadata);
-    outputsMetadata = getOutputsMetadata(modelMetadata);
-    minVersion = modelMetadata.minParserVersion();
-  }
-
-  /** Gets the count of input tensors with metadata in the metadata FlatBuffer. */
-  int getInputTensorCount() {
-    return inputsMetadata.size();
-  }
-
-  /**
-   * Gets the metadata for the input tensor specified by {@code inputIndex}.
-   *
-   * @param inputIndex The index of the desired intput tensor.
-   * @throws IllegalArgumentException if the inputIndex specified is invalid.
-   */
-  @Nullable
-  TensorMetadata getInputTensorMetadata(int inputIndex) {
-    checkArgument(
-        inputIndex >= 0 && inputIndex < inputsMetadata.size(),
-        "The inputIndex specified is invalid.");
-    return inputsMetadata.get(inputIndex);
-  }
-
-  /**
-   * Gets the minimum parser version of the metadata. It can be {@code null} if the version is not
-   * populated.
-   */
-  @Nullable
-  String getMininumParserVersion() {
-    return minVersion;
-  }
-
-  /** Gets the root handler for the model metadata. */
-  ModelMetadata getModelMetadata() {
-    return modelMetadata;
-  }
-
-  /** Gets the count of output tensors with metadata in the metadata FlatBuffer. */
-  int getOutputTensorCount() {
-    return outputsMetadata.size();
-  }
-
-  /**
-   * Gets the metadata for the output tensor specified by {@code outputIndex}.
-   *
-   * @param outputIndex The index of the desired output tensor.
-   * @throws IllegalArgumentException if the outputIndex specified is invalid.
-   */
-  @Nullable
-  TensorMetadata getOutputTensorMetadata(int outputIndex) {
-    checkArgument(
-        outputIndex >= 0 && outputIndex < outputsMetadata.size(),
-        "The outputIndex specified is invalid.");
-    return outputsMetadata.get(outputIndex);
-  }
-
-  /**
-   * Verifies if the buffer is a valid TFLite metadata flatbuffer.
-   *
-   * @param buffer the TFLite metadata flatbuffer
-   * @throws NullPointerException if {@code buffer} is null.
-   * @throws IllegalArgumentException if {@code buffer} does not contain the expected identifier
-   */
-  private static void assertTFLiteMetadata(ByteBuffer buffer) {
-    checkNotNull(buffer, "Metadata flatbuffer cannot be null.");
-    checkArgument(
-        ModelMetadata.ModelMetadataBufferHasIdentifier(buffer),
-        "The identifier of the metadata is invalid. The buffer may not be a valid TFLite metadata"
-            + " flatbuffer.");
-  }
-
-  /** Gets metadata for all input tensors. */
-  private static List getInputsMetadata(ModelMetadata modelMetadata) {
-    SubGraphMetadata subgraphMetadata = modelMetadata.subgraphMetadata(0);
-    int tensorNum = subgraphMetadata.inputTensorMetadataLength();
-    ArrayList inputsMetadata = new ArrayList<>(tensorNum);
-    for (int i = 0; i < tensorNum; i++) {
-      inputsMetadata.add(subgraphMetadata.inputTensorMetadata(i));
-    }
-    return Collections.unmodifiableList(inputsMetadata);
-  }
-
-  /** Gets metadata for all output tensors. */
-  private static List getOutputsMetadata(ModelMetadata modelMetadata) {
-    SubGraphMetadata subgraphMetadata = modelMetadata.subgraphMetadata(0);
-    int tensorNum = subgraphMetadata.outputTensorMetadataLength();
-    ArrayList outputsMetadata = new ArrayList<>(tensorNum);
-    for (int i = 0; i < tensorNum; i++) {
-      outputsMetadata.add(subgraphMetadata.outputTensorMetadata(i));
-    }
-    return Collections.unmodifiableList(outputsMetadata);
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/Preconditions.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/Preconditions.java
deleted file mode 100644
index c2f20fbaacd..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/Preconditions.java
+++ /dev/null
@@ -1,184 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import org.checkerframework.checker.nullness.qual.Nullable;
-
-/** Static error checking util methods. */
-final class Preconditions {
-  /**
-   * Ensures that an object reference passed as a parameter to the calling method is not null.
-   *
-   * @param reference an object reference
-   * @return the non-null reference that was validated
-   * @throws NullPointerException if {@code reference} is null
-   */
-  public static  T checkNotNull(T reference) {
-    if (reference == null) {
-      throw new NullPointerException("The object reference is null.");
-    }
-    return reference;
-  }
-
-  /**
-   * Ensures that an object reference passed as a parameter to the calling method is not null.
-   *
-   * @param reference an object reference
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @return the non-null reference that was validated
-   * @throws NullPointerException if {@code reference} is null
-   */
-  public static  T checkNotNull(T reference, @Nullable Object errorMessage) {
-    if (reference == null) {
-      throw new NullPointerException(String.valueOf(errorMessage));
-    }
-    return reference;
-  }
-
-  /**
-   * Ensures that the given String is not empty and not null.
-   *
-   * @param string the String to test
-   * @return the non-null non-empty String that was validated
-   * @throws IllegalArgumentException if {@code string} is null or empty
-   */
-  public static String checkNotEmpty(String string) {
-    if (string == null || string.length() == 0) {
-      throw new IllegalArgumentException("Given String is empty or null.");
-    }
-    return string;
-  }
-
-  /**
-   * Ensures that the given String is not empty and not null.
-   *
-   * @param string the String to test
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @return the non-null non-empty String that was validated
-   * @throws IllegalArgumentException if {@code string} is null or empty
-   */
-  public static String checkNotEmpty(String string, Object errorMessage) {
-    if (string == null || string.length() == 0) {
-      throw new IllegalArgumentException(String.valueOf(errorMessage));
-    }
-    return string;
-  }
-
-  /**
-   * Ensures the truth of an expression involving one or more parameters to the calling method.
-   *
-   * @param expression a boolean expression.
-   * @throws IllegalArgumentException if {@code expression} is false.
-   */
-  public static void checkArgument(boolean expression) {
-    if (!expression) {
-      throw new IllegalArgumentException();
-    }
-  }
-
-  /**
-   * Ensures the truth of an expression involving one or more parameters to the calling method.
-   *
-   * @param expression a boolean expression.
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}.
-   * @throws IllegalArgumentException if {@code expression} is false.
-   */
-  public static void checkArgument(boolean expression, @Nullable Object errorMessage) {
-    if (!expression) {
-      throw new IllegalArgumentException(String.valueOf(errorMessage));
-    }
-  }
-
-  /**
-   * Ensures that {@code index} specifies a valid element in an array, list or string of size
-   * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive.
-   *
-   * @param index a user-supplied index identifying an element of an array, list or string
-   * @param size the size of that array, list or string
-   * @return the value of {@code index}
-   * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size}
-   * @throws IllegalArgumentException if {@code size} is negative
-   */
-  public static int checkElementIndex(int index, int size) {
-    return checkElementIndex(index, size, "index");
-  }
-
-  /**
-   * Ensures that {@code index} specifies a valid element in an array, list or string of size
-   * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive.
-   *
-   * @param index a user-supplied index identifying an element of an array, list or string
-   * @param size the size of that array, list or string
-   * @param desc the text to use to describe this index in an error message
-   * @return the value of {@code index}
-   * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size}
-   * @throws IllegalArgumentException if {@code size} is negative
-   */
-  public static int checkElementIndex(int index, int size, @Nullable String desc) {
-    // Carefully optimized for execution by hotspot (explanatory comment above)
-    if (index < 0 || index >= size) {
-      throw new IndexOutOfBoundsException(badElementIndex(index, size, desc));
-    }
-    return index;
-  }
-
-  /**
-   * Ensures the truth of an expression involving the state of the calling instance, but not
-   * involving any parameters to the calling method.
-   *
-   * @param expression a boolean expression
-   * @throws IllegalStateException if {@code expression} is false
-   * @see Verify#verify Verify.verify()
-   */
-  public static void checkState(boolean expression) {
-    if (!expression) {
-      throw new IllegalStateException();
-    }
-  }
-
-  /**
-   * Ensures the truth of an expression involving the state of the calling instance, but not
-   * involving any parameters to the calling method.
-   *
-   * @param expression a boolean expression
-   * @param errorMessage the exception message to use if the check fails; will be converted to a
-   *     string using {@link String#valueOf(Object)}
-   * @throws IllegalStateException if {@code expression} is false
-   * @see Verify#verify Verify.verify()
-   */
-  public static void checkState(boolean expression, @Nullable Object errorMessage) {
-    if (!expression) {
-      throw new IllegalStateException(String.valueOf(errorMessage));
-    }
-  }
-
-  private static String badElementIndex(int index, int size, @Nullable String desc) {
-    if (index < 0) {
-      return String.format("%s (%s) must not be negative", desc, index);
-    } else if (size < 0) {
-      throw new IllegalArgumentException("negative size: " + size);
-    } else { // index >= size
-      return String.format("%s (%s) must be less than size (%s)", desc, index, size);
-    }
-  }
-
-  private Preconditions() {
-    throw new AssertionError("Preconditions is Uninstantiable.");
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/SeekableByteChannelCompat.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/SeekableByteChannelCompat.java
deleted file mode 100644
index c655786755b..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/SeekableByteChannelCompat.java
+++ /dev/null
@@ -1,107 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import java.io.IOException;
-import java.nio.ByteBuffer;
-import java.nio.channels.Channel;
-
-/**
- * A byte channel that maintains a current position and allows the position to be changed.
- * {@link SeekableByteChannelCompat} is compatible with {@link
- * java.nio.channels.SeekableByteChannel}.
- *
- * 
{@link java.nio.channels.SeekableByteChannel} is not available in Android API 23 and under.
- * Therefore, {@link SeekableByteChannelCompat} is introduced here to make the interfaces used in
- * the MetadtaExtractor library consistent with the common used Java libraries.
- */
-interface SeekableByteChannelCompat extends Channel {
-  /**
-   * Reads a sequence of bytes from this channel into the given buffer.
-   *
-   * @param dst The buffer into which bytes are to be transferred
-   * @return The number of bytes read, possibly zero, or -1 if the channel has reached
-   *     end-of-stream
-   * @throws NonReadableChannelException If this channel was not opened for reading
-   * @throws ClosedChannelException If this channel is closed
-   * @throws AsynchronousCloseException If another thread closes this channel while the read
-   *     operation is in progress
-   * @throws ClosedByInterruptException If another thread interrupts the current thread while the
-   *     read operation is in progress, thereby closing the channel and setting the current thread's
-   *     interrupt status
-   * @throws IOException If some other I/O error occurs
-   */
-  int read(ByteBuffer dst) throws IOException;
-
-  /**
-   * Writes a sequence of bytes to this channel from the given buffer.
-   *
-   * @param src The buffer from which bytes are to be retrieved
-   * @return The number of bytes written, possibly zero
-   * @throws NonWritableChannelException If this channel was not opened for writing
-   * @throws ClosedChannelException If this channel is closed
-   * @throws AsynchronousCloseException If another thread closes this channel while the write
-   *     operation is in progress
-   * @throws ClosedByInterruptException If another thread interrupts the current thread while the
-   *     write operation is in progress, thereby closing the channel and setting the current
-   *     thread's interrupt status
-   * @throws IOException If some other I/O error occurs
-   */
-  int write(ByteBuffer src) throws IOException;
-
-  /**
-   * Returns this channel's position.
-   *
-   * @return This channel's position, a non-negative integer counting the number of bytes from the
-   *     beginning of the entity to the current position
-   * @throws ClosedChannelException If this channel is closed
-   * @throws IOException If some other I/O error occurs
-   */
-  long position() throws IOException;
-
-  /**
-   * Sets this channel's position.
-   *
-   * @param newPosition The new position, a non-negative integer counting the number of bytes from
-   *     the beginning of the entity
-   * @return This channel
-   * @throws ClosedChannelException If this channel is closed
-   * @throws IllegalArgumentException If the new position is negative
-   * @throws IOException If some other I/O error occurs
-   */
-  SeekableByteChannelCompat position(long newPosition) throws IOException;
-
-  /**
-   * Returns the current size of entity to which this channel is connected.
-   *
-   * @return The current size, measured in bytes
-   * @throws ClosedChannelException If this channel is closed
-   * @throws IOException If some other I/O error occurs
-   */
-  long size() throws IOException;
-
-  /**
-   * Truncates the entity, to which this channel is connected, to the given size.
-   *
-   * @param size The new size, a non-negative byte count
-   * @return This channel
-   * @throws NonWritableChannelException If this channel was not opened for writing
-   * @throws ClosedChannelException If this channel is closed
-   * @throws IllegalArgumentException If the new size is negative
-   * @throws IOException If some other I/O error occurs
-   */
-  SeekableByteChannelCompat truncate(long size) throws IOException;
-}
diff --git a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ZipFile.java b/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ZipFile.java
deleted file mode 100644
index f055d7dcd7e..00000000000
--- a/tensorflow/lite/experimental/support/metadata/java/src/java/org/tensorflow/lite/support/metadata/ZipFile.java
+++ /dev/null
@@ -1,427 +0,0 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-
-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
-
-    http://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.
-==============================================================================*/
-
-package org.tensorflow.lite.support.metadata;
-
-import static org.tensorflow.lite.support.metadata.Preconditions.checkArgument;
-import static org.tensorflow.lite.support.metadata.Preconditions.checkNotNull;
-
-import java.io.Closeable;
-import java.io.EOFException;
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
-import java.nio.charset.Charset;
-import java.util.ArrayList;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.zip.ZipException;
-
-/**
- * Reads uncompressed files from the TFLite model, a zip file.
- *
- * 
TODO(b/150237111): add a link to the webpage of MetadataPopulator once it's available.
- *
- * 
A TFLite model file becomes a zip file when it contains associated files. The associated files
- * can be packed to a TFLite model file using the MetadataPopulator. The associated files are not
- * compressed when being added to the model file.
- *
- * 
{@link ZipFile} does not support Zip64 format, because TFLite models are much smaller than the
- * size limit for Zip64, which is 4GB.
- */
-final class ZipFile implements Closeable {
-  /** Maps String to list of ZipEntrys, name -> actual entries. */
-  private final Map> nameMap;
-
-  /** The actual data source. */
-  private final ByteBufferChannel archive;
-
-  /**
-   * Opens the given {@link ByteBufferChannel} for reading, assuming "UTF8" for file names. {@link
-   * ZipFile} does not synchronized over the buffer that is passed into it.
-   *
-   * @param channel the archive
-   * @throws IOException if an error occurs while creating this {@link ZipFile}
-   * @throws ZipException if the channel is not a zip archive
-   * @throws NullPointerException if the archive is null
-   */
-  public static ZipFile createFrom(ByteBufferChannel channel) throws IOException {
-    checkNotNull(channel);
-    ZipParser zipParser = new ZipParser(channel);
-    Map> nameMap = zipParser.parseEntries();
-    return new ZipFile(channel, nameMap);
-  }
-
-  @Override
-  public void close() {
-    archive.close();
-  }
-
-  /**
-   * Exposes the raw stream of the archive entry.
-   *
-   * 
Since the associated files will not be compressed when being packed to the zip file, the raw
-   * stream represents the non-compressed files.
-   *
-   * 
WARNING: The returned {@link InputStream}, is not thread-safe. If multiple
-   * threads concurrently reading from the returned {@link InputStream}, it must be synchronized
-   * externally.
-   *
-   * @param name name of the entry to get the stream for
-   * @return the raw input stream containing data
-   * @throws IllegalArgumentException if the specified file does not exist in the zip file
-   */
-  public InputStream getRawInputStream(String name) {
-    checkArgument(
-        nameMap.containsKey(name),
-        String.format("The file, %s, does not exist in the zip file.", name));
-
-    List entriesWithTheSameName = nameMap.get(name);
-    ZipEntry entry = entriesWithTheSameName.get(0);
-    long start = entry.getDataOffset();
-    long remaining = entry.getSize();
-    return new BoundedInputStream(archive, start, remaining);
-  }
-
-  private ZipFile(ByteBufferChannel channel, Map> nameMap) {
-    archive = channel;
-    this.nameMap = nameMap;
-  }
-
-  /* Parses a Zip archive and gets the information for each {@link ZipEntry}. */
-  private static class ZipParser {
-    private final ByteBufferChannel archive;
-
-    // Cached buffers that will only be used locally in the class to reduce garbage collection.
-    private final ByteBuffer longBuffer =
-        ByteBuffer.allocate(ZipConstants.LONG_BYTE_SIZE).order(ByteOrder.LITTLE_ENDIAN);
-    private final ByteBuffer intBuffer =
-        ByteBuffer.allocate(ZipConstants.INT_BYTE_SIZE).order(ByteOrder.LITTLE_ENDIAN);
-    private final ByteBuffer shortBuffer =
-        ByteBuffer.allocate(ZipConstants.SHORT_BYTE_SIZE).order(ByteOrder.LITTLE_ENDIAN);
-
-    private ZipParser(ByteBufferChannel archive) {
-      this.archive = archive;
-    }
-
-    /**
-     * Parses the underlying {@code archive} and returns the information as a list of {@link
-     * ZipEntry}.
-     */
-    private Map> parseEntries() throws IOException {
-      List entries = parseCentralDirectory();
-      return parseLocalFileHeaderData(entries);
-    }
-
-    /**
-     * Checks if the current position contains a central file header signature, {@link
-     * ZipConstants#CENSIG}.
-     */
-    private boolean foundCentralFileheaderSignature() {
-      long signature = (long) getInt();
-      return signature == ZipConstants.CENSIG;
-    }
-
-    /**
-     * Gets the value as a Java int from two bytes starting at the current position of the archive.
-     */
-    private int getShort() {
-      shortBuffer.rewind();
-      archive.read(shortBuffer);
-      shortBuffer.flip();
-      return (int) shortBuffer.getShort();
-    }
-
-    /**
-     * Gets the value as a Java long from four bytes starting at the current position of the
-     * archive.
-     */
-    private int getInt() {
-      intBuffer.rewind();
-      archive.read(intBuffer);
-      intBuffer.flip();
-      return intBuffer.getInt();
-    }
-
-    /**
-     * Gets the value as a Java long from four bytes starting at the current position of the
-     * archive.
-     */
-    private long getLong() {
-      longBuffer.rewind();
-      archive.read(longBuffer);
-      longBuffer.flip();
-      return longBuffer.getLong();
-    }
-
-    /**
-     * Positions the archive at the start of the central directory.
-     *
-     * 
First, it searches for the signature of the "end of central directory record", {@link
-     * ZipConstants#ENDSIG}. Position the stream at the start of the "end of central directory
-     * record". The zip file are created without archive comments, thus {@link ZipConstants#ENDSIG}
-     * should appear exactly at {@link ZipConstants#ENDHDR} from the end of the zip file.
-     *
-     * 
Then, parse the "end of central dir record" and position the archive at the start of the
-     * central directory.
-     */
-    private void locateCentralDirectory() throws IOException {
-      if (archive.size() < ZipConstants.ENDHDR) {
-        throw new ZipException("The archive is not a ZIP archive.");
-      }
-
-      // Positions the archive at the start of the "end of central directory record".
-      long offsetRecord = archive.size() - ZipConstants.ENDHDR;
-      archive.position(offsetRecord);
-
-      // Checks for the signature, {@link ZipConstants#ENDSIG}.
-      long endSig = getLong();
-      if (endSig != ZipConstants.ENDSIG) {
-        throw new ZipException("The archive is not a ZIP archive.");
-      }
-
-      // Positions the archive at the “offset of central directory”.
-      skipBytes(ZipConstants.ENDOFF - ZipConstants.ENDSUB);
-      // Gets the offset to central directory
-      long offsetDirectory = getInt();
-      // Goes to the central directory.
-      archive.position(offsetDirectory);
-    }
-
-    /**
-     * Reads the central directory of the given archive and populates the internal tables with
-     * {@link ZipEntry} instances.
-     */
-    private List parseCentralDirectory() throws IOException {
-      /** List of entries in the order they appear inside the central directory. */
-      List entries = new ArrayList<>();
-      locateCentralDirectory();
-
-      while (foundCentralFileheaderSignature()) {
-        ZipEntry entry = parseCentralDirectoryEntry();
-        entries.add(entry);
-      }
-
-      return entries;
-    }
-
-    /**
-     * Reads an individual entry of the central directory, creats an ZipEntry from it and adds it to
-     * the global maps.
-     */
-    private ZipEntry parseCentralDirectoryEntry() throws IOException {
-      // Positions the archive at the "compressed size" and read the value.
-      skipBytes(ZipConstants.CENSIZ - ZipConstants.CENVEM);
-      long compressSize = getInt();
-
-      // Positions the archive at the "filename length" and read the value.
-      skipBytes(ZipConstants.CENNAM - ZipConstants.CENLEN);
-      int fileNameLen = getShort();
-
-      // Reads the extra field length and the comment length.
-      int extraLen = getShort();
-      int commentLen = getShort();
-
-      // Positions the archive at the "local file header offset" and read the value.
-      skipBytes(ZipConstants.CENOFF - ZipConstants.CENDSK);
-      long localHeaderOffset = getInt();
-
-      // Reads the file name.
-      byte[] fileNameBuf = new byte[fileNameLen];
-      archive.read(ByteBuffer.wrap(fileNameBuf));
-      String fileName = new String(fileNameBuf, Charset.forName("UTF-8"));
-
-      // Skips the extra field and the comment.
-      skipBytes(extraLen + commentLen);
-
-      ZipEntry entry = new ZipEntry();
-      entry.setSize(compressSize);
-      entry.setLocalHeaderOffset(localHeaderOffset);
-      entry.setName(fileName);
-
-      return entry;
-    }
-
-    /** Walks through all recorded entries and records the offsets for the entry data. */
-    private Map> parseLocalFileHeaderData(List entries) {
-      /** Maps String to list of ZipEntrys, name -> actual entries. */
-      Map> nameMap = new LinkedHashMap<>();
-
-      for (ZipEntry entry : entries) {
-        long offset = entry.getLocalHeaderOffset();
-        archive.position(offset + ZipConstants.LOCNAM);
-
-        // Gets the data offset of this entry.
-        int fileNameLen = getShort();
-        int extraFieldLen = getShort();
-        long dataOffset =
-            offset
-                + ZipConstants.LOCEXT
-                + ZipConstants.SHORT_BYTE_SIZE
-                + fileNameLen
-                + extraFieldLen;
-        entry.setDataOffset(dataOffset);
-
-        // Puts the entry into the nameMap.
-        String name = entry.getName();
-        List entriesWithTheSameName;
-        if (nameMap.containsKey(name)) {
-          entriesWithTheSameName = nameMap.get(name);
-        } else {
-          entriesWithTheSameName = new ArrayList<>();
-          nameMap.put(name, entriesWithTheSameName);
-        }
-        entriesWithTheSameName.add(entry);
-      }
-
-      return nameMap;
-    }
-
-    /** Skips the given number of bytes or throws an EOFException if skipping failed. */
-    private void skipBytes(int count) throws IOException {
-      long currentPosition = archive.position();
-      long newPosition = currentPosition + count;
-      if (newPosition > archive.size()) {
-        throw new EOFException();
-      }
-      archive.position(newPosition);
-    }
-  }
-
-  /** Stores the data offset and the size of an entry in the archive. */
-  private static class ZipEntry {
-
-    private String name;
-    private long dataOffset = -1;
-    private long size = -1;
-    private long localHeaderOffset = -1;
-
-    public long getSize() {
-      return size;
-    }
-
-    public long getDataOffset() {
-      return dataOffset;
-    }
-
-    public String getName() {
-      return name;
-    }
-
-    public long getLocalHeaderOffset() {
-      return localHeaderOffset;
-    }
-
-    public void setSize(long size) {
-      this.size = size;
-    }
-
-    public void setDataOffset(long dataOffset) {
-      this.dataOffset = dataOffset;
-    }
-
-    public void setName(String name) {
-      this.name = name;
-    }
-
-    public void setLocalHeaderOffset(long localHeaderOffset) {
-      this.localHeaderOffset = localHeaderOffset;
-    }
-  }
-
-  /**
-   * Various constants for this {@link ZipFile}.
-   *
-   * 
Referenced from {@link java.util.zip.ZipConstants}.
-   */
-  private static class ZipConstants {
-    /** length of Java short in bytes. */
-    static final int SHORT_BYTE_SIZE = Short.SIZE / 8;
-
-    /** length of Java int in bytes. */
-    static final int INT_BYTE_SIZE = Integer.SIZE / 8;
-
-    /** length of Java long in bytes. */
-    static final int LONG_BYTE_SIZE = Long.SIZE / 8;
-
-    /*
-     * Header signatures
-     */
-    static final long LOCSIG = 0x04034b50L; // "PK\003\004"
-    static final long EXTSIG = 0x08074b50L; // "PK\007\008"
-    static final long CENSIG = 0x02014b50L; // "PK\001\002"
-    static final long ENDSIG = 0x06054b50L; // "PK\005\006"
-
-    /*
-     * Header sizes in bytes (including signatures)
-     */
-    static final int LOCHDR = 30; // LOC header size
-    static final int EXTHDR = 16; // EXT header size
-    static final int CENHDR = 46; // CEN header size
-    static final int ENDHDR = 22; // END header size
-
-    /*
-     * Local file (LOC) header field offsets
-     */
-    static final int LOCVER = 4; // version needed to extract
-    static final int LOCFLG = 6; // general purpose bit flag
-    static final int LOCHOW = 8; // compression method
-    static final int LOCTIM = 10; // modification time
-    static final int LOCCRC = 14; // uncompressed file crc-32 value
-    static final int LOCSIZ = 18; // compressed size
-    static final int LOCLEN = 22; // uncompressed size
-    static final int LOCNAM = 26; // filename length
-    static final int LOCEXT = 28; // extra field length
-
-    /*
-     * Extra local (EXT) header field offsets
-     */
-    static final int EXTCRC = 4; // uncompressed file crc-32 value
-    static final int EXTSIZ = 8; // compressed size
-    static final int EXTLEN = 12; // uncompressed size
-
-    /*
-     * Central directory (CEN) header field offsets
-     */
-    static final int CENVEM = 4; // version made by
-    static final int CENVER = 6; // version needed to extract
-    static final int CENFLG = 8; // encrypt, decrypt flags
-    static final int CENHOW = 10; // compression method
-    static final int CENTIM = 12; // modification time
-    static final int CENCRC = 16; // uncompressed file crc-32 value
-    static final int CENSIZ = 20; // compressed size
-    static final int CENLEN = 24; // uncompressed size
-    static final int CENNAM = 28; // filename length
-    static final int CENEXT = 30; // extra field length
-    static final int CENCOM = 32; // comment length
-    static final int CENDSK = 34; // disk number start
-    static final int CENATT = 36; // internal file attributes
-    static final int CENATX = 38; // external file attributes
-    static final int CENOFF = 42; // LOC header offset
-
-    /*
-     * End of central directory (END) header field offsets
-     */
-    static final int ENDSUB = 8; // number of entries on this disk
-    static final int ENDTOT = 10; // total number of entries
-    static final int ENDSIZ = 12; // central directory size in bytes
-    static final int ENDOFF = 16; // offset of first CEN header
-    static final int ENDCOM = 20; // zip file comment length
-
-    private ZipConstants() {}
-  }
-}
diff --git a/tensorflow/lite/experimental/support/metadata/metadata.py b/tensorflow/lite/experimental/support/metadata/metadata.py
deleted file mode 100644
index b3d8d28806b..00000000000
--- a/tensorflow/lite/experimental/support/metadata/metadata.py
+++ /dev/null
@@ -1,615 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"""TensorFlow Lite metadata tools."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import copy
-import os
-import shutil
-import tempfile
-import warnings
-import zipfile
-
-from flatbuffers.python import flatbuffers
-from tensorflow.lite.experimental.support.metadata import metadata_schema_py_generated as _metadata_fb
-from tensorflow.lite.experimental.support.metadata import schema_py_generated as _schema_fb
-from tensorflow.lite.experimental.support.metadata.cc.python import _pywrap_metadata_version
-from tensorflow.lite.experimental.support.metadata.flatbuffers_lib import _pywrap_flatbuffers
-from tensorflow.python.platform import resource_loader
-
-_FLATC_TFLITE_METADATA_SCHEMA_FILE = resource_loader.get_path_to_datafile(
-    "metadata_schema.fbs")
-
-
-# TODO(b/141467403): add delete method for associated files.
-class MetadataPopulator(object):
-  """Packs metadata and associated files into TensorFlow Lite model file.
-
-  MetadataPopulator can be used to populate metadata and model associated files
-  into a model file or a model buffer (in bytearray). It can also help to
-  inspect list of files that have been packed into the model or are supposed to
-  be packed into the model.
-
-  The metadata file (or buffer) should be generated based on the metadata
-  schema:
-  third_party/tensorflow/lite/schema/metadata_schema.fbs
-
-  Example usage:
-  Populate matadata and label file into an image classifier model.
-
-  First, based on metadata_schema.fbs, generate the metadata for this image
-  classifer model using Flatbuffers API. Attach the label file onto the ouput
-  tensor (the tensor of probabilities) in the metadata.
-
-  Then, pack the metadata and label file into the model as follows.
-
-    ```python
-    # Populating a metadata file (or a metadta buffer) and associated files to
-    a model file:
-    populator = MetadataPopulator.with_model_file(model_file)
-    # For metadata buffer (bytearray read from the metadata file), use:
-    # populator.load_metadata_buffer(metadata_buf)
-    populator.load_metadata_file(metadata_file)
-    populator.load_associated_files([label.txt])
-    populator.populate()
-
-    # Populating a metadata file (or a metadta buffer) and associated files to
-    a model buffer:
-    populator = MetadataPopulator.with_model_buffer(model_buf)
-    populator.load_metadata_file(metadata_file)
-    populator.load_associated_files([label.txt])
-    populator.populate()
-    # Writing the updated model buffer into a file.
-    updated_model_buf = populator.get_model_buffer()
-    with open("updated_model.tflite", "wb") as f:
-      f.write(updated_model_buf)
-    ```
-
-  Note that existing metadata buffer (if applied) will be overridden by the new
-  metadata buffer.
-  """
-  # As Zip API is used to concatenate associated files after tflite model file,
-  # the populating operation is developed based on a model file. For in-memory
-  # model buffer, we create a tempfile to serve the populating operation.
-  # Creating the deleting such a tempfile is handled by the class,
-  # _MetadataPopulatorWithBuffer.
-
-  METADATA_FIELD_NAME = "TFLITE_METADATA"
-  TFLITE_FILE_IDENTIFIER = b"TFL3"
-  METADATA_FILE_IDENTIFIER = b"M001"
-
-  def __init__(self, model_file):
-    """Constructor for MetadataPopulator.
-
-    Args:
-      model_file: valid path to a TensorFlow Lite model file.
-
-    Raises:
-      IOError: File not found.
-      ValueError: the model does not have the expected flatbuffer identifer.
-    """
-    _assert_model_file_identifier(model_file)
-    self._model_file = model_file
-    self._metadata_buf = None
-    self._associated_files = set()
-
-  @classmethod
-  def with_model_file(cls, model_file):
-    """Creates a MetadataPopulator object that populates data to a model file.
-
-    Args:
-      model_file: valid path to a TensorFlow Lite model file.
-
-    Returns:
-      MetadataPopulator object.
-
-    Raises:
-      IOError: File not found.
-      ValueError: the model does not have the expected flatbuffer identifer.
-    """
-    return cls(model_file)
-
-  # TODO(b/141468993): investigate if type check can be applied to model_buf for
-  # FB.
-  @classmethod
-  def with_model_buffer(cls, model_buf):
-    """Creates a MetadataPopulator object that populates data to a model buffer.
-
-    Args:
-      model_buf: TensorFlow Lite model buffer in bytearray.
-
-    Returns:
-      A MetadataPopulator(_MetadataPopulatorWithBuffer) object.
-
-    Raises:
-      ValueError: the model does not have the expected flatbuffer identifer.
-    """
-    return _MetadataPopulatorWithBuffer(model_buf)
-
-  def get_model_buffer(self):
-    """Gets the buffer of the model with packed metadata and associated files.
-
-    Returns:
-      Model buffer (in bytearray).
-    """
-    with open(self._model_file, "rb") as f:
-      return f.read()
-
-  def get_packed_associated_file_list(self):
-    """Gets a list of associated files packed to the model file.
-
-    Returns:
-      List of packed associated files.
-    """
-    if not zipfile.is_zipfile(self._model_file):
-      return []
-
-    with zipfile.ZipFile(self._model_file, "r") as zf:
-      return zf.namelist()
-
-  def get_recorded_associated_file_list(self):
-    """Gets a list of associated files recorded in metadata of the model file.
-
-    Associated files may be attached to a model, a subgraph, or an input/output
-    tensor.
-
-    Returns:
-      List of recorded associated files.
-    """
-    recorded_files = []
-
-    if not self._metadata_buf:
-      return recorded_files
-
-    metadata = _metadata_fb.ModelMetadata.GetRootAsModelMetadata(
-        self._metadata_buf, 0)
-
-    # Add associated files attached to ModelMetadata
-    self._get_associated_files_from_metadata_struct(metadata, recorded_files)
-
-    # Add associated files attached to each SubgraphMetadata
-    for j in range(metadata.SubgraphMetadataLength()):
-      subgraph = metadata.SubgraphMetadata(j)
-      self._get_associated_files_from_metadata_struct(subgraph, recorded_files)
-
-      # Add associated files attached to each input tensor
-      for k in range(subgraph.InputTensorMetadataLength()):
-        tensor = subgraph.InputTensorMetadata(k)
-        self._get_associated_files_from_metadata_struct(tensor, recorded_files)
-
-      # Add associated files attached to each output tensor
-      for k in range(subgraph.OutputTensorMetadataLength()):
-        tensor = subgraph.OutputTensorMetadata(k)
-        self._get_associated_files_from_metadata_struct(tensor, recorded_files)
-
-    return recorded_files
-
-  def load_associated_files(self, associated_files):
-    """Loads associated files that to be concatenated after the model file.
-
-    Args:
-      associated_files: list of file paths.
-
-    Raises:
-      IOError:
-        File not found.
-    """
-    for af in associated_files:
-      _assert_exist(af)
-      self._associated_files.add(af)
-
-  def load_metadata_buffer(self, metadata_buf):
-    """Loads the metadata buffer (in bytearray) to be populated.
-
-    Args:
-      metadata_buf: metadata buffer (in bytearray) to be populated.
-
-    Raises:
-      ValueError: The metadata to be populated is empty.
-      ValueError: The metadata does not have the expected flatbuffer identifer.
-      ValueError: Error occurs when getting the minimum metadata parser version.
-    """
-    if not metadata_buf:
-      raise ValueError("The metadata to be populated is empty.")
-
-    _assert_metadata_buffer_identifier(metadata_buf)
-
-    # Gets the minimum metadata parser version of the metadata_buf.
-    min_version = _pywrap_metadata_version.GetMinimumMetadataParserVersion(
-        bytes(metadata_buf))
-
-    # Inserts in the minimum metadata parser version into the metadata_buf.
-    metadata = _metadata_fb.ModelMetadataT.InitFromObj(
-        _metadata_fb.ModelMetadata.GetRootAsModelMetadata(metadata_buf, 0))
-    metadata.minParserVersion = min_version
-
-    b = flatbuffers.Builder(0)
-    b.Finish(metadata.Pack(b), self.METADATA_FILE_IDENTIFIER)
-    metadata_buf_with_version = b.Output()
-
-    self._metadata_buf = metadata_buf_with_version
-
-  def load_metadata_file(self, metadata_file):
-    """Loads the metadata file to be populated.
-
-    Args:
-      metadata_file: path to the metadata file to be populated.
-
-    Raises:
-      IOError: File not found.
-      ValueError: The metadata does not have the expected flatbuffer identifer.
-    """
-    _assert_exist(metadata_file)
-    with open(metadata_file, "rb") as f:
-      metadata_buf = f.read()
-    self.load_metadata_buffer(bytearray(metadata_buf))
-
-  def populate(self):
-    """Populates loaded metadata and associated files into the model file."""
-    self._assert_validate()
-    self._populate_metadata_buffer()
-    self._populate_associated_files()
-
-  def _assert_validate(self):
-    """Validates the metadata and associated files to be populated.
-
-    Raises:
-      ValueError:
-        File is recorded in the metadata, but is not going to be populated.
-        File has already been packed.
-    """
-    # Gets files that are recorded in metadata.
-    recorded_files = self.get_recorded_associated_file_list()
-
-    # Gets files that have been packed to self._model_file.
-    packed_files = self.get_packed_associated_file_list()
-
-    # Gets the file name of those associated files to be populated.
-    to_be_populated_files = []
-    for af in self._associated_files:
-      to_be_populated_files.append(os.path.basename(af))
-
-    # Checks all files recorded in the metadata will be populated.
-    for rf in recorded_files:
-      if rf not in to_be_populated_files and rf not in packed_files:
-        raise ValueError("File, '{0}', is recorded in the metadata, but has "
-                         "not been loaded into the populator.".format(rf))
-
-    for f in to_be_populated_files:
-      if f in packed_files:
-        raise ValueError("File, '{0}', has already been packed.".format(f))
-
-      if f not in recorded_files:
-        warnings.warn(
-            "File, '{0}', does not exsit in the metadata. But packing it to "
-            "tflite model is still allowed.".format(f))
-
-  def _copy_archived_files(self, src_zip, dst_zip, file_list):
-    """Copy archieved files in file_list from src_zip ro dst_zip."""
-
-    if not zipfile.is_zipfile(src_zip):
-      raise ValueError("File, '{0}', is not a zipfile.".format(src_zip))
-
-    with zipfile.ZipFile(src_zip,
-                         "r") as src_zf, zipfile.ZipFile(dst_zip,
-                                                         "a") as dst_zf:
-      src_list = src_zf.namelist()
-      for f in file_list:
-        if f not in src_list:
-          raise ValueError(
-              "File, '{0}', does not exist in the zipfile, {1}.".format(
-                  f, src_zip))
-        file_buffer = src_zf.read(f)
-        dst_zf.writestr(f, file_buffer)
-
-  def _get_associated_files_from_metadata_struct(self, file_holder, file_list):
-    for j in range(file_holder.AssociatedFilesLength()):
-      file_list.append(file_holder.AssociatedFiles(j).Name().decode("utf-8"))
-
-  def _populate_associated_files(self):
-    """Concatenates associated files after TensorFlow Lite model file.
-
-    If the MetadataPopulator object is created using the method,
-    with_model_file(model_file), the model file will be updated.
-    """
-    # Opens up the model file in "appending" mode.
-    # If self._model_file already has pack files, zipfile will concatenate
-    # addition files after self._model_file. For example, suppose we have
-    # self._model_file = old_tflite_file | label1.txt | label2.txt
-    # Then after trigger populate() to add label3.txt, self._model_file becomes
-    # self._model_file = old_tflite_file | label1.txt | label2.txt | label3.txt
-    with zipfile.ZipFile(self._model_file, "a") as zf:
-      for af in self._associated_files:
-        filename = os.path.basename(af)
-        zf.write(af, filename)
-
-  def _populate_metadata_buffer(self):
-    """Populates the metadata buffer (in bytearray) into the model file.
-
-    Inserts metadata_buf into the metadata field of schema.Model. If the
-    MetadataPopulator object is created using the method,
-    with_model_file(model_file), the model file will be updated.
-
-    Existing metadata buffer (if applied) will be overridden by the new metadata
-    buffer.
-    """
-
-    with open(self._model_file, "rb") as f:
-      model_buf = f.read()
-
-    model = _schema_fb.ModelT.InitFromObj(
-        _schema_fb.Model.GetRootAsModel(model_buf, 0))
-    buffer_field = _schema_fb.BufferT()
-    buffer_field.data = self._metadata_buf
-
-    is_populated = False
-    if not model.metadata:
-      model.metadata = []
-    else:
-      # Check if metadata has already been populated.
-      for meta in model.metadata:
-        if meta.name.decode("utf-8") == self.METADATA_FIELD_NAME:
-          is_populated = True
-          model.buffers[meta.buffer] = buffer_field
-
-    if not is_populated:
-      if not model.buffers:
-        model.buffers = []
-      model.buffers.append(buffer_field)
-      # Creates a new metadata field.
-      metadata_field = _schema_fb.MetadataT()
-      metadata_field.name = self.METADATA_FIELD_NAME
-      metadata_field.buffer = len(model.buffers) - 1
-      model.metadata.append(metadata_field)
-
-    # Packs model back to a flatbuffer binaray file.
-    b = flatbuffers.Builder(0)
-    b.Finish(model.Pack(b), self.TFLITE_FILE_IDENTIFIER)
-    model_buf = b.Output()
-
-    # Saves the updated model buffer to model file.
-    # Gets files that have been packed to self._model_file.
-    packed_files = self.get_packed_associated_file_list()
-    if packed_files:
-      # Writes the updated model buffer and associated files into a new model
-      # file. Then overwrites the original model file.
-      with tempfile.NamedTemporaryFile() as temp:
-        new_file = temp.name
-      with open(new_file, "wb") as f:
-        f.write(model_buf)
-      self._copy_archived_files(self._model_file, new_file, packed_files)
-      shutil.copy(new_file, self._model_file)
-      os.remove(new_file)
-    else:
-      with open(self._model_file, "wb") as f:
-        f.write(model_buf)
-
-
-class _MetadataPopulatorWithBuffer(MetadataPopulator):
-  """Subclass of MetadtaPopulator that populates metadata to a model buffer.
-
-  This class is used to populate metadata into a in-memory model buffer. As we
-  use Zip API to concatenate associated files after tflite model file, the
-  populating operation is developed based on a model file. For in-memory model
-  buffer, we create a tempfile to serve the populating operation. This class is
-  then used to generate this tempfile, and delete the file when the
-  MetadataPopulator object is deleted.
-  """
-
-  def __init__(self, model_buf):
-    """Constructor for _MetadataPopulatorWithBuffer.
-
-    Args:
-      model_buf: TensorFlow Lite model buffer in bytearray.
-
-    Raises:
-      ValueError: model_buf is empty.
-      ValueError: model_buf does not have the expected flatbuffer identifer.
-    """
-    if not model_buf:
-      raise ValueError("model_buf cannot be empty.")
-
-    with tempfile.NamedTemporaryFile() as temp:
-      model_file = temp.name
-
-    with open(model_file, "wb") as f:
-      f.write(model_buf)
-
-    MetadataPopulator.__init__(self, model_file)
-
-  def __del__(self):
-    """Destructor of _MetadataPopulatorWithBuffer.
-
-    Deletes the tempfile.
-    """
-    if os.path.exists(self._model_file):
-      os.remove(self._model_file)
-
-
-class MetadataDisplayer(object):
-  """Displays metadata and associated file info in human-readable format."""
-
-  def __init__(self, model_file, metadata_file, associated_file_list):
-    """Constructor for MetadataDisplayer.
-
-    Args:
-      model_file: valid path to the model file.
-      metadata_file: valid path to the metadata file.
-      associated_file_list: list of associate files in the model file.
-    """
-    _assert_model_file_identifier(model_file)
-    _assert_metadata_file_identifier(metadata_file)
-    self._model_file = model_file
-    self._metadata_file = metadata_file
-    self._associated_file_list = associated_file_list
-
-  @classmethod
-  def with_model_file(cls, model_file):
-    """Creates a MetadataDisplayer object for the model file.
-
-    Args:
-      model_file: valid path to a TensorFlow Lite model file.
-
-    Returns:
-      MetadataDisplayer object.
-
-    Raises:
-      IOError: File not found.
-      ValueError: The model does not have metadata.
-    """
-    _assert_exist(model_file)
-    metadata_file = cls._save_temporary_metadata_file(model_file)
-    associated_file_list = cls._parse_packed_associted_file_list(model_file)
-    return cls(model_file, metadata_file, associated_file_list)
-
-  @classmethod
-  def with_model_buffer(cls, model_buffer):
-    """Creates a MetadataDisplayer object for a file buffer.
-
-    Args:
-      model_buffer: TensorFlow Lite model buffer in bytearray.
-
-    Returns:
-      MetadataDisplayer object.
-    """
-    if not model_buffer:
-      raise ValueError("model_buffer cannot be empty.")
-
-    with tempfile.NamedTemporaryFile() as temp:
-      model_file = temp.name
-
-    with open(model_file, "wb") as f:
-      f.write(model_buffer)
-    return cls.with_model_file(model_file)
-
-  def get_metadata_json(self):
-    """Converts the metadata into a json string."""
-    opt = _pywrap_flatbuffers.IDLOptions()
-    opt.strict_json = True
-    parser = _pywrap_flatbuffers.Parser(opt)
-    with open(_FLATC_TFLITE_METADATA_SCHEMA_FILE) as f:
-      metadata_schema_content = f.read()
-    with open(self._metadata_file, "rb") as f:
-      metadata_file_content = f.read()
-    if not parser.parse(metadata_schema_content):
-      raise ValueError("Cannot parse metadata schema. Reason: " + parser.error)
-    with open(self._metadata_file, "rb") as f:
-      metadata_file_content = f.read()
-    return _pywrap_flatbuffers.generate_text(parser, metadata_file_content)
-
-  def get_packed_associated_file_list(self):
-    """Returns a list of associated files that are packed in the model.
-
-    Returns:
-      A name list of associated files.
-    """
-    return copy.deepcopy(self._associated_file_list)
-
-  @staticmethod
-  def _save_temporary_metadata_file(model_file):
-    """Saves the metadata in the model file to a temporary file.
-
-    Args:
-      model_file: valid path to the model file.
-
-    Returns:
-      Path to the metadata temporary file.
-
-    Raises:
-      ValueError: The model does not have metadata.
-    """
-    with open(model_file, "rb") as f:
-      model_buf = f.read()
-
-    tflite_model = _schema_fb.Model.GetRootAsModel(model_buf, 0)
-
-    # Gets metadata from the model file.
-    for i in range(tflite_model.MetadataLength()):
-      meta = tflite_model.Metadata(i)
-      if meta.Name().decode("utf-8") == MetadataPopulator.METADATA_FIELD_NAME:
-        buffer_index = meta.Buffer()
-        metadata = tflite_model.Buffers(buffer_index)
-        metadata_buf = metadata.DataAsNumpy().tobytes()
-        # Creates a temporary file to store the metadata.
-        with tempfile.NamedTemporaryFile() as temp:
-          metadata_file = temp.name
-        # Saves the metadata into the temporary file.
-        with open(metadata_file, "wb") as f:
-          f.write(metadata_buf)
-          return metadata_file
-
-    raise ValueError("The model does not have metadata.")
-
-  @staticmethod
-  def _parse_packed_associted_file_list(model_file):
-    """Gets a list of associated files packed to the model file.
-
-    Args:
-      model_file: valid path to the model file.
-
-    Returns:
-      List of packed associated files.
-    """
-    if not zipfile.is_zipfile(model_file):
-      return []
-
-    with zipfile.ZipFile(model_file, "r") as zf:
-      return zf.namelist()
-
-  def __del__(self):
-    """Destructor of MetadataDisplayer.
-
-    Deletes the tempfile.
-    """
-    if os.path.exists(self._metadata_file):
-      os.remove(self._metadata_file)
-
-
-def _assert_exist(filename):
-  """Checks if a file exists."""
-  if not os.path.exists(filename):
-    raise IOError("File, '{0}', does not exist.".format(filename))
-
-
-def _assert_model_file_identifier(model_file):
-  """Checks if a model file has the expected TFLite schema identifier."""
-  _assert_exist(model_file)
-  with open(model_file, "rb") as f:
-    model_buf = f.read()
-
-  if not _schema_fb.Model.ModelBufferHasIdentifier(model_buf, 0):
-    raise ValueError(
-        "The model provided does not have the expected identifier, and "
-        "may not be a valid TFLite model.")
-
-
-def _assert_metadata_file_identifier(metadata_file):
-  """Checks if a metadata file has the expected Metadata schema identifier."""
-  _assert_exist(metadata_file)
-  with open(metadata_file, "rb") as f:
-    metadata_buf = f.read()
-  _assert_metadata_buffer_identifier(metadata_buf)
-
-
-def _assert_metadata_buffer_identifier(metadata_buf):
-  """Checks if a metadata buffer has the expected Metadata schema identifier."""
-  if not _metadata_fb.ModelMetadata.ModelMetadataBufferHasIdentifier(
-      metadata_buf, 0):
-    raise ValueError(
-        "The metadata buffer does not have the expected identifier, and may not"
-        " be a valid TFLite Metadata.")
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template b/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
deleted file mode 100644
index a41ac06969c..00000000000
--- a/tensorflow/lite/experimental/support/metadata/metadata_parser.py.template
+++ /dev/null
@@ -1,26 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"""Information about the metadata parser that this python library depends on."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
-class MetadataParser(object):
-  """Information about the metadata parser."""
-
-  # The version of the metadata parser.
-  VERSION = "{LATEST_METADATA_PARSER_VERSION}"
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py b/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py
deleted file mode 100644
index 3b1d19278cd..00000000000
--- a/tensorflow/lite/experimental/support/metadata/metadata_parser_test.py
+++ /dev/null
@@ -1,38 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"""Tests for tensorflow.lite.experimental.support.metadata.metadata_parser."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import re
-
-from tensorflow.lite.experimental.support.metadata import metadata_parser
-from tensorflow.python.framework import test_util
-from tensorflow.python.platform import test
-
-
-class MetadataParserTest(test_util.TensorFlowTestCase):
-
-  def test_version_wellFormedSemanticVersion(self):
-    # Validates that the version is well-formed (x.y.z).
-    self.assertTrue(
-        re.match('[0-9]+\\.[0-9]+\\.[0-9]+',
-                 metadata_parser.MetadataParser.VERSION))
-
-
-if __name__ == '__main__':
-  test.main()
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs b/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs
deleted file mode 100644
index a88225f1960..00000000000
--- a/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs
+++ /dev/null
@@ -1,570 +0,0 @@
-// Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-//
-// 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
-//
-//     http://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.
-
-namespace tflite;
-
-// TFLite metadata contains both human readable and machine readable information
-// about what the model does and how to use the model. It can be used as a
-// README file, which elaborates the details of the model, each input/ouput
-// tensor, and each associated file.
-//
-// An important use case of TFLite metadata is the TFLite codegen tool, which
-// automatically generates the model interface based on the properties of the
-// model and the tensors. The model interface provides high-level APIs to
-// interact with the model, such as preprocessing the input data and running
-// inferences.
-//
-// Entries marked with "" are used in TFLite codegen tool to
-// generate the model interface. It is recommended to fill in at least those
-// enties to boost the codegen performance.
-
-// The Metadata schema is versioned by the Semantic versioning number, such as
-// MAJOR.MINOR.PATCH. It tracks the schema changes according to the rules below:
-//  * Bump up the MAJOR number when making potentially backwards incompatible
-//    changes. It must be incremented if the new changes break the backwards
-//    compatibility. It may also include minor and patch level changes as
-//    needed. The true backwards compatibility is indicated by the file
-//    identifier.
-//  * Bump up the MINOR number when making backwards compatible updates for
-//    major features, such as supporting new content types or adding new
-//    processing units.
-//  * Bump up the PATCH number when making small backwards compatible changes,
-//    such as adding a new fields or deprecating certain fields (not deleting
-//    them).
-//
-// ModelMetadata.min_parser_version indicates the minimum necessary metadata
-// parser version to fully understand all fields in a given metadata flatbuffer.
-//
-// New fields and types will have associated comments with the schema version
-// for which they were added.
-//
-// LINT.IfChange
-// Schema Semantic version: 1.0.1
-// LINT.ThenChange(//tensorflow/lite/experimental/\
-//.    support/metadata/java/src/java/org/tensorflow/lite/support/metadata/\
-//.    MetadataParser.java)
-
-// This indicates the flatbuffer compatibility. The number will bump up when a
-// break change is applied to the schema, such as removing fields or adding new
-// fields to the middle of a table.
-file_identifier "M001";
-
-// History:
-// 1.0.1 - Added VOCABULARY type to AssociatedFileType.
-
-// File extension of any written files.
-file_extension "tflitemeta";
-
-// LINT.IfChange
-enum AssociatedFileType : byte {
-  UNKNOWN = 0,
-
-  // Files such as readme.txt.
-  DESCRIPTIONS = 1,
-
-  // Contains labels that annotate certain axis of the tensor. For example,
-  // the label file in image classification. Those labels annotate the
-  // the output tensor, such that each value in the output tensor is the
-  // probability of that corresponding category specified by the label.
-  //
-  // :
-  // If an output tensor has an associated file as TENSOR_AXIS_LABELS, return
-  // the output as a mapping between the labels and probability in the model
-  // interface.
-  // If multiple files of the same type are present, the first one is used by
-  // default; additional ones are to be distinguished from one another by their
-  // specified locale.
-  TENSOR_AXIS_LABELS = 2,
-
-  // Contains labels that tensor values correspond to. For example, in
-  // the object detection model, one of the output tensors is the detected
-  // classes. And each value in the tensor refers to the index of label in the
-  // category label file.
-  //
-  // :
-  // If an output tensor has an associated file as TENSOR_VALUE_LABELS, convert
-  // the tensor values into labels, and return a list of string as the output.
-  // If multiple files of the same type are present, the first one is used by
-  // default; additional ones are to be distinguished from one another by their
-  // specified locale.
-  TENSOR_VALUE_LABELS = 3,
-
-  // Contains sigmoid-based score calibration parameters, formatted as CSV.
-  // Lines contain for each index of an output tensor the scale, slope, offset
-  // and (optional) min_score parameters to be used for sigmoid fitting (in this
-  // order and in `strtof`-compatible [1] format).
-  // A line may be left empty to default calibrated scores for this index to
-  // default_score.
-  // In summary, each line should thus contain 0, 3 or 4 comma-separated values.
-  //
-  // See documentation for ScoreCalibrationOptions for details.
-  //
-  // [1]: https://en.cppreference.com/w/c/string/byte/strtof
-  TENSOR_AXIS_SCORE_CALIBRATION = 4,
-
-  // Contains a list of unique words (characters separated by "\n" or in lines)
-  // that help to convert natural language words to embedding vectors.
-  // Added in: 1.0.1
-  VOCABULARY = 5,
-}
-
-table AssociatedFile {
-  // Name of this file. Need to be exact the same as the name of the actual file
-  // packed into the TFLite model as a zip file.
-  //
-  // :
-  // Locates to the actual file in the TFLite model.
-  name:string;
-
-  // A description of what the file is.
-  description:string;
-
-  // Type of the associated file. There may be special pre/post processing for
-  // some types. For example in image classification, a label file of the output
-  // will be used to convert object index into string.
-  //
-  // :
-  // Determines how to process the corresponding tensor.
-  type:AssociatedFileType;
-
-  // An optional locale for this associated file (if applicable). It is
-  // recommended to use an ISO 639-1 letter code (e.g. "en" for English),
-  // optionally completed by a two letter region code (e.g. "en-US" for US
-  // English and "en-CA" for Canadian English).
-  // Leverage this in order to specify e.g multiple label files translated in
-  // different languages.
-  locale:string;
-}
-
-// The basic content type for all tensors.
-//
-// :
-// Input feature tensors:
-// 1. Generates the method to load data from a TensorBuffer.
-// 2. Creates the preprocessing logic. The default processing pipeline is:
-// [NormalizeOp, QuantizeOp].
-// Output feature tensors:
-// 1. Generates the method to return the output data to a TensorBuffer.
-// 2. Creates the post-processing logic. The default processing pipeline is:
-// [DeQuantizeOp].
-table FeatureProperties {
-}
-
-// The type of color space of an image.
-enum ColorSpaceType : byte {
-  UNKNOWN = 0,
-  RGB = 1,
-  GRAYSCALE = 2,
-}
-
-table ImageSize {
-  width:uint;
-  height:uint;
-}
-
-// The properties for image tensors.
-//
-// :
-// Input image tensors:
-// 1. Generates the method to load an image from a TensorImage.
-// 2. Creates the preprocessing logic. The default processing pipeline is:
-// [ResizeOp, NormalizeOp, QuantizeOp].
-// Output image tensors:
-// 1. Generates the method to return the output data to a TensorImage.
-// 2. Creates the post-processing logic. The default processing pipeline is:
-// [DeQuantizeOp].
-table ImageProperties {
-  // The color space of the image.
-  //
-  // :
-  // Determines how to convert the color space of a given image from users.
-  color_space:ColorSpaceType;
-
-  // Indicates the default value of image width and height if the tensor shape
-  // is dynamic. For fixed-size tensor, this size will be consistent with the
-  // expected size.
-  default_size:ImageSize;
-}
-
-// The properties for tensors representing bounding boxes.
-//
-// :
-// Input image tensors: NA.
-// Output image tensors: parses the values into a data stucture that represents
-// bounding boxes. For example, in the generated wrapper for Android, it returns
-// the output as android.graphics.Rect objects.
-enum BoundingBoxType : byte {
-  UNKNOWN = 0,
-  // Represents the bounding box by using the combination of boundaries,
-  // {left, top, right, bottom}.
-  // The default order is {left, top, right, bottom}. Other orders can be
-  // indicated by BoundingBoxProperties.index.
-  BOUNDARIES = 1,
-
-  // Represents the bounding box by using the upper_left corner, width and
-  // height.
-  // The default order is {upper_left_x, upper_left_y, width, height}. Other
-  // orders can be indicated by BoundingBoxProperties.index.
-  UPPER_LEFT = 2,
-
-  // Represents the bounding box by using the center of the box, width and
-  // height. The default order is {center_x, center_y, width, height}. Other
-  // orders can be indicated by BoundingBoxProperties.index.
-  CENTER = 3,
-
-}
-
-enum CoordinateType : byte {
-  // The coordinates are float values from 0 to 1.
-  RATIO = 0,
-  // The coordinates are integers.
-  PIXEL = 1,
-}
-
-table BoundingBoxProperties {
-  // Denotes the order of the elements defined in each bounding box type. An
-  // empty index array represent the default order of each bounding box type.
-  // For example, to denote the default order of BOUNDARIES, {left, top, right,
-  // bottom}, the index should be {0, 1, 2, 3}. To denote the order {left,
-  // right, top, bottom}, the order should be {0, 2, 1, 3}.
-  //
-  // The index array can be applied to all bounding box types to adjust the
-  // order of their corresponding underlying elements.
-  //
-  // :
-  // Indicates how to parse the bounding box values.
-  index:[uint];
-
-  // :
-  // Indicates how to parse the bounding box values.
-  type:BoundingBoxType;
-
-  // :
-  // Indicates how to convert the bounding box back to the original image in
-  // pixels.
-  coordinate_type:CoordinateType;
-}
-
-union ContentProperties {
-  FeatureProperties,
-  ImageProperties,
-  BoundingBoxProperties,
-}
-
-table ValueRange {
-  min:int;
-  max:int;
-}
-
-table Content {
-  // The properties that the content may have, indicating the type of the
-  // Content.
-  //
-  // :
-  // Indicates how to process the tensor.
-  content_properties:ContentProperties;
-
-  // The range of dimensions that the content corresponds to. A NULL
-  // "range" indicates that the content uses up all dimensions,
-  // except the batch axis if applied.
-  //
-  // Here are all the possible situations of how a tensor is composed.
-  // Case 1: The tensor is a single object, such as an image.
-  // For example, the input of an image classifier
-  // (https://www.tensorflow.org/lite/models/image_classification/overview),
-  // a tensor of shape [1, 224, 224, 3]. Dimensions 1 to 3 correspond to the
-  // image. Since dimension 0 is a batch axis, which can be ignored,
-  // "range" can be left as NULL.
-  //
-  // Case 2: The tensor contains multiple instances of the same object.
-  // For example, the output tensor of detected bounding boxes of an object
-  // detection model
-  // (https://www.tensorflow.org/lite/models/object_detection/overview).
-  // The tensor shape is [1, 10, 4]. Here is the what the three dimensions
-  // represent for:
-  // dimension 0: the batch axis.
-  // dimension 1: the 10 objects detected with the highest confidence.
-  // dimension 2: the bounding boxes of the 10 detected objects.
-  // The tensor is essentially 10 bounding boxes. In this case,
-  // "range" should be {min=2; max=2;}.
-  // Another example is the pose estimation model
-  // (https://www.tensorflow.org/lite/models/pose_estimation/overview).
-  // The output tensor of heatmaps is in the shape of [1, 9, 9, 17].
-  // Here is the what the four dimensions represent for:
-  // dimension 0: the batch axis.
-  // dimension 1/2: the heatmap image.
-  // dimension 3: 17 body parts of a person.
-  // Even though the last axis is body part, the real content of this tensor is
-  // the heatmap. "range" should be [min=1; max=2].
-  //
-  // Case 3: The tensor contains multiple different objects. (Not supported by
-  // Content at this point).
-  // Sometimes a tensor may contain multiple different objects, thus different
-  // contents. It is very common for regression models. For example, a model
-  // to predict the fuel efficiency
-  // (https://www.tensorflow.org/tutorials/keras/regression).
-  // The input tensor has shape [1, 9], consisting of 9 features, such as
-  // "Cylinders", "Displacement", "Weight", etc. In this case, dimension 1
-  // contains 9 different contents. However, since these sub-dimension objects
-  // barely need to be specifically processed, their contents are not recorded
-  // in the metadata. Through, the name of each dimension can be set through
-  // TensorMetadata.dimension_names.
-  //
-  // Note that if it is not case 3, a tensor can only have one content type.
-  //
-  // :
-  // Case 1: return a processed single object of certain content type.
-  // Case 2: return a list of processed objects of certain content type. The
-  // generated model interface have API to random access those objects from
-  // the output.
-  range:ValueRange;
-}
-
-// Parameters that are used when normalizing the tensor.
-table NormalizationOptions{
-  // mean and std are normalization parameters. Tensor values are normalized
-  // on a per-channel basis, by the formula
-  //   (x - mean) / std.
-  // If there is only one value in mean or std, we'll propogate the value to
-  // all channels.
-  //
-  // Quantized models share the same normalization parameters as their
-  // corresponding float models. For example, an image input tensor may have
-  // the normalization parameter of
-  //   mean = 127.5f and std = 127.5f.
-  // The image value will be normalized from [0, 255] to [-1, 1].
-  // Then, for quantized models, the image data should be further quantized
-  // according to the quantization parameters. In the case of uint8, the image
-  // data will be scaled back to [0, 255], while for int8, the image data will
-  // be scaled to [-128, 127].
-  //
-  // Both the normalization parameters and quantization parameters can be
-  // retrieved through the metadata extractor library.
-  // TODO(b/156644598): add link for the metadata extractor library.
-
-  // Per-channel mean of the possible values used in normalization.
-  //
-  // :
-  // Apply normalization to input tensors accordingly.
-  mean:[float];
-
-  // Per-channel standard dev. of the possible values used in normalization.
-  //
-  // :
-  // Apply normalization to input tensors accordingly.
-  std:[float];
-}
-
-// The different possible score transforms to apply to uncalibrated scores
-// before applying score calibration.
-enum ScoreTransformationType : byte {
-  // Identity function: g(x) = x.
-  IDENTITY = 0,
-  // Log function: g(x) = log(x).
-  LOG = 1,
-  // Inverse logistic function: g(x) = log(x) - log(1-x).
-  INVERSE_LOGISTIC = 2,
-}
-
-// Options to perform score calibration on an output tensor through sigmoid
-// functions. One of the main purposes of score calibration is to make scores
-// across classes comparable, so that a common threshold can be used for all
-// output classes. This is meant for models producing class predictions as
-// output, e.g. image classification or detection models.
-//
-// For each index in the output tensor, this applies:
-// * `f(x) = scale / (1 + e^-(slope*g(x)+offset))` if `x > min_score` or if no
-//   `min_score` has been specified,
-// * `f(x) = default_score` otherwise or if no scale, slope and offset have been
-//   specified.
-// Where:
-// * scale, slope, offset and (optional) min_score are index-specific parameters
-// * g(x) is an index-independent transform among those defined in
-//   ScoreTransformationType
-// * default_score is an index-independent parameter.
-// An AssociatedFile with type TANSOR_AXIS_SCORE_CALIBRATION specifying the
-// index-specific parameters must be associated with the corresponding
-// TensorMetadata for score calibration be applied.
-table ScoreCalibrationOptions {
-  // The function to use for transforming the uncalibrated score before
-  // applying score calibration.
-  score_transformation:ScoreTransformationType;
-
-  // The default calibrated score to apply if the uncalibrated score is
-  // below min_score or if no parameters were specified for a given index.
-  default_score:float;
-}
-
-// Performs thresholding on output tensor values, in order to filter out
-// low-confidence results.
-table ScoreThresholdingOptions {
-  // The recommended global threshold below which results are considered
-  // low-confidence and should be filtered out.
-  global_score_threshold:float;
-}
-
-// Options that are used when processing the tensor.
-union ProcessUnitOptions {
-  NormalizationOptions,
-  ScoreCalibrationOptions,
-  ScoreThresholdingOptions,
-}
-
-// A process unit that is used to process the tensor out-of-graph.
-table ProcessUnit {
-  options:ProcessUnitOptions;
-}
-
-
-// Statistics to describe a tensor.
-table Stats {
-  // Max and min are not currently used in tflite.support codegen. They mainly
-  // serve as references for users to better understand the model. They can also
-  // be used to validate model pre/post processing results.
-  // If there is only one value in max or min, we'll propogate the value to
-  // all channels.
-
-  // Per-channel maximum value of the tensor.
-  max:[float];
-
-  // Per-channel minimum value of the tensor.
-  min:[float];
-}
-
-// Detailed information of an input or output tensor.
-table TensorMetadata {
-  // Name of the tensor.
-  //
-  // :
-  // The name of this tensor in the generated model interface.
-  name:string;
-
-  // A description of the tensor.
-  description:string;
-
-  // A list of names of the dimensions in this tensor. The length of
-  // dimension_names need to match the number of dimensions in this tensor.
-  //
-  // :
-  // The name of each dimension in the generated model interface. See "Case 2"
-  // in the comments of Content.range.
-  dimension_names:[string];
-
-  // The content that represents this tensor.
-  //
-  // :
-  // Determines how to process this tensor. See each item in ContentProperties
-  // for the default process units that will be applied to the tensor.
-  content:Content;
-
-  // The process units that are used to process the tensor out-of-graph.
-  //
-  // :
-  // Contains the parameters of the default processing pipeline for each content
-  // type, such as the normalization parameters in all content types. See the
-  // items under ContentProperties for the details of the default processing
-  // pipeline.
-  process_units:[ProcessUnit];
-
-  // The statistics of the tensor values.
-  stats:Stats;
-
-  // A list of associated files of this tensor.
-  //
-  // :
-  // Contains processing parameters of this tensor, such as normalization.
-  associated_files:[AssociatedFile];
-}
-
-table SubGraphMetadata {
-  // Name of the subgraph.
-  //
-  // Note that, since TFLite only support one subgraph at this moment, the
-  // Codegen tool will use the name in ModelMetadata in the generated model
-  // interface.
-  name:string;
-
-  // A description explains details about what the subgraph does.
-  description:string;
-
-  // Metadata of all input tensors used in this subgraph. It matches extactly
-  // the input tensors specified by `SubGraph.inputs` in the TFLite
-  // schema.fbs file[2]. The number of `TensorMetadata` in the array should
-  // equal to the number of indices in `SubGraph.inputs`.
-  //
-  // [2]: tensorflow/lite/schema/schema.fbs
-  // :
-  // Determines how to process the inputs.
-  input_tensor_metadata:[TensorMetadata];
-
-  // Metadata of all output tensors used in this subgraph. It matches extactly
-  // the output tensors specified by `SubGraph.outputs` in the TFLite
-  // schema.fbs file[2]. The number of `TensorMetadata` in the array should
-  // equal to the number of indices in `SubGraph.outputs`.
-  //
-  // :
-  // Determines how to process the outputs.
-  output_tensor_metadata:[TensorMetadata];
-
-  // A list of associated files of this subgraph.
-  associated_files:[AssociatedFile];
-}
-
-table ModelMetadata {
-  // Name of the model.
-  //
-  // :
-  // The name of the model in the generated model interface.
-  name:string;
-
-  // Model description in schema.
-  description:string;
-
-  // Version of the model that specified by model creators.
-  version:string;
-
-  // Noted that, the minimum required TFLite runtime version that the model is
-  // compatible with, has already been added as a metadata entry in tflite
-  // schema. We'll decide later if we want to move it here, and keep it with
-  // other metadata entries.
-
-  // Metadata of all the subgraphs of the model. The 0th is assumed to be the
-  // main subgraph.
-  //
-  // :
-  // Determines how to process the inputs and outputs.
-  subgraph_metadata:[SubGraphMetadata];
-
-  // The person who creates this model.
-  author:string;
-
-  // Licenses that may apply to this model.
-  license:string;
-
-  // A list of associated files of this model.
-  associated_files:[AssociatedFile];
-
-  // The minimum metadata parser version that can fully understand the fields in
-  // the metadata flatbuffer. The version is effectively the largest version
-  // number among the versions of all the fields populated and the smallest
-  // compatible version indicated by the file identifier.
-  //
-  // This field is automaticaly populated by the MetadataPopulator when
-  // the metadata is populated into a TFLite model.
-  min_parser_version:string;
-}
-// LINT.ThenChange(//tensorflow/lite/experimental/\
-//     support/metadata/cc/metadata_version.cc)
-
-root_type ModelMetadata;
diff --git a/tensorflow/lite/experimental/support/metadata/metadata_test.py b/tensorflow/lite/experimental/support/metadata/metadata_test.py
deleted file mode 100644
index 28395041746..00000000000
--- a/tensorflow/lite/experimental/support/metadata/metadata_test.py
+++ /dev/null
@@ -1,484 +0,0 @@
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# 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
-#
-#     http://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.
-# ==============================================================================
-"""Tests for tensorflow.lite.experimental.support.metadata.metadata."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
-import six
-
-from flatbuffers.python import flatbuffers
-from tensorflow.lite.experimental.support.metadata import metadata as _metadata
-from tensorflow.lite.experimental.support.metadata import metadata_schema_py_generated as _metadata_fb
-from tensorflow.lite.experimental.support.metadata import schema_py_generated as _schema_fb
-from tensorflow.python.framework import test_util
-from tensorflow.python.platform import resource_loader
-from tensorflow.python.platform import test
-
-
-class MetadataTest(test_util.TensorFlowTestCase):
-
-  def setUp(self):
-    super(MetadataTest, self).setUp()
-    self._invalid_model_buf = None
-    self._invalid_file = "not_existed_file"
-    self._empty_model_buf = self._create_empty_model_buf()
-    self._empty_model_file = self.create_tempfile().full_path
-    with open(self._empty_model_file, "wb") as f:
-      f.write(self._empty_model_buf)
-    self._model_file = self._create_model_file_with_metadata_and_buf_fields()
-    self._metadata_file = self._create_metadata_file()
-    self._metadata_file_with_version = self._create_metadata_file_with_version(
-        self._metadata_file, "1.0.0")
-    self._file1 = self.create_tempfile("file1").full_path
-    self._file2 = self.create_tempfile("file2").full_path
-    self._file3 = self.create_tempfile("file3").full_path
-
-  def _create_empty_model_buf(self):
-    model = _schema_fb.ModelT()
-    model_builder = flatbuffers.Builder(0)
-    model_builder.Finish(
-        model.Pack(model_builder),
-        _metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER)
-    return model_builder.Output()
-
-  def _create_model_file_with_metadata_and_buf_fields(self):
-    metadata_field = _schema_fb.MetadataT()
-    metadata_field.name = "meta"
-    buffer_field = _schema_fb.BufferT()
-    model = _schema_fb.ModelT()
-    model.metadata = [metadata_field, metadata_field]
-    model.buffers = [buffer_field, buffer_field, buffer_field]
-    model_builder = flatbuffers.Builder(0)
-    model_builder.Finish(
-        model.Pack(model_builder),
-        _metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER)
-
-    mnodel_file = self.create_tempfile().full_path
-    with open(mnodel_file, "wb") as f:
-      f.write(model_builder.Output())
-
-    return mnodel_file
-
-  def _create_metadata_file(self):
-    associated_file1 = _metadata_fb.AssociatedFileT()
-    associated_file1.name = b"file1"
-    associated_file2 = _metadata_fb.AssociatedFileT()
-    associated_file2.name = b"file2"
-    self.expected_recorded_files = [
-        six.ensure_str(associated_file1.name),
-        six.ensure_str(associated_file2.name)
-    ]
-
-    output_meta = _metadata_fb.TensorMetadataT()
-    output_meta.associatedFiles = [associated_file2]
-    subgraph = _metadata_fb.SubGraphMetadataT()
-    subgraph.outputTensorMetadata = [output_meta]
-
-    model_meta = _metadata_fb.ModelMetadataT()
-    model_meta.name = "Mobilenet_quantized"
-    model_meta.associatedFiles = [associated_file1]
-    model_meta.subgraphMetadata = [subgraph]
-    b = flatbuffers.Builder(0)
-    b.Finish(
-        model_meta.Pack(b),
-        _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
-
-    metadata_file = self.create_tempfile().full_path
-    with open(metadata_file, "wb") as f:
-      f.write(b.Output())
-    return metadata_file
-
-  def _create_model_buffer_with_wrong_identifier(self):
-    wrong_identifier = b"widn"
-    model = _schema_fb.ModelT()
-    model_builder = flatbuffers.Builder(0)
-    model_builder.Finish(model.Pack(model_builder), wrong_identifier)
-    return model_builder.Output()
-
-  def _create_metadata_buffer_with_wrong_identifier(self):
-    # Creates a metadata with wrong identifier
-    wrong_identifier = b"widn"
-    metadata = _metadata_fb.ModelMetadataT()
-    metadata_builder = flatbuffers.Builder(0)
-    metadata_builder.Finish(metadata.Pack(metadata_builder), wrong_identifier)
-    return metadata_builder.Output()
-
-  def _populate_metadata_with_identifier(self, model_buf, metadata_buf,
-                                         identifier):
-    # For testing purposes only. MetadataPopulator cannot populate metadata with
-    # wrong identifiers.
-    model = _schema_fb.ModelT.InitFromObj(
-        _schema_fb.Model.GetRootAsModel(model_buf, 0))
-    buffer_field = _schema_fb.BufferT()
-    buffer_field.data = metadata_buf
-    model.buffers = [buffer_field]
-    # Creates a new metadata field.
-    metadata_field = _schema_fb.MetadataT()
-    metadata_field.name = _metadata.MetadataPopulator.METADATA_FIELD_NAME
-    metadata_field.buffer = len(model.buffers) - 1
-    model.metadata = [metadata_field]
-    b = flatbuffers.Builder(0)
-    b.Finish(model.Pack(b), identifier)
-    return b.Output()
-
-  def _create_metadata_file_with_version(self, metadata_file, min_version):
-    # Creates a new metadata file with the specified min_version for testing
-    # purposes.
-    with open(metadata_file, "rb") as f:
-      metadata_buf = bytearray(f.read())
-
-    metadata = _metadata_fb.ModelMetadataT.InitFromObj(
-        _metadata_fb.ModelMetadata.GetRootAsModelMetadata(metadata_buf, 0))
-    metadata.minParserVersion = min_version
-
-    b = flatbuffers.Builder(0)
-    b.Finish(
-        metadata.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
-
-    metadata_file_with_version = self.create_tempfile().full_path
-    with open(metadata_file_with_version, "wb") as f:
-      f.write(b.Output())
-    return metadata_file_with_version
-
-
-class MetadataPopulatorTest(MetadataTest):
-
-  def testToValidModelFile(self):
-    populator = _metadata.MetadataPopulator.with_model_file(
-        self._empty_model_file)
-    self.assertIsInstance(populator, _metadata.MetadataPopulator)
-
-  def testToInvalidModelFile(self):
-    with self.assertRaises(IOError) as error:
-      _metadata.MetadataPopulator.with_model_file(self._invalid_file)
-    self.assertEqual("File, '{0}', does not exist.".format(self._invalid_file),
-                     str(error.exception))
-
-  def testToValidModelBuffer(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    self.assertIsInstance(populator, _metadata.MetadataPopulator)
-
-  def testToInvalidModelBuffer(self):
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataPopulator.with_model_buffer(self._invalid_model_buf)
-    self.assertEqual("model_buf cannot be empty.", str(error.exception))
-
-  def testToModelBufferWithWrongIdentifier(self):
-    model_buf = self._create_model_buffer_with_wrong_identifier()
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataPopulator.with_model_buffer(model_buf)
-    self.assertEqual(
-        "The model provided does not have the expected identifier, and "
-        "may not be a valid TFLite model.", str(error.exception))
-
-  def testSinglePopulateAssociatedFile(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    populator.load_associated_files([self._file1])
-    populator.populate()
-
-    packed_files = populator.get_packed_associated_file_list()
-    expected_packed_files = [os.path.basename(self._file1)]
-    self.assertEqual(set(packed_files), set(expected_packed_files))
-
-  def testRepeatedPopulateAssociatedFile(self):
-    populator = _metadata.MetadataPopulator.with_model_file(
-        self._empty_model_file)
-    populator.load_associated_files([self._file1, self._file2])
-    # Loads file2 multiple times.
-    populator.load_associated_files([self._file2])
-    populator.populate()
-
-    packed_files = populator.get_packed_associated_file_list()
-    expected_packed_files = [
-        os.path.basename(self._file1),
-        os.path.basename(self._file2)
-    ]
-    self.assertEqual(len(packed_files), 2)
-    self.assertEqual(set(packed_files), set(expected_packed_files))
-
-    # Check if the model buffer read from file is the same as that read from
-    # get_model_buffer().
-    with open(self._empty_model_file, "rb") as f:
-      model_buf_from_file = f.read()
-    model_buf_from_getter = populator.get_model_buffer()
-    self.assertEqual(model_buf_from_file, model_buf_from_getter)
-
-  def testPopulateInvalidAssociatedFile(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    with self.assertRaises(IOError) as error:
-      populator.load_associated_files([self._invalid_file])
-    self.assertEqual("File, '{0}', does not exist.".format(self._invalid_file),
-                     str(error.exception))
-
-  def testPopulatePackedAssociatedFile(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    populator.load_associated_files([self._file1])
-    populator.populate()
-    with self.assertRaises(ValueError) as error:
-      populator.load_associated_files([self._file1])
-      populator.populate()
-    self.assertEqual(
-        "File, '{0}', has already been packed.".format(
-            os.path.basename(self._file1)), str(error.exception))
-
-  def testGetPackedAssociatedFileList(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    packed_files = populator.get_packed_associated_file_list()
-    self.assertEqual(packed_files, [])
-
-  def testPopulateMetadataFileToEmptyModelFile(self):
-    populator = _metadata.MetadataPopulator.with_model_file(
-        self._empty_model_file)
-    populator.load_metadata_file(self._metadata_file)
-    populator.load_associated_files([self._file1, self._file2])
-    populator.populate()
-
-    with open(self._empty_model_file, "rb") as f:
-      model_buf_from_file = f.read()
-    model = _schema_fb.Model.GetRootAsModel(model_buf_from_file, 0)
-    metadata_field = model.Metadata(0)
-    self.assertEqual(
-        six.ensure_str(metadata_field.Name()),
-        six.ensure_str(_metadata.MetadataPopulator.METADATA_FIELD_NAME))
-
-    buffer_index = metadata_field.Buffer()
-    buffer_data = model.Buffers(buffer_index)
-    metadata_buf_np = buffer_data.DataAsNumpy()
-    metadata_buf = metadata_buf_np.tobytes()
-    with open(self._metadata_file_with_version, "rb") as f:
-      expected_metadata_buf = bytearray(f.read())
-    self.assertEqual(metadata_buf, expected_metadata_buf)
-
-    recorded_files = populator.get_recorded_associated_file_list()
-    self.assertEqual(set(recorded_files), set(self.expected_recorded_files))
-
-    # Up to now, we've proved the correctness of the model buffer that read from
-    # file. Then we'll test if get_model_buffer() gives the same model buffer.
-    model_buf_from_getter = populator.get_model_buffer()
-    self.assertEqual(model_buf_from_file, model_buf_from_getter)
-
-  def testPopulateMetadataFileWithoutAssociatedFiles(self):
-    populator = _metadata.MetadataPopulator.with_model_file(
-        self._empty_model_file)
-    populator.load_metadata_file(self._metadata_file)
-    populator.load_associated_files([self._file1])
-    # Suppose to populate self._file2, because it is recorded in the metadta.
-    with self.assertRaises(ValueError) as error:
-      populator.populate()
-    self.assertEqual(("File, '{0}', is recorded in the metadata, but has "
-                      "not been loaded into the populator.").format(
-                          os.path.basename(self._file2)), str(error.exception))
-
-  def testPopulateMetadataBufferWithWrongIdentifier(self):
-    metadata_buf = self._create_metadata_buffer_with_wrong_identifier()
-    populator = _metadata.MetadataPopulator.with_model_file(self._model_file)
-    with self.assertRaises(ValueError) as error:
-      populator.load_metadata_buffer(metadata_buf)
-    self.assertEqual(
-        "The metadata buffer does not have the expected identifier, and may not"
-        " be a valid TFLite Metadata.", str(error.exception))
-
-  def _assert_golden_metadata(self, model_file):
-    with open(model_file, "rb") as f:
-      model_buf_from_file = f.read()
-    model = _schema_fb.Model.GetRootAsModel(model_buf_from_file, 0)
-    # There are two elements in model.Metadata array before the population.
-    # Metadata should be packed to the third element in the array.
-    metadata_field = model.Metadata(2)
-    self.assertEqual(
-        six.ensure_str(metadata_field.Name()),
-        six.ensure_str(_metadata.MetadataPopulator.METADATA_FIELD_NAME))
-
-    buffer_index = metadata_field.Buffer()
-    buffer_data = model.Buffers(buffer_index)
-    metadata_buf_np = buffer_data.DataAsNumpy()
-    metadata_buf = metadata_buf_np.tobytes()
-    with open(self._metadata_file_with_version, "rb") as f:
-      expected_metadata_buf = bytearray(f.read())
-    self.assertEqual(metadata_buf, expected_metadata_buf)
-
-  def testPopulateMetadataFileToModelWithMetadataAndAssociatedFiles(self):
-    # First, creates a dummy metadata. Populates it and the associated files
-    # into the model.
-    model_meta = _metadata_fb.ModelMetadataT()
-    model_meta.name = "Mobilenet_quantized"
-    b = flatbuffers.Builder(0)
-    b.Finish(
-        model_meta.Pack(b),
-        _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
-    metadata_buf = b.Output()
-
-    populator1 = _metadata.MetadataPopulator.with_model_file(self._model_file)
-    populator1.load_metadata_buffer(metadata_buf)
-    populator1.load_associated_files([self._file1, self._file2])
-    populator1.populate()
-
-    # Then, populates the metadata again.
-    populator2 = _metadata.MetadataPopulator.with_model_file(self._model_file)
-    populator2.load_metadata_file(self._metadata_file)
-    populator2.populate()
-
-    # Tests if the metadata is populated correctly.
-    self._assert_golden_metadata(self._model_file)
-
-  def testPopulateMetadataFileToModelFileWithMetadataAndBufFields(self):
-    populator = _metadata.MetadataPopulator.with_model_file(self._model_file)
-    populator.load_metadata_file(self._metadata_file)
-    populator.load_associated_files([self._file1, self._file2])
-    populator.populate()
-
-    # Tests if the metadata is populated correctly.
-    self._assert_golden_metadata(self._model_file)
-
-    recorded_files = populator.get_recorded_associated_file_list()
-    self.assertEqual(set(recorded_files), set(self.expected_recorded_files))
-
-    # Up to now, we've proved the correctness of the model buffer that read from
-    # file. Then we'll test if get_model_buffer() gives the same model buffer.
-    with open(self._model_file, "rb") as f:
-      model_buf_from_file = f.read()
-    model_buf_from_getter = populator.get_model_buffer()
-    self.assertEqual(model_buf_from_file, model_buf_from_getter)
-
-  def testPopulateInvalidMetadataFile(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    with self.assertRaises(IOError) as error:
-      populator.load_metadata_file(self._invalid_file)
-    self.assertEqual("File, '{0}', does not exist.".format(self._invalid_file),
-                     str(error.exception))
-
-  def testPopulateInvalidMetadataBuffer(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    with self.assertRaises(ValueError) as error:
-      populator.load_metadata_buffer([])
-    self.assertEqual("The metadata to be populated is empty.",
-                     str(error.exception))
-
-  def testGetModelBufferBeforePopulatingData(self):
-    populator = _metadata.MetadataPopulator.with_model_buffer(
-        self._empty_model_buf)
-    model_buf = populator.get_model_buffer()
-    expected_model_buf = self._empty_model_buf
-    self.assertEqual(model_buf, expected_model_buf)
-
-
-class MetadataDisplayerTest(MetadataTest):
-
-  def setUp(self):
-    super(MetadataDisplayerTest, self).setUp()
-    self._model_file = self._create_model_with_metadata_and_associated_files()
-
-  def _create_model_with_metadata_and_associated_files(self):
-    model_buf = self._create_empty_model_buf()
-    model_file = self.create_tempfile().full_path
-    with open(model_file, "wb") as f:
-      f.write(model_buf)
-
-    populator = _metadata.MetadataPopulator.with_model_file(model_file)
-    populator.load_metadata_file(self._metadata_file)
-    populator.load_associated_files([self._file1, self._file2])
-    populator.populate()
-    return model_file
-
-  def test_load_model_buffer_metadataBufferWithWrongIdentifier_throwsException(
-      self):
-    model_buf = self._create_model_buffer_with_wrong_identifier()
-    metadata_buf = self._create_metadata_buffer_with_wrong_identifier()
-    model_buf = self._populate_metadata_with_identifier(
-        model_buf, metadata_buf,
-        _metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER)
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataDisplayer.with_model_buffer(model_buf)
-    self.assertEqual(
-        "The metadata buffer does not have the expected identifier, and may not"
-        " be a valid TFLite Metadata.", str(error.exception))
-
-  def test_load_model_buffer_modelBufferWithWrongIdentifier_throwsException(
-      self):
-    model_buf = self._create_model_buffer_with_wrong_identifier()
-    metadata_file = self._create_metadata_file()
-    wrong_identifier = b"widn"
-    with open(metadata_file, "rb") as f:
-      metadata_buf = bytearray(f.read())
-    model_buf = self._populate_metadata_with_identifier(model_buf, metadata_buf,
-                                                        wrong_identifier)
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataDisplayer.with_model_buffer(model_buf)
-    self.assertEqual(
-        "The model provided does not have the expected identifier, and "
-        "may not be a valid TFLite model.", str(error.exception))
-
-  def test_load_model_file_invalidModelFile_throwsException(self):
-    with self.assertRaises(IOError) as error:
-      _metadata.MetadataDisplayer.with_model_file(self._invalid_file)
-    self.assertEqual("File, '{0}', does not exist.".format(self._invalid_file),
-                     str(error.exception))
-
-  def test_load_model_file_modelWithoutMetadata_throwsException(self):
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataDisplayer.with_model_file(self._empty_model_file)
-    self.assertEqual("The model does not have metadata.", str(error.exception))
-
-  def test_load_model_file_modelWithMetadata(self):
-    displayer = _metadata.MetadataDisplayer.with_model_file(self._model_file)
-    self.assertIsInstance(displayer, _metadata.MetadataDisplayer)
-
-  def test_load_model_buffer_modelWithOutMetadata_throwsException(self):
-    with self.assertRaises(ValueError) as error:
-      _metadata.MetadataDisplayer.with_model_buffer(
-          self._create_empty_model_buf())
-    self.assertEqual("The model does not have metadata.", str(error.exception))
-
-  def test_load_model_buffer_modelWithMetadata(self):
-    displayer = _metadata.MetadataDisplayer.with_model_buffer(
-        open(self._model_file, "rb").read())
-    self.assertIsInstance(displayer, _metadata.MetadataDisplayer)
-
-  def test_get_metadata_json_modelWithMetadata(self):
-    displayer = _metadata.MetadataDisplayer.with_model_file(self._model_file)
-    actual = displayer.get_metadata_json()
-
-    # Verifies the generated json file.
-    golden_json_file_path = resource_loader.get_path_to_datafile(
-        "testdata/golden_json.json")
-    with open(golden_json_file_path, "r") as f:
-      expected = f.read()
-    self.assertEqual(actual, expected)
-
-  def test_get_packed_associated_file_list_modelWithMetadata(self):
-    displayer = _metadata.MetadataDisplayer.with_model_file(self._model_file)
-    packed_files = displayer.get_packed_associated_file_list()
-
-    expected_packed_files = [
-        os.path.basename(self._file1),
-        os.path.basename(self._file2)
-    ]
-    self.assertEqual(len(packed_files), 2)
-    self.assertEqual(set(packed_files), set(expected_packed_files))
-
-
-if __name__ == "__main__":
-  test.main()
diff --git a/tensorflow/lite/experimental/support/metadata/testdata/golden_json.json b/tensorflow/lite/experimental/support/metadata/testdata/golden_json.json
deleted file mode 100644
index 9ff5581fbff..00000000000
--- a/tensorflow/lite/experimental/support/metadata/testdata/golden_json.json
+++ /dev/null
@@ -1,22 +0,0 @@
-{
-  "name": "Mobilenet_quantized",
-  "subgraph_metadata": [
-    {
-      "output_tensor_metadata": [
-        {
-          "associated_files": [
-            {
-              "name": "file2"
-            }
-          ]
-        }
-      ]
-    }
-  ],
-  "associated_files": [
-    {
-      "name": "file1"
-    }
-  ],
-  "min_parser_version": "1.0.0"
-}

From 50eb68909561ac674ab905f52f02828fd71d05a3 Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Tue, 30 Jun 2020 19:58:57 -0700
Subject: [PATCH 1370/1390] More closely mimic the tf.Tensor api from
 KerasTensor, and mark tf op layers as being loaded from the configs in
 savedmodels. This fixes issues with more subtle usages of Keras functional
 models when we enable KerasTensors.

(E.g. `is_tensor(keras_tensor)` will now return True, like how it does for functional models in head. This is important in code that relies heavily on op layers).

PiperOrigin-RevId: 319148736
Change-Id: I209583f049dfe08b12ee51fd396e6fe21075aee0
---
 .../python/keras/engine/keras_tensor.py       | 35 +++++++++++++++++++
 tensorflow/python/keras/layers/core.py        |  3 ++
 tensorflow/python/keras/losses.py             |  5 +--
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/keras/engine/keras_tensor.py b/tensorflow/python/keras/engine/keras_tensor.py
index 50d377c5292..98560aa8e46 100644
--- a/tensorflow/python/keras/engine/keras_tensor.py
+++ b/tensorflow/python/keras/engine/keras_tensor.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import type_spec as type_spec_module
 from tensorflow.python.ops import array_ops
@@ -122,6 +123,35 @@ class KerasTensor(object):
   def get_shape(self):
     return self.shape
 
+  def __len__(self):
+    raise TypeError('Symbolic Functional model inputs/outputs in Keras do not '
+                    'implement `__len__`. You may be '
+                    'seeing this error if you are passing it '
+                    'to a TF API that does not support dispatching to Keras '
+                    'lambda layers.')
+
+  @property
+  def is_tensor_like(self):
+    return True
+
+  def set_shape(self, shape):
+    """Updates the shape of this KerasTensor. Mimics `tf.Tensor.set_shape()`."""
+    if not isinstance(shape, tensor_shape.TensorShape):
+      shape = tensor_shape.TensorShape(shape)
+    if shape.dims is not None:
+      dim_list = [dim.value for dim in shape.dims]
+      for dim in range(len(dim_list)):
+        if dim_list[dim] is None and self.shape.dims is not None:
+          dim_list[dim] = self.shape.dims[dim]
+      shape = tensor_shape.TensorShape(dim_list)
+    if not self.shape.is_compatible_with(shape):
+      raise ValueError(
+          "Keras Intermediate Value's shape %s is not"
+          "compatible with supplied shape %s" %
+          (self.shape, shape))
+    else:
+      self._type_spec._shape = shape  # pylint: disable=protected-access
+
   @property
   def dtype(self):
     """Returns the `dtype` of elements in the tensor."""
@@ -168,6 +198,11 @@ class KerasTensor(object):
     for operator in ops.Tensor.OVERLOADABLE_OPERATORS:
       cls._overload_operator(operator)
 
+    # We include `experimental_ref` for versions of TensorFlow that
+    # still include the deprecated method in Tensors.
+    if hasattr(ops.Tensor, 'experimental_ref'):
+      cls._overload_operator('experimental_ref')
+
   @classmethod
   def _overload_operator(cls, operator):  # pylint: disable=invalid-name
     """Overload an operator with the same overloading as `ops.Tensor`.
diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py
index e64a1c27bcf..292c85560b4 100644
--- a/tensorflow/python/keras/layers/core.py
+++ b/tensorflow/python/keras/layers/core.py
@@ -1298,6 +1298,9 @@ class TFOpLambda(Layer):
       return self._call_wrapper(*args, **kwargs)
     self.call = tf_decorator.make_decorator(function, _call_wrapper)
 
+    # Do not individually trace op layers in the SavedModel.
+    self._must_restore_from_config = True
+
     super(TFOpLambda, self).__init__(**kwargs)
 
     # Warning on every invocation will be quite irksome in Eager mode.
diff --git a/tensorflow/python/keras/losses.py b/tensorflow/python/keras/losses.py
index d1f406fc40f..9b162560120 100644
--- a/tensorflow/python/keras/losses.py
+++ b/tensorflow/python/keras/losses.py
@@ -247,8 +247,9 @@ class LossFunctionWrapper(Loss):
       Loss values per sample.
     """
     if tensor_util.is_tensor(y_pred) and tensor_util.is_tensor(y_true):
-      y_pred, y_true = tf_losses_util.squeeze_or_expand_dimensions(
-          y_pred, y_true)
+      if not K.is_keras_tensor(y_pred) and not K.is_keras_tensor(y_true):
+        y_pred, y_true = tf_losses_util.squeeze_or_expand_dimensions(
+            y_pred, y_true)
     ag_fn = autograph.tf_convert(self.fn, ag_ctx.control_status_ctx())
     return ag_fn(y_true, y_pred, **self._fn_kwargs)
 

From 3a65a5e05b91f52327f2eff2d74d058f0f351197 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 30 Jun 2020 21:00:14 -0700
Subject: [PATCH 1371/1390] More closely mimic the tf.Tensor api from
 KerasTensor, and mark tf op layers as being loaded from the configs in
 savedmodels. This fixes issues with more subtle usages of Keras functional
 models when we enable KerasTensors.

(E.g. `is_tensor(keras_tensor)` will now return True, like how it does for functional models in head. This is important in code that relies heavily on op layers).

PiperOrigin-RevId: 319155010
Change-Id: Iabb06e84b5b2a0b0385a76b6982cd649e1d16816
---
 .../python/keras/engine/keras_tensor.py       | 35 -------------------
 tensorflow/python/keras/layers/core.py        |  3 --
 tensorflow/python/keras/losses.py             |  5 ++-
 3 files changed, 2 insertions(+), 41 deletions(-)

diff --git a/tensorflow/python/keras/engine/keras_tensor.py b/tensorflow/python/keras/engine/keras_tensor.py
index 98560aa8e46..50d377c5292 100644
--- a/tensorflow/python/keras/engine/keras_tensor.py
+++ b/tensorflow/python/keras/engine/keras_tensor.py
@@ -21,7 +21,6 @@ from __future__ import print_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import type_spec as type_spec_module
 from tensorflow.python.ops import array_ops
@@ -123,35 +122,6 @@ class KerasTensor(object):
   def get_shape(self):
     return self.shape
 
-  def __len__(self):
-    raise TypeError('Symbolic Functional model inputs/outputs in Keras do not '
-                    'implement `__len__`. You may be '
-                    'seeing this error if you are passing it '
-                    'to a TF API that does not support dispatching to Keras '
-                    'lambda layers.')
-
-  @property
-  def is_tensor_like(self):
-    return True
-
-  def set_shape(self, shape):
-    """Updates the shape of this KerasTensor. Mimics `tf.Tensor.set_shape()`."""
-    if not isinstance(shape, tensor_shape.TensorShape):
-      shape = tensor_shape.TensorShape(shape)
-    if shape.dims is not None:
-      dim_list = [dim.value for dim in shape.dims]
-      for dim in range(len(dim_list)):
-        if dim_list[dim] is None and self.shape.dims is not None:
-          dim_list[dim] = self.shape.dims[dim]
-      shape = tensor_shape.TensorShape(dim_list)
-    if not self.shape.is_compatible_with(shape):
-      raise ValueError(
-          "Keras Intermediate Value's shape %s is not"
-          "compatible with supplied shape %s" %
-          (self.shape, shape))
-    else:
-      self._type_spec._shape = shape  # pylint: disable=protected-access
-
   @property
   def dtype(self):
     """Returns the `dtype` of elements in the tensor."""
@@ -198,11 +168,6 @@ class KerasTensor(object):
     for operator in ops.Tensor.OVERLOADABLE_OPERATORS:
       cls._overload_operator(operator)
 
-    # We include `experimental_ref` for versions of TensorFlow that
-    # still include the deprecated method in Tensors.
-    if hasattr(ops.Tensor, 'experimental_ref'):
-      cls._overload_operator('experimental_ref')
-
   @classmethod
   def _overload_operator(cls, operator):  # pylint: disable=invalid-name
     """Overload an operator with the same overloading as `ops.Tensor`.
diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py
index 292c85560b4..e64a1c27bcf 100644
--- a/tensorflow/python/keras/layers/core.py
+++ b/tensorflow/python/keras/layers/core.py
@@ -1298,9 +1298,6 @@ class TFOpLambda(Layer):
       return self._call_wrapper(*args, **kwargs)
     self.call = tf_decorator.make_decorator(function, _call_wrapper)
 
-    # Do not individually trace op layers in the SavedModel.
-    self._must_restore_from_config = True
-
     super(TFOpLambda, self).__init__(**kwargs)
 
     # Warning on every invocation will be quite irksome in Eager mode.
diff --git a/tensorflow/python/keras/losses.py b/tensorflow/python/keras/losses.py
index 9b162560120..d1f406fc40f 100644
--- a/tensorflow/python/keras/losses.py
+++ b/tensorflow/python/keras/losses.py
@@ -247,9 +247,8 @@ class LossFunctionWrapper(Loss):
       Loss values per sample.
     """
     if tensor_util.is_tensor(y_pred) and tensor_util.is_tensor(y_true):
-      if not K.is_keras_tensor(y_pred) and not K.is_keras_tensor(y_true):
-        y_pred, y_true = tf_losses_util.squeeze_or_expand_dimensions(
-            y_pred, y_true)
+      y_pred, y_true = tf_losses_util.squeeze_or_expand_dimensions(
+          y_pred, y_true)
     ag_fn = autograph.tf_convert(self.fn, ag_ctx.control_status_ctx())
     return ag_fn(y_true, y_pred, **self._fn_kwargs)
 

From 116092e99a146375bb9edd7462272bf6981ed2b9 Mon Sep 17 00:00:00 2001
From: Jiho Choi 
Date: Tue, 30 Jun 2020 21:09:57 -0700
Subject: [PATCH 1372/1390] Add element tracing for map_and_batch.

PiperOrigin-RevId: 319156094
Change-Id: I68b8289071f41fbf4f039000df13f70003ec5b56
---
 .../core/kernels/data/experimental/BUILD      |  2 ++
 .../experimental/map_and_batch_dataset_op.cc  | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index dc188071814..8457cfa6145 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -319,6 +319,8 @@ tf_kernel_library(
         "//tensorflow/core/kernels/data:dataset_utils",
         "//tensorflow/core/kernels/data:name_utils",
         "//tensorflow/core/kernels/data:stats_utils",
+        "//tensorflow/core/profiler/lib:traceme",
+        "//tensorflow/core/profiler/lib:traceme_encode",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
index 0cf85a58985..fdc63bdb913 100644
--- a/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
@@ -36,6 +36,8 @@ limitations under the License.
 #include "tensorflow/core/platform/cpu_info.h"
 #include "tensorflow/core/platform/stringprintf.h"
 #include "tensorflow/core/platform/tracing.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
+#include "tensorflow/core/profiler/lib/traceme_encode.h"
 
 namespace tensorflow {
 namespace data {
@@ -242,6 +244,10 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
         batch_results_.pop_front();
         cond_var_->notify_all();
       }
+      profiler::TraceMe traceme([&] {
+        return profiler::TraceMeEncode("MapAndBatchConsume",
+                                       {{"element_id", result->id}});
+      });
       return ProcessResult(ctx, result, out_tensors, end_of_sequence);
     }
 
@@ -314,7 +320,7 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
     // BatchResult encapsulates the output batch, as well as ancillary
     // metadata required to execute the fused map-and-batch operation.
     struct BatchResult {
-      explicit BatchResult(int64 batch_size) {
+      explicit BatchResult(int64 batch_size, int64 id = -1) : id(id) {
         end_of_input = false;
         num_calls = batch_size;
         num_elements = 0;
@@ -348,6 +354,7 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
       int64 status_offset TF_GUARDED_BY(mu);
       // Counts the number of outstanding calls for this batch.
       int64 num_calls;  // access guarded by owner's mutex
+      int64 id = -1;
     };
 
     void CallCompleted(const std::shared_ptr& ctx,
@@ -370,6 +377,10 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
     void CallFunction(std::shared_ptr ctx,
                       const std::shared_ptr& result, int64 offset)
         TF_LOCKS_EXCLUDED(*mu_) {
+      profiler::TraceMe traceme([&] {
+        return profiler::TraceMeEncode("MapAndBatchProduce",
+                                       {{"element_id", result->id}});
+      });
       // Get the next input element.
       std::vector input_element;
       bool end_of_input = false;
@@ -583,6 +594,8 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
                 (batch_results_.size() == max_batch_results_ &&
                  call_counter_ % dataset()->batch_size_ == 0));
       };
+      // Counts the total number of batches to use as an id of BatchResult.
+      int64 num_total_batches = 1;
       while (true) {
         {
           mutex_lock l(*mu_);
@@ -607,8 +620,8 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
 
           while (!busy()) {
             if (call_counter_ % dataset()->batch_size_ == 0) {
-              batch_results_.push_back(
-                  std::make_shared(dataset()->batch_size_));
+              batch_results_.push_back(std::make_shared(
+                  dataset()->batch_size_, num_total_batches++));
             }
             int64 offset = call_counter_++ % dataset()->batch_size_;
             new_calls.emplace_back(batch_results_.back(), offset);

From 848de750514ac608d7c16cbf52fc4cdb5daa5b5e Mon Sep 17 00:00:00 2001
From: Henry Tan 
Date: Tue, 30 Jun 2020 21:24:20 -0700
Subject: [PATCH 1373/1390] Add LIBTFTPU macro qualifier to the definition of
 the TpuCompilationCacheMetrics.

PiperOrigin-RevId: 319157570
Change-Id: I0ee8c7036ec37eaba0cf661e5448c5cbc599d6df
---
 tensorflow/core/tpu/kernels/tpu_compilation_cache_metrics.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/core/tpu/kernels/tpu_compilation_cache_metrics.cc b/tensorflow/core/tpu/kernels/tpu_compilation_cache_metrics.cc
index ba4e2ccff93..7c8ac937fe2 100644
--- a/tensorflow/core/tpu/kernels/tpu_compilation_cache_metrics.cc
+++ b/tensorflow/core/tpu/kernels/tpu_compilation_cache_metrics.cc
@@ -17,6 +17,9 @@ limitations under the License.
 namespace tensorflow {
 namespace tpu {
 
+// TODO(henrytan): remove this once `TpuCompilationCache` migration to OSS is
+// completed.
+#if defined(LIBTFTPU)
 /* static */
 void TpuCompilationCacheMetrics::IncrementCacheLookupCount(
     bool is_cache_hit, absl::string_view session_name) {
@@ -27,6 +30,7 @@ void TpuCompilationCacheMetrics::IncrementCacheLookupCount(
 void TpuCompilationCacheMetrics::SetCacheEntryCount(int64 count) {
   // A placeholder for tracking metrics.
 }
+#endif  // LIBTFTPU
 
 }  // namespace tpu
 }  // namespace tensorflow

From 2db7754d4c5bcf878dc7f1d4ce85cfc0ec41486b Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Tue, 30 Jun 2020 21:59:43 -0700
Subject: [PATCH 1374/1390] Add microbenchmarks to check the overhead of a
 small model made up entirely out of 20 small op layers.

PiperOrigin-RevId: 319161658
Change-Id: Id715b7b883c23fe5b54d2cafcfe7641724be09d3
---
 .../benchmarks/eager_microbenchmarks_test.py      | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorflow/python/keras/benchmarks/eager_microbenchmarks_test.py b/tensorflow/python/keras/benchmarks/eager_microbenchmarks_test.py
index aa33618fbe7..2d58e984d67 100644
--- a/tensorflow/python/keras/benchmarks/eager_microbenchmarks_test.py
+++ b/tensorflow/python/keras/benchmarks/eager_microbenchmarks_test.py
@@ -99,6 +99,21 @@ class MicroBenchmarksBase(test.Benchmark):
 
     self._run(fn, 10000)
 
+  def benchmark_op_layer_call_overhead(self):
+    model_input = tf.keras.Input(shape=(1,))
+    model_output = model_input
+    x = tf.convert_to_tensor([[1.1]])
+
+    for _ in range(20):
+      model_output = tf.multiply(model_output, x)
+    model = tf.keras.Model(inputs=model_input, outputs=model_output)
+
+    def fn():
+      model(x)  # pylint: disable=not-callable
+
+    fn()
+    self._run(fn, 100)
+
   def benchmark_model_predict_tensorlike_overhead(self):
 
     class OnlyOverheadLayer(tf.keras.layers.Layer):

From ac9969a2762d4068edef7e451b53256d2adb2692 Mon Sep 17 00:00:00 2001
From: Feng Liu 
Date: Tue, 30 Jun 2020 22:30:04 -0700
Subject: [PATCH 1375/1390] Add a trait verifier to the same operands and
 results scales trait

PiperOrigin-RevId: 319164800
Change-Id: I48ffe33160d5dbbdc549a4d71294979a13841df3
---
 .../compiler/mlir/lite/flatbuffer_import.cc   | 56 ++++++++++++++++++-
 tensorflow/compiler/mlir/lite/ir/tfl_ops.td   |  9 +++
 .../mlir/lite/quantization/quantization.td    | 23 ++++++--
 .../lite/quantization/quantization_driver.cc  |  4 +-
 .../lite/quantization/quantization_traits.h   | 18 +++---
 .../lite/quantization/quantization_utils.cc   | 56 ++++++++++++++++++-
 .../lite/quantization/quantization_utils.h    |  2 +-
 .../tests/flatbuffer2mlir/quantization.mlir   | 16 ++++++
 tensorflow/compiler/mlir/lite/tests/ops.mlir  |  8 +++
 9 files changed, 169 insertions(+), 23 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
index 4b3658adbc2..867c7560344 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc
@@ -724,6 +724,59 @@ StatusOr> PruneSubgraph(
   return visited;
 }
 
+// We want to adjust the func op according to some cross ops information.
+static StatusOr PostProcessFuncOp(FuncOp func) {
+  OpBuilder builder(func);
+  // When a quantized constant is imported, its quantization parameter is set
+  // to be narrow range. Here revert to be the fully range if the user doesn't
+  // require narrow range.
+  func.walk([&](tfl::QConstOp cst) {
+    Value value = cst.getResult();
+    // This is a quantized constant, so it only has one use.
+    assert(value.hasOneUse() && "QConst has only one use.");
+    auto& use = *value.getUses().begin();
+    Operation* user = use.getOwner();
+    if (user->isKnownTerminator()) return;
+    auto affine_user = llvm::dyn_cast(user);
+    if (affine_user &&
+        affine_user.GetAffineOperandIndex() == use.getOperandNumber() &&
+        affine_user.RequiredNarrowRangeAffineOperand())
+      return;
+    auto qtype = mlir::quant::UniformQuantizedType::getQuantizedElementType(
+        value.getType());
+    // Only the 8-bit constants are imported with narrow range.
+    if (!qtype || qtype.getStorageTypeIntegralWidth() != 8) return;
+    mlir::quant::QuantizedType new_qtype;
+    if (auto per_axis =
+            qtype.dyn_cast()) {
+      new_qtype = mlir::quant::UniformQuantizedPerAxisType::get(
+          per_axis.getFlags(), per_axis.getStorageType(),
+          per_axis.getExpressedType(), per_axis.getScales(),
+          per_axis.getZeroPoints(), per_axis.getQuantizedDimension(),
+          per_axis.getStorageTypeMin() - 1, per_axis.getStorageTypeMax());
+    } else if (auto per_tensor =
+                   qtype.dyn_cast()) {
+      new_qtype = mlir::quant::UniformQuantizedType::get(
+          per_tensor.getFlags(), per_tensor.getStorageType(),
+          per_tensor.getExpressedType(), per_tensor.getScale(),
+          per_tensor.getZeroPoint(), per_tensor.getStorageTypeMin() - 1,
+          per_tensor.getStorageTypeMax());
+    } else {
+      return;
+    }
+    auto new_output_type = new_qtype.castFromExpressedType(
+        mlir::quant::UniformQuantizedType::castToExpressedType(
+            value.getType()));
+    builder.setInsertionPointAfter(cst);
+    auto new_op = builder.create(
+        cst.getLoc(), new_output_type, mlir::TypeAttr::get(new_output_type),
+        cst.valueAttr());
+    cst.replaceAllUsesWith(new_op.getResult());
+    cst.erase();
+  });
+  return func;
+}
+
 // Build a FuncOp from a tflite SubGraph
 // The op_names are a mapping from indexes into the TFLite operators array to
 // the operator name MLIR expects (tfl.foo_op). The buffers are directly taken
@@ -978,7 +1031,7 @@ StatusOr ConvertSubgraph(
 
   op_builder.create(base_loc, return_operands);
 
-  return func;
+  return PostProcessFuncOp(func);
 }
 
 // TFLite subgraphs do not necessarily have names, though MLIR functions must
@@ -1059,6 +1112,5 @@ OwningModuleRef tflite::FlatBufferToMlir(
     }
     module.push_back(func_or_error.ConsumeValueOrDie());
   }
-
   return OwningModuleRef(module);
 }
diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
index d109e425cae..f02e761cab5 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
@@ -739,6 +739,15 @@ def TFL_ConcatenationOp : TFL_Op<"concatenation",
   let hasFolder = 1;
 
   let verifier = [{ return Verify(*this); }];
+
+  let extraClassDeclaration = [{
+    // SameScalesOpInterface:
+    bool RequiredSameOperandsAndResultsScale(bool sign, int bit_width) {
+      // uint8 doesn't require same operands and results scales.
+      bool is_uint8 = !sign && (bit_width == 8);
+      return !is_uint8;
+    }
+  }];
 }
 
 def TFL_ConstOp : Op {
+  let description = [{
+    Interface for ops potentially have same operands and results scales.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      [{Returns whether same operands and results scales are required.}],
+      "bool", "RequiredSameOperandsAndResultsScale",
+      (ins "bool":$sign, "int":$bit_width), [{}], [{return true;}]
+    >,
+  ];
+
+  let verify = [{
+    return quant::VerifySameScales($_op);
+  }];
+}
+
 // Specify this trait if the op has a fixed output value range.
 class FixedResultScale : NativeOpTrait::Impl")>;
 
-// Specify this trait if the op requires same inputs and outputs quantization
-// scales.
-def SameOperandsAndResultsScale : NativeOpTrait<
-  "quant::SameOperandsAndResultsScale">;
-
 // Specify this trait if the bias-th input of the op is a bias input, which
 // needs a scale based on the scales of op1 and op2.
 class AccumulatorUniformScale : NativeOpTrait<
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
index 9b63290a10b..c1cf4354c04 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_driver.cc
@@ -652,7 +652,7 @@ void QuantizationDriver::PreprocessConstantOps() {
       }
 
       if (biases.find(operand_num) == biases.end() &&
-          !user->hasTrait()) {
+          !llvm::dyn_cast(user)) {
         // Needs to scan the content to get the quantiztion parameters if there
         // are no quantization parameters (FakeQuant ops).
         weights_.insert(cst);
@@ -764,7 +764,7 @@ bool QuantizationDriver::PropagateParams() {
       continue;
     }
 
-    if (op->hasTrait()) {
+    if (llvm::isa(op)) {
       auto params = GetQuantParamsForSameScaleConstraint(op);
       // The quantization parameters haven't been propagated to any operands
       // or results. Skip this node for now.
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_traits.h b/tensorflow/compiler/mlir/lite/quantization/quantization_traits.h
index 693f692c61a..d5eea94e848 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_traits.h
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_traits.h
@@ -20,11 +20,18 @@ limitations under the License.
 
 #include "mlir/Dialect/Quant/QuantTypes.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
 
 using QuantizedType = mlir::quant::QuantizedType;
 using UniformQuantizedType = mlir::quant::UniformQuantizedType;
 
 namespace mlir {
+namespace quant {
+// Verify that the op satisfies the same operands and results scales
+// constraints. Note that this constraint can only be applied on some
+// storage types of the op.
+LogicalResult VerifySameScales(Operation* op);
+}  // namespace quant
 
 // This includes the interface class definition. It couldn't be in a namespace
 // because the table gen doesn't emit the namespace when it is used.
@@ -40,17 +47,6 @@ struct QuantizationSpecTraitBase : public TraitBase {
   static bool IsQuantizable() { return true; }
 };
 
-// This class provides the API for TFL ops that requires same input and output
-// scale as the quantization results. This is used as a trait like this:
-//
-//   class TransposeOp
-//       : public Op {
-//
-template 
-class SameOperandsAndResultsScale
-    : public QuantizationSpecTraitBase {};
-
 // This class provides the API for TFL ops that has a fixed output value range.
 // This is used as a trait like this:
 //
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
index a0392583f36..dbfbe451d37 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/StandardTypes.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
 
 namespace mlir {
 
@@ -475,7 +476,7 @@ bool RemoveRedundantStatsOps(mlir::FuncOp func,
       // We don't propagate this parameter down if it has multiple operands.
       // We want to use the result parameter scales instead.
 
-      if (user->hasTrait() &&
+      if (llvm::dyn_cast(user) &&
           !PreferResultScale(user)) {
         for (Value res : user->getResults()) {
           if (res.hasOneUse()) {
@@ -506,7 +507,7 @@ bool RemoveRedundantStatsOps(mlir::FuncOp func,
     all_stats_ops.pop_back();
 
     if (auto def = stats_op.arg().getDefiningOp()) {
-      if (def->hasTrait()) {
+      if (llvm::dyn_cast(def)) {
         for (auto input : def->getOperands()) {
           if (auto next_stats = llvm::dyn_cast_or_null(
                   input.getDefiningOp())) {
@@ -529,5 +530,56 @@ bool RemoveRedundantStatsOps(mlir::FuncOp func,
   // Returns false if the steps finish without errors.
   return false;
 }
+
+LogicalResult VerifySameScales(Operation* op) {
+  auto same_scale_op = llvm::cast(op);
+
+  llvm::SmallVector collected_quant_params;
+  for (auto input : op->getOperands()) {
+    auto quant_params =
+        UniformQuantizedType::getQuantizedElementType(input.getType());
+    // Skip non-quantizable operands.
+    if (quant_params) {
+      collected_quant_params.push_back(quant_params);
+    }
+  }
+
+  for (auto output : op->getResults()) {
+    auto quant_params =
+        UniformQuantizedType::getQuantizedElementType(output.getType());
+    // Skip non-quantizable results.
+    if (quant_params) {
+      collected_quant_params.push_back(quant_params);
+    }
+  }
+
+  if (collected_quant_params.size() <= 1) return success();
+  for (int i = 1; i < collected_quant_params.size(); i++) {
+    auto expected_params = collected_quant_params[0];
+    auto compared_paras = collected_quant_params[i];
+    // Same quantization parameters are always ok.
+    if (expected_params == compared_paras) continue;
+    // If the quantization parameters are not the same, as long as it has the
+    // same storage type and the op interface doesn't require same scale
+    // constraint for this storage type, it is still ok.
+    if ((expected_params.isSigned() == compared_paras.isSigned() &&
+         expected_params.getStorageTypeIntegralWidth() ==
+             compared_paras.getStorageTypeIntegralWidth()) &&
+        !same_scale_op.RequiredSameOperandsAndResultsScale(
+            expected_params.isSigned(),
+            expected_params.getStorageTypeIntegralWidth()))
+      continue;
+
+    std::string err_msg =
+        "quantization parameters violate the same scale constraint: ";
+    llvm::raw_string_ostream os(err_msg);
+    collected_quant_params[0].print(os);
+    os << " vs. ";
+    collected_quant_params[i].print(os);
+    os.flush();
+    return op->emitOpError(err_msg);
+  }
+  return success();
+}
 }  // namespace quant
 }  // namespace mlir
diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.h b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.h
index ad99b1c58d2..35c930281d0 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.h
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.h
@@ -387,7 +387,7 @@ struct FoldTrivalRequantizeOp : public OpRewritePattern {
     Operation* def = pre_quantized.getDefiningOp();
     if (!def) return failure();
     if (llvm::isa(def) ||
-        def->hasTrait() ||
+        llvm::isa(def) ||
         def->hasTrait()) {
       return failure();
     }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
index 22943b55f66..f5de214a692 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
@@ -1,5 +1,6 @@
 // RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s
 
+// CHECK-LABEL: main
 func @main(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x1001xf32> {
 // CHECK:   %{{.*}} = "tfl.quantize"(%{{.*}}) {qtype = tensor<1x224x224x3x!quant.uniform>} : (tensor<1x224x224x3xf32>) -> tensor<1x224x224x3x!quant.uniform>
 // The float values here doesn't match exactly because double -> float -> double is lossy
@@ -17,3 +18,18 @@ func @main(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x1001xf32> {
   %6 = "tfl.dequantize"(%5) : (tensor<1x1001x!quant.uniform>) -> tensor<1x1001xf32>
   return %6 : tensor<1x1001xf32>
 }
+
+// CHECK-LABEL: quantized_constant
+func @quantized_constant(%arg0: tensor<1x2xf32>) -> tensor<2x2xf32> {
+  %1 = "tfl.quantize"(%arg0) {qtype = tensor<1x2x!quant.uniform>, volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform>
+  %cst = "tfl.pseudo_qconst"() {qtype = tensor<1x2x!quant.uniform>, value = dense<-76> : tensor<1x2xi8>} : () -> tensor<1x2x!quant.uniform>
+  %2 = "tfl.concatenation"(%1, %cst) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2x!quant.uniform>, tensor<1x2x!quant.uniform>) -> tensor<2x2x!quant.uniform>
+  %3 = "tfl.dequantize"(%2) : (tensor<2x2x!quant.uniform>) -> tensor<2x2xf32>
+  return %3 : tensor<2x2xf32>
+
+// CHECK-NEXT: %[[Q:.*]] = "tfl.quantize"(%arg0) {qtype = tensor<1x2x!quant.uniform>} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform>
+// CHECK-NEXT: %[[CST:.*]] = "tfl.pseudo_qconst"() {qtype = tensor<1x2x!quant.uniform>, value = dense<-76> : tensor<1x2xi8>} : () -> tensor<1x2x!quant.uniform>
+// CHECK-NEXT: %[[CONCAT:.*]] = "tfl.concatenation"(%[[Q]], %[[CST]]) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2x!quant.uniform>, tensor<1x2x!quant.uniform>) -> tensor<2x2x!quant.uniform>
+// CHECK-NEXT: %[[DQ:.*]] = "tfl.dequantize"(%[[CONCAT]]) : (tensor<2x2x!quant.uniform>) -> tensor<2x2xf32>
+// CHECK-NEXT: return %[[DQ]] : tensor<2x2xf32>
+}
diff --git a/tensorflow/compiler/mlir/lite/tests/ops.mlir b/tensorflow/compiler/mlir/lite/tests/ops.mlir
index f1742538935..5f434e954c8 100644
--- a/tensorflow/compiler/mlir/lite/tests/ops.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/ops.mlir
@@ -1310,6 +1310,14 @@ func @testConcatInvalidOperandDimSizeComparedToPrevInput(%arg0: tensor<1x2xi32>,
 
 // -----
 
+func @testConcatInvalidScales(%arg0: tensor<*x!quant.uniform>, %arg1: tensor<*x!quant.uniform>) -> tensor<*x!quant.uniform> {
+  // expected-error @+1 {{'tfl.concatenation' op quantization parameters violate the same scale constraint: !quant.uniform vs. !quant.uniform}}
+  %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 3 : i32, fused_activation_function = "NONE"} : (tensor<*x!quant.uniform>, tensor<*x!quant.uniform>) -> tensor<*x!quant.uniform>
+  return %0 : tensor<*x!quant.uniform>
+}
+
+// -----
+
 func @testConcatBenignUnrankedOperand(%arg0: tensor<*xi32>, %arg1: tensor<1x2xi32>) -> tensor<2x2xi32> {
   %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<*xi32>, tensor<1x2xi32>) -> tensor<2x2xi32>
   return %0 : tensor<2x2xi32>

From 3252cc128c6bc1d4855d657d3eb4481b3a0e7879 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Tue, 30 Jun 2020 22:49:18 -0700
Subject: [PATCH 1376/1390] Internal change

PiperOrigin-RevId: 319166670
Change-Id: Ie098bec86a9eb6da20f51dd1eef9569cf3ca11f1
---
 tensorflow/python/keras/saving/hdf5_format.py | 35 +------------------
 1 file changed, 1 insertion(+), 34 deletions(-)

diff --git a/tensorflow/python/keras/saving/hdf5_format.py b/tensorflow/python/keras/saving/hdf5_format.py
index 6c93e518741..3aa4fe1245a 100644
--- a/tensorflow/python/keras/saving/hdf5_format.py
+++ b/tensorflow/python/keras/saving/hdf5_format.py
@@ -33,7 +33,6 @@ from tensorflow.python.keras.saving.saved_model import json_utils
 from tensorflow.python.keras.utils import conv_utils
 from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
 from tensorflow.python.ops import variables as variables_module
-from tensorflow.python.platform import gfile
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import serialization
 from tensorflow.python.util.lazy_loader import LazyLoader
@@ -55,24 +54,7 @@ sequential_lib = LazyLoader(
 # pylint:enable=g-inconsistent-quotes
 
 
-# create lock file
-def create_lockfile(filepath):
-  lockfile_path = filepath + '.lock'
-
-  f = gfile.GFile(lockfile_path, 'w')
-  f.write(str(os.getpid()))
-  f.close()
-
-  return lockfile_path
-
-
-def check_lockfile(filepath):
-  lockfile_path = filepath + '.lock'
-  return gfile.Exists(lockfile_path)
-
-
-def save_model_to_hdf5(model, filepath, overwrite=True, \
-                       lockfile=True, include_optimizer=True):
+def save_model_to_hdf5(model, filepath, overwrite=True, include_optimizer=True):
   """Saves a model to a HDF5 file.
 
   The saved model contains:
@@ -92,9 +74,6 @@ def save_model_to_hdf5(model, filepath, overwrite=True, \
       overwrite: Whether we should overwrite any existing
           model at the target location, or instead
           ask the user with a manual prompt.
-      lockfile: Create a lockfile before saving the model
-          file to prevent from reading, while saving
-          is not done.
       include_optimizer: If True, save optimizer's state together.
 
   Raises:
@@ -121,10 +100,6 @@ def save_model_to_hdf5(model, filepath, overwrite=True, \
       if not proceed:
         return
 
-    # create lock file
-    if lockfile:
-      lockfile_path = create_lockfile(filepath)
-
     f = h5py.File(filepath, mode='w')
     opened_new_file = True
   else:
@@ -155,10 +130,6 @@ def save_model_to_hdf5(model, filepath, overwrite=True, \
     if opened_new_file:
       f.close()
 
-      # remove lock file
-      if lockfile:
-        gfile.Remove(lockfile_path)
-
 
 def load_model_from_hdf5(filepath, custom_objects=None, compile=True):  # pylint: disable=redefined-builtin
   """Loads a model saved via `save_model_to_hdf5`.
@@ -193,10 +164,6 @@ def load_model_from_hdf5(filepath, custom_objects=None, compile=True):  # pylint
 
   opened_new_file = not isinstance(filepath, h5py.File)
   if opened_new_file:
-    # check if lock file exist
-    if check_lockfile(filepath):
-      raise ValueError('Cannot read from file at this time.')
-
     f = h5py.File(filepath, mode='r')
   else:
     f = filepath

From 6d0cf63bb2cce6986ea6940d732f2942b53dfdeb Mon Sep 17 00:00:00 2001
From: Jay Shi 
Date: Tue, 30 Jun 2020 23:06:13 -0700
Subject: [PATCH 1377/1390] [tf.data] Avoid calling `CollectTunableParameters`
 function repeatedly.

PiperOrigin-RevId: 319168243
Change-Id: I8f3f0389820250bd0f6f9727f4ad10b6cf5c8ea9
---
 tensorflow/core/framework/model.cc | 8 ++++----
 tensorflow/core/framework/model.h  | 5 ++++-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/framework/model.cc b/tensorflow/core/framework/model.cc
index 198d2f6574c..6dcaf8ecac2 100644
--- a/tensorflow/core/framework/model.cc
+++ b/tensorflow/core/framework/model.cc
@@ -1320,14 +1320,14 @@ Model::CollectTunableParameters(std::shared_ptr node) {
 }
 
 absl::flat_hash_map>
-Model::CollectEssentialParallelism(std::shared_ptr node) {
+Model::CollectEssentialParallelism(
+    std::shared_ptr node,
+    const absl::flat_hash_map>& parameters) {
   // Parallelism parameter is considered to be essential if the corresponding
   // transformations's processing time is greater than essential rate times the
   // average transformation self processing time.
   constexpr double kEssentialRate = 0.3L;
 
-  absl::flat_hash_map> parameters;
-  node->CollectTunableParameters(¶meters);
   absl::flat_hash_map processing_times;
   double processing_time = node->TotalProcessingTime(&processing_times);
   double uniform_share =
@@ -1350,7 +1350,7 @@ void Model::OptimizeGradientDescent(int64 cpu_budget, int64 ram_budget) {
   }
   VLOG(2) << "Starting optimization of tunable parameters with GradientDescent";
   auto parameters = CollectTunableParameters(snapshot);
-  auto essential_parameters = CollectEssentialParallelism(snapshot);
+  auto essential_parameters = CollectEssentialParallelism(snapshot, parameters);
   // We add the number of model's buffered bytes because it is excluded from the
   // memory budget, but it is included in the maximum number of buffered bytes.
   ram_budget += TotalBufferedBytes(snapshot);
diff --git a/tensorflow/core/framework/model.h b/tensorflow/core/framework/model.h
index 82814ed1353..e8d78756192 100644
--- a/tensorflow/core/framework/model.h
+++ b/tensorflow/core/framework/model.h
@@ -628,7 +628,10 @@ class Model {
   // relative to other transformations. The collected parameters are returned
   // as a mapping from a (unique) node name to a parallelism parameter.
   absl::flat_hash_map>
-  CollectEssentialParallelism(std::shared_ptr node);
+  CollectEssentialParallelism(
+      std::shared_ptr node,
+      const absl::flat_hash_map>&
+          parameters);
 
   // This optimization algorithm starts by setting all tunable parallelism
   // parameters to the minimum value. It then repeatedly identifies the

From c4ae3c7a0430adb6970fb440ecd9fbe5da9c0461 Mon Sep 17 00:00:00 2001
From: Terry Heo 
Date: Tue, 30 Jun 2020 23:52:28 -0700
Subject: [PATCH 1378/1390] Add an external delegate which is initialized from
 a shared library

This CL introduces the following API to manage an external delegate.
- TfLiteExternalDelegateCreate()
- TfLiteExternalDelegateOptionsDefault()
- TfLiteExternalDelegateDelete()
Also refactored the ExternalDelegateProvider to adopt this change.

PiperOrigin-RevId: 319171849
Change-Id: Ibccbb3a70aa91dfcf96745e13342cdbb2e77d832
---
 tensorflow/lite/delegates/external/BUILD      |  35 +++
 .../delegates/external/external_delegate.cc   | 243 ++++++++++++++++++
 .../delegates/external/external_delegate.h    |  53 ++++
 tensorflow/lite/tools/delegates/BUILD         |   1 +
 .../delegates/external_delegate_provider.cc   | 103 ++------
 5 files changed, 352 insertions(+), 83 deletions(-)
 create mode 100644 tensorflow/lite/delegates/external/BUILD
 create mode 100644 tensorflow/lite/delegates/external/external_delegate.cc
 create mode 100644 tensorflow/lite/delegates/external/external_delegate.h

diff --git a/tensorflow/lite/delegates/external/BUILD b/tensorflow/lite/delegates/external/BUILD
new file mode 100644
index 00000000000..ca23f95122f
--- /dev/null
+++ b/tensorflow/lite/delegates/external/BUILD
@@ -0,0 +1,35 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+
+package(
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+cc_library(
+    name = "external_delegate",
+    srcs = ["external_delegate.cc"],
+    hdrs = ["external_delegate.h"],
+    deps = [
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:common",
+    ],
+)
+
+exports_files([
+    "external_delegate.h",
+])
diff --git a/tensorflow/lite/delegates/external/external_delegate.cc b/tensorflow/lite/delegates/external/external_delegate.cc
new file mode 100644
index 00000000000..5df158942f2
--- /dev/null
+++ b/tensorflow/lite/delegates/external/external_delegate.cc
@@ -0,0 +1,243 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+#include "tensorflow/lite/delegates/external/external_delegate.h"
+
+#include 
+#include 
+
+#if defined(_WIN32)
+#include 
+#else
+#include 
+#endif
+
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+namespace {
+
+// Library Support construct to handle dynamic library operations
+#if defined(_WIN32)
+struct LibSupport {
+  static void* Load(const char* lib) { return LoadLibrary(lib); }
+
+  static void* GetSymbol(void* handle, const char* symbol) {
+    return (void*)GetProcAddress((HMODULE)handle, symbol);
+  }
+
+  static int UnLoad(void* handle) { return FreeLibrary((HMODULE)handle); }
+};
+#else
+struct LibSupport {
+  static void* Load(const char* lib) {
+    return dlopen(lib, RTLD_LAZY | RTLD_LOCAL);
+  }
+
+  static void* GetSymbol(void* handle, const char* symbol) {
+    return dlsym(handle, symbol);
+  }
+
+  static int UnLoad(void* handle) { return dlclose(handle); }
+};
+#endif
+
+// External delegate library construct
+struct ExternalLib {
+  using CreateDelegatePtr = std::add_pointer::type;
+  using DestroyDelegatePtr = std::add_pointer::type;
+
+  // Open a given delegate library and load the create/destroy symbols
+  bool load(const std::string library) {
+    void* handle = LibSupport::Load(library.c_str());
+    if (handle == nullptr) {
+      TFLITE_LOG(TFLITE_LOG_INFO, "Unable to load external delegate from : %s",
+                 library.c_str());
+    } else {
+      create = reinterpret_cast(
+          LibSupport::GetSymbol(handle, "tflite_plugin_create_delegate"));
+      destroy = reinterpret_cast(
+          LibSupport::GetSymbol(handle, "tflite_plugin_destroy_delegate"));
+      return create && destroy;
+    }
+    return false;
+  }
+
+  CreateDelegatePtr create{nullptr};
+  DestroyDelegatePtr destroy{nullptr};
+};
+
+// An ExternalDelegateWrapper is responsibile to manage a TFLite delegate
+// initialized from a shared library. It creates a delegate from the given
+// option and storages it to external_delegate_ member variable. On the
+// destruction, it conducts necessary clean up process.
+class ExternalDelegateWrapper {
+ public:
+  explicit ExternalDelegateWrapper(
+      const TfLiteExternalDelegateOptions* options);
+  ~ExternalDelegateWrapper();
+
+  // Return a TfLiteDelegate which is created from
+  // tflite_plugin_create_delegate() of an external delegate logic.
+  TfLiteDelegate* tflite_external_delegate() { return external_delegate_; }
+
+  // Return a TfLiteDelegate which is convertibile to this class.
+  TfLiteDelegate* tflite_wrapper_delegate() { return &wrapper_delegate_; }
+
+ private:
+  ExternalLib external_lib_;
+
+  // external delegate instance owned by external delegate logic.
+  // It's created by "tflite_plugin_destroy_delegate()" function in the external
+  // delegate logic And it should be released by
+  // "tflite_plugin_destroy_delegate()" function.
+  TfLiteDelegate* external_delegate_;
+
+  // TfLiteDelegate representation of this ExternalDelegateWrapper object.
+  TfLiteDelegate wrapper_delegate_;
+};
+
+// Converts the given TfLiteDelegate to an ExternalDelegateWrapper instance.
+inline ExternalDelegateWrapper* GetExternalDelegateWrapper(
+    TfLiteDelegate* delegate) {
+  return reinterpret_cast(delegate->data_);
+}
+
+// Relay Prepare() call to the associated external TfLiteDelegate object.
+TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->Prepare(context, external_delegate);
+}
+
+// Relay CopyFromBufferHandle() call to the associated external TfLiteDelegate
+// object.
+TfLiteStatus DelegateCopyFromBufferHandle(TfLiteContext* context,
+                                          struct TfLiteDelegate* delegate,
+                                          TfLiteBufferHandle buffer_handle,
+                                          TfLiteTensor* tensor) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->CopyFromBufferHandle(context, delegate,
+                                                 buffer_handle, tensor);
+}
+
+// Relay CopyToBufferHandle() call to the associated external TfLiteDelegate
+// object.
+TfLiteStatus DelegateCopyToBufferHandle(TfLiteContext* context,
+                                        struct TfLiteDelegate* delegate,
+                                        TfLiteBufferHandle buffer_handle,
+                                        TfLiteTensor* tensor) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->CopyToBufferHandle(context, delegate, buffer_handle,
+                                               tensor);
+}
+
+// Relay FreeBufferHandle() call to the associated external TfLiteDelegate
+// object.
+void DelegateFreeBufferHandle(TfLiteContext* context,
+                              struct TfLiteDelegate* delegate,
+                              TfLiteBufferHandle* handle) {
+  auto external_delegate_wrapper = GetExternalDelegateWrapper(delegate);
+  TfLiteDelegate* external_delegate =
+      external_delegate_wrapper->tflite_external_delegate();
+  return external_delegate->FreeBufferHandle(context, delegate, handle);
+}
+
+ExternalDelegateWrapper::ExternalDelegateWrapper(
+    const TfLiteExternalDelegateOptions* options) {
+  external_delegate_ = nullptr;
+  if (external_lib_.load(options->lib_path)) {
+    std::vector ckeys, cvalues;
+    for (int i = 0; i < options->count; i++) {
+      ckeys.push_back(options->keys[i]);
+      cvalues.push_back(options->values[i]);
+    }
+
+    external_delegate_ = external_lib_.create(ckeys.data(), cvalues.data(),
+                                              ckeys.size(), nullptr);
+    if (external_delegate_) {
+      wrapper_delegate_ = {
+          .data_ = reinterpret_cast(this),
+          .Prepare = DelegatePrepare,
+          .CopyFromBufferHandle = nullptr,
+          .CopyToBufferHandle = nullptr,
+          .FreeBufferHandle = nullptr,
+          .flags = external_delegate_->flags,
+      };
+      if (external_delegate_->CopyFromBufferHandle) {
+        wrapper_delegate_.CopyFromBufferHandle = DelegateCopyFromBufferHandle;
+      }
+      if (external_delegate_->CopyToBufferHandle) {
+        wrapper_delegate_.CopyToBufferHandle = DelegateCopyToBufferHandle;
+      }
+      if (external_delegate_->FreeBufferHandle) {
+        wrapper_delegate_.FreeBufferHandle = DelegateFreeBufferHandle;
+      }
+    }
+  }
+}
+
+ExternalDelegateWrapper::~ExternalDelegateWrapper() {
+  if (external_delegate_ != nullptr) {
+    external_lib_.destroy(external_delegate_);
+  }
+}
+
+}  // namespace
+}  // namespace tflite
+
+// TfLiteExternalDelegateOptionsInsert adds key/value to the given
+// TfLiteExternalDelegateOptions instance.
+TfLiteStatus TfLiteExternalDelegateOptionsInsert(
+    TfLiteExternalDelegateOptions* options, const char* key,
+    const char* value) {
+  if (options->count >= kMaxOptions) {
+    return kTfLiteError;
+  }
+  options->keys[options->count] = key;
+  options->values[options->count] = value;
+  options->count++;
+  return kTfLiteOk;
+}
+
+TfLiteExternalDelegateOptions TfLiteExternalDelegateOptionsDefault(
+    const char* lib_path) {
+  TfLiteExternalDelegateOptions options = {
+      .lib_path = lib_path,
+      .count = 0,
+      .insert = TfLiteExternalDelegateOptionsInsert,
+  };
+  return options;
+}
+
+TfLiteDelegate* TfLiteExternalDelegateCreate(
+    const TfLiteExternalDelegateOptions* options) {
+  auto* external_delegate_wrapper =
+      new tflite::ExternalDelegateWrapper(options);
+  if (external_delegate_wrapper) {
+    return external_delegate_wrapper->tflite_wrapper_delegate();
+  }
+  return nullptr;
+}
+
+void TfLiteExternalDelegateDelete(TfLiteDelegate* delegate) {
+  delete tflite::GetExternalDelegateWrapper(delegate);
+}
diff --git a/tensorflow/lite/delegates/external/external_delegate.h b/tensorflow/lite/delegates/external/external_delegate.h
new file mode 100644
index 00000000000..774c0f07db3
--- /dev/null
+++ b/tensorflow/lite/delegates/external/external_delegate.h
@@ -0,0 +1,53 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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
+
+    http://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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_
+#define TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_
+
+#include "tensorflow/lite/c/common.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+// TfLiteExternalDelegateOptions is a structure of key/value options to create
+// an external delegate.
+const int kMaxOptions = 256;
+typedef struct TfLiteExternalDelegateOptions {
+  const char* lib_path;
+  int count;
+  const char* keys[kMaxOptions];
+  const char* values[kMaxOptions];
+  TfLiteStatus (*insert)(TfLiteExternalDelegateOptions* options,
+                         const char* key, const char* value);
+} TfLiteExternalDelegateOptions;
+
+// Populates TfLiteExternalDelegateOptions with the given shared library path.
+TfLiteExternalDelegateOptions TfLiteExternalDelegateOptionsDefault(
+    const char* lib_path);
+
+// Creates a new delegate instance that need to be destroyed with
+// `TfLiteExternalDelegateDelete` when delegate is no longer used by TFLite.
+TfLiteDelegate* TfLiteExternalDelegateCreate(
+    const TfLiteExternalDelegateOptions* options);
+
+// Destroys a delegate created with `TfLiteExternalDelegateCreate` call.
+void TfLiteExternalDelegateDelete(TfLiteDelegate* delegate);
+
+#ifdef __cplusplus
+}
+#endif  // __cplusplus
+
+#endif  // TENSORFLOW_LITE_DELEGATES_EXTERNAL_EXTERNAL_DELEGATE_H_
diff --git a/tensorflow/lite/tools/delegates/BUILD b/tensorflow/lite/tools/delegates/BUILD
index 93b918d37b1..def41cc6c69 100644
--- a/tensorflow/lite/tools/delegates/BUILD
+++ b/tensorflow/lite/tools/delegates/BUILD
@@ -151,6 +151,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":delegate_provider_hdr",
+        "//tensorflow/lite/delegates/external:external_delegate",
     ],
     alwayslink = 1,
 )
diff --git a/tensorflow/lite/tools/delegates/external_delegate_provider.cc b/tensorflow/lite/tools/delegates/external_delegate_provider.cc
index 193860820b1..4f7dfeb8646 100644
--- a/tensorflow/lite/tools/delegates/external_delegate_provider.cc
+++ b/tensorflow/lite/tools/delegates/external_delegate_provider.cc
@@ -12,45 +12,15 @@ 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.
 ==============================================================================*/
-#include "tensorflow/lite/tools/delegates/delegate_provider.h"
-
-#if defined(_WIN32)
-#include 
-#else
-#include 
-#endif
 
 #include 
-#include 
 #include 
 
+#include "tensorflow/lite/delegates/external/external_delegate.h"
+#include "tensorflow/lite/tools/delegates/delegate_provider.h"
+
 namespace tflite {
 namespace tools {
-namespace {
-// Library Support construct to handle dynamic library operations
-#if defined(_WIN32)
-struct LibSupport {
-  static void* Load(const char* lib) { return LoadLibrary(lib); }
-
-  static void* GetSymbol(void* handle, const char* symbol) {
-    return (void*)GetProcAddress((HMODULE)handle, symbol);
-  }
-
-  static int UnLoad(void* handle) { return FreeLibrary((HMODULE)handle); }
-};
-#else
-struct LibSupport {
-  static void* Load(const char* lib) {
-    return dlopen(lib, RTLD_LAZY | RTLD_LOCAL);
-  }
-
-  static void* GetSymbol(void* handle, const char* symbol) {
-    return dlsym(handle, symbol);
-  }
-
-  static int UnLoad(void* handle) { return dlclose(handle); }
-};
-#endif
 
 // Split a given string to a vector of string using a delimiter character
 std::vector SplitString(const std::string& str, char delimiter) {
@@ -63,32 +33,6 @@ std::vector SplitString(const std::string& str, char delimiter) {
   return tokens;
 }
 
-// External delegate library construct
-struct ExternalLib {
-  using CreateDelegatePtr = std::add_pointer::type;
-  using DestroyDelegatePtr = std::add_pointer::type;
-
-  // Open a given delegate library and load the create/destroy symbols
-  bool load(const std::string library) {
-    void* handle = LibSupport::Load(library.c_str());
-    if (handle == nullptr) {
-      TFLITE_LOG(INFO) << "Unable to load external delegate from : " << library;
-    } else {
-      create = reinterpret_cast(
-          LibSupport::GetSymbol(handle, "tflite_plugin_create_delegate"));
-      destroy = reinterpret_cast(
-          LibSupport::GetSymbol(handle, "tflite_plugin_destroy_delegate"));
-      return create && destroy;
-    }
-    return false;
-  }
-
-  CreateDelegatePtr create{nullptr};
-  DestroyDelegatePtr destroy{nullptr};
-};
-}  // namespace
 
 // External delegate provider used to dynamically load delegate libraries
 // Note: Assumes the lifetime of the provider exceeds the usage scope of
@@ -136,35 +80,28 @@ TfLiteDelegatePtr ExternalDelegateProvider::CreateTfLiteDelegate(
   TfLiteDelegatePtr delegate(nullptr, [](TfLiteDelegate*) {});
   std::string lib_path = params.Get("external_delegate_path");
   if (!lib_path.empty()) {
-    ExternalLib delegate_lib;
-    if (delegate_lib.load(lib_path)) {
-      // Parse delegate options
-      const std::vector options = SplitString(
-          params.Get("external_delegate_options"), ';');
-      std::vector keys, values;
-      for (const auto& option : options) {
-        auto key_value = SplitString(option, ':');
-        if (key_value.size() == 2) {
-          values.push_back(std::move(key_value[1]));
-          keys.push_back(std::move(key_value[0]));
-        }
-      }
+    auto delegate_options =
+        TfLiteExternalDelegateOptionsDefault(lib_path.c_str());
 
-      const size_t num_options = keys.size();
-      std::vector ckeys, cvalues;
-      for (int i = 0; i < num_options; ++i) {
-        ckeys.push_back(keys[i].c_str());
-        cvalues.push_back(values[i].c_str());
+    // Parse delegate options
+    const std::vector options =
+        SplitString(params.Get("external_delegate_options"), ';');
+    std::vector keys, values;
+    for (const auto& option : options) {
+      auto key_value = SplitString(option, ':');
+      if (key_value.size() == 2) {
+        delegate_options.insert(&delegate_options, key_value[0].c_str(),
+                                key_value[1].c_str());
       }
-
-      // Create delegate
-      delegate =
-          TfLiteDelegatePtr(delegate_lib.create(ckeys.data(), cvalues.data(),
-                                                num_options, nullptr),
-                            delegate_lib.destroy);
     }
+
+    auto external_delegate = TfLiteExternalDelegateCreate(&delegate_options);
+    return TfLiteDelegatePtr(external_delegate, [](TfLiteDelegate* delegate) {
+      TfLiteExternalDelegateDelete(delegate);
+    });
   }
   return delegate;
 }
+
 }  // namespace tools
 }  // namespace tflite

From d8c49c2fdeed114470e2773967dddb932e9a9fb3 Mon Sep 17 00:00:00 2001
From: Xinyi Wang 
Date: Wed, 1 Jul 2020 00:10:24 -0700
Subject: [PATCH 1379/1390] Add multi_worker_mirrored_strategy to
 moving_averages_test

PiperOrigin-RevId: 319173462
Change-Id: I940527e1c205599324c2bfb2fcdef9c24076d1d9
---
 .../python/distribute/moving_averages_test.py | 110 +++++++++---------
 1 file changed, 56 insertions(+), 54 deletions(-)

diff --git a/tensorflow/python/distribute/moving_averages_test.py b/tensorflow/python/distribute/moving_averages_test.py
index 2d0ab80f1f2..577a6c1168f 100644
--- a/tensorflow/python/distribute/moving_averages_test.py
+++ b/tensorflow/python/distribute/moving_averages_test.py
@@ -20,14 +20,13 @@ from __future__ import print_function
 
 from absl.testing import parameterized
 
+from tensorflow.python.distribute import collective_all_reduce_strategy
 from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import strategy_combinations
 from tensorflow.python.distribute import tpu_strategy
 from tensorflow.python.eager import def_function
 from tensorflow.python.eager import test
 from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.training import moving_averages
 
@@ -38,6 +37,10 @@ all_distributions = [
     strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
     strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
     strategy_combinations.tpu_strategy,
+    strategy_combinations.multi_worker_mirrored_2x1_cpu,
+    strategy_combinations.multi_worker_mirrored_2x1_gpu,
+    strategy_combinations.multi_worker_mirrored_2x2_gpu,
+    strategy_combinations.multi_worker_mirrored_4x1_cpu,
 ]
 
 all_combinations = combinations.combine(
@@ -62,11 +65,11 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):
           var, val, decay, zero_debias=False)
       return var, assign
 
-    with distribution.scope(), self.cached_session() as sess:
+    with distribution.scope():
       var, assign = distribution.extended.call_for_each_replica(replica_fn)
-      variables.global_variables_initializer().run()
-      self.assertAllClose([10.0, 11.0], var)
-      sess.run(distribution.experimental_local_results(assign))
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([10.0, 11.0], self.evaluate(var))
+      self.evaluate(distribution.experimental_local_results(assign))
       # Mean of val across calls to replica_fn().
       average_val = [1.0 + 0.5 * (replica_id[0] - 1),
                      2.0 - 0.5 * (replica_id[0] - 1)]
@@ -74,7 +77,7 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):
       self.assertAllClose(
           [10.0 * 0.25 + average_val[0] * val_weight,
            11.0 * 0.25 + average_val[1] * val_weight],
-          var.eval())
+          self.evaluate(var))
 
   @combinations.generate(all_combinations)
   def testReplicaMode(self, distribution):
@@ -88,19 +91,19 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):
       assign = moving_averages.assign_moving_average(var, val, decay)
       return var, assign.op
 
-    with distribution.scope(), self.cached_session() as sess:
+    with distribution.scope():
       var, assign_op = distribution.extended.call_for_each_replica(replica_fn)
-      variables.global_variables_initializer().run()
-      self.assertAllClose([0.0, 0.0], var)
-      sess.run(distribution.experimental_local_results(assign_op))
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([0.0, 0.0], self.evaluate(var))
+      self.evaluate(distribution.experimental_local_results(assign_op))
       # Mean of val across calls to replica_fn().
       average_val = [1.0 + 0.5 * (replica_id[0] - 1),
                      2.0 - 0.5 * (replica_id[0] - 1)]
-      self.assertAllClose(average_val, var)
+      self.assertAllClose(average_val, self.evaluate(var))
 
   @combinations.generate(all_combinations)
   def testCrossDeviceWithoutZeroDebias(self, distribution):
-    with distribution.scope(), self.cached_session() as sess:
+    with distribution.scope():
       var = variables.Variable([10.0, 11.0])
       val = constant_op.constant([1.0, 2.0])
       decay = 0.25
@@ -109,45 +112,38 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):
       assign = moving_averages.assign_moving_average(
           var, val, decay, zero_debias=False)
 
-      variables.global_variables_initializer().run()
-      self.assertAllClose([10.0, 11.0], var)
-      sess.run(assign)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([10.0, 11.0], self.evaluate(var))
+      self.evaluate(assign)
       average_val = [1.0, 2.0]
       val_weight = 1.0 - 0.25
       self.assertAllClose(
           [10.0 * 0.25 + average_val[0] * val_weight,
            11.0 * 0.25 + average_val[1] * val_weight],
-          var.eval())
+          self.evaluate(var))
       # Also try assign.op.
-      sess.run(assign.op)
+      self.evaluate(assign.op)
       orig_weight = 0.25 * 0.25
       val_weight = 1.0 - orig_weight
       self.assertAllClose(
           [10.0 * orig_weight + average_val[0] * val_weight,
            11.0 * orig_weight + average_val[1] * val_weight],
-          var.eval())
+          self.evaluate(var))
 
   @combinations.generate(all_combinations)
   def testCrossDevice(self, distribution):
-    with distribution.scope(), self.cached_session() as sess:
+    with distribution.scope():
       var = variables.Variable([0.0, 0.0])
-      val = array_ops.placeholder(dtypes.float32)
+      val = variables.Variable([1.0, 2.0])
       decay = 0.25
       # NOTE(josh11b): We currently generate an error if val is a PerReplica
       # value.
       assign = moving_averages.assign_moving_average(var, val, decay)
 
-      variables.global_variables_initializer().run()
-      self.assertAllClose([0.0, 0.0], var)
-      sess.run(assign, feed_dict={val: [1.0, 2.0]})
-      self.assertAllClose([1.0, 2.0], var)
-
-      # Also try assign.op.
-      sess.run(assign.op, feed_dict={val: [10.0, 0.0]})
-      self.assertAllClose(
-          [(1.0 * 0.25 + 10.0) / (1.0 * 0.25 + 1.0),
-           (2.0 * 0.25 + 0.0) / (1.0 * 0.25 + 1.0)],
-          var.eval())
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([0.0, 0.0], self.evaluate(var))
+      self.evaluate(assign)
+      self.assertAllClose([1.0, 2.0], self.evaluate(var))
 
   @combinations.generate(all_combinations_eager)
   def testUpdateContext(self, distribution, use_function):
@@ -179,14 +175,14 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase):
           var, val, decay, zero_debias=False)
       return var, assign
 
-    with distribution.scope(), self.cached_session() as sess:
+    with distribution.scope():
       var, assign = distribution.extended.call_for_each_replica(replica_fn)
-      variables.global_variables_initializer().run()
-      self.assertAllClose([10.0, 11.0], var)
-      sess.run(distribution.experimental_local_results(assign))
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([10.0, 11.0], self.evaluate(var))
+      self.evaluate(distribution.experimental_local_results(assign))
       self.assertAllClose(
           [10 * 0.25 + 1. * (1 - 0.25), 11 * 0.25 + 2. * (1 - 0.25)],
-          var.eval())
+          self.evaluate(var))
 
 
 class ExponentialMovingAverageTest(test.TestCase, parameterized.TestCase):
@@ -196,6 +192,10 @@ class ExponentialMovingAverageTest(test.TestCase, parameterized.TestCase):
     if not use_function and isinstance(
         distribution, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1)):
       self.skipTest("TPUStrategy doesn't support pure eager execution.")
+    if isinstance(distribution,
+                  collective_all_reduce_strategy.CollectiveAllReduceStrategy):
+      self.skipTest("b/160194267: Cannot do variable.assign([0.5]) in replica "
+                    "context with MultiWorkerMirroredStrategy.")
     with distribution.scope():
       w = variables.Variable([1.0],
                              name="w",
@@ -255,33 +255,35 @@ class ExponentialMovingAverageTest(test.TestCase, parameterized.TestCase):
                   (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1)):
       self.skipTest("b/139550827: Cannot do variable.assign in replica context "
                     "of TPUStrategy")
+    if isinstance(distribution,
+                  collective_all_reduce_strategy.CollectiveAllReduceStrategy):
+      self.skipTest("b/160194267: Cannot do variable.assign([0.5]) in replica "
+                    "context with MultiWorkerMirroredStrategy.")
     with distribution.scope():
       w_assign, w_apply, ema_w = distribution.run(
           self._ema_replica_fn_graph)
     self.assertEqual(ema_w.name, "w/ExponentialMovingAverage:0")
-    with self.cached_session():
-      variables.global_variables_initializer().run()
-      self.evaluate(distribution.experimental_local_results(w_apply))
-      self.evaluate(distribution.experimental_local_results(w_assign))
-      self.evaluate(distribution.experimental_local_results(w_apply))
-      self.assertAllClose(
-          self.evaluate(distribution.experimental_local_results(ema_w))[0],
-          [0.89999998])
+    self.evaluate(variables.global_variables_initializer())
+    self.evaluate(distribution.experimental_local_results(w_apply))
+    self.evaluate(distribution.experimental_local_results(w_assign))
+    self.evaluate(distribution.experimental_local_results(w_apply))
+    self.assertAllClose(
+        self.evaluate(distribution.experimental_local_results(ema_w))[0],
+        [0.89999998])
 
   @combinations.generate(all_combinations)
   def testCrossReplicaContextGraph(self, distribution):
     with distribution.scope():
       w_assign, w_apply, ema_w = self._ema_replica_fn_graph()
     self.assertEqual(ema_w.name, "w/ExponentialMovingAverage:0")
-    with self.cached_session():
-      variables.global_variables_initializer().run()
-      self.evaluate(distribution.experimental_local_results(w_apply))
-      self.evaluate(distribution.experimental_local_results(w_assign))
-      self.evaluate(distribution.experimental_local_results(w_apply))
-      self.assertAllClose(
-          self.evaluate(distribution.experimental_local_results(ema_w))[0],
-          [0.89999998])
+    self.evaluate(variables.global_variables_initializer())
+    self.evaluate(distribution.experimental_local_results(w_apply))
+    self.evaluate(distribution.experimental_local_results(w_assign))
+    self.evaluate(distribution.experimental_local_results(w_apply))
+    self.assertAllClose(
+        self.evaluate(distribution.experimental_local_results(ema_w))[0],
+        [0.89999998])
 
 
 if __name__ == "__main__":
-  test.main()
+  combinations.main()

From cd0a2e6c1fbbff32f9b56e28634d831ccfbcb11d Mon Sep 17 00:00:00 2001
From: Tomer Kaftan 
Date: Wed, 1 Jul 2020 01:20:38 -0700
Subject: [PATCH 1380/1390] Enable dispatching for the publicly exposed
 `tf.convert_to_tensor` api.

This CL avoids enabling dispatch for the internal `convert_to_tensor` that api symbols use, so that the current version of dispatching continues to work.

This change is useful because user code that expects tensor-like objects often starts of with `tf.convert_to_tensor`, but those very same objects may be relying on dispatching to function like tensors (e.g. KerasTensors). Actually registering a tensor conversion would break the dispatching logic they utilize.

This will probably not be needed in future designs for dispatching

PiperOrigin-RevId: 319181371
Change-Id: I7788ef9c63bcacafd17b353808f58df921d1ecc8
---
 tensorflow/python/framework/ops.py      | 34 ++++++++++++++++++++-----
 tensorflow/python/util/dispatch_test.py | 23 +++++++++++++++--
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index 000e3bb87a0..8633665fe29 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -67,6 +67,7 @@ from tensorflow.python.types import internal
 from tensorflow.python.util import compat
 from tensorflow.python.util import decorator_utils
 from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
 from tensorflow.python.util import function_utils
 from tensorflow.python.util import lock_util
 from tensorflow.python.util import memory
@@ -1257,11 +1258,13 @@ EagerTensor = pywrap_tfe.TFE_Py_InitEagerTensor(_EagerTensorBase)
 
 
 @tf_export(v1=["convert_to_tensor"])
-def convert_to_tensor_v1(value,
-                         dtype=None,
-                         name=None,
-                         preferred_dtype=None,
-                         dtype_hint=None):
+@dispatch.add_dispatch_support
+def convert_to_tensor_v1_with_dispatch(
+    value,
+    dtype=None,
+    name=None,
+    preferred_dtype=None,
+    dtype_hint=None):
   """Converts the given `value` to a `Tensor`.
 
   This function converts Python objects of various types to `Tensor`
@@ -1311,13 +1314,26 @@ def convert_to_tensor_v1(value,
     RuntimeError: If a registered conversion function returns an invalid value.
     ValueError: If the `value` is a tensor not of given `dtype` in graph mode.
   """
+  return convert_to_tensor_v1(value, dtype=dtype, name=name,
+                              preferred_dtype=preferred_dtype,
+                              dtype_hint=dtype_hint)
+
+
+def convert_to_tensor_v1(value,
+                         dtype=None,
+                         name=None,
+                         preferred_dtype=None,
+                         dtype_hint=None):
+  """Converts the given `value` to a `Tensor` (with the TF1 API)."""
   preferred_dtype = deprecation.deprecated_argument_lookup(
       "dtype_hint", dtype_hint, "preferred_dtype", preferred_dtype)
   return convert_to_tensor_v2(value, dtype, preferred_dtype, name)
 
 
 @tf_export("convert_to_tensor", v1=[])
-def convert_to_tensor_v2(value, dtype=None, dtype_hint=None, name=None):
+@dispatch.add_dispatch_support
+def convert_to_tensor_v2_with_dispatch(
+    value, dtype=None, dtype_hint=None, name=None):
   """Converts the given `value` to a `Tensor`.
 
   This function converts Python objects of various types to `Tensor`
@@ -1378,6 +1394,12 @@ def convert_to_tensor_v2(value, dtype=None, dtype_hint=None, name=None):
     RuntimeError: If a registered conversion function returns an invalid value.
     ValueError: If the `value` is a tensor not of given `dtype` in graph mode.
   """
+  return convert_to_tensor_v2(
+      value, dtype=dtype, dtype_hint=dtype_hint, name=name)
+
+
+def convert_to_tensor_v2(value, dtype=None, dtype_hint=None, name=None):
+  """Converts the given `value` to a `Tensor`."""
   return convert_to_tensor(
       value=value,
       dtype=dtype,
diff --git a/tensorflow/python/util/dispatch_test.py b/tensorflow/python/util/dispatch_test.py
index 49026a754e4..f6074ac415d 100644
--- a/tensorflow/python/util/dispatch_test.py
+++ b/tensorflow/python/util/dispatch_test.py
@@ -27,6 +27,7 @@ from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util import deprecation
 from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import get_canonical_name_for_symbol
 from tensorflow.python.util.tf_export import tf_export
 
 
@@ -76,7 +77,8 @@ class TensorTracerOpDispatcher(dispatch.GlobalOpDispatcher):
             any(self.is_tensor_tracer_arg(x) for x in kwargs.values())):
       return self.NOT_SUPPORTED
 
-    return TensorTracer(op.__name__, args, kwargs)
+    symbol_name = get_canonical_name_for_symbol(op)
+    return TensorTracer(symbol_name, args, kwargs)
 
   def is_tensor_tracer_arg(self, value):
     if isinstance(value, TensorTracer):
@@ -183,12 +185,29 @@ class DispatchTest(test_util.TensorFlowTestCase):
       y = TensorTracer("y")
       trace = math_ops.reduce_sum(math_ops.add(math_ops.abs(x), y), axis=3)
       self.assertEqual(
-          str(trace), "reduce_sum(add(name=None, x=abs(x), y=y), axis=3)")
+          str(trace),
+          "math.reduce_sum(math.add(name=None, x=math.abs(x), y=y), axis=3)")
 
     finally:
       # Clean up.
       dispatch._GLOBAL_DISPATCHERS = original_global_dispatchers
 
+  def testGlobalDispatcherConvertToTensor(self):
+    original_global_dispatchers = dispatch._GLOBAL_DISPATCHERS
+    try:
+      TensorTracerOpDispatcher().register()
+
+      x = TensorTracer("x")
+      y = TensorTracer("y")
+      trace = math_ops.add(math_ops.abs(
+          ops.convert_to_tensor_v2_with_dispatch(x)), y)
+      self.assertEqual(
+          str(trace),
+          "math.add(name=None, x=math.abs(convert_to_tensor(x)), y=y)")
+
+    finally:
+      # Clean up.
+      dispatch._GLOBAL_DISPATCHERS = original_global_dispatchers
 
 if __name__ == "__main__":
   googletest.main()

From a5d8f188bcc4d5ffe9d5bfcc0fe1de411fb214d9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 1 Jul 2020 02:01:38 -0700
Subject: [PATCH 1381/1390] Update GraphDef version to 449.

PiperOrigin-RevId: 319185220
Change-Id: I76ff23cadc4e0d9435c63a69929a1d8d993146a6
---
 tensorflow/core/public/version.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 9d183076374..308b29f9ed5 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -108,7 +108,7 @@ limitations under the License.
 
 #define TF_GRAPH_DEF_VERSION_MIN_PRODUCER 0
 #define TF_GRAPH_DEF_VERSION_MIN_CONSUMER 0
-#define TF_GRAPH_DEF_VERSION 448  // Updated: 2020/6/30
+#define TF_GRAPH_DEF_VERSION 449  // Updated: 2020/7/1
 
 // Checkpoint compatibility versions (the versions field in SavedSliceMeta).
 //

From 6fc54250a5f3142cb1108dd53d3b7f291ba2e4f0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 1 Jul 2020 02:01:39 -0700
Subject: [PATCH 1382/1390] compat: Update forward compatibility horizon to
 2020-07-01

PiperOrigin-RevId: 319185221
Change-Id: I3a9b75bf05016c3abcead58de570d0600ba6a4cf
---
 tensorflow/python/compat/compat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/compat/compat.py b/tensorflow/python/compat/compat.py
index f1935b80ed5..80a087b0cfe 100644
--- a/tensorflow/python/compat/compat.py
+++ b/tensorflow/python/compat/compat.py
@@ -33,7 +33,7 @@ from tensorflow.python.util.tf_export import tf_export
 # This value changes every day with an automatic CL. It can be modified in code
 # via `forward_compatibility_horizon()` or with the environment variable
 # TF_FORWARD_COMPATIBILITY_DELTA_DAYS, which is added to the compatibility date.
-_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 6, 30)
+_FORWARD_COMPATIBILITY_HORIZON = datetime.date(2020, 7, 1)
 _FORWARD_COMPATIBILITY_DELTA_DAYS_VAR_NAME = "TF_FORWARD_COMPATIBILITY_DELTA_DAYS"
 _FORWARD_COMPATIBILITY_DATE_NUMBER = None
 

From fc49cbb2adf0b69f843f7bf904978648bffbe268 Mon Sep 17 00:00:00 2001
From: Taehee Jeong 
Date: Wed, 1 Jul 2020 02:26:51 -0700
Subject: [PATCH 1383/1390] Fix bug on referencing invalid reference to a
 tensor.

Adding tensors to the TfLiteContext will give TfLiteContext.tensors a new address, so the old references should be updated with new one.

PiperOrigin-RevId: 319188033
Change-Id: I1538d6260236f7cf5a710621d6af330a8639f443
---
 tensorflow/lite/delegates/gpu/common/object_reader.cc | 6 +++++-
 tensorflow/lite/delegates/utils.cc                    | 2 +-
 tensorflow/lite/delegates/utils.h                     | 3 ++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/common/object_reader.cc b/tensorflow/lite/delegates/gpu/common/object_reader.cc
index 55a0aea01a1..f299232a8a7 100644
--- a/tensorflow/lite/delegates/gpu/common/object_reader.cc
+++ b/tensorflow/lite/delegates/gpu/common/object_reader.cc
@@ -37,7 +37,7 @@ absl::Status ObjectReader::ReadNonConstantTensor(
   }
 
   if (tensor_to_value->find(tensor_idx) == tensor_to_value->end()) {
-    const TfLiteTensor& tflite_tensor = context->tensors[tensor_idx];
+    TfLiteTensor& tflite_tensor = context->tensors[tensor_idx];
     if (tflite::IsConstantTensor(&tflite_tensor)) {
       return absl::InvalidArgumentError(absl::StrCat(
           "ReadNonConstantTensor: value is a constant tensor: ", tensor_idx));
@@ -58,6 +58,7 @@ absl::Status ObjectReader::ReadNonConstantTensor(
                 &fp_tensor_index) != kTfLiteOk) {
           return absl::InternalError("Could not add new tensor to graph");
         }
+
         // Remember this tensor for later.
         (*quant_conversion_map)[fp_tensor_index] = tensor_idx;
         (*quant_conversion_map)[tensor_idx] = fp_tensor_index;
@@ -67,6 +68,9 @@ absl::Status ObjectReader::ReadNonConstantTensor(
             ConvertTfLiteTensorToTensorRef(*fp_tflite_tensor, &value->tensor));
         value->tensor.ref = fp_tensor_index;
         value->quant_params.emplace();
+        // tflite_tensor from the outer scope is invalidated due to calling
+        // CreateNewTensorWithDifferentType
+        tflite_tensor = context->tensors[tensor_idx];
         RETURN_IF_ERROR(
             PopulateQuantParams(tflite_tensor, &value->quant_params.value()));
         (*tensor_to_value)[fp_tensor_index] = value;
diff --git a/tensorflow/lite/delegates/utils.cc b/tensorflow/lite/delegates/utils.cc
index 873cadc180f..289586c5346 100644
--- a/tensorflow/lite/delegates/utils.cc
+++ b/tensorflow/lite/delegates/utils.cc
@@ -29,8 +29,8 @@ TfLiteStatus CreateNewTensorWithDifferentType(TfLiteContext* context,
                                               TfLiteType new_type,
                                               TfLiteTensor** new_tensor,
                                               int* new_tensor_index) {
-  const TfLiteTensor& original_tensor = context->tensors[original_tensor_index];
   TF_LITE_ENSURE_STATUS(context->AddTensors(context, 1, new_tensor_index));
+  const TfLiteTensor& original_tensor = context->tensors[original_tensor_index];
   *new_tensor = &context->tensors[*new_tensor_index];
   (*new_tensor)->type = new_type;
   (*new_tensor)->allocation_type = kTfLiteArenaRw;
diff --git a/tensorflow/lite/delegates/utils.h b/tensorflow/lite/delegates/utils.h
index 12684fcb84a..a9fb67316fc 100644
--- a/tensorflow/lite/delegates/utils.h
+++ b/tensorflow/lite/delegates/utils.h
@@ -33,7 +33,8 @@ namespace tflite {
 namespace delegates {
 
 // Creates a new Read/Write tensor having the same shape as the original, but
-// with a different type.
+// with a different type. Note that this might void existing references to
+// tensors.
 TfLiteStatus CreateNewTensorWithDifferentType(TfLiteContext* context,
                                               const int original_tensor_index,
                                               TfLiteType new_type,

From 4fba4cbfdc098b4192622654fa74848c45755278 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 1 Jul 2020 02:49:04 -0700
Subject: [PATCH 1384/1390] Modify InitializeTableFromDatasetOp to be async and
 add the table init op to the table initializers collection

PiperOrigin-RevId: 319190503
Change-Id: Ibd94b6f194e839fc11e37f1153c43533462bc265
---
 tensorflow/core/kernels/BUILD                 |  1 +
 .../core/kernels/lookup_table_init_op.cc      | 21 +++---
 tensorflow/core/kernels/lookup_util.cc        | 71 +++++++++++--------
 tensorflow/core/kernels/lookup_util.h         |  7 +-
 .../python/kernel_tests/lookup_ops_test.py    | 12 ++++
 tensorflow/python/ops/lookup_ops.py           |  1 +
 6 files changed, 71 insertions(+), 42 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index bd540baa65a..908deb06d0e 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -417,6 +417,7 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "//tensorflow/core/framework:op_requires",
     ],
 )
 
diff --git a/tensorflow/core/kernels/lookup_table_init_op.cc b/tensorflow/core/kernels/lookup_table_init_op.cc
index 49744cea59e..7bffb5ac547 100644
--- a/tensorflow/core/kernels/lookup_table_init_op.cc
+++ b/tensorflow/core/kernels/lookup_table_init_op.cc
@@ -164,24 +164,29 @@ REGISTER_KERNEL_BUILDER(
     Name("InitializeTableFromTextFileV2").Device(DEVICE_CPU),
     InitializeTableFromTextFileOp);
 
-class InitializeTableFromDatasetOp : public OpKernel {
+class InitializeTableFromDatasetOp : public AsyncOpKernel {
  public:
   explicit InitializeTableFromDatasetOp(OpKernelConstruction* ctx)
-      : OpKernel(ctx) {}
+      : AsyncOpKernel(ctx),
+        background_worker_(ctx->env(), "initialize_table_from_dataset") {}
 
-  void Compute(OpKernelContext* ctx) override {
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
     lookup::InitializableLookupTable* table;
-    OP_REQUIRES_OK(ctx,
-                   GetInitializableLookupTable("table_handle", ctx, &table));
+    OP_REQUIRES_OK_ASYNC(
+        ctx, GetInitializableLookupTable("table_handle", ctx, &table), done);
     core::ScopedUnref unref_me(table);
     DatasetBase* dataset;
-    OP_REQUIRES_OK(ctx, GetDatasetFromVariantTensor(ctx->input(1), &dataset));
-    OP_REQUIRES_OK(ctx,
-                   lookup::InitializeTableFromDataset(ctx, dataset, table));
+    OP_REQUIRES_OK_ASYNC(
+        ctx, GetDatasetFromVariantTensor(ctx->input(1), &dataset), done);
+    background_worker_.Schedule([ctx, dataset, table, done]() {
+      lookup::InitializeTableFromDataset(ctx, dataset, table, done);
+    });
   }
 
  private:
   TF_DISALLOW_COPY_AND_ASSIGN(InitializeTableFromDatasetOp);
+
+  data::BackgroundWorker background_worker_;
 };
 
 REGISTER_KERNEL_BUILDER(Name("InitializeTableFromDataset").Device(DEVICE_CPU),
diff --git a/tensorflow/core/kernels/lookup_util.cc b/tensorflow/core/kernels/lookup_util.cc
index ef9bca4475c..802e1ff5c35 100644
--- a/tensorflow/core/kernels/lookup_util.cc
+++ b/tensorflow/core/kernels/lookup_util.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/function_handle_cache.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -451,46 +452,54 @@ class DatasetIterator : public InitializableLookupTable::InitTableIterator {
   Status status_;
 };
 
-Status InitializeTableFromDataset(OpKernelContext* ctx,
-                                  data::DatasetBase* dataset,
-                                  InitializableLookupTable* table) {
+void InitializeTableFromDataset(OpKernelContext* ctx,
+                                data::DatasetBase* dataset,
+                                InitializableLookupTable* table,
+                                AsyncOpKernel::DoneCallback done) {
   // Assert that the dataset types match up to that expected in the table.
   const auto& dataset_types = dataset->output_dtypes();
-  if (dataset_types.size() != 2) {
-    return errors::InvalidArgument("Dataset should have two output types only");
-  }
-  if (dataset_types[0] != table->key_dtype()) {
-    return errors::InvalidArgument("Key dtype expected: ", table->key_dtype(),
-                                   " but obtained: ", dataset_types[0],
-                                   " from the dataset");
-  }
-  if (dataset_types[1] != table->value_dtype()) {
-    return errors::InvalidArgument(
-        "Value dtype expected: ", table->value_dtype(),
-        " but obtained: ", dataset_types[1], " from the dataset");
-  }
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_types.size() == 2,
+      errors::InvalidArgument("Dataset should have two output types only"),
+      done);
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_types[0] == table->key_dtype(),
+      errors::InvalidArgument("Key dtype expected: ", table->key_dtype(),
+                              " but obtained: ", dataset_types[0],
+                              " from the dataset"),
+      done);
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_types[1] == table->value_dtype(),
+      errors::InvalidArgument("Value dtype expected: ", table->value_dtype(),
+                              " but obtained: ", dataset_types[1],
+                              " from the dataset"),
+      done);
   // Assert that the dataset output shapes are scalars.
   const auto& dataset_shapes = dataset->output_shapes();
-  if (dataset_shapes.size() != 2) {
-    return errors::InvalidArgument(
-        "Dataset should have two output shapes only");
-  }
-  if (!dataset_shapes[0].IsCompatibleWith(PartialTensorShape({}))) {
-    return errors::InvalidArgument("Expected scalar for key. Obtained: ",
-                                   dataset_shapes[0].DebugString());
-  }
-  if (!dataset_shapes[1].IsCompatibleWith(PartialTensorShape({}))) {
-    return errors::InvalidArgument("Expected scalar for key. Obtained: ",
-                                   dataset_shapes[1].DebugString());
-  }
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_shapes.size() == 2,
+      errors::InvalidArgument("Dataset should have two output shapes only"),
+      done);
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_shapes[0].IsCompatibleWith(PartialTensorShape({})),
+      errors::InvalidArgument("Expected scalar for key. Obtained: ",
+                              dataset_shapes[0].DebugString()),
+      done);
+  OP_REQUIRES_ASYNC(
+      ctx, dataset_shapes[1].IsCompatibleWith(PartialTensorShape({})),
+      errors::InvalidArgument("Expected scalar for key. Obtained: ",
+                              dataset_shapes[1].DebugString()),
+      done);
   DatasetIterator iter(dataset);
-  TF_RETURN_IF_ERROR(iter.Init(ctx));
+  OP_REQUIRES_OK_ASYNC(ctx, iter.Init(ctx), done);
   Status s = table->Initialize(iter);
   if (errors::IsFailedPrecondition(s) && table->is_initialized()) {
     LOG(INFO) << "Table already initialized from dataset.";
-    return Status::OK();
+    done();
+    return;
   }
-  return s;
+  ctx->SetStatus(s);
+  done();
 }
 
 }  // namespace lookup
diff --git a/tensorflow/core/kernels/lookup_util.h b/tensorflow/core/kernels/lookup_util.h
index 97893d0c17a..7e53ed5db51 100644
--- a/tensorflow/core/kernels/lookup_util.h
+++ b/tensorflow/core/kernels/lookup_util.h
@@ -58,9 +58,10 @@ Status InitializeTableFromTextFile(const string& filename, int64 vocab_size,
 
 // Initializes `table` from `dataset` by iterating over it. Caller retains
 // ownership of `dataset`.
-Status InitializeTableFromDataset(OpKernelContext* ctx,
-                                  data::DatasetBase* dataset,
-                                  InitializableLookupTable* table);
+void InitializeTableFromDataset(OpKernelContext* ctx,
+                                data::DatasetBase* dataset,
+                                InitializableLookupTable* table,
+                                AsyncOpKernel::DoneCallback done);
 
 }  // namespace lookup
 }  // namespace tensorflow
diff --git a/tensorflow/python/kernel_tests/lookup_ops_test.py b/tensorflow/python/kernel_tests/lookup_ops_test.py
index 9b237b258a9..59afb2c27ab 100644
--- a/tensorflow/python/kernel_tests/lookup_ops_test.py
+++ b/tensorflow/python/kernel_tests/lookup_ops_test.py
@@ -565,6 +565,18 @@ class DatasetInitializerTest(BaseLookupTableTest):
     result = self.evaluate(output)
     self.assertAllEqual([1, 2, -1], result)
 
+  def test_compatibility(self):
+    with ops.Graph().as_default():
+      keys = dataset_ops.Dataset.range(100)
+      values = dataset_ops.Dataset.range(100).map(string_ops.as_string)
+      ds = dataset_ops.Dataset.zip((keys, values))
+      init = lookup_ops.DatasetInitializer(ds)
+      table = self.getHashTable()(init, default_value="")
+      output = table.lookup(constant_op.constant([0, 2, 5], dtypes.int64))
+      self.evaluate(lookup_ops.tables_initializer())
+      result = self.evaluate(output)
+    self.assertAllEqual(["0", "2", "5"], result)
+
 
 class InitializeTableFromFileOpTest(BaseLookupTableTest):
 
diff --git a/tensorflow/python/ops/lookup_ops.py b/tensorflow/python/ops/lookup_ops.py
index 96f3cf91499..87b8aaa30bd 100644
--- a/tensorflow/python/ops/lookup_ops.py
+++ b/tensorflow/python/ops/lookup_ops.py
@@ -468,6 +468,7 @@ class DatasetInitializer(TableInitializerBase):
     _check_table_dtypes(table, self._key_dtype, self._value_dtype)
     init_op = gen_lookup_ops.initialize_table_from_dataset(
         table.resource_handle, self.dataset._variant_tensor)  # pylint: disable=protected-access
+    ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op)
     return init_op
 
 

From 3bdb3a75bd17f0f7f15d1b492f38464a70190974 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" 
Date: Wed, 1 Jul 2020 03:26:11 -0700
Subject: [PATCH 1385/1390] Reword comment in eager runtime.

PiperOrigin-RevId: 319194534
Change-Id: I9cc602c9eefa9e647ffa74db01109db2e90a9509
---
 tensorflow/core/common_runtime/eager/context.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h
index 68f618adbec..5a14ebdfda7 100644
--- a/tensorflow/core/common_runtime/eager/context.h
+++ b/tensorflow/core/common_runtime/eager/context.h
@@ -377,8 +377,8 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
   // class/struct.
   //
   // Enables the eager context to communicate with remote devices. When
-  // initializing with this method, this context will be the master context,
-  // which will kill all its slaves in shutdown.
+  // initializing with this method, this context will be the primary context,
+  // which will kill all its remote contexts in shutdown.
   //
   // - server: A ServerInterface that exports the tensorflow.WorkerService.
   // Note that this class expects the server to already have been started.

From 8802516b56e190cba5846f7b7dfca7a0902bcf03 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev 
Date: Wed, 1 Jul 2020 03:52:55 -0700
Subject: [PATCH 1386/1390] Integrate LLVM at
 https://github.com/llvm/llvm-project/commit/2501e86acda2

PiperOrigin-RevId: 319196952
Change-Id: I078a64a0b84eb9cd8f3c5d277ad30c943b58fd1c
---
 tensorflow/compiler/mlir/xla/tests/BUILD           |  4 ++++
 .../mlir/xla/transforms/lhlo_legalize_to_llvm.cc   |  6 ++++--
 tensorflow/workspace.bzl                           |  4 ++--
 third_party/llvm/llvm.autogenerated.BUILD          | 14 +++++++++++++-
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/tests/BUILD b/tensorflow/compiler/mlir/xla/tests/BUILD
index 87d26728f3d..1ad83b5ea4a 100644
--- a/tensorflow/compiler/mlir/xla/tests/BUILD
+++ b/tensorflow/compiler/mlir/xla/tests/BUILD
@@ -6,6 +6,10 @@ package(licenses = ["notice"])
 glob_lit_tests(
     data = [":test_utilities"],
     driver = "@llvm-project//mlir:run_lit.sh",
+    exclude = [
+        # TODO(b/160227541): Re-enable LHLO->LLVM lowering.
+        "lhlo-legalize-to-llvm.mlir",
+    ],
     test_file_exts = ["mlir"],
 )
 
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
index 99d2c08aa98..4be175b8afa 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
@@ -128,8 +128,10 @@ struct DynamicMemRefCastOpConverter
 
 void PopulateLhloToLLVMConversionPatterns(LLVMTypeConverter *converter,
                                           OwningRewritePatternList *patterns) {
-  patterns->insert(
-      *converter);
+  // TODO(b/160227541): Re-enable LHLO->LLVM lowering.
+  //  patterns->insert(
+  //     *converter);
 }
 
 }  // namespace xla_lhlo
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 9435ef96bd1..4f1ff3e7bc6 100755
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -710,8 +710,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
     )
 
     # Check out LLVM and MLIR from llvm-project.
-    LLVM_COMMIT = "e34523c87c3f1cfabcf741568dede026bbb12d3a"
-    LLVM_SHA256 = "04e82e8fa5d492dc4c298d538e48e461c711bed3d58ab6f4dbc8aa735765ac4b"
+    LLVM_COMMIT = "2501e86acda2905e50012f7e9fc1942517c1237d"
+    LLVM_SHA256 = "3428d4f4806c80745be4a035167ee5b32a67533a85db61d937c56b28a26bfbf2"
     LLVM_URLS = [
         "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
         "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT),
diff --git a/third_party/llvm/llvm.autogenerated.BUILD b/third_party/llvm/llvm.autogenerated.BUILD
index 85efc0db65e..6e53745166d 100644
--- a/third_party/llvm/llvm.autogenerated.BUILD
+++ b/third_party/llvm/llvm.autogenerated.BUILD
@@ -687,7 +687,18 @@ cc_library(
 
 gentbl(
     name = "omp_gen",
-    tbl_outs = [("--gen-directive-decls", "include/llvm/Frontend/OpenMP/OMP.h.inc")],
+    tbl_outs = [("--gen-directive-decl", "include/llvm/Frontend/OpenMP/OMP.h.inc")],
+    tblgen = ":llvm-tblgen",
+    td_file = "include/llvm/Frontend/OpenMP/OMP.td",
+    td_srcs = glob([
+        "include/llvm/Frontend/OpenMP/*.td",
+        "include/llvm/Frontend/Directive/*.td",
+    ]),
+)
+
+gentbl(
+    name = "omp_gen_impl",
+    tbl_outs = [("--gen-directive-impl", "include/llvm/Frontend/OpenMP/OMP.cpp.inc")],
     tblgen = ":llvm-tblgen",
     td_file = "include/llvm/Frontend/OpenMP/OMP.td",
     td_srcs = glob([
@@ -2092,6 +2103,7 @@ cc_library(
         ":TransformUtils",
         ":config",
         ":omp_gen",
+        ":omp_gen_impl",
     ],
 )
 

From b09a0f8b1417b1f1522e777feb6f602e847ab189 Mon Sep 17 00:00:00 2001
From: Chao Mei 
Date: Wed, 1 Jul 2020 03:54:17 -0700
Subject: [PATCH 1387/1390] Fix the compilation error by not using designated
 initializers.

PiperOrigin-RevId: 319197094
Change-Id: Ibe90c15087a341668383b2862578e8d50608508d
---
 .../lite/delegates/external/external_delegate.cc     | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/delegates/external/external_delegate.cc b/tensorflow/lite/delegates/external/external_delegate.cc
index 5df158942f2..0ebfb62421c 100644
--- a/tensorflow/lite/delegates/external/external_delegate.cc
+++ b/tensorflow/lite/delegates/external/external_delegate.cc
@@ -220,11 +220,13 @@ TfLiteStatus TfLiteExternalDelegateOptionsInsert(
 
 TfLiteExternalDelegateOptions TfLiteExternalDelegateOptionsDefault(
     const char* lib_path) {
-  TfLiteExternalDelegateOptions options = {
-      .lib_path = lib_path,
-      .count = 0,
-      .insert = TfLiteExternalDelegateOptionsInsert,
-  };
+  // As 'keys' and 'values' don't need to be set here, using designated
+  // initializers may cause a compiling error as "non-trivial designated
+  // initializers not supported" by some compiler.
+  TfLiteExternalDelegateOptions options;
+  options.lib_path = lib_path;
+  options.count = 0;
+  options.insert = TfLiteExternalDelegateOptionsInsert;
   return options;
 }
 

From 0e1ba4145096be86dfa1f32ebd6e0c11156491b1 Mon Sep 17 00:00:00 2001
From: Christian Sigg 
Date: Wed, 1 Jul 2020 04:44:37 -0700
Subject: [PATCH 1388/1390] Bump minimum supported cuDNN version to 7.3.

Remove code to support older versions.

PiperOrigin-RevId: 319202156
Change-Id: I297d7950a93b6b802d07ae0f89232f26a5f2592e
---
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 246 +-------------------
 1 file changed, 9 insertions(+), 237 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index a97850bd8d5..28902b65722 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -61,7 +61,7 @@ PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuDnnPlugin);
 
 namespace {
 
-static_assert(CUDNN_VERSION >= 6000, "cuDNN needs to be version 6.0 or higher");
+static_assert(CUDNN_VERSION >= 7300, "cuDNN needs to be version 7.3 or higher");
 
 // Exits the program if 'expr' doesn't return CUDNN_STATUS_SUCCESS.
 #define CHECK_CUDNN_OK(expr) CHECK_EQ(expr, CUDNN_STATUS_SUCCESS)
@@ -115,12 +115,10 @@ std::string ToString(cudnnStatus_t status) {
       return "CUDNN_STATUS_LICENSE_ERROR";
     case CUDNN_STATUS_RUNTIME_PREREQUISITE_MISSING:
       return "CUDNN_STATUS_RUNTIME_PREREQUISITE_MISSING";
-#if CUDNN_VERSION >= 7000
     case CUDNN_STATUS_RUNTIME_IN_PROGRESS:
       return "CUDNN_STATUS_RUNTIME_IN_PROGRESS";
     case CUDNN_STATUS_RUNTIME_FP_OVERFLOW:
       return "CUDNN_STATUS_RUNTIME_FP_OVERFLOW";
-#endif
     default:
       return absl::StrCat("(status),
                           ">");
@@ -309,12 +307,11 @@ port::Status CudnnSupport::Init() {
       const std::string error = absl::StrCat(
           "Loaded runtime CuDNN library: ", loaded_version.ToString(),
           " but source was compiled with: ", source_version.ToString(),
-          ".  CuDNN library major and minor version needs to match or have "
-          "higher minor version in case of CuDNN 7.0 or later version. If "
-          "using a binary install, upgrade your CuDNN library.  If building "
-          "from sources, make sure the library loaded at runtime is "
-          "compatible "
-          "with the version specified during compile configuration.");
+          ".  CuDNN library needs to have matching major version and equal or "
+          "higher minor version. If using a binary install, upgrade your CuDNN "
+          "library.  If building from sources, make sure the library loaded at "
+          "runtime is compatible with the version specified during compile "
+          "configuration.");
       LOG(ERROR) << error;
       cudnnDestroy(cudnn_handle);
       return port::Status(port::error::INTERNAL, error);
@@ -359,13 +356,11 @@ struct TensorDescriptorDeleter {
     CHECK_CUDNN_OK(cudnnDestroyTensorDescriptor(descriptor));
   }
 };
-#if CUDNN_VERSION >= 7201
 struct RNNDataDescriptorDeleter {
   void operator()(cudnnRNNDataDescriptor_t descriptor) const {
     CHECK_CUDNN_OK(cudnnDestroyRNNDataDescriptor(descriptor));
   }
 };
-#endif
 struct FilterDescriptorDeleter {
   void operator()(cudnnFilterDescriptor_t descriptor) const {
     CHECK_CUDNN_OK(cudnnDestroyFilterDescriptor(descriptor));
@@ -418,10 +413,8 @@ struct CtcLossDescriptorDeleter {
 // RAII wrappers for cuDNN types.
 using TensorDescriptor =
     std::unique_ptr;
-#if CUDNN_VERSION >= 7201
 using RNNDataDescriptor =
     std::unique_ptr;
-#endif
 using FilterDescriptor =
     std::unique_ptr;
 using ConvolutionDescriptor =
@@ -447,13 +440,11 @@ TensorDescriptor CreateTensorDescriptor() {
   CHECK_CUDNN_OK(cudnnCreateTensorDescriptor(&result));
   return TensorDescriptor(result);
 }
-#if CUDNN_VERSION >= 7201
 RNNDataDescriptor CreateRNNDataDescriptor() {
   cudnnRNNDataDescriptor_t result;
   CHECK_CUDNN_OK(cudnnCreateRNNDataDescriptor(&result));
   return RNNDataDescriptor(result);
 }
-#endif
 FilterDescriptor CreateFilterDescriptor() {
   cudnnFilterDescriptor_t result;
   CHECK_CUDNN_OK(cudnnCreateFilterDescriptor(&result));
@@ -718,7 +709,6 @@ class CudnnConvolutionDescriptor {
   }
 
   void set_use_tensor_op_math(bool use_tensor_op_math) {
-#if CUDNN_VERSION >= 7000
     cudnnMathType_t math_type =
 #if CUDNN_VERSION >= 8000
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH);
@@ -726,7 +716,6 @@ class CudnnConvolutionDescriptor {
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
 #endif
     CHECK_CUDNN_OK(cudnnSetConvolutionMathType(handle_.get(), math_type));
-#endif
   }
 
   cudnnConvolutionDescriptor_t handle() const { return handle_.get(); }
@@ -749,9 +738,7 @@ static bool IsTensorMathOpSet(const CudnnConvolutionDescriptor& conv) {
 #endif
 }
 
-static bool TensorOpMathAvailable(int cc_major) {
-  return cc_major >= 7 && CUDNN_VERSION >= 7000;
-}
+static bool TensorOpMathAvailable(int cc_major) { return cc_major >= 7; }
 
 static bool IsTensorMathAllowed(Stream* stream, dnn::DataType input_type) {
   int cc_major, cc_minor;
@@ -862,11 +849,9 @@ class CudnnActivationDescriptor {
     double relu_ceiling = 0.0;
     cudnnActivationMode_t mode;
     switch (activation_mode) {
-#if CUDNN_VERSION >= 7100
       case dnn::ActivationMode::kNone:
         mode = CUDNN_ACTIVATION_IDENTITY;
         break;
-#endif
       case dnn::ActivationMode::kRelu6:
         relu_ceiling = 6.0;
         mode = CUDNN_ACTIVATION_CLIPPED_RELU;
@@ -1113,26 +1098,18 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
         /*direction=*/direction_mode, /*mode=*/rnn_mode, /*algo=*/rnn_algo,
         /*dataType=*/compute_type));
     if (use_projection) {
-#if CUDNN_VERSION >= 7101
       RETURN_IF_CUDNN_ERROR(cudnnSetRNNProjectionLayers(
           cudnn.handle(), /*rnnDesc=*/rnn_desc.get(),
           /*recProjSize=*/hidden_size, /*outProjSize=*/0));
-#else
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "No supported cudnnSetRNNProjectionLayers when "
-                          "CUDNN_VERSION < 7.1.1");
-#endif
     }
 
     // TODO: For now, we only use cudnnRNN**Ex API to process padded inputs.
     // But in the future if these APIs are used to process full length arrays,
     // we need to distinguish when to set it.
-#if CUDNN_VERSION >= 7201
     if (use_padded_io) {
       RETURN_IF_CUDNN_ERROR(
           cudnnSetRNNPaddingMode(rnn_desc.get(), CUDNN_RNN_PADDED_IO_ENABLED));
     }
-#endif
 
     port::StatusOr rnn_plan_wrapper;
     PersistentRnnPlan rnn_plan;
@@ -1155,7 +1132,6 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
                             cudnn, input_size, data_type, rnn_desc.get(),
                             rnn_mode, direction_mode, num_layers));
 
-#if CUDNN_VERSION >= 7000
     // Require explicit algorithm config to enable tensor cores. Some configs
     // return CUDNN_NOT_SUPPORTED when tensor ops are enabled (which is against
     // the idiom that enabling tensor ops is only a hint: see nvbugs/2172799).
@@ -1169,7 +1145,7 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
     if (algorithm_config.algorithm().has_value()) {
       use_tensor_ops = algorithm_config.algorithm()->tensor_ops_enabled();
     } else {
-      use_tensor_ops = CUDNN_VERSION >= 7201 && allow_tensor_ops;
+      use_tensor_ops = allow_tensor_ops;
     }
 
     if (use_tensor_ops && !allow_tensor_ops) {
@@ -1184,7 +1160,6 @@ class CudnnRnnDescriptor : public dnn::RnnDescriptor {
     math_type = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
 #endif
     CHECK_CUDNN_OK(cudnnSetRNNMatrixMathType(rnn_desc.get(), math_type));
-#endif  // CUDNN_VERSION >= 7000
 
     return CudnnRnnDescriptor(cudnn, std::move(rnn_desc), std::move(rnn_plan),
                               num_layers, hidden_size, input_size, cell_size,
@@ -1281,7 +1256,6 @@ port::Status CheckAndFetchProjectionWeights(
     const TensorDescriptor& input_desc, const FilterDescriptor& filter_desc,
     const FilterDescriptor& region_desc_handle,
     dnn::RnnDescriptor::ParamsRegions* weights) {
-#if CUDNN_VERSION >= 7101
   int hidden_size_v;
   int num_layers_v;
   cudnnDropoutDescriptor_t dropout_desc;
@@ -1345,7 +1319,6 @@ port::Status CheckAndFetchProjectionWeights(
                                                size};
     weights->push_back(region);
   }
-#endif  // CUDNN_VERSION >= 7101
   return port::Status::OK();
 }
 
@@ -1452,18 +1425,14 @@ class CudnnRnnSequenceTensorDescriptor
   CudnnRnnSequenceTensorDescriptor(GpuExecutor* parent, int max_seq_length,
                                    int batch_size, int data_size,
                                    cudnnDataType_t data_type,
-#if CUDNN_VERSION >= 7201
                                    RNNDataDescriptor data_handle,
-#endif
                                    TensorDescriptor handle)
       : max_seq_length_(max_seq_length),
         batch_size_(batch_size),
         data_size_(data_size),
         data_type_(data_type),
         handle_(std::move(handle)),
-#if CUDNN_VERSION >= 7201
         rnn_data_handle_(std::move(data_handle)),
-#endif
         handles_(max_seq_length, handle_.get()) {
   }
 
@@ -1484,9 +1453,7 @@ class CudnnRnnSequenceTensorDescriptor
         /*strideA=*/strides));
     return CudnnRnnSequenceTensorDescriptor(parent, max_seq_length, batch_size,
                                             data_size, data_type,
-#if CUDNN_VERSION >= 7201
                                             nullptr,
-#endif
                                             std::move(tensor_desc));
   }
 
@@ -1494,7 +1461,6 @@ class CudnnRnnSequenceTensorDescriptor
       GpuExecutor* parent, int max_seq_length, int batch_size, int data_size,
       const absl::Span& seq_lengths, bool time_major,
       cudnnDataType_t data_type) {
-#if CUDNN_VERSION >= 7201
     CHECK_GT(max_seq_length, 0);
     int dims[] = {batch_size, data_size, 1};
     int strides[] = {dims[1] * dims[2], dims[2], 1};
@@ -1522,29 +1488,18 @@ class CudnnRnnSequenceTensorDescriptor
     return CudnnRnnSequenceTensorDescriptor(
         parent, max_seq_length, batch_size, data_size, data_type,
         std::move(data_desc), std::move(tensor_desc));
-#else
-    return port::Status(port::error::INVALID_ARGUMENT,
-                        "No supported cudnnSetRNNDataDescriptor when "
-                        "CUDNN_VERSION < 7.2.1");
-#endif
   }
 
   const cudnnTensorDescriptor_t* handles() const { return handles_.data(); }
-#if CUDNN_VERSION >= 7201
   const cudnnRNNDataDescriptor_t data_handle() const {
     return rnn_data_handle_.get();
   }
-#endif
 
   int max_seq_length() const { return max_seq_length_; }
   int batch_size() const { return batch_size_; }
   int data_size() const { return data_size_; }
   bool is_var_seq_lengths() const {
-#if CUDNN_VERSION >= 7201
     return rnn_data_handle_ != nullptr;
-#else
-    return false;
-#endif
   }
 
  private:
@@ -1553,9 +1508,7 @@ class CudnnRnnSequenceTensorDescriptor
   int data_size_;
   cudnnDataType_t data_type_;
   TensorDescriptor handle_;
-#if CUDNN_VERSION >= 7201
   RNNDataDescriptor rnn_data_handle_;
-#endif
   std::vector handles_;  // Copies of handle_.
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnRnnSequenceTensorDescriptor);
 };
@@ -1816,7 +1769,6 @@ port::Status CudnnSupport::DoRnnForwardImpl(
 
   if (!is_training) {
     if (input_desc.is_var_seq_lengths()) {
-#if CUDNN_VERSION >= 7201
       RETURN_IF_CUDNN_ERROR(cudnnRNNForwardInferenceEx(
           /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
           /*xDesc=*/input_desc.data_handle(), /*x=*/input_data.opaque(),
@@ -1831,11 +1783,6 @@ port::Status CudnnSupport::DoRnnForwardImpl(
           nullptr,
           /*workspace=*/workspace.opaque(),
           /*workSpaceSizeInBytes=*/workspace.size()));
-#else
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "No supported cudnnRNNForwardInferenceEx when "
-                          "CUDNN_VERSION < 7.2.1");
-#endif
     } else {
       RETURN_IF_CUDNN_ERROR(cudnnRNNForwardInference(
           /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
@@ -1852,7 +1799,6 @@ port::Status CudnnSupport::DoRnnForwardImpl(
     }
   } else {
     if (input_desc.is_var_seq_lengths()) {
-#if CUDNN_VERSION >= 7201
       // cudnnSetRNNPaddingMode(rnn_desc.handle(), CUDNN_RNN_PADDED_IO_ENABLED);
       RETURN_IF_CUDNN_ERROR(cudnnRNNForwardTrainingEx(
           /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
@@ -1870,11 +1816,6 @@ port::Status CudnnSupport::DoRnnForwardImpl(
           /*workSpaceSizeInBytes=*/workspace.size(),
           /*reserveSpace=*/reserve_space.opaque(),
           /*reserveSpaceSizeInBytes=*/reserve_space.size()));
-#else
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "No supported cudnnRNNForwardTrainingEx when "
-                          "CUDNN_VERSION < 7.2.1");
-#endif
     } else {
       RETURN_IF_CUDNN_ERROR(cudnnRNNForwardTraining(
           /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
@@ -1958,7 +1899,6 @@ port::Status CudnnSupport::DoRnnBackwardImpl(
   }
 
   if (input_desc.is_var_seq_lengths()) {
-#if CUDNN_VERSION >= 7201
     RETURN_IF_CUDNN_ERROR(cudnnRNNBackwardDataEx(
         /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
         /*yDesc=*/output_desc.data_handle(), /*y=*/output_data.opaque(),
@@ -1981,11 +1921,6 @@ port::Status CudnnSupport::DoRnnBackwardImpl(
         /*workSpaceSizeInBytes=*/workspace.size(),
         /*reserveSpace=*/reserve_space_data->opaque(),
         /*reserveSpaceSizeInBytes=*/reserve_space_data->size()));
-#else
-    return port::Status(port::error::INVALID_ARGUMENT,
-                        "No supported cudnnRNNBackwardDataEx when "
-                        "CUDNN_VERSION < 7.2.1");
-#endif
   } else {
     RETURN_IF_CUDNN_ERROR(cudnnRNNBackwardData(
         /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
@@ -2015,7 +1950,6 @@ port::Status CudnnSupport::DoRnnBackwardImpl(
     // Clear the dw to zeros.
     stream->ThenMemZero(params_backprop_data, params_backprop_data->size());
     if (input_desc.is_var_seq_lengths()) {
-#if CUDNN_VERSION >= 7201
       RETURN_IF_CUDNN_ERROR(cudnnRNNBackwardWeightsEx(
           /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
           /*xDesc=*/input_desc.data_handle(), /*x=*/input_data.opaque(),
@@ -2028,11 +1962,6 @@ port::Status CudnnSupport::DoRnnBackwardImpl(
           /*dw=*/params_backprop_data->opaque(),
           /*reserveSpace=*/reserve_space_data->opaque(),
           /*reserveSpaceSizeInBytes=*/reserve_space_data->size()));
-#else
-      return port::Status(port::error::INVALID_ARGUMENT,
-                          "No supported cudnnRNNBackwardWeightsEx when "
-                          "CUDNN_VERSION < 7.2.1");
-#endif
     } else {
       // make the backward weight call
       RETURN_IF_CUDNN_ERROR(cudnnRNNBackwardWeights(
@@ -2935,7 +2864,7 @@ class CudnnEnvVar {
 // algorithm through an env-var "TF_ENABLE_FFT_TILING_FORWARD=1".
 struct FftTilingForward {
   static constexpr const char* kName = "TF_ENABLE_FFT_TILING_FORWARD";
-  static constexpr bool kDefaultFlag = CUDNN_VERSION >= 7000;
+  static constexpr bool kDefaultFlag = true;
 };
 
 // A helper struct to decide whether to enable the WINOGRAD_NONFUSED algorithms.
@@ -3014,34 +2943,6 @@ dnn::DataType GetConvAccumulatorType(dnn::DataType data_type) {
       LOG(FATAL) << "Invalid DNN data type: " << static_cast(data_type);
   }
 }
-
-// Determines whether we can safely perform a winograd non-fused convolution for
-// the given input and output shapes.  This works around b/68264959, an integer
-// overflow in cuDNNv5 and cuDNNv6.
-#if CUDNN_VERSION >= 7000
-bool ShouldIncludeWinogradNonfusedAlgo(const dnn::BatchDescriptor&,
-                                       const dnn::BatchDescriptor&) {
-  return true;
-}
-#else
-bool ShouldIncludeWinogradNonfusedAlgo(
-    const dnn::BatchDescriptor& input_desc,
-    const dnn::BatchDescriptor& output_desc) {
-  int64 batch = input_desc.count();
-  int64 in_depths = input_desc.feature_map_count();
-  int64 in_rows = input_desc.height();
-  int64 in_cols = input_desc.ndims() == 1 ? 1 : input_desc.width();
-  int64 out_depths = output_desc.feature_map_count();
-
-  int64 total_size = port::MathUtil::CeilOfRatio(batch, int64{16}) *
-                     std::max(in_depths, out_depths) * in_cols * in_rows *
-                     sizeof(float);
-
-  const int64 threshold = 1L << 31;
-  return total_size < threshold;
-}
-#endif
-
 }  // namespace
 
 port::Status CudnnSupport::DoPrepareForConvolution(
@@ -3147,41 +3048,6 @@ port::Status CudnnSupport::DoConvolve(
   }
 
   const auto get_fwd_bugs = [&]() -> port::Status {
-    // Report an error if we might be hitting a cuDNN bug that accesses illegal
-    // memory. See nvbugs/2138754, b/80018418.
-    if (CUDNN_VERSION < 7300) {
-      if (algorithm_desc.algo_id() != CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING) {
-        return port::Status::OK();
-      }
-      if (input_descriptor.ndims() < 3) {
-        return port::Status::OK();
-      }
-      // Checks that a*b is within the valid range (as provided by NVIDIA).
-      const auto check_sizes = [](size_t a, size_t b) {
-        if ((a * b * 4608 - 1) >> 31 == 0) {
-          return port::Status::OK();
-        }
-        return port::Status(
-            port::error::FAILED_PRECONDITION,
-            "This configuration potentially accesses illegal memory.");
-      };
-      SE_RETURN_IF_ERROR(check_sizes(input_descriptor.feature_map_count(),
-                                     output_descriptor.feature_map_count()));
-      SE_RETURN_IF_ERROR(check_sizes(input_descriptor.count(),
-                                     input_descriptor.feature_map_count()));
-      SE_RETURN_IF_ERROR(check_sizes(input_descriptor.count(),
-                                     output_descriptor.feature_map_count()));
-      return port::Status::OK();
-    }
-    if (algorithm_desc.algo_id() ==
-            CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED &&
-        !ShouldIncludeWinogradNonfusedAlgo(input_descriptor,
-                                           output_descriptor)) {
-      return port::Status(
-          port::error::FAILED_PRECONDITION,
-          "This configuration has potential integer overflow in "
-          "cuDNNv5 and cuDNNv6. See b/68264959.");
-    }
     if (CUDNN_VERSION < 8000) {
       if (algorithm_desc.algo_id() ==
               CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM &&
@@ -3196,91 +3062,10 @@ port::Status CudnnSupport::DoConvolve(
   };
 
   auto get_bwd_data_bugs = [&]() -> port::Status {
-    if (algorithm_desc.algo_id() ==
-            CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED &&
-        !ShouldIncludeWinogradNonfusedAlgo(input_descriptor,
-                                           output_descriptor)) {
-      return port::Status(
-          port::error::FAILED_PRECONDITION,
-          "This configuration has potential integer overflow in "
-          "cuDNNv5 and cuDNNv6. See b/68264959.");
-    }
-
-    // Cudnn 7.1.4 has a bug if the workspace of the following convolution is
-    // not zero-initialized, nvbugs/2254619.
-    if (CUDNN_VERSION >= 7000 && CUDNN_VERSION < 7300 &&
-        algorithm_desc.algo_id() == CUDNN_CONVOLUTION_BWD_DATA_ALGO_1 &&
-        cudnn_type == CUDNN_DATA_HALF && algorithm_desc.tensor_ops_enabled() &&
-        input_descriptor.layout() == dnn::DataLayout::kBatchYXDepth &&
-        filter_descriptor.layout() == dnn::FilterLayout::kOutputInputYX &&
-        output_descriptor.layout() == dnn::DataLayout::kBatchDepthYX &&
-        (convolution_descriptor.vertical_filter_stride() > 1 ||
-         convolution_descriptor.horizontal_filter_stride() > 1)) {
-      stream->ThenMemZero(&scratch_memory, scratch_memory.size());
-    }
     return port::Status::OK();
   };
 
   const auto get_bwd_filter_bugs = [&]() -> port::Status {
-    // Report an error if we might be hitting a cuDNN bug that produces
-    // incorrect results. See nvbugs/2072856
-    if (CUDNN_VERSION < 7300) {
-      SE_RETURN_IF_ERROR([&] {
-        if (algorithm_desc.algo_id() !=
-            CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING) {
-          return port::Status::OK();
-        }
-        if (output_descriptor.height() > 1 && output_descriptor.width() > 1) {
-          return port::Status::OK();
-        }
-        int convolution_size = output_descriptor.height() > 1
-                                   ? filter_descriptor.input_filter_height()
-                                   : filter_descriptor.input_filter_width();
-        if (convolution_size <= 32) {
-          return port::Status::OK();
-        }
-        cudnnConvolutionMode_t convolution_mode;
-        cudnnDataType_t compute_type;
-        RETURN_IF_CUDNN_ERROR(cudnnGetConvolutionNdDescriptor(
-            conv.handle(), 0, nullptr, nullptr, nullptr, nullptr,
-            &convolution_mode, &compute_type));
-        if (convolution_mode != CUDNN_CONVOLUTION) {
-          return port::Status::OK();
-        }
-        return port::Status(
-            port::error::FAILED_PRECONDITION,
-            "This configuration potentially produces incorrect results.");
-      }());
-    }
-
-    if (algorithm_desc.algo_id() ==
-            CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED &&
-        !ShouldIncludeWinogradNonfusedAlgo(input_descriptor,
-                                           output_descriptor)) {
-      return port::Status(
-          port::error::FAILED_PRECONDITION,
-          "This configuration has potential integer overflow in "
-          "cuDNNv5 and cuDNNv6. See b/68264959.");
-    }
-
-    // Zero out the result buffer for strided conv backward filter for NHWC
-    // layouts. cuDNN 7.1.4 and 7.2 has non-determinisic bug if the buffer is
-    // not zeroed.
-    //
-    // This wrong result caused by the bug is very flaky. It needs to be run for
-    // up to 20 times to produce a mismatch.
-    //
-    // See nvbugs/2379553.
-    if (CUDNN_VERSION >= 7100 && CUDNN_VERSION < 7300 &&
-        algorithm_desc.algo_id() == CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1 &&
-        cudnn_type == CUDNN_DATA_HALF &&
-        input_descriptor.layout() == dnn::DataLayout::kBatchYXDepth &&
-        filter_descriptor.layout() == dnn::FilterLayout::kOutputYXInput &&
-        output_descriptor.layout() == dnn::DataLayout::kBatchYXDepth &&
-        (convolution_descriptor.vertical_filter_stride() > 1 ||
-         convolution_descriptor.horizontal_filter_stride() > 1)) {
-      stream->ThenMemZero(&filter_data, filter_data.size());
-    }
     return port::Status::OK();
   };
 
@@ -3439,13 +3224,6 @@ port::Status CudnnSupport::DoFusedConvolveImpl(
           << "\noutput_nd.handle() = " << output_nd.handle()
           << "\noutput_data->opaque() = " << output_data->opaque();
 
-  if (algo_desc.algo_id() == CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED &&
-      !ShouldIncludeWinogradNonfusedAlgo(conv_input_descriptor,
-                                         output_descriptor)) {
-    return port::Status(port::error::FAILED_PRECONDITION,
-                        "This configuration has potential integer overflow in "
-                        "cuDNNv5 and cuDNNv6. See around b/68264959.");
-  }
   if (IsTensorMathOpSet(conv) != algo_desc.tensor_ops_enabled()) {
     return port::Status(port::error::FAILED_PRECONDITION,
                         "Tensor op math type in dnn::AlgorithmDesc does not "
@@ -3531,9 +3309,7 @@ bool CudnnSupport::GetRnnAlgorithms(
   out_algorithms->clear();
   for (auto i : algo_types) {
     out_algorithms->push_back({i, /*use_tensor_ops=*/false});
-#if CUDNN_VERSION >= 7100
     out_algorithms->push_back({i, /*use_tensor_ops=*/true});
-#endif
   }
   return true;
 }
@@ -3672,11 +3448,9 @@ port::Status CudnnSupport::DoBatchNormalizationForwardImpl(
   CudnnTensorDescriptor scale_offset_descriptor(
       scale_offset_desc, ToCudnnDataType(scale_data_type));
   cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL;
-#if CUDNN_VERSION >= 7000
   if (BatchnormSpatialPersistentEnabled() && is_training) {
     mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
   }
-#endif
   float one = 1.0;
   float zero = 0.0;
   auto cudnn = cudnn_->GetHandle(parent_, stream);
@@ -3855,11 +3629,9 @@ port::Status CudnnSupport::DoBatchNormalizationBackwardImpl(
   CudnnTensorDescriptor scale_offset_descriptor(
       scale_offset_desc, static_cast(cudnn_scale_type));
   cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL;
-#if CUDNN_VERSION >= 7000
   if (BatchnormSpatialPersistentEnabled()) {
     mode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
   }
-#endif
   float one = 1.0;
   float zero = 0.0;
 

From 55aaf2124d6bdfd30b3e874a399b11c1a10a0e0e Mon Sep 17 00:00:00 2001
From: Nat Jeffries 
Date: Wed, 1 Jul 2020 05:17:47 -0700
Subject: [PATCH 1389/1390] Enable int8 input and int16 output for cmsis-nn
 softmax.

PiperOrigin-RevId: 319205673
Change-Id: Ibdf817b78266801e8a86bd20d0215fd6edf409d3
---
 .../lite/micro/kernels/cmsis-nn/softmax.cc    | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/tensorflow/lite/micro/kernels/cmsis-nn/softmax.cc b/tensorflow/lite/micro/kernels/cmsis-nn/softmax.cc
index 51a48ec8a93..2db8caba243 100644
--- a/tensorflow/lite/micro/kernels/cmsis-nn/softmax.cc
+++ b/tensorflow/lite/micro/kernels/cmsis-nn/softmax.cc
@@ -36,8 +36,15 @@ TfLiteStatus CalculateSoftmaxParams(TfLiteContext* context,
       TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
     } else {
       TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt8);
-      TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt8);
-      TF_LITE_ENSURE_EQ(context, output->params.zero_point, -128);
+      if (output->type == kTfLiteInt16) {
+        TF_LITE_ENSURE_EQ(context, output->params.zero_point, -32768);
+        // NOTE: Current int16 softmax output does not require symmetric scaling
+        // - so no need to verify scale here.
+      } else {
+        TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt8);
+        TF_LITE_ENSURE_EQ(context, output->params.zero_point, -128);
+        TF_LITE_ENSURE(context, output->params.scale == 1.f / 256);
+      }
     }
     TF_LITE_ENSURE(context, (output->params.scale == 1.f / 256) ||
                                 (output->params.scale == 1.f / 255));
@@ -90,17 +97,23 @@ void SoftmaxQuantized(const TfLiteTensor* input, TfLiteTensor* output,
                                    GetTensorData(input), output_shape,
                                    GetTensorData(output));
   } else {
-    const unsigned int num_dims = NumDimensions(input);
+    if (output->type == kTfLiteInt16) {
+      tflite::reference_ops::Softmax(
+          op_data, GetTensorShape(input), GetTensorData(input),
+          GetTensorShape(output), GetTensorData(output));
+    } else {
+      const unsigned int num_dims = NumDimensions(input);
 
-    const int trailing_dim = input_shape.DimensionsCount() - 1;
-    const int outer_size =
-        MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
-    const int depth =
-        MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
+      const int trailing_dim = input_shape.DimensionsCount() - 1;
+      const int outer_size =
+          MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape);
+      const int depth =
+          MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim);
 
-    arm_softmax_s8(GetTensorData(input), outer_size, depth,
-                   op_data.input_multiplier, op_data.input_left_shift,
-                   op_data.diff_min, GetTensorData(output));
+      arm_softmax_s8(GetTensorData(input), outer_size, depth,
+                     op_data.input_multiplier, op_data.input_left_shift,
+                     op_data.diff_min, GetTensorData(output));
+    }
   }
 }
 

From ed7033c7fc2787aa50fae345fc1be4030608b54f Mon Sep 17 00:00:00 2001
From: Alexander Belyaev 
Date: Wed, 1 Jul 2020 05:39:45 -0700
Subject: [PATCH 1390/1390] [MLIR][XLA] Re-enable LHLO->LLVM pass.

PiperOrigin-RevId: 319207613
Change-Id: If9f5b8c8d72eee419f92952901c84b01212aa34e
---
 tensorflow/compiler/mlir/xla/tests/BUILD                 | 4 ----
 .../mlir/xla/transforms/lhlo_legalize_to_llvm.cc         | 9 ++++-----
 .../mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc    | 3 ++-
 tensorflow/compiler/mlir/xla/transforms/rewriters.h      | 4 +++-
 4 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/mlir/xla/tests/BUILD b/tensorflow/compiler/mlir/xla/tests/BUILD
index 1ad83b5ea4a..87d26728f3d 100644
--- a/tensorflow/compiler/mlir/xla/tests/BUILD
+++ b/tensorflow/compiler/mlir/xla/tests/BUILD
@@ -6,10 +6,6 @@ package(licenses = ["notice"])
 glob_lit_tests(
     data = [":test_utilities"],
     driver = "@llvm-project//mlir:run_lit.sh",
-    exclude = [
-        # TODO(b/160227541): Re-enable LHLO->LLVM lowering.
-        "lhlo-legalize-to-llvm.mlir",
-    ],
     test_file_exts = ["mlir"],
 )
 
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
index 4be175b8afa..9f7f0fe6108 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm.cc
@@ -126,12 +126,11 @@ struct DynamicMemRefCastOpConverter
 
 }  // namespace
 
-void PopulateLhloToLLVMConversionPatterns(LLVMTypeConverter *converter,
+void PopulateLhloToLLVMConversionPatterns(const LowerToLLVMOptions &options,
+                                          LLVMTypeConverter *converter,
                                           OwningRewritePatternList *patterns) {
-  // TODO(b/160227541): Re-enable LHLO->LLVM lowering.
-  //  patterns->insert(
-  //     *converter);
+  patterns->insert(
+      *converter, options);
 }
 
 }  // namespace xla_lhlo
diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
index 63265c4a7e7..03a4f7320fe 100644
--- a/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_legalize_to_llvm_pass.cc
@@ -36,7 +36,8 @@ class TestLhloToLLVMPass
     OwningRewritePatternList patterns;
     LLVMTypeConverter converter(m.getContext());
     populateStdToLLVMConversionPatterns(converter, patterns);
-    PopulateLhloToLLVMConversionPatterns(&converter, &patterns);
+    PopulateLhloToLLVMConversionPatterns(
+        LowerToLLVMOptions::getDefaultOptions(), &converter, &patterns);
 
     ConversionTarget target(getContext());
     target.addLegalDialect();
diff --git a/tensorflow/compiler/mlir/xla/transforms/rewriters.h b/tensorflow/compiler/mlir/xla/transforms/rewriters.h
index 7303b87be75..8b25576cd14 100644
--- a/tensorflow/compiler/mlir/xla/transforms/rewriters.h
+++ b/tensorflow/compiler/mlir/xla/transforms/rewriters.h
@@ -24,6 +24,7 @@ limitations under the License.
 
 namespace mlir {
 class LLVMTypeConverter;
+class LowerToLLVMOptions;
 class OwningRewritePatternList;
 class BufferAssignmentPlacer;
 namespace xla_hlo {
@@ -77,7 +78,8 @@ void PopulateUnfuseBatchNormPatterns(MLIRContext *context,
 namespace xla_lhlo {
 
 /// Collect a set of patterns to convert from the LHLO dialect to LLVM.
-void PopulateLhloToLLVMConversionPatterns(LLVMTypeConverter *converter,
+void PopulateLhloToLLVMConversionPatterns(const LowerToLLVMOptions &options,
+                                          LLVMTypeConverter *converter,
                                           OwningRewritePatternList *patterns);
 
 }  // namespace xla_lhlo